From 34156235687cc2d6a9bf11782844c889886b6f48 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 06:05:44 -0700 Subject: [PATCH] fix(proxy): keep masked mappings and odd keys safe in Responses stream frames Only read-only dict views and pydantic models nested in a hook chunk or error event expand into JSON objects. Other mappings, such as httpx.Headers, keep their own string form so masked values stay masked. Keys JSON cannot take become strings, and a chunk JSON still cannot encode goes out as the plain data frame it was before. --- .../common_utils/responses_stream_errors.py | 31 ++++++++++-- litellm/proxy/proxy_server.py | 7 ++- .../proxy_server/test_streaming_helpers.py | 47 +++++++++++++++++++ 3 files changed, 77 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/common_utils/responses_stream_errors.py b/litellm/proxy/common_utils/responses_stream_errors.py index ac4854f4014..137b1221d7d 100644 --- a/litellm/proxy/common_utils/responses_stream_errors.py +++ b/litellm/proxy/common_utils/responses_stream_errors.py @@ -74,14 +74,37 @@ def _is_mapping( return isinstance(value, Mapping) -def jsonable_stream_value(value: object) -> object: - if _is_mapping(value): - return {**value} +def _is_mapping_proxy( + value: object, +) -> TypeIs[MappingProxyType[object, object]]: # guard-ok: isinstance decides it, keys and values stay object + return isinstance(value, MappingProxyType) + + +def _json_key(key: object) -> str | int | float | bool | None: + return key if key is None or isinstance(key, str | int | float | bool) else str(key) + + +def _with_json_keys(value: Mapping[object, object]) -> Mapping[str | int | float | bool | None, object]: + return {_json_key(key): item for key, item in value.items()} + + +def _jsonable_stream_value(value: object) -> object: + if _is_mapping_proxy(value): + return _with_json_keys(value) if isinstance(value, BaseModel): return value.model_dump() return str(value) +def responses_chunk_json(chunk: object) -> str | None: + if not _is_mapping(chunk): + return None + try: + return json.dumps(_with_json_keys(chunk), default=_jsonable_stream_value) + except (TypeError, ValueError): + return None + + def sse_error_object(data: str | bytes) -> Mapping[str, object] | None: try: return _ErrorPayload.model_validate_json(data).error_object() @@ -206,7 +229,7 @@ class ResponsesStreamErrorState: "sequence_number": self.sequence_number + 1, "error": error, }, - default=jsonable_stream_value, + default=_jsonable_stream_value, ) return f"event: error\ndata: {payload}\n\n" diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 2ddb7a9241b..7acdd582658 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -539,7 +539,7 @@ from litellm.proxy.common_utils.periodic_reload_schedule import ( ) from litellm.proxy.common_utils.proxy_state import ProxyState from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob -from litellm.proxy.common_utils.responses_stream_errors import ResponsesStreamErrorState, jsonable_stream_value +from litellm.proxy.common_utils.responses_stream_errors import ResponsesStreamErrorState, responses_chunk_json from litellm.proxy.common_utils.scheduled_job_stagger import ( apply_scheduled_job_stagger, attach_job_timing_logger, @@ -10203,9 +10203,8 @@ def _format_streaming_sse_chunk(chunk: object) -> str | bytes: def _format_responses_sse_chunk(chunk: object) -> str | bytes: - if isinstance(chunk, Mapping): - return f"data: {json.dumps(chunk, default=jsonable_stream_value)}\n\n" - return _format_streaming_sse_chunk(chunk=chunk) + data: Final = responses_chunk_json(chunk) + return _format_streaming_sse_chunk(chunk=chunk) if data is None else f"data: {data}\n\n" _SSE_FRAME_DELIMITERS: Final = ("\r\n\r\n", "\n\n", "\r\r") diff --git a/tests/unit/proxy/proxy_server/test_streaming_helpers.py b/tests/unit/proxy/proxy_server/test_streaming_helpers.py index a188326c685..05bde512608 100644 --- a/tests/unit/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/unit/proxy/proxy_server/test_streaming_helpers.py @@ -219,6 +219,53 @@ def test_responses_error_event_encodes_nested_error_fields_as_json_objects() -> payload: Final = json.loads(frame.split("data: ", 1)[1]) assert payload["error"]["provider_specific_fields"] == {"upstream": {"type": "overloaded_error"}} + +_UPSTREAM_HEADERS: Final = httpx.Headers({"authorization": "Bearer hook-canary", "x-request-id": "req_1"}) + + +def _responses_frame_payload(nested: object) -> Mapping[str, object]: + chunk_frame: Final = _format_responses_sse_chunk( + chunk=MappingProxyType({"type": "response.output_text.delta", "nested": nested}) + ) + error_frame: Final = ResponsesStreamErrorState().format_error({"message": "blocked", "nested": nested}) + assert isinstance(chunk_frame, str) + chunk_payload: Final = json.loads(chunk_frame.removeprefix("data: ")) + error_payload: Final = json.loads(error_frame.split("data: ", 1)[1]) + assert chunk_payload["nested"] == error_payload["error"]["nested"] + return chunk_payload + + +def test_responses_stream_frames_send_nested_header_mappings_in_their_masked_string_form() -> None: + assert _responses_frame_payload(_UPSTREAM_HEADERS)["nested"] == str(_UPSTREAM_HEADERS) + + +def test_responses_stream_frames_turn_mapping_keys_json_cannot_take_into_strings() -> None: + assert _responses_frame_payload(MappingProxyType({("a", "b"): 1}))["nested"] == {str(("a", "b")): 1} + + +def test_responses_sse_chunk_turns_top_level_keys_json_cannot_take_into_strings() -> None: + frame: Final = _format_responses_sse_chunk( + chunk=MappingProxyType({"type": "response.output_text.delta", ("a", "b"): "x"}) + ) + + assert isinstance(frame, str) + assert json.loads(frame.removeprefix("data: ")) == {"type": "response.output_text.delta", str(("a", "b")): "x"} + + +def _self_referencing_chunk() -> Mapping[str, object]: + chunk: Final[dict[str, object]] = {"type": "response.output_text.delta"} + chunk["self"] = chunk # mutable-ok: a self-reference needs the dict to exist before it can hold itself + return chunk + + +@pytest.mark.parametrize( + "chunk", + [{"type": "response.output_text.delta", "meta": {("a", "b"): 1}}, _self_referencing_chunk()], + ids=["nested_dict_tuple_key", "self_reference"], +) +def test_responses_sse_chunk_json_cannot_encode_streams_as_a_plain_data_frame(chunk: Mapping[str, object]) -> None: + assert _format_responses_sse_chunk(chunk=chunk) == _format_streaming_sse_chunk(chunk=chunk) + @pytest.mark.asyncio @pytest.mark.parametrize( ("guardrail_error", "error_object"),