From ca873a1327928e8cc1812d19ae892e10c7399a50 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 05:26:01 -0700 Subject: [PATCH] fix(proxy): keep cursor chat framing and JSON-encode Responses hook errors --- litellm/proxy/common_request_processing.py | 20 ++++-- .../common_utils/responses_stream_errors.py | 9 ++- litellm/proxy/proxy_server.py | 8 ++- .../proxy/response_api_endpoints/endpoints.py | 4 +- .../response_polling/background_streaming.py | 4 +- .../proxy_server/test_streaming_helpers.py | 27 ++++++-- .../response_api_endpoints/test_endpoints.py | 68 +++++++++++++------ 7 files changed, 99 insertions(+), 41 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 27164ba0c83..a920b42b9d0 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1250,6 +1250,16 @@ RESPONSES_FIRST_CHUNK_FRAMING: Final = FirstChunkSseFraming( ) +@dataclass(frozen=True, slots=True) +class SseFraming: + first_chunk: FirstChunkSseFraming + late: LateSseFraming + + +OPENAI_SSE_FRAMING: Final = SseFraming(first_chunk=OPENAI_FIRST_CHUNK_FRAMING, late=OPENAI_LATE_SSE_FRAMING) +RESPONSES_SSE_FRAMING: Final = SseFraming(first_chunk=RESPONSES_FIRST_CHUNK_FRAMING, late=RESPONSES_LATE_SSE_FRAMING) + + 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}) @@ -2945,7 +2955,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, + sse_framing: SseFraming = OPENAI_SSE_FRAMING, ) -> object: """Run the request, sending SSE keepalives while the upstream is still silent. @@ -2986,10 +2996,11 @@ class ProxyBaseLLMRequestProcessing: is_streaming_request=is_streaming_request, contents=contents, skip_pre_call_logic=skip_pre_call_logic, + first_chunk_framing=sse_framing.first_chunk, ), ping_interval_seconds=ttft_keepalive_interval(self.data, llm_router), on_late_failure=_audit_late_failure, - late_framing=late_framing, + late_framing=sse_framing.late, ) async def _process_llm_request( @@ -3013,6 +3024,7 @@ class ProxyBaseLLMRequestProcessing: is_streaming_request: bool | None = False, contents: list[object] | None = None, # Add contents parameter skip_pre_call_logic: bool = False, + first_chunk_framing: FirstChunkSseFraming = OPENAI_FIRST_CHUNK_FRAMING, ) -> Any: """ Common request processing logic for both chat completions and responses API endpoints @@ -3306,9 +3318,7 @@ class ProxyBaseLLMRequestProcessing: request=request, refresh_headers=refresh_stream_headers, general_settings=general_settings, - framing=( - RESPONSES_FIRST_CHUNK_FRAMING if route_type == "aresponses" else OPENAI_FIRST_CHUNK_FRAMING - ), + framing=first_chunk_framing, ) ### CALL HOOKS ### - modify outgoing data diff --git a/litellm/proxy/common_utils/responses_stream_errors.py b/litellm/proxy/common_utils/responses_stream_errors.py index 4a806c32284..0024c210fab 100644 --- a/litellm/proxy/common_utils/responses_stream_errors.py +++ b/litellm/proxy/common_utils/responses_stream_errors.py @@ -61,19 +61,22 @@ class _FailureDetails(LiteLLMBaseModel): class _ErrorPayload(LiteLLMBaseModel): model_config = ConfigDict(frozen=True) - error: Mapping[str, object] + error: Mapping[str, object] | str + + def error_object(self) -> Mapping[str, object]: + return {"message": self.error} if isinstance(self.error, str) else self.error def sse_error_object(data: str | bytes) -> Mapping[str, object] | None: try: - return _ErrorPayload.model_validate_json(data).error + return _ErrorPayload.model_validate_json(data).error_object() except ValidationError: return None def event_error_object(event: Mapping[str, object]) -> Mapping[str, object] | None: try: - return _ErrorPayload.model_validate(event).error or None + return _ErrorPayload.model_validate(event).error_object() or None except ValidationError: return None diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 0377ed809ad..a6648e47de6 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -10202,6 +10202,12 @@ def _format_streaming_sse_chunk(chunk: object) -> str | bytes: return f"data: {chunk}\n\n" +def _format_responses_sse_chunk(chunk: object) -> str | bytes: + if isinstance(chunk, Mapping): + return f"data: {json.dumps(chunk, default=str)}\n\n" + return _format_streaming_sse_chunk(chunk=chunk) + + _SSE_FRAME_DELIMITERS: Final = ("\r\n\r\n", "\n\n", "\r\r") _OPENAI_STREAM_DONE_FRAME: Final = "data: [DONE]\n\n" _MAX_RAW_SSE_BUFFER_CHARS: Final = 8 * 1024 * 1024 @@ -10602,7 +10608,7 @@ async def async_data_generator( if not raw_passthrough: try: if error_state is not None: - yield error_state.mark_emitted(_format_streaming_sse_chunk(chunk=chunk)) + yield error_state.mark_emitted(_format_responses_sse_chunk(chunk=chunk)) else: yield _format_streaming_sse_chunk(chunk=chunk) except Exception as e: diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 1fd487ea076..948a32d4449 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -52,7 +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, + RESPONSES_SSE_FRAMING, ProxyBaseLLMRequestProcessing, cost_headers, create_response, @@ -438,7 +438,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, + sse_framing=RESPONSES_SSE_FRAMING, model=None, user_model=user_model, user_temperature=user_temperature, diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py index cd91bfa204f..99c988317bc 100644 --- a/litellm/proxy/response_polling/background_streaming.py +++ b/litellm/proxy/response_polling/background_streaming.py @@ -21,7 +21,7 @@ from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth -from litellm.proxy.common_request_processing import RESPONSES_LATE_SSE_FRAMING, ProxyBaseLLMRequestProcessing +from litellm.proxy.common_request_processing import RESPONSES_SSE_FRAMING, ProxyBaseLLMRequestProcessing from litellm.proxy.common_utils.responses_stream_errors import ( event_error_object, responses_error_fields, @@ -188,7 +188,7 @@ async def background_streaming_task( user_api_base=user_api_base, version=version, skip_pre_call_logic=True, - late_framing=RESPONSES_LATE_SSE_FRAMING, + sse_framing=RESPONSES_SSE_FRAMING, ) # Process streaming response following OpenAI events format diff --git a/tests/unit/proxy/proxy_server/test_streaming_helpers.py b/tests/unit/proxy/proxy_server/test_streaming_helpers.py index 93eb8fa226b..b3a6e7b1fbf 100644 --- a/tests/unit/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/unit/proxy/proxy_server/test_streaming_helpers.py @@ -17,7 +17,7 @@ from __future__ import annotations import asyncio import json -from collections.abc import AsyncIterator, Callable +from collections.abc import AsyncIterator, Callable, Mapping from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock @@ -148,7 +148,8 @@ async def test_responses_stream_dispatches_named_events_without_a_chat_done_sent @pytest.mark.asyncio -async def test_responses_stream_sends_no_failure_event_after_an_upstream_error_event() -> None: +@pytest.mark.parametrize("as_mapping", [False, True]) +async def test_responses_stream_sends_no_failure_event_after_an_upstream_error_event(as_mapping: bool) -> None: response: Final = ResponsesAPIResponse(id="resp_error_then_raise", created_at=1, model="local-model", output=[]) created: Final = ResponseCreatedEvent.model_validate( {"type": "response.created", "sequence_number": 0, "response": response} @@ -157,9 +158,9 @@ async def test_responses_stream_sends_no_failure_event_after_an_upstream_error_e {"type": "error", "sequence_number": 1, "code": "server_error", "message": "upstream failed", "param": None} ) - async def upstream() -> AsyncIterator[BaseModel]: + async def upstream() -> AsyncIterator[BaseModel | Mapping[str, object]]: yield created - yield error_event + yield error_event.model_dump() if as_mapping else error_event raise litellm.APIError(status_code=500, message="Stream close failed", llm_provider="openai", model="local-model") frames: Final = tuple( @@ -172,14 +173,26 @@ async def test_responses_stream_sends_no_failure_event_after_an_upstream_error_e ) assert tuple(frame.split("\n", 1)[0] for frame in frames) == ("event: response.created", "event: error") + assert json.loads(frames[-1].split("data: ", 1)[1])["message"] == "upstream failed" @pytest.mark.asyncio -async def test_a_guardrail_error_string_mid_responses_stream_ends_it_as_a_named_error_event() -> None: +@pytest.mark.parametrize( + ("guardrail_error", "error_object"), + [ + ( + {"message": "response withheld", "type": "guardrail_scan_error", "code": 500}, + {"message": "response withheld", "type": "guardrail_scan_error", "code": 500}, + ), + ("response withheld", {"message": "response withheld"}), + ], +) +async def test_a_guardrail_error_string_mid_responses_stream_ends_it_as_a_named_error_event( + guardrail_error: Mapping[str, object] | str, error_object: Mapping[str, object] +) -> None: response: Final = ResponsesAPIResponse(id="resp_hook_error", created_at=1, model="local-model", output=[]) created: Final = ResponseCreatedEvent.model_validate( {"type": "response.created", "sequence_number": 0, "response": response} ) - guardrail_error: Final = {"message": "response withheld", "type": "guardrail_scan_error", "code": 500} async def upstream() -> AsyncIterator[object]: yield created @@ -203,7 +216,7 @@ async def test_a_guardrail_error_string_mid_responses_stream_ends_it_as_a_named_ "message": "response withheld", "param": None, "sequence_number": 1, - "error": guardrail_error, + "error": error_object, } diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index abf5933bd35..d7d6321dd2b 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -181,14 +181,16 @@ async def test_streaming_upstream_errors_keep_the_client_protocol( class _FirstFrameHook(CustomLogger): - def __init__(self, frame: str | bytes) -> None: + def __init__(self, first: str | bytes | Exception) -> None: super().__init__() - self.frame: Final = frame + self.first: Final = first async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: object, response: object, request_data: object ) -> AsyncIterator[str | bytes]: - yield self.frame + if isinstance(self.first, Exception): + raise self.first + yield self.first _DATA_ERROR_FRAME: Final = 'data: {"error": {"message": "denied", "type": "guardrail_scan_error", "code": 403}}\n\n' @@ -197,22 +199,9 @@ _NAMED_ERROR_FRAME: Final = ( ) -@pytest.mark.asyncio -@pytest.mark.parametrize( - "path,first_frame,expected_status", - [ - ("/v1/responses", _DATA_ERROR_FRAME, 403), - ("/v1/chat/completions", _DATA_ERROR_FRAME, 403), - ("/v1/chat/completions", _NAMED_ERROR_FRAME, 200), - ("/v1/responses", _NAMED_ERROR_FRAME, 529), - ("/v1/responses", _NAMED_ERROR_FRAME.replace(b"event: error", b"event: response.created"), 200), - ("/v1/chat/completions", _DATA_ERROR_FRAME + "data: [DONE]\n\n", 200), - ("/v1/responses", _DATA_ERROR_FRAME + "data: [DONE]\n\n", 200), - ], -) -async def test_only_a_lone_first_error_frame_answers_with_its_status( - monkeypatch: pytest.MonkeyPatch, path: str, first_frame: str | bytes, expected_status: int -) -> None: +async def _post_through_first_frame_hook( + monkeypatch: pytest.MonkeyPatch, path: str, first: str | bytes | Exception +) -> httpx.Response: import litellm.proxy.proxy_server as ps from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -230,14 +219,14 @@ async def test_only_a_lone_first_error_frame_answers_with_its_status( ) monkeypatch.setattr(ps, "llm_router", router) monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - monkeypatch.setattr(litellm, "callbacks", [_FirstFrameHook(first_frame)]) + monkeypatch.setattr(litellm, "callbacks", [_FirstFrameHook(first)]) 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")).respond( 200, content="data: [DONE]\n\n", headers={"Content-Type": "text/event-stream"} ) async with httpx.AsyncClient(transport=httpx.ASGITransport(app), base_url="http://testserver") as client: - result: Final = await client.post( + return await client.post( path, json={ "model": model, @@ -246,6 +235,26 @@ async def test_only_a_lone_first_error_frame_answers_with_its_status( }, ) + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "path,first_frame,expected_status", + [ + ("/v1/responses", _DATA_ERROR_FRAME, 403), + ("/v1/chat/completions", _DATA_ERROR_FRAME, 403), + ("/v1/chat/completions", _NAMED_ERROR_FRAME, 200), + ("/cursor/chat/completions", _NAMED_ERROR_FRAME, 200), + ("/v1/responses", _NAMED_ERROR_FRAME, 529), + ("/v1/responses", _NAMED_ERROR_FRAME.replace(b"event: error", b"event: response.created"), 200), + ("/v1/chat/completions", _DATA_ERROR_FRAME + "data: [DONE]\n\n", 200), + ("/v1/responses", _DATA_ERROR_FRAME + "data: [DONE]\n\n", 200), + ], +) +async def test_only_a_lone_first_error_frame_answers_with_its_status( + monkeypatch: pytest.MonkeyPatch, path: str, first_frame: str | bytes, expected_status: int +) -> None: + result: Final = await _post_through_first_frame_hook(monkeypatch, path, first_frame) + assert result.status_code == expected_status, result.text if expected_status == 200: assert result.headers["content-type"].startswith("text/event-stream"), result.text @@ -253,6 +262,23 @@ async def test_only_a_lone_first_error_frame_answers_with_its_status( assert result.json()["error"]["message"] == "denied" +@pytest.mark.asyncio +@pytest.mark.parametrize("path", ["/cursor/chat/completions", "/v1/chat/completions"]) +async def test_a_hook_failure_before_the_first_chunk_answers_in_chat_completions_shape( + monkeypatch: pytest.MonkeyPatch, path: str +) -> None: + message: Final = "rate limited before the first chunk" + result: Final = await _post_through_first_frame_hook( + monkeypatch, path, HTTPException(status_code=429, detail=message) + ) + frames: Final = tuple(frame for frame in result.text.split("\n\n") if frame) + + assert result.status_code == 429, result.text + assert frames[-1] == "data: [DONE]", result.text + assert message in json.loads(frames[0].removeprefix("data: "))["error"]["message"] + assert len(frames) == 2, result.text + + @pytest.mark.parametrize("path", ["/v1/responses", "/v1/chat/completions"]) def test_a_failure_after_keepalive_pings_keeps_the_client_protocol(monkeypatch: pytest.MonkeyPatch, path: str) -> None: async def request_and_check() -> None: