Don't treat capacity phrases in model output as provider failures

backend/app/routes/proxy.py

@@ -354,62 +354,63 @@ async def aiter_bytes(self):
 _SUCCESS_MARKERS = ("message_stop",)
 
 
-def _is_pre_output_failure(lowered: str) -> bool:
-    """Return True if the lowered SSE text shows a failure with no output."""
+_SSE_FRAME_SEPARATOR = re.compile(rb"\r?\n\r?\n")
+
+
+def _is_error_text(lowered: str) -> bool:
+    """Return True if lowered response text is a provider failure.
+
+    Model output can contain any phrase, so failure markers only count in text that carries no output event.
+    """
     has_failed = any(m in lowered for m in _FAIL_MARKERS)
     has_output = any(m in lowered for m in _OUTPUT_DELTA_MARKERS)
     return has_failed and not has_output
 
 
-def _is_provider_error_frame(frame: bytes) -> bool:
-    """Suppress Anthropic error/capacity frames, including midstream failures."""
-    lowered = frame.decode(errors="replace").lower()
-    has_error = any(marker in lowered for marker in _FAIL_MARKERS)
-    has_output = any(marker in lowered for marker in _OUTPUT_DELTA_MARKERS)
-    return (has_error and not has_output) or any(phrase in lowered for phrase in _CAPACITY_PHRASES)
-
+def _is_capacity_text(lowered: str) -> bool:
+    return _is_error_text(lowered) and any(p in lowered for p in _CAPACITY_PHRASES)
 
-_SYNTHETIC_MESSAGE_STOP = b'event: message_stop\ndata: {"type":"message_stop"}\n\n'
 
+def _is_error_body(raw: bytes) -> bool:
+    """Return True if a non-streaming 200 body is a provider failure rather than a message."""
+    try:
+        payload = json.loads(raw)
+    except (json.JSONDecodeError, UnicodeDecodeError, ValueError):
+        payload = None
+    if isinstance(payload, dict) and payload.get("type") == "message":
+        return False
+    return _is_error_text(raw.decode(errors="replace").lower())
 
-def _sanitize_sse_payload(raw: bytes) -> bytes:
-    """Remove provider failure frames and terminate a buffered stream cleanly."""
-    output = bytearray()
-    remaining = raw
-    while remaining:
-        match = re.search(rb"\r?\n\r?\n", remaining)
-        if match is None:
-            if not _is_provider_error_frame(remaining):
-                output.extend(remaining)
-            break
-        end = match.end()
-        frame = remaining[: match.start()]
-        if _is_provider_error_frame(frame):
-            output.extend(_SYNTHETIC_MESSAGE_STOP)
-            return bytes(output)
-        output.extend(remaining[:end])
-        remaining = remaining[end:]
-    return bytes(output)
 
+def _first_error_frame(raw: bytes, *, complete: bool) -> Optional[str]:
+    """Return the first SSE frame in ``raw`` that is a provider failure, lowered.
 
-async def _iter_sanitized_sse(response: httpx.Response):
-    """Stream SSE frames without ever forwarding provider error events."""
+    A trailing frame without its separator is only inspected once ``complete`` says no more bytes will arrive.
+    """
+    frames = _SSE_FRAME_SEPARATOR.split(raw)
+    if not complete:
+        frames = frames[:-1]
+    for frame in frames:
+        lowered = frame.decode(errors="replace").lower()
+        if _is_error_text(lowered):
+            return lowered
+    return None
+
+
+async def _iter_sse_frames(response: httpx.Response):
+    """Re-chunk an SSE stream into whole frames so each can be rewritten independently."""
     buffer = bytearray()
     async for chunk in response.aiter_bytes():
         buffer.extend(chunk)
         while True:
-            match = re.search(rb"\r?\n\r?\n", buffer)
+            match = _SSE_FRAME_SEPARATOR.search(buffer)
             if match is None:
                 break
             end = match.end()
-            frame = bytes(buffer[: match.start()])
-            separator = match.group(0)
+            frame = bytes(buffer[:end])
             del buffer[:end]
-            if _is_provider_error_frame(frame):
-                yield _SYNTHETIC_MESSAGE_STOP
-                return
-            yield bytes(frame) + separator
-    if buffer and not _is_provider_error_frame(bytes(buffer)):
+            yield frame
+    if buffer:
         yield bytes(buffer)
 
 
@@ -426,19 +427,18 @@ async def _prepare_candidate(candidate: httpx.Response):
     reasonable confidence the upstream will produce output.  This function
     buffers SSE frames and keeps reading until one of three outcomes:
 
-    1. An **output-delta** event is seen  -> the model is generating; start
-       streaming (return ``capacity_error=False``).
-    2. A **response.failed** / capacity phrase is seen *without* any output
-       delta  -> the request failed before generating; the caller can safely
-       retry on the next account (``capacity_error=True``).
+    1. A complete frame is a provider error (``"type":"error"`` /
+       ``response.failed``) -> nothing has reached the client, so the caller
+       retries on the next account.  Capacity phrases inside that error frame
+       mark it as a capacity error (longer cooldown).
+    2. An **output-delta** or completion event is seen -> the model is
+       generating; start streaming.
     3. A safety bound is reached (64 KB or 90 s total prefetch time) without
        either signal -> treat this attempt as failed and rotate accounts.
 
-    The original bug forwarded buffered keepalives after a prefix timeout;
-    the provider could then emit a late failure/capacity event directly to the
-    CLI, ending the user's run instead of triggering account failover. We now
-    hold the stream until real output/completion and rotate on any pre-output
-    timeout or terminal failure.
+    Phrases inside model output never count as failures.  Once streaming has
+    started, a provider error frame is forwarded to the client unchanged, so
+    the client sees the failure instead of a truncated response.
     """
     if candidate.status_code != 200:
         return candidate
@@ -451,33 +451,23 @@ async def _prepare_candidate(candidate: httpx.Response):
 
     if not is_sse:
         raw = await candidate.aread()
-        lowered = raw.decode(errors="replace").lower()
 
         if raw.lstrip().startswith((b"event:", b"data:")):
             is_sse = True
         else:
-            candidate.extensions["proxy_capacity_error"] = any(p in lowered for p in _CAPACITY_PHRASES)
-            candidate.extensions["proxy_pre_output_failure"] = _is_pre_output_failure(lowered)
+            is_error = _is_error_body(raw)
+            candidate.extensions["proxy_capacity_error"] = is_error and _is_capacity_text(raw.decode(errors="replace").lower())
+            candidate.extensions["proxy_pre_output_failure"] = is_error
             return candidate
 
     if is_sse and hasattr(candidate, "_content") and candidate._content:
-        lowered = candidate._content.decode(errors="replace").lower()
-        is_capacity_error = any(p in lowered for p in _CAPACITY_PHRASES)
-        is_pre_output_failure = _is_pre_output_failure(lowered)
-        if is_capacity_error:
-            return _PrefetchedResponse(
-                candidate,
-                candidate._content,
-                _empty_aiter(),
-                True,
-                is_pre_output_failure,
-            )
+        error_frame = _first_error_frame(candidate._content, complete=True)
         return _PrefetchedResponse(
             candidate,
             candidate._content,
             _empty_aiter(),
-            False,
-            is_pre_output_failure,
+            error_frame is not None and _is_capacity_text(error_frame),
+            error_frame is not None,
         )
 
     iterator = candidate.aiter_bytes().__aiter__()
@@ -506,31 +496,22 @@ async def _prepare_candidate(candidate: httpx.Response):
             logger.warning("Upstream produced no output within the SSE prefetch timeout; treating as capacity unavailable")
             return _PrefetchedResponse(candidate, bytes(prefix), iterator, False, True)
         prefix.extend(chunk)
-        lowered = bytes(prefix).decode(errors="replace").lower()
 
-        # Explicit capacity phrases -> always retry.
-        if any(p in lowered for p in _CAPACITY_PHRASES):
-            return _PrefetchedResponse(
-                candidate,
-                bytes(prefix),
-                iterator,
-                True,
-                _is_pre_output_failure(lowered),
-            )
-
-        has_output = any(m in lowered for m in _OUTPUT_DELTA_MARKERS)
-        completed = any(m in lowered for m in _SUCCESS_MARKERS)
+        # Nothing has reached the client yet, so a terminal upstream failure is
+        # safe to retry on another account. The caller applies the cooldown
+        # that prevents every concurrent request from piling onto this account.
+        error_frame = _first_error_frame(bytes(prefix), complete=False)
+        if error_frame is not None:
+            return _PrefetchedResponse(candidate, bytes(prefix), iterator, _is_capacity_text(error_frame), True)
 
         # Real output is flowing -> response is healthy, start streaming.
-        if has_output or completed:
+        lowered = bytes(prefix).decode(errors="replace").lower()
+        if any(m in lowered for m in _OUTPUT_DELTA_MARKERS + _SUCCESS_MARKERS):
             return _PrefetchedResponse(candidate, bytes(prefix), iterator, False, False)
 
-        # Nothing has reached the client yet, so a terminal upstream failure is
-        # safe to retry on another account. The caller applies the short cooldown
-        # that prevents every concurrent request from piling onto this account.
-        if _is_pre_output_failure(lowered):
-            return _PrefetchedResponse(candidate, bytes(prefix), iterator, False, True)
-
+    error_frame = _first_error_frame(bytes(prefix), complete=True)
+    if error_frame is not None:
+        return _PrefetchedResponse(candidate, bytes(prefix), iterator, _is_capacity_text(error_frame), True)
     logger.warning("Upstream filled the SSE prefetch buffer without producing output; treating as capacity unavailable")
     return _PrefetchedResponse(candidate, bytes(prefix), iterator, False, True)
 
@@ -1217,8 +1198,6 @@ async def proxy_messages(
     if resp.status_code >= 400 or "text/event-stream" not in content_type:
         raw = await resp.aread()
         await resp.aclose()
-        if "text/event-stream" in content_type or raw.lstrip().startswith((b"event:", b"data:")):
-            raw = _sanitize_sse_payload(raw)
         raw = _restore_requested_model_in_error(raw, resp.status_code, requested_model, request_model)
 
         if not is_count_tokens:
@@ -1258,7 +1237,7 @@ async def proxy_messages(
 
     async def stream_body():
         try:
-            async for chunk in _iter_sanitized_sse(resp):
+            async for chunk in _iter_sse_frames(resp):
                 accumulator.feed(chunk)
                 yield _restore_requested_model_in_sse_frame(chunk, requested_model)
         finally:

backend/tests/test_api.py

@@ -403,6 +403,108 @@ def test_proxy_fails_over_on_streamed_capacity_event(client, seed_account, make_
     assert len(route.calls) == 2
 
 
+@respx.mock
+def test_model_text_mentioning_capacity_is_not_a_provider_failure(client, seed_account, make_user):
+    from app.utils.models.api import AccountStatus
+    from app.utils.postgres import AccountDb
+    from app.utils.postgres.base import SessionFactory
+
+    first_id = seed_account("chatty-first")
+    seed_account("chatty-second")
+    key = make_user("chatty-user")
+    text = 'The cluster is at capacity; try a different model. Logged \\"type\\":\\"error\\" and response.failed.'
+    sse = (
+        'event: message_start\ndata: {"type":"message_start","message":{"model":"claude-haiku-4-5","usage":{"input_tokens":5}}}\n\n'
+        f'event: content_block_delta\ndata: {{"type":"content_block_delta","index":0,"delta":{{"type":"text_delta","text":"{text}"}}}}\n\n'
+        'event: message_delta\ndata: {"type":"message_delta","usage":{"output_tokens":9}}\n\n'
+        'event: message_stop\ndata: {"type":"message_stop"}\n\n'
+    )
+    message = {
+        "type": "message",
+        "model": "claude-haiku-4-5",
+        "content": [{"type": "text", "text": "The cluster is at capacity; response.failed was logged."}],
+        "usage": {"input_tokens": 1, "output_tokens": 1},
+    }
+
+    respx.route(host="testserver").pass_through()
+    route = respx.post(ANTHROPIC_MESSAGES).mock(
+        side_effect=[
+            httpx.Response(200, headers={"content-type": "text/event-stream"}, text=sse),
+            httpx.Response(200, json=message),
+        ]
+    )
+    headers = {"Authorization": f"Bearer {key}"}
+
+    streamed = client.post("/api/v1/messages", headers=headers, json={"model": "claude-haiku-4-5", "stream": True})
+    assert streamed.status_code == 200
+    assert "try a different model" in streamed.text
+    assert streamed.text.count("message_stop") == 2  # event name + payload of the one real stop frame
+
+    plain = client.post("/api/v1/messages", headers=headers, json={"model": "claude-haiku-4-5"})
+    assert plain.status_code == 200
+    assert plain.json()["content"][0]["text"] == message["content"][0]["text"]
+
+    assert len(route.calls) == 2
+    with SessionFactory() as db:
+        assert db.get(AccountDb, first_id).status == AccountStatus.ACTIVE
+
+
+def _streaming_response(*frames: str) -> httpx.Response:
+    async def body():
+        for frame in frames:
+            yield frame.encode()
+
+    return httpx.Response(200, headers={"content-type": "text/event-stream"}, content=body())
+
+
+def test_prefetch_starts_stream_when_output_mentions_capacity():
+    import asyncio
+
+    from app.routes.proxy import _prepare_candidate
+
+    prepared = asyncio.run(
+        _prepare_candidate(
+            _streaming_response(
+                'event: message_start\ndata: {"type":"message_start","message":{}}\n\n'
+                'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"text":"temporarily overloaded"}}\n\n',
+            )
+        )
+    )
+    assert prepared._proxy_capacity_error is False
+    assert prepared._proxy_pre_output_failure is False
+
+
+def test_prefetch_retries_an_error_frame_before_anything_is_sent():
+    import asyncio
+
+    from app.routes.proxy import _prepare_candidate
+
+    prepared = asyncio.run(
+        _prepare_candidate(
+            _streaming_response(
+                'event: message_start\ndata: {"type":"message_start","message":{}}\n\n'
+                'event: error\ndata: {"type":"error","error":{"message":"Selected model is at capacity."}}\n\n',
+            )
+        )
+    )
+    assert prepared._proxy_capacity_error is True
+    assert prepared._proxy_pre_output_failure is True
+
+
+def test_midstream_error_frame_is_forwarded_not_replaced_with_a_stop():
+    import asyncio
+
+    from app.routes.proxy import _iter_sse_frames
+
+    delta = 'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"text":"partial"}}\n\n'
+    error = 'event: error\ndata: {"type":"error","error":{"type":"overloaded_error","message":"Overloaded"}}\n\n'
+
+    async def collect():
+        return [frame async for frame in _iter_sse_frames(_streaming_response(delta[:20], delta[20:] + error))]
+
+    assert asyncio.run(collect()) == [delta.encode(), error.encode()]
+
+
 @respx.mock
 def test_refresh_quota_reports_rate_limit_cleanly(client, admin_headers, seed_account, monkeypatch):
     # A persistently rate-limited probe must return a friendly, actionable message -- not a raw httpx 429 string.