diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 24a36a7f2dc..abc99a3be23 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -4,6 +4,7 @@ import json import logging import math from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Mapping, Sequence +from dataclasses import dataclass from datetime import datetime from functools import lru_cache from itertools import chain @@ -114,6 +115,7 @@ from litellm.proxy.common_utils.openai_error_payload import ( openai_error_param, openai_error_type, ) +from litellm.proxy.common_utils.responses_stream_errors import ResponsesStreamErrorState from litellm.proxy.common_utils.sse_keepalive import ( SSE_COMMENT_PING_BYTES, SSE_STREAM_START_TAIL, @@ -907,11 +909,18 @@ async def _enforce_guardrail_added_tag_budgets( ) from e +def _sse_data_payload(event_line: str | bytes) -> str | None: + text: Final = event_line.decode("utf-8") if isinstance(event_line, bytes) else event_line + return next( + (line.removeprefix("data: ").strip() for line in text.splitlines() if line.startswith("data: ")), + None, + ) + + async def _parse_event_data_for_error(event_line: str | bytes) -> int | None: """Parses an event line and returns an error code if present, else None.""" - event_line = event_line.decode("utf-8") if isinstance(event_line, bytes) else event_line - if event_line.startswith("data: "): - json_str: Final = event_line[len("data: ") :].strip() + json_str: Final = _sse_data_payload(event_line) + if json_str is not None: if not json_str or json_str == "[DONE]": # handle empty data or [DONE] message return None try: @@ -952,8 +961,6 @@ def _extract_error_from_sse_chunk(event_line: str | bytes) -> dict: Returns: Error dictionary in OpenAI API format """ - event_line = event_line.decode("utf-8") if isinstance(event_line, bytes) else event_line - # Default error format default_error: Final = { "message": "Unknown error", @@ -962,8 +969,8 @@ def _extract_error_from_sse_chunk(event_line: str | bytes) -> dict: "code": "500", } - if event_line.startswith("data: "): - json_str: Final = event_line[len("data: ") :].strip() + json_str: Final = _sse_data_payload(event_line) + if json_str is not None: if not json_str or json_str == "[DONE]": return default_error @@ -1171,6 +1178,39 @@ def _sse_error_frames(error_obj: Mapping[str, object]) -> tuple[str, str]: return f"data: {json.dumps({'error': error_obj}, default=str)}\n\n", "data: [DONE]\n\n" +@dataclass(frozen=True, slots=True) +class LateSseFraming: + """How a response the keepalive already opened ends once its upstream fails late.""" + + failure: Callable[[Exception], tuple[str, ...]] + body: Callable[[Response], tuple[bytes, ...]] + + +def _openai_late_failure_frames(exc: Exception) -> tuple[str, ...]: + _, error_obj = sse_error_payload(exc) + return _sse_error_frames(error_obj) + + +def _openai_late_body_frames(produced: Response) -> tuple[bytes, ...]: + return b"data: " + (bytes(produced.body) or b"{}") + b"\n\n", b"data: [DONE]\n\n" + + +def _responses_late_failure_frames(exc: Exception) -> tuple[str, ...]: + frame: Final = ResponsesStreamErrorState().format_failure(exc) + return () if frame is None else (frame,) + + +def _responses_late_body_frames(produced: Response) -> tuple[bytes, ...]: + frame: Final = ResponsesStreamErrorState().format_error_body(bytes(produced.body), produced.status_code) + return () if frame is None else (frame.encode(),) + + +OPENAI_LATE_SSE_FRAMING: Final = LateSseFraming(failure=_openai_late_failure_frames, body=_openai_late_body_frames) +RESPONSES_LATE_SSE_FRAMING: Final = LateSseFraming( + failure=_responses_late_failure_frames, body=_responses_late_body_frames +) + + def _sse_stream_headers(headers: Mapping[str, str]) -> Mapping[str, str]: """`headers` plus the two that stop reverse proxies from buffering SSE (issue #28384).""" return MappingProxyType({**headers, **_TTFT_KEEPALIVE_HEADERS}) @@ -1371,13 +1411,13 @@ async def _aclose_late_response(produced: Response) -> None: verbose_proxy_logger.debug("error closing relayed streaming generator: %s", exc) -async def _relay_late_response(produced: Response) -> AsyncGenerator[bytes, None]: +async def _relay_late_response(produced: Response, framing: LateSseFraming) -> AsyncGenerator[bytes, None]: """Replay a Response that was built after a keepalive had already opened the wire.""" if not isinstance(produced, StreamingResponse): # The status line is already on the wire, so a non-streaming body, an error # body included, can only reach the client as an SSE frame. - yield b"data: " + (bytes(produced.body) or b"{}") + b"\n\n" - yield b"data: [DONE]\n\n" + for frame in framing.body(produced): + yield frame return try: @@ -1419,6 +1459,7 @@ async def open_sse_before_first_byte( ping_interval_seconds: float | str | None, media_type: str = "text/event-stream", on_late_failure: Callable[[Exception], Awaitable[HTTPException | None]] | None = None, + late_framing: LateSseFraming = OPENAI_LATE_SSE_FRAMING, ) -> _LateResponseT | StreamingResponse: """Write SSE keepalive comments while the upstream LLM call is still in flight. @@ -1464,11 +1505,10 @@ async def open_sse_before_first_byte( # would never fire and the failure would go unaudited. The hook # also gets to sanitize what reaches the client, by returning or # raising a replacement, so its answer decides the frame. - _, error_obj = sse_error_payload(await _sanitized_late_failure(exc, on_late_failure)) - for frame in _sse_error_frames(error_obj): + for frame in late_framing.failure(await _sanitized_late_failure(exc, on_late_failure)): yield frame.encode() return - async for chunk in _relay_late_response(produced): + async for chunk in _relay_late_response(produced, late_framing): yield chunk finally: if not produce_task.done(): @@ -2812,6 +2852,7 @@ class ProxyBaseLLMRequestProcessing: is_streaming_request: bool | None = False, contents: list[object] | None = None, skip_pre_call_logic: bool = False, + late_framing: LateSseFraming = OPENAI_LATE_SSE_FRAMING, ) -> object: """Run the request, sending SSE keepalives while the upstream is still silent. @@ -2855,6 +2896,7 @@ class ProxyBaseLLMRequestProcessing: ), ping_interval_seconds=ttft_keepalive_interval(self.data, llm_router), on_late_failure=_audit_late_failure, + late_framing=late_framing, ) async def _process_llm_request( diff --git a/litellm/proxy/common_utils/responses_stream_errors.py b/litellm/proxy/common_utils/responses_stream_errors.py index 6e563b2b257..312bf472594 100644 --- a/litellm/proxy/common_utils/responses_stream_errors.py +++ b/litellm/proxy/common_utils/responses_stream_errors.py @@ -57,6 +57,12 @@ class _FailureDetails(LiteLLMBaseModel): return cls.model_validate({field: getattr(exception, field, None) for field in cls.model_fields}) +class _ErrorBody(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + error: _FailureDetails = _FailureDetails() + + def _original_failure(exception: Exception) -> Exception: current = exception # rebind-ok: the recursion gate requires iterative wrapper traversal while isinstance(current, MidStreamFallbackError) and current.original_exception is not None: @@ -137,10 +143,16 @@ class ResponsesStreamErrorState: return prefix.encode("utf-8") + frame if isinstance(frame, bytes) else prefix + frame def format_failure(self, exception: Exception) -> str | None: + original: Final = _original_failure(exception) + return self._format_failed(_failure_details(original), str(original)) + + def format_error_body(self, body: bytes, status_code: int) -> str | None: + error: Final = _ErrorBody.model_validate_json(body or b"{}").error + return self._format_failed(error.model_copy(update={"status_code": status_code}), body.decode()) + + def _format_failed(self, details: _FailureDetails, fallback_message: str) -> str | None: if self.terminal_emitted: return None - original: Final = _original_failure(exception) - details: Final = _failure_details(original) response: Final = ResponsesAPIResponse.model_validate( MappingProxyType( { @@ -153,7 +165,7 @@ class ResponsesStreamErrorState: "error": MappingProxyType( { "code": _response_error_code(details), - "message": redact_internal_details_from_client_message(details.message or str(original)), + "message": redact_internal_details_from_client_message(details.message or fallback_message), } ), } diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index a5a287fb42e..15832d81b7f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -10620,7 +10620,7 @@ async def async_data_generator( if error_message is not None: yield error_message - # OpenAI-compatible streams terminate with data: [DONE]; Google GenAI (?alt=sse) does not. + # Chat Completions streams terminate with data: [DONE]; Responses and Google GenAI (?alt=sse) streams do not. if ends_with_openai_done: yield _OPENAI_STREAM_DONE_FRAME except (asyncio.CancelledError, GeneratorExit): diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index aa5fe472c66..1e16c16f33a 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -52,6 +52,7 @@ from litellm.proxy.auth.user_api_key_auth import ( user_api_key_auth_websocket, ) from litellm.proxy.common_request_processing import ( + RESPONSES_LATE_SSE_FRAMING, ProxyBaseLLMRequestProcessing, cost_headers, create_response, @@ -435,6 +436,7 @@ async def responses_api( if requests_background(_PROXY_SERVER_REQUEST_ADAPTER.validate_python(data)) else native_data_generator ), + late_framing=RESPONSES_LATE_SSE_FRAMING, model=None, user_model=user_model, user_temperature=user_temperature, diff --git a/tests/integration/observability/test_guardrail_effects.py b/tests/integration/observability/test_guardrail_effects.py index 8f83287696c..c1ed810a428 100644 --- a/tests/integration/observability/test_guardrail_effects.py +++ b/tests/integration/observability/test_guardrail_effects.py @@ -1197,9 +1197,7 @@ def test_responses_pre_call_denial_streams_sse_with_typed_message_item(gateway: response.headers["content-type"], response.text, ) - lines: Final = tuple(line for line in response.text.split("\n") if line.startswith("data: ")) - assert lines[-1] == "data: [DONE]", response.text - events: Final = tuple(json.loads(line.removeprefix("data: ")) for line in lines[:-1]) + events: Final = _blocked_stream_events(response.text) kinds: Final = tuple(event["type"] for event in events) assert tuple(kind for kind in kinds if kind != "response.output_text.delta") == ( "response.created", @@ -1286,9 +1284,13 @@ def _responses_output_denial_config(tmp_path: Path, identity: str, model: str) - def _blocked_stream_events(text: str) -> tuple[dict[str, object], ...]: - lines: Final = tuple(line for line in text.split("\n") if line.startswith("data: ")) - assert lines[-1] == "data: [DONE]", text - return tuple(json.loads(line.removeprefix("data: ")) for line in lines[:-1]) + assert "data: [DONE]" not in text, text + lines: Final = tuple(line for line in text.split("\n") if line.startswith(("event: ", "data: "))) + events: Final[tuple[dict[str, object], ...]] = tuple( + json.loads(line.removeprefix("data: ")) for line in lines[1::2] + ) + assert lines[::2] == tuple(f"event: {event['type']}" for event in events), text + return events def _dead_api_base() -> str: @@ -2786,7 +2788,7 @@ def test_client_disconnect_after_a_finished_responses_tool_call_scans_the_text_a with _disconnect_rig(gateway, tmp_path, shape="text_then_tool_call") as rig: try: received: Final = _stream_and_close( - rig, "/v1/responses", _tool_call_request(rig, "/v1/responses"), "response.output_item.done" + rig, "/v1/responses", _tool_call_request(rig, "/v1/responses"), '"type":"response.output_item.done"' ) assert rig.token() in received, received scans: Final = eventually( diff --git a/tests/integration/providers/test_bedrock_mantle_claude_chat_chaos.py b/tests/integration/providers/test_bedrock_mantle_claude_chat_chaos.py index bcbf6d0f6f3..2325b705790 100644 --- a/tests/integration/providers/test_bedrock_mantle_claude_chat_chaos.py +++ b/tests/integration/providers/test_bedrock_mantle_claude_chat_chaos.py @@ -261,7 +261,7 @@ async def test_slow_native_streams_reach_every_caller_whole_and_are_logged_once( assert len(served) == 10 for item in served: _assert_answered_with_its_own_marker(item) - assert item.text.rstrip().endswith("data: [DONE]"), item.text + assert item.text.rstrip().endswith("data: [DONE]") == (item.call.endpoint == "chat"), item.text _assert_native_requests_for(wire.drain(), calls) assert _spend_status_by_call(served) == {item.call_id: "success" for item in served} diff --git a/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py b/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py index d4f9533fcdc..9c73fcc3b6e 100644 --- a/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py @@ -846,7 +846,12 @@ def test_responses_api_stream_on_a_mantle_claude_id_bridges_to_the_native_stream "POST", "/v1/responses", {"model": model, "input": _prompt(marker), "stream": True} ) assert response.status_code == 200, response.text - events: Final = _sse_objects(response.text) + assert "data: [DONE]" not in response.text, response.text + events: Final = tuple( + _JSON_OBJECT.validate_json(line.removeprefix("data: ")) + for line in response.text.splitlines() + if line.startswith("data: ") + ) kinds: Final = [event["type"] for event in events] assert (kinds[0], kinds[-1]) == ("response.created", "response.completed"), kinds assert "".join( diff --git a/tests/integration/streaming/test_chat_stream_guardrail_error_events.py b/tests/integration/streaming/test_chat_stream_guardrail_error_events.py index a8ebe510b8e..b2390de7737 100644 --- a/tests/integration/streaming/test_chat_stream_guardrail_error_events.py +++ b/tests/integration/streaming/test_chat_stream_guardrail_error_events.py @@ -482,8 +482,7 @@ def _responses_body(marker: str, *directives: str) -> dict[str, JsonValue]: def _failed_response_error(response: httpx.Response) -> dict[str, JsonValue]: assert response.status_code == 200, response.text - lines: Final = _data_lines(response.text) - assert lines[-1] == "[DONE]" and lines.count("[DONE]") == 1, response.text + assert "[DONE]" not in _data_lines(response.text), response.text events: Final = _payloads(response.text) kinds: Final = tuple(string_value(event["type"]) for event in events) assert kinds[0] == "response.created" and kinds[-1] == "response.failed", response.text diff --git a/tests/integration/streaming/test_responses_bridge_stream_errors.py b/tests/integration/streaming/test_responses_bridge_stream_errors.py index 8a5655be0f8..6b0e35f9c56 100644 --- a/tests/integration/streaming/test_responses_bridge_stream_errors.py +++ b/tests/integration/streaming/test_responses_bridge_stream_errors.py @@ -474,9 +474,10 @@ def test_native_responses_stream_forwards_the_failed_response_on_both_legs(gatew assert response.status_code == 200, response.text frames: Final = data_frames(response.text) assert [frame["type"] for frame in frames] == ["response.created", "response.failed"], response.text + assert [name for name, _ in sse_events(response.text)] == ["response.created", "response.failed"], response.text failed: Final = object_value(frames[-1]["response"]) assert failed["status"] == "failed", failed assert object_value(failed["error"])["code"] == "rate_limit_exceeded", failed - assert response.text.rstrip().endswith("data: [DONE]"), response.text + assert "data: [DONE]" not in response.text, response.text assert len(wire.drain()) == 1 assert spend_row(response.headers["x-litellm-call-id"])["status"] == "failure" diff --git a/tests/integration/streaming/test_responses_top_level_error_events.py b/tests/integration/streaming/test_responses_top_level_error_events.py index 63bdadceebc..6d361492dd6 100644 --- a/tests/integration/streaming/test_responses_top_level_error_events.py +++ b/tests/integration/streaming/test_responses_top_level_error_events.py @@ -256,8 +256,10 @@ def _responses_frames(gateway: Gateway, stream: tuple[bytes, ...]) -> tuple[http ) _assert_one_attempt(served.wire, marker) assert served.response.status_code == 200, served.response.text - assert served.response.text.rstrip().endswith("data: [DONE]"), served.response.text + assert "data: [DONE]" not in served.response.text, served.response.text frames: Final = data_frames(served.response.text) + names: Final = tuple(name for name, _ in sse_events(served.response.text)) + assert names == tuple(item["type"] for item in frames), served.response.text assert frames[-1]["type"] == "response.failed", served.response.text return served.response, object_value(object_value(frames[-1]["response"])["error"]) diff --git a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py index 0b7f32575c5..8aa887bfd78 100644 --- a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -3783,7 +3783,7 @@ class TestBuildBlockSseChunks: def _payloads(self, chunks): import json - payloads = [ + payloads: Final = [ json.loads(next(line[6:] for line in chunk.decode().splitlines() if line.startswith("data: "))) for chunk in chunks ] diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 0bf14f59b1a..5945352c61e 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -177,6 +177,67 @@ async def test_streaming_upstream_errors_keep_the_client_protocol( assert "error" in events[-1] +@pytest.mark.asyncio +@pytest.mark.parametrize("path", ["/v1/responses", "/v1/chat/completions"]) +async def test_a_failure_after_keepalive_pings_keeps_the_client_protocol(monkeypatch: pytest.MonkeyPatch, path: str): + import litellm.proxy.proxy_server as ps + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + model: Final = "gpt-6-astra" + message: Final = "Upstream is rate limiting this key" + upstream_url: Final = "https://streaming.example/v1" + is_chat: Final = path == "/v1/chat/completions" + + async def slow_rate_limit(_: httpx.Request) -> httpx.Response: + await asyncio.sleep(0.3) + return httpx.Response(429, json={"error": {"message": message, "type": "rate_limit_error", "code": None}}) + + router: Final = litellm.Router( + model_list=[ + { + "model_name": model, + "litellm_params": { + "model": "openai/" + model, + "api_base": upstream_url, + "api_key": "fixture-key", + "max_retries": 0, + }, + } + ], + num_retries=0, + ) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "sse_keepalive_ping_interval_seconds", 0.05) + monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, _auth_override) + with respx.mock as transport: + transport.post(upstream_url + ("/chat/completions" if is_chat else "/responses")).mock( + side_effect=slow_rate_limit + ) + async with httpx.AsyncClient(transport=httpx.ASGITransport(app), base_url="http://testserver") as client: + result: Final = await client.post( + path, + json={ + "model": model, + "stream": True, + **({"messages": [{"role": "user", "content": "hello"}]} if is_chat else {"input": "hello"}), + }, + ) + frames: Final = tuple(frame for frame in result.text.split("\n\n") if frame) + + assert result.status_code == 200, result.text + assert frames[0] == ": ping", result.text + if is_chat: + assert frames[-1] == "data: [DONE]", result.text + assert message in json.loads(frames[-2].removeprefix("data: "))["error"]["message"] + else: + assert "data: [DONE]" not in result.text + assert frames[-1].startswith("event: response.failed\ndata: "), result.text + failed_error: Final = json.loads(frames[-1].split("data: ", 1)[1])["response"]["error"] + assert failed_error["code"] == "rate_limit_exceeded" + assert message in failed_error["message"] + + _COMPACTION_ITEM: Final = { "type": "compaction", "encrypted_content": json.dumps({"type": "compaction", "content": "the cat is named Biscuit"}), diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 52816dc50a8..310b526365f 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -50,6 +50,7 @@ from litellm.proxy.common_request_processing import ( _has_attribute_error_in_chain, include_guardrail_response_requested, is_azure_model_router_request, + RESPONSES_LATE_SSE_FRAMING, open_sse_before_first_byte, resolve_litellm_call_id, ttft_keepalive_interval, @@ -2161,6 +2162,18 @@ class TestCommonRequestProcessingHelpers: assert body["error"]["code"] == 403 assert body["error"]["message"] == "forbidden" + async def test_a_responses_stream_whose_first_event_is_a_named_error_answers_as_json(self): + error: Final = {"type": "guardrail_error", "code": "400", "message": "blocked", "param": None} + + async def responses_stream(): + yield f"event: error\ndata: {json.dumps({'type': 'error', 'sequence_number': 0, 'error': error})}\n\n" + + response: Final = await create_response(responses_stream(), "text/event-stream", {}) + + assert isinstance(response, JSONResponse) + assert response.status_code == 400 + assert json.loads(response.body) == {"error": error} + async def test_create_streaming_response_first_chunk_not_error(self): async def mock_generator(): yield 'data: {"content": "first part"}\n\n' @@ -10378,6 +10391,43 @@ async def test_ttft_keepalive_relays_a_late_non_streaming_body_as_an_sse_frame() assert collected[-1] == b"data: [DONE]\n\n" +async def _slow_rate_limit_failure(): + await asyncio.sleep(0.2) + raise HTTPException(status_code=429, detail="rate limited") + + +async def _slow_guardrail_block_body(): + await asyncio.sleep(0.2) + return JSONResponse(status_code=400, content={"error": {"message": "blocked", "type": "guardrail_error", "code": "400"}}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "produce,expected_code,expected_message", + [ + (_slow_rate_limit_failure, "rate_limit_exceeded", "rate limited"), + (_slow_guardrail_block_body, "invalid_request_error", "blocked"), + ], +) +async def test_ttft_keepalive_ends_a_late_responses_failure_with_response_failed_and_no_done( + produce, expected_code, expected_message +): + response: Final = await open_sse_before_first_byte( + produce(), ping_interval_seconds=0.05, late_framing=RESPONSES_LATE_SSE_FRAMING + ) + collected: Final = await _drain(response) + + assert collected[0] == TTFT_PING + assert b"data: [DONE]\n\n" not in collected + event_line, data_line = collected[-1].decode().strip().split("\n") + failed: Final = json.loads(data_line.removeprefix("data: ")) + assert event_line == "event: response.failed" + assert failed["type"] == "response.failed" + assert failed["response"]["status"] == "failed" + assert failed["response"]["error"]["code"] == expected_code + assert expected_message in failed["response"]["error"]["message"] + + @pytest.mark.asyncio async def test_ttft_keepalive_closes_the_upstream_stream_it_relayed(): """Starlette never calls the produced response, so its own cleanup never runs