diff --git a/backend/app/routes/proxy.py b/backend/app/routes/proxy.py index dfc898d..97c7972 100644 --- a/backend/app/routes/proxy.py +++ b/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: diff --git a/backend/tests/test_api.py b/backend/tests/test_api.py index a789b7c..f832732 100644 --- a/backend/tests/test_api.py +++ b/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.