From 6b5ed50633348efa3e7c9b416e2a99d5a959a00a Mon Sep 17 00:00:00 2001 From: HootRock Date: Fri, 23 Oct 2026 07:42:05 +0800 Subject: [PATCH] fix(proxy): preserve Responses SSE event framing --- .../guardrail_translation/handler.py | 2 +- .../common_utils/responses_stream_errors.py | 5 +- litellm/proxy/proxy_server.py | 2 +- .../proxy/response_api_endpoints/endpoints.py | 1 - ...test_openai_responses_guardrail_handler.py | 9 ++- .../proxy_server/test_streaming_helpers.py | 62 +++++++++++++++++-- .../response_api_endpoints/test_endpoints.py | 14 +++-- .../proxy/test_response_model_sanitization.py | 5 +- 8 files changed, 85 insertions(+), 15 deletions(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 3ea126c5a2a..d08806b19cb 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -1789,7 +1789,6 @@ class OpenAIResponsesHandler(BaseTranslation): new output item under the same response id, and close with a ``response.completed`` carrying only the replacement item. - The proxy's data generator appends ``data: [DONE]`` itself. """ events: Final = ( self._block_continuation_events(exc, responses_so_far or ()) @@ -1797,6 +1796,7 @@ class OpenAIResponsesHandler(BaseTranslation): else self._standalone_block_events(exc) ) return tuple( + f"event: {event.type.value if isinstance(event.type, ResponsesAPIStreamEvents) else event.type}\n" f"data: {event.model_dump_json(exclude_none=True, exclude_unset=True, serialize_as_any=True)}\n\n".encode() for event in events ) diff --git a/litellm/proxy/common_utils/responses_stream_errors.py b/litellm/proxy/common_utils/responses_stream_errors.py index 3409b4b6657..6e563b2b257 100644 --- a/litellm/proxy/common_utils/responses_stream_errors.py +++ b/litellm/proxy/common_utils/responses_stream_errors.py @@ -131,7 +131,10 @@ class ResponsesStreamErrorState: self.created_at = event.response.created_at if event.type in ("response.completed", "response.failed", "response.incomplete"): self.terminal_emitted = True - return frame + if not event.type or "\n" in event.type or "\r" in event.type: + return frame + prefix: Final = f"event: {event.type}\n" + return prefix.encode("utf-8") + frame if isinstance(frame, bytes) else prefix + frame def format_failure(self, exception: Exception) -> str | None: if self.terminal_emitted: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f632317d9a8..6aff6601ac4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -10390,7 +10390,7 @@ async def async_data_generator( stream_completed = False client_disconnected = False error_state: Final = ResponsesStreamErrorState() if responses_stream_errors else None - ends_with_openai_done: Final = not request_data.get("_litellm_skip_openai_stream_done") + ends_with_openai_done: Final = not responses_stream_errors and not request_data.get("_litellm_skip_openai_stream_done") needs_iterator_wrap: Final = proxy_logging_obj.needs_iterator_wrap() stream_iterator: Final = _object_stream( proxy_logging_obj.async_post_call_streaming_iterator_hook( diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 312c5124ca5..8a13ce876ea 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -493,7 +493,6 @@ async def responses_api( async def _blocked_stream() -> AsyncGenerator[str, None]: for chunk in block_chunks: yield chunk.decode() - yield "data: [DONE]\n\n" return await create_response(generator=_blocked_stream(), media_type="text/event-stream", headers={}) return build_blocked_response(e) 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 2f10fd6fd0d..23dbf3906d6 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 @@ -3766,7 +3766,14 @@ class TestBuildBlockSseChunks: def _payloads(self, chunks): import json - return [json.loads(chunk.decode().removeprefix("data: ").strip()) for chunk in chunks] + payloads = [ + json.loads(next(line[6:] for line in chunk.decode().splitlines() if line.startswith("data: "))) + for chunk in chunks + ] + assert [chunk.decode().splitlines()[0] for chunk in chunks] == [ + f"event: {payload['type']}" for payload in payloads + ] + return payloads def test_standalone_block_emits_complete_synthetic_stream(self): handler = OpenAIResponsesHandler() diff --git a/tests/unit/proxy/proxy_server/test_streaming_helpers.py b/tests/unit/proxy/proxy_server/test_streaming_helpers.py index 9a3516d20d6..015269801b8 100644 --- a/tests/unit/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/unit/proxy/proxy_server/test_streaming_helpers.py @@ -50,9 +50,12 @@ from litellm.proxy.proxy_server import ( select_data_generator, ) from litellm.types.llms.openai import ( + ErrorEvent, + OutputTextDeltaEvent, ResponseCompletedEvent, ResponseCreatedEvent, ResponseFailedEvent, + ResponseIncompleteEvent, ResponsesAPIResponse, ) from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Usage @@ -92,6 +95,56 @@ async def _async_iter_raises(exc: Exception): raise exc +@pytest.mark.asyncio +@pytest.mark.parametrize("terminal_type", ["response.completed", "response.failed", "response.incomplete", "error"]) +async def test_responses_stream_dispatches_named_events_without_a_chat_done_sentinel(terminal_type: str) -> None: + response: Final = ResponsesAPIResponse(id="resp_framing", created_at=1, model="local-model", output=[]) + created: Final = ResponseCreatedEvent.model_validate( + {"type": "response.created", "sequence_number": 0, "response": response} + ) + delta: Final = OutputTextDeltaEvent.model_validate( + { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": "msg_framing", + "output_index": 0, + "content_index": 0, + "delta": "hello", + } + ) + terminal_models: Final = { + "response.completed": ResponseCompletedEvent, + "response.failed": ResponseFailedEvent, + "response.incomplete": ResponseIncompleteEvent, + "error": ErrorEvent, + } + terminal: Final = terminal_models[terminal_type].model_validate( + {"type": terminal_type, "sequence_number": 2, "response": response} + ) + + async def upstream() -> AsyncIterator[BaseModel]: + for event in (created, delta, terminal): + yield event + + frames: Final = tuple( + [ + frame.decode() if isinstance(frame, bytes) else frame + async for frame in select_data_generator( + response=upstream(), user_api_key_dict=_user_auth(), request_data={}, responses_stream_errors=True + ) + ] + ) + assert len(frames) == 3 + for frame, event_type in zip( + frames, ("response.created", "response.output_text.delta", terminal_type), strict=True + ): + assert frame.startswith(f"event: {event_type}\n") + payload: Final = json.loads( + next(line.removeprefix("data: ") for line in frame.splitlines() if line.startswith("data: ")) + ) + assert payload["type"] == event_type + assert json.loads(frames[1].split("data: ", 1)[1])["delta"] == "hello" + class _FakeStream: def __init__(self, chunks, hidden_params=None): self._chunks = chunks @@ -1132,14 +1185,15 @@ async def test_responses_stream_failure_sends_the_failure_hook_replacement(monke ] ) failure: Final = ResponseFailedEvent.model_validate_json( - next(line.removeprefix("data: ") for line in frames[-2].splitlines() if line.startswith("data: ")) + next(line.removeprefix("data: ") for line in frames[-1].splitlines() if line.startswith("data: ")) ) assert _BLOCK_MESSAGE not in "".join(frames) assert failure.response.error is not None assert failure.response.error["code"] == "permission_error" assert "Request blocked by policy" in failure.response.error["message"] - assert frames[-1] == "data: [DONE]\n\n" + assert frames[-1].startswith("event: response.failed\n") + assert "data: [DONE]\n\n" not in frames @pytest.mark.asyncio @@ -1335,8 +1389,8 @@ async def test_responses_stream_keeps_tool_deltas_and_only_emits_a_valid_termina for frame in event_frames ) - assert decoded[-1] == "data: [DONE]\n\n" - assert len(decoded) == len(event_frames) + 1 + assert len(decoded) == len(event_frames) + assert "data: [DONE]\n\n" not in decoded assert payloads[0]["response"]["id"] == "resp_visible" assert payloads[1] == tool_delta.model_dump() assert len(payloads) == 3 diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 985aee061bc..0bf14f59b1a 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -142,8 +142,9 @@ async def test_streaming_upstream_errors_keep_the_client_protocol( assert result.status_code == 200, result.text assert message in result.text if path == "/v1/responses": - assert frames[-1] == "data: [DONE]", result.text - assert frames[-2].startswith("event: response.failed\n"), result.text + assert all("data: [DONE]" not in frame for frame in frames), result.text + assert frames[-1].startswith("event: response.failed\n"), result.text + assert [frame.splitlines()[0] for frame in frames] == [f"event: {event['type']}" for event in events] if partial: assert [event["type"] for event in events] == [ "response.created", @@ -4862,9 +4863,11 @@ class TestGuardrailBlockedResponsesShape: assert response.status_code == 200, response.text assert response.headers["content-type"].startswith("text/event-stream") frames = _sse_data_frames(response.text) - assert frames[-1] == "[DONE]" - events = [json.loads(frame) for frame in frames[:-1]] + assert "[DONE]" not in frames + events = [json.loads(frame) for frame in frames] types = [event["type"] for event in events] + assert types[-1] == "response.completed" + assert [line[7:] for line in response.text.splitlines() if line.startswith("event: ")] == types assert "response.created" in types completed = [event for event in events if event["type"] == "response.completed"] assert len(completed) == 1 @@ -4891,7 +4894,8 @@ class TestGuardrailBlockedResponsesShape: assert response.status_code == 200, response.text frames = _sse_data_frames(response.text) - completed = [json.loads(frame) for frame in frames[:-1] if json.loads(frame)["type"] == "response.completed"] + assert "[DONE]" not in frames + completed = [json.loads(frame) for frame in frames if json.loads(frame)["type"] == "response.completed"] usage = completed[0]["response"]["usage"] assert usage["input_tokens"] == 14 assert usage["output_tokens"] == 20 diff --git a/tests/unit/proxy/test_response_model_sanitization.py b/tests/unit/proxy/test_response_model_sanitization.py index f52736453c1..6f3c5d8430b 100644 --- a/tests/unit/proxy/test_response_model_sanitization.py +++ b/tests/unit/proxy/test_response_model_sanitization.py @@ -763,7 +763,10 @@ def test_restamp_streaming_chunk_rewrites_nested_model_in_guardrail_block_frames )[0] for frame in frames ] - events = [json.loads(frame.decode().removeprefix("data: ")) for frame in client_frames] + events = [ + json.loads(next(line[6:] for line in frame.decode().splitlines() if line.startswith("data: "))) + for frame in client_frames + ] assert all(frame.endswith(b"\n\n") for frame in client_frames) assert events[-1]["type"] == "response.completed"