mirror of
https://github.com/Sea-Haven-Industries/open-swe.git
synced 2026-10-07 16:19:09 +00:00
Merge a094ab099b into 0546085672
This commit is contained in:
seahaven-openswe[bot]
GitHub
commit
30abaf600e
5 changed files with 357 additions and 48 deletions
|
|
@ -1,22 +1,40 @@
|
|||
"""Middleware that falls back to a secondary model when the primary fails transiently.
|
||||
"""Middleware that retries model calls across a primary and fallback provider.
|
||||
|
||||
Wraps the model call. When the primary model raises a transient provider error
|
||||
(5xx, 429, connection/timeout), the same request is retried once against the
|
||||
configured fallback model. The fallback is bound to tools by the agent factory
|
||||
on the second call, so swapping ``request.model`` is sufficient.
|
||||
Wraps the model call. When a model raises a transient provider error (5xx,
|
||||
429, connection/timeout), the request is retried, alternating between the
|
||||
primary and the configured fallback model with exponential backoff between
|
||||
attempts. The fallback is bound to tools by the agent factory on each call,
|
||||
so swapping ``request.model`` is sufficient.
|
||||
|
||||
Why alternate with backoff instead of failing over once: both providers can
|
||||
be routed through the same LLM Gateway, so a gateway outage takes out the
|
||||
"cross-provider" fallback too. A single immediate failover cannot ride out
|
||||
even a short shared outage (the gateway's 502 page literally says "try again
|
||||
in 30 seconds"), and an unprotected fallback call crashes the whole run.
|
||||
Alternating with a backoff schedule that reaches past 30s lets a long-running
|
||||
agent run survive multi-minute provider or gateway blips.
|
||||
|
||||
Bidirectional: if the primary is Anthropic the fallback is typically OpenAI,
|
||||
and vice versa. The middleware itself is provider-agnostic — it inspects the
|
||||
exception type/status code to decide whether to fall over.
|
||||
exception type/status code to decide whether an attempt is retryable.
|
||||
|
||||
If every attempt fails, the middleware either raises the last error or (by
|
||||
default) returns a terminal ``AIMessage`` explaining the outage, so the run
|
||||
ends with a visible message in Slack/GitHub instead of an abrupt crash. The
|
||||
turn's progress is checkpointed, so the user can retrigger to continue.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from collections.abc import Awaitable, Callable
|
||||
import random
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable, Sequence
|
||||
from typing import Any
|
||||
|
||||
import anthropic
|
||||
import httpx
|
||||
import openai
|
||||
from botocore.exceptions import ClientError
|
||||
from langchain.agents.middleware import AgentMiddleware
|
||||
|
|
@ -37,6 +55,7 @@ _TRANSIENT_EXCEPTIONS: tuple[type[BaseException], ...] = (
|
|||
openai.APITimeoutError,
|
||||
openai.RateLimitError,
|
||||
openai.InternalServerError,
|
||||
httpx.TransportError,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -48,6 +67,23 @@ _RETRYABLE_BEDROCK_ERROR_CODES = {
|
|||
}
|
||||
|
||||
|
||||
# Seconds slept before each retry attempt (attempt 0 is the initial call).
|
||||
# The first failover is immediate: a provider-specific outage should not delay
|
||||
# the cross-provider retry. Later delays grow past the ~30s the gateway's 502
|
||||
# page asks for. Each attempt additionally benefits from the SDK's own
|
||||
# ``max_retries`` backoff, so worst-case wall time before giving up is a few
|
||||
# minutes — acceptable for a long-running agent, far better than crashing.
|
||||
DEFAULT_BACKOFF_SCHEDULE: tuple[float, ...] = (0.0, 5.0, 15.0, 30.0, 45.0)
|
||||
|
||||
MODEL_OUTAGE_MESSAGE = (
|
||||
"I wasn't able to reach the language model providers after several retries "
|
||||
"(both the primary and fallback models returned transient errors, e.g. "
|
||||
"502/503/overloaded). This is a temporary provider or gateway outage, not a "
|
||||
"problem with your task. My progress so far has been saved — please retrigger "
|
||||
"the run in a few minutes to continue."
|
||||
)
|
||||
|
||||
|
||||
def _should_fallback(exc: BaseException) -> bool:
|
||||
if isinstance(exc, _TRANSIENT_EXCEPTIONS):
|
||||
return True
|
||||
|
|
@ -117,52 +153,132 @@ def _provider_access_error_message(exc: BaseException) -> str | None:
|
|||
|
||||
|
||||
class ModelFallbackMiddleware(AgentMiddleware):
|
||||
"""Retry the model call against a fallback provider on transient errors."""
|
||||
"""Retry the model call across primary and fallback providers on transient errors.
|
||||
|
||||
def __init__(self, fallback_model: BaseChatModel) -> None:
|
||||
Args:
|
||||
fallback_model: Cross-provider model used on odd-numbered attempts.
|
||||
backoff_schedule: Seconds slept before each retry. ``len(schedule) + 1``
|
||||
is the total number of attempts. Delays get ±25% jitter.
|
||||
surface_outage_message: When all attempts fail, return a terminal
|
||||
``AIMessage`` describing the outage instead of raising, so the run
|
||||
ends gracefully with a user-visible message rather than a crash.
|
||||
Set to ``False`` to re-raise the last error (e.g. if platform-level
|
||||
alerting keys off failed runs).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
fallback_model: BaseChatModel,
|
||||
*,
|
||||
backoff_schedule: Sequence[float] = DEFAULT_BACKOFF_SCHEDULE,
|
||||
surface_outage_message: bool = True,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self._fallback_model = fallback_model
|
||||
self._backoff_schedule = tuple(backoff_schedule)
|
||||
self._surface_outage_message = surface_outage_message
|
||||
|
||||
def _fallback_name(self) -> str:
|
||||
return (
|
||||
getattr(self._fallback_model, "model_name", None)
|
||||
or getattr(self._fallback_model, "model", None)
|
||||
or "fallback"
|
||||
)
|
||||
|
||||
def _jittered_delay(self, delay: float) -> float:
|
||||
return delay + random.uniform(0, delay * 0.25) if delay > 0 else 0.0
|
||||
|
||||
def _log_retry(
|
||||
self, exc_type: str, use_fallback: bool, attempt: int, total: int, delay: float
|
||||
) -> None:
|
||||
logger.warning(
|
||||
"Model call failed transiently (%s) on %s model "
|
||||
"(attempt %d/%d); retrying %s model in %.1fs",
|
||||
exc_type,
|
||||
"fallback" if use_fallback else "primary",
|
||||
attempt + 1,
|
||||
total,
|
||||
"primary" if use_fallback else f"fallback ({self._fallback_name()})",
|
||||
delay,
|
||||
)
|
||||
|
||||
def _resolve_request(self, request: ModelRequest, use_fallback: bool) -> ModelRequest:
|
||||
return request.override(model=self._fallback_model) if use_fallback else request
|
||||
|
||||
def wrap_model_call(
|
||||
self,
|
||||
request: ModelRequest,
|
||||
handler: Callable[[ModelRequest], ModelResponse],
|
||||
) -> ModelCallResult:
|
||||
try:
|
||||
return handler(request)
|
||||
except Exception as exc:
|
||||
access_error_message = _provider_access_error_message(exc)
|
||||
if access_error_message is not None:
|
||||
logger.warning("Model access error surfaced to user: %s", type(exc).__name__)
|
||||
return AIMessage(content=access_error_message)
|
||||
if not _should_fallback(exc):
|
||||
raise
|
||||
logger.warning(
|
||||
"Primary model failed (%s); falling back to %s",
|
||||
type(exc).__name__,
|
||||
getattr(self._fallback_model, "model_name", None)
|
||||
or getattr(self._fallback_model, "model", "fallback"),
|
||||
)
|
||||
return handler(request.override(model=self._fallback_model))
|
||||
total_attempts = len(self._backoff_schedule) + 1
|
||||
last_exc: BaseException | None = None
|
||||
|
||||
for attempt in range(total_attempts):
|
||||
use_fallback = attempt % 2 == 1
|
||||
attempt_request = self._resolve_request(request, use_fallback)
|
||||
try:
|
||||
return handler(attempt_request)
|
||||
except Exception as exc:
|
||||
access_error_message = _provider_access_error_message(exc)
|
||||
if access_error_message is not None:
|
||||
logger.warning("Model access error surfaced to user: %s", type(exc).__name__)
|
||||
return AIMessage(content=access_error_message)
|
||||
if not _should_fallback(exc):
|
||||
raise
|
||||
last_exc = exc
|
||||
if attempt + 1 >= total_attempts:
|
||||
break
|
||||
delay = self._jittered_delay(self._backoff_schedule[attempt])
|
||||
self._log_retry(type(exc).__name__, use_fallback, attempt, total_attempts, delay)
|
||||
if delay > 0:
|
||||
time.sleep(delay)
|
||||
|
||||
assert last_exc is not None
|
||||
logger.error(
|
||||
"Model call failed after %d attempts across primary and fallback (%s): %s",
|
||||
total_attempts,
|
||||
self._fallback_name(),
|
||||
last_exc,
|
||||
)
|
||||
if self._surface_outage_message:
|
||||
return AIMessage(content=MODEL_OUTAGE_MESSAGE)
|
||||
raise last_exc
|
||||
|
||||
async def awrap_model_call(
|
||||
self,
|
||||
request: ModelRequest,
|
||||
handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
|
||||
) -> Any:
|
||||
try:
|
||||
return await handler(request)
|
||||
except Exception as exc:
|
||||
access_error_message = _provider_access_error_message(exc)
|
||||
if access_error_message is not None:
|
||||
logger.warning("Model access error surfaced to user: %s", type(exc).__name__)
|
||||
return AIMessage(content=access_error_message)
|
||||
if not _should_fallback(exc):
|
||||
raise
|
||||
logger.warning(
|
||||
"Primary model failed (%s); falling back to %s",
|
||||
type(exc).__name__,
|
||||
getattr(self._fallback_model, "model_name", None)
|
||||
or getattr(self._fallback_model, "model", "fallback"),
|
||||
)
|
||||
return await handler(request.override(model=self._fallback_model))
|
||||
total_attempts = len(self._backoff_schedule) + 1
|
||||
last_exc: BaseException | None = None
|
||||
|
||||
for attempt in range(total_attempts):
|
||||
use_fallback = attempt % 2 == 1
|
||||
attempt_request = self._resolve_request(request, use_fallback)
|
||||
try:
|
||||
return await handler(attempt_request)
|
||||
except Exception as exc:
|
||||
access_error_message = _provider_access_error_message(exc)
|
||||
if access_error_message is not None:
|
||||
logger.warning("Model access error surfaced to user: %s", type(exc).__name__)
|
||||
return AIMessage(content=access_error_message)
|
||||
if not _should_fallback(exc):
|
||||
raise
|
||||
last_exc = exc
|
||||
if attempt + 1 >= total_attempts:
|
||||
break
|
||||
delay = self._jittered_delay(self._backoff_schedule[attempt])
|
||||
self._log_retry(type(exc).__name__, use_fallback, attempt, total_attempts, delay)
|
||||
if delay > 0:
|
||||
await asyncio.sleep(delay)
|
||||
|
||||
assert last_exc is not None
|
||||
logger.error(
|
||||
"Model call failed after %d attempts across primary and fallback (%s): %s",
|
||||
total_attempts,
|
||||
self._fallback_name(),
|
||||
last_exc,
|
||||
)
|
||||
if self._surface_outage_message:
|
||||
return AIMessage(content=MODEL_OUTAGE_MESSAGE)
|
||||
raise last_exc
|
||||
|
|
|
|||
121
agent/utils/ttl_cache.py
Normal file
121
agent/utils/ttl_cache.py
Normal file
|
|
@ -0,0 +1,121 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import TypeVar, cast
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
_CACHE: dict[str, tuple[object, float]] = {}
|
||||
_LOCKS: dict[tuple[str, int], asyncio.Lock] = {}
|
||||
_REFRESH_TASKS: dict[tuple[str, int], asyncio.Task[None]] = {}
|
||||
|
||||
|
||||
def _now() -> float:
|
||||
return asyncio.get_running_loop().time()
|
||||
|
||||
|
||||
def _lock_for(key: str) -> asyncio.Lock:
|
||||
loop = asyncio.get_running_loop()
|
||||
lock_key = (key, id(loop))
|
||||
lock = _LOCKS.get(lock_key)
|
||||
if lock is None:
|
||||
lock = asyncio.Lock()
|
||||
_LOCKS[lock_key] = lock
|
||||
return lock
|
||||
|
||||
|
||||
async def cached(key: str, ttl_seconds: float, loader: Callable[[], Awaitable[T]]) -> T:
|
||||
now = _now()
|
||||
entry = _CACHE.get(key)
|
||||
if entry is not None:
|
||||
value, expires_at = entry
|
||||
if expires_at > now:
|
||||
return cast(T, value)
|
||||
|
||||
async with _lock_for(key):
|
||||
now = _now()
|
||||
entry = _CACHE.get(key)
|
||||
if entry is not None:
|
||||
value, expires_at = entry
|
||||
if expires_at > now:
|
||||
return cast(T, value)
|
||||
|
||||
stale = entry[0] if entry is not None else None
|
||||
has_stale = entry is not None
|
||||
try:
|
||||
value = await loader()
|
||||
except Exception:
|
||||
if has_stale:
|
||||
logger.warning(
|
||||
"TTL cache loader failed for %s; serving stale value", key, exc_info=True
|
||||
)
|
||||
return cast(T, stale)
|
||||
raise
|
||||
_CACHE[key] = (value, now + ttl_seconds)
|
||||
return value
|
||||
|
||||
|
||||
async def _refresh_stale_entry(
|
||||
key: str, ttl_seconds: float, loader: Callable[[], Awaitable[object]]
|
||||
) -> None:
|
||||
async with _lock_for(key):
|
||||
try:
|
||||
value = await loader()
|
||||
except Exception:
|
||||
logger.warning("TTL cache background refresh failed for %s", key, exc_info=True)
|
||||
return
|
||||
_CACHE[key] = (value, _now() + ttl_seconds)
|
||||
|
||||
|
||||
def _schedule_refresh(
|
||||
key: str, ttl_seconds: float, loader: Callable[[], Awaitable[object]]
|
||||
) -> None:
|
||||
loop = asyncio.get_running_loop()
|
||||
task_key = (key, id(loop))
|
||||
existing = _REFRESH_TASKS.get(task_key)
|
||||
if existing is not None and not existing.done():
|
||||
return
|
||||
|
||||
task = asyncio.create_task(_refresh_stale_entry(key, ttl_seconds, loader))
|
||||
_REFRESH_TASKS[task_key] = task
|
||||
task.add_done_callback(lambda _task: _REFRESH_TASKS.pop(task_key, None))
|
||||
|
||||
|
||||
async def cached_stale_while_revalidate(
|
||||
key: str,
|
||||
ttl_seconds: float,
|
||||
loader: Callable[[], Awaitable[T]],
|
||||
) -> T:
|
||||
now = _now()
|
||||
entry = _CACHE.get(key)
|
||||
if entry is not None:
|
||||
value, expires_at = entry
|
||||
if expires_at > now:
|
||||
return cast(T, value)
|
||||
_schedule_refresh(key, ttl_seconds, loader)
|
||||
return cast(T, value)
|
||||
return await cached(key, ttl_seconds, loader)
|
||||
|
||||
|
||||
def set_cached(key: str, value: object, ttl_seconds: float) -> None:
|
||||
try:
|
||||
expires_at = _now() + ttl_seconds
|
||||
except RuntimeError:
|
||||
expires_at = float("inf")
|
||||
_CACHE[key] = (value, expires_at)
|
||||
|
||||
|
||||
def invalidate(key: str) -> None:
|
||||
_CACHE.pop(key, None)
|
||||
|
||||
|
||||
def clear() -> None:
|
||||
_CACHE.clear()
|
||||
_LOCKS.clear()
|
||||
for task in _REFRESH_TASKS.values():
|
||||
task.cancel()
|
||||
_REFRESH_TASKS.clear()
|
||||
|
|
@ -62,8 +62,8 @@
|
|||
{"sha": "13b40113", "pr": 1666, "subject": "chore(deps): bump cryptography from 48.0.1 to 49.0.0 in the major group (#1666)", "disposition": "wont-merge", "reason": "already in dev — same 48->49 cryptography bump dev did via #105", "branch": "", "local_sha": null, "updated": "2026-07-08T20:14:42Z"}
|
||||
{"sha": "fbc6de85", "pr": 1667, "subject": "chore(deps): bump fireworks-ai from 1.2.0a75 to 1.2.0a86 (#1667)", "disposition": "deferred", "reason": "conflicts on pick — dev diverged to fireworks-ai a85; a86 still wanted, needs manual bump + uv lock", "branch": "deps", "local_sha": null, "updated": "2026-07-08T23:17:35Z"}
|
||||
{"sha": "290d0fee", "pr": 1669, "subject": "chore(deps): update langgraph-cli[inmem] requirement (#1669)", "disposition": "wont-merge", "reason": "already in dev — langgraph-cli[inmem] at 0.4.30", "branch": "", "local_sha": null, "updated": "2026-07-08T20:14:42Z"}
|
||||
{"sha": "c9f6dd86", "pr": 1694, "subject": "fix: fall back on model stream transport errors (#1694)", "disposition": "deferred", "reason": "adds httpx.TransportError to transient set; small conflict w/ dev's diverged Bedrock fallback", "branch": "model-fallback", "local_sha": null, "updated": "2026-07-08T20:14:42Z"}
|
||||
{"sha": "c9a9a7cd", "pr": 1695, "subject": "fix: retry model fallback exhaustion (#1695)", "disposition": "deferred", "reason": "alternating retry+backoff rewrite; reconcile by hand w/ dev's Bedrock + sync wrap path", "branch": "model-fallback", "local_sha": null, "updated": "2026-07-08T20:14:42Z"}
|
||||
{"sha": "c9f6dd86", "pr": 1694, "subject": "fix: fall back on model stream transport errors (#1694)", "disposition": "landed", "reason": "adds httpx.TransportError to transient set; small conflict w/ dev's diverged Bedrock fallback", "branch": "model-fallback", "local_sha": null, "updated": "2026-07-09T18:59:20Z"}
|
||||
{"sha": "c9a9a7cd", "pr": 1695, "subject": "fix: retry model fallback exhaustion (#1695)", "disposition": "landed", "reason": "alternating retry+backoff rewrite; reconcile by hand w/ dev's Bedrock + sync wrap path", "branch": "model-fallback", "local_sha": null, "updated": "2026-07-09T18:59:20Z"}
|
||||
{"sha": "e5dbc788", "pr": 1696, "subject": "fix: Harden durable agent runs (#1696)", "disposition": "landed", "reason": "large durable-run hardening; 3 new modules dev lacks; rewrites fork dispatch/completion", "branch": "durable-dispatch", "local_sha": null, "updated": "2026-07-09T17:22:53Z"}
|
||||
{"sha": "52fe2916", "pr": 1698, "subject": "feat: add PR review link route (#1698)", "disposition": "landed", "reason": "cherry-picked (-x) in #127 (PR review link route)", "branch": "reviewer-misc", "local_sha": null, "updated": "2026-07-08T22:58:21Z"}
|
||||
{"sha": "5f7f5fbd", "pr": 1697, "subject": "fix: Reduce graph import and loader startup latency (#1697)", "disposition": "landed", "reason": "import-hygiene refactor; cross-cutting, references many deferred upstream-only modules", "branch": "durable-dispatch", "local_sha": null, "updated": "2026-07-09T17:22:53Z"}
|
||||
|
|
|
|||
|
|
@ -46,6 +46,8 @@ Rows key on the **upstream SHA** (stable across local cherry-picks). Deferred ro
|
|||
| `4f913198` | #1647 | widen split review diffs | Landed | already present in dev; empty pick confirmed in #127 | reviewer-misc |
|
||||
| `20f63e8c` | #1646 | install missing deps before verification | Landed | already in dev via #81 (upstream-sync); ledger was stale | prompt-tweaks |
|
||||
| `2f237b53` | #1626 | fall back to vision model for image threads | Landed | ported (adapted to Bedrock/Fireworks vision) in #128 | gateway-routing |
|
||||
| `c9f6dd86` | #1694 | fix: fall back on model stream transport errors (#1694) | Landed | adds httpx.TransportError to transient set; small conflict w/ dev's diverged Bedrock fallback | model-fallback |
|
||||
| `c9a9a7cd` | #1695 | fix: retry model fallback exhaustion (#1695) | Landed | alternating retry+backoff rewrite; reconcile by hand w/ dev's Bedrock + sync wrap path | model-fallback |
|
||||
| `e5dbc788` | #1696 | fix: Harden durable agent runs (#1696) | Landed | large durable-run hardening; 3 new modules dev lacks; rewrites fork dispatch/completion | durable-dispatch |
|
||||
| `52fe2916` | #1698 | feat: add PR review link route (#1698) | Landed | cherry-picked (-x) in #127 (PR review link route) | reviewer-misc |
|
||||
| `5f7f5fbd` | #1697 | fix: Reduce graph import and loader startup latency (#1697) | Landed | import-hygiene refactor; cross-cutting, references many deferred upstream-only modules | durable-dispatch |
|
||||
|
|
@ -85,8 +87,6 @@ Rows key on the **upstream SHA** (stable across local cherry-picks). Deferred ro
|
|||
| `4f8bc2dd` | #1692 | refactor: simplify open-swe agent sandbox lifecycle (#1692) | Deferred | FLAG-HUMAN: structural rewrite of ensure_sandbox_for_thread (drops __creating__ 4-case sentinel) | sandbox-refactor |
|
||||
| `48217b68` | #1489 | feat(open-swe): add E2B sandbox provider (#1489) | Deferred | additive E2B provider; separable but ships on the async sandbox.py base | sandbox-refactor |
|
||||
| `fbc6de85` | #1667 | chore(deps): bump fireworks-ai from 1.2.0a75 to 1.2.0a86 (#1667) | Deferred | conflicts on pick — dev diverged to fireworks-ai a85; a86 still wanted, needs manual bump + uv lock | deps |
|
||||
| `c9f6dd86` | #1694 | fix: fall back on model stream transport errors (#1694) | Deferred | adds httpx.TransportError to transient set; small conflict w/ dev's diverged Bedrock fallback | model-fallback |
|
||||
| `c9a9a7cd` | #1695 | fix: retry model fallback exhaustion (#1695) | Deferred | alternating retry+backoff rewrite; reconcile by hand w/ dev's Bedrock + sync wrap path | model-fallback |
|
||||
| `67abf5b0` | #1659 | fix: surface attributed PR creation failures (#1659) | Deferred | PR-attribution-failure guard (new mw, safe imports); heavy conflict on diverged open_pull_request.py | pr-attribution |
|
||||
| `c75cbb1f` | #1677 | feat: re-add Fable 5 with an admin toggle to disable it (#1677) | Deferred | FLAG-HUMAN: re-adds Fable 5 via anthropic: — contradicts dev's deliberate hide (#1483) + Bedrock migration (#62); wont-merge candidate | fable-admin-toggle |
|
||||
| `304032fa` | #1680 | chore: clarify question answering prompt (#1680) | Deferred | reword Slack info-only answer guidance; conflicts w/ fork's customized Slack prompt | prompt-tweaks |
|
||||
|
|
|
|||
|
|
@ -73,6 +73,12 @@ class TestShouldFallback:
|
|||
exc = anthropic.BadRequestError("bad", response=response, body={})
|
||||
assert _should_fallback(exc) is False
|
||||
|
||||
def test_httpx_remote_protocol_error_falls_back(self) -> None:
|
||||
exc = httpx.RemoteProtocolError(
|
||||
"peer closed connection without sending complete message body (incomplete chunked read)"
|
||||
)
|
||||
assert _should_fallback(exc) is True
|
||||
|
||||
def test_value_error_does_not_fall_back(self) -> None:
|
||||
assert _should_fallback(ValueError("nope")) is False
|
||||
|
||||
|
|
@ -100,6 +106,31 @@ class TestModelFallbackMiddleware:
|
|||
request.override.assert_called_once_with(model=fallback_model)
|
||||
assert calls[1] is request.override.return_value
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_falls_over_on_stream_transport_error(self) -> None:
|
||||
fallback_model = MagicMock(name="fallback_model")
|
||||
middleware = ModelFallbackMiddleware(fallback_model)
|
||||
|
||||
calls: list[object] = []
|
||||
good_response = MagicMock(result=[AIMessage(content="ok from fallback")])
|
||||
|
||||
async def handler(req: object) -> object:
|
||||
calls.append(req)
|
||||
if len(calls) == 1:
|
||||
raise httpx.RemoteProtocolError(
|
||||
"peer closed connection without sending complete message body "
|
||||
"(incomplete chunked read)"
|
||||
)
|
||||
return good_response
|
||||
|
||||
request = _make_request()
|
||||
result = await middleware.awrap_model_call(request, handler)
|
||||
|
||||
assert result is good_response
|
||||
assert len(calls) == 2
|
||||
request.override.assert_called_once_with(model=fallback_model)
|
||||
assert calls[1] is request.override.return_value
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_propagates_non_transient_error(self) -> None:
|
||||
middleware = ModelFallbackMiddleware(MagicMock())
|
||||
|
|
@ -128,9 +159,50 @@ class TestModelFallbackMiddleware:
|
|||
assert "data retention enabled" in result.text
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_does_not_double_fall_back(self) -> None:
|
||||
"""If the fallback also fails transiently, the error propagates."""
|
||||
middleware = ModelFallbackMiddleware(MagicMock())
|
||||
async def test_async_retries_primary_after_fallback_failure(self) -> None:
|
||||
"""If the fallback also fails transiently, retry the primary instead of crashing."""
|
||||
fallback_model = MagicMock(name="fallback_model")
|
||||
middleware = ModelFallbackMiddleware(fallback_model, backoff_schedule=(0.0, 0.0, 0.0))
|
||||
calls: list[object] = []
|
||||
good_response = MagicMock(result=[AIMessage(content="ok from primary retry")])
|
||||
|
||||
async def handler(req: object) -> object:
|
||||
calls.append(req)
|
||||
if len(calls) <= 2: # primary fails, then fallback fails
|
||||
raise _openai_5xx()
|
||||
return good_response
|
||||
|
||||
request = _make_request()
|
||||
result = await middleware.awrap_model_call(request, handler)
|
||||
|
||||
assert result is good_response
|
||||
assert len(calls) == 3
|
||||
# Attempts alternate primary -> fallback -> primary.
|
||||
assert calls[0] is request
|
||||
assert calls[1] is request.override.return_value
|
||||
assert calls[2] is request
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_exhaustion_returns_outage_message(self) -> None:
|
||||
"""After exhausting all attempts, the run ends with a visible message, not a crash."""
|
||||
middleware = ModelFallbackMiddleware(MagicMock(), backoff_schedule=(0.0, 0.0))
|
||||
calls: list[object] = []
|
||||
|
||||
async def handler(req: object) -> object:
|
||||
calls.append(req)
|
||||
raise _openai_5xx()
|
||||
|
||||
result = await middleware.awrap_model_call(_make_request(), handler)
|
||||
|
||||
assert len(calls) == 3
|
||||
assert isinstance(result, AIMessage)
|
||||
assert "retrigger" in result.text
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_exhaustion_raises_when_message_disabled(self) -> None:
|
||||
middleware = ModelFallbackMiddleware(
|
||||
MagicMock(), backoff_schedule=(0.0,), surface_outage_message=False
|
||||
)
|
||||
calls: list[object] = []
|
||||
|
||||
async def handler(req: object) -> object:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue