From b9e838282e1d3b7da4ecd03333b46cd7bbb091aa Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 2 Nov 2026 12:03:21 -0800 Subject: [PATCH] fix(proxy): skip containers already visited when checking hook chunks for JSON --- .../common_utils/responses_stream_errors.py | 24 +++++++++++++++---- .../llm_translation/test_containers_e2e.py | 2 +- .../proxy_server/test_streaming_helpers.py | 8 +++++++ 3 files changed, 29 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/common_utils/responses_stream_errors.py b/litellm/proxy/common_utils/responses_stream_errors.py index d3de5e18cea..2bb17281f3d 100644 --- a/litellm/proxy/common_utils/responses_stream_errors.py +++ b/litellm/proxy/common_utils/responses_stream_errors.py @@ -3,6 +3,7 @@ import re import time from collections.abc import Iterator, Mapping, Sequence from contextlib import suppress +from dataclasses import dataclass from enum import Enum from http import HTTPStatus from itertools import accumulate, chain, groupby, islice, repeat, takewhile @@ -116,13 +117,28 @@ def _children(value: object) -> tuple[object, ...]: return tuple(value) if _is_sequence(value) else () -def _next_level(level: tuple[object, ...], _: None) -> tuple[object, ...]: - return tuple(chain.from_iterable(_children(value) for value in level)) +@dataclass(frozen=True, slots=True) +class _PlainJsonLevel: + values: tuple[object, ...] + seen: frozenset[int] + + +def _next_level(level: _PlainJsonLevel, _: None) -> _PlainJsonLevel: + unseen: Final = {id(value): value for value in level.values if id(value) not in level.seen} + return _PlainJsonLevel( + values=tuple(chain.from_iterable(map(_children, unseen.values()))), + seen=level.seen.union(unseen), + ) + + +def _settles(level: _PlainJsonLevel) -> bool: + return not level.values or not all(map(_is_plain_node, level.values)) def _is_plain_json(chunk: object) -> bool: - levels: Final = islice(accumulate(repeat(None), _next_level, initial=(chunk,)), _MAX_PLAIN_JSON_DEPTH) - return next((not level for level in levels if not level or not all(map(_is_plain_node, level))), False) + first: Final = _PlainJsonLevel(values=(chunk,), seen=frozenset()) + levels: Final = islice(accumulate(repeat(None), _next_level, initial=first), _MAX_PLAIN_JSON_DEPTH) + return next((not level.values for level in levels if _settles(level)), False) def responses_chunk_data(chunk: object) -> str | None: diff --git a/tests/e2e/llm_translation/test_containers_e2e.py b/tests/e2e/llm_translation/test_containers_e2e.py index 0a494d11108..95dc68341e8 100644 --- a/tests/e2e/llm_translation/test_containers_e2e.py +++ b/tests/e2e/llm_translation/test_containers_e2e.py @@ -32,7 +32,7 @@ second regression, since the global-credential fallback then reaches the container anyway. The streaming cell repeats the flow with ``stream=True`` and uploads right -after the last event. The OpenAI SDK closes the connection at ``[DONE]``, so an +after the last event. The OpenAI SDK closes the connection after the last frame, so an ownership row written after the stream is cancelled with the body task and every follow-up container call 403s (LIT-8612); the row has to land before the ``response.completed`` frame goes out. diff --git a/tests/unit/proxy/proxy_server/test_streaming_helpers.py b/tests/unit/proxy/proxy_server/test_streaming_helpers.py index e55e195e6cc..ec09b1d4b28 100644 --- a/tests/unit/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/unit/proxy/proxy_server/test_streaming_helpers.py @@ -293,6 +293,12 @@ def _self_referencing_chunk() -> Mapping[str, object]: return chunk +def _chunk_holding_itself_twice() -> Mapping[str, object]: + chunk: Final[dict[str, object]] = {"type": "response.output_text.delta"} + chunk.update(left=chunk, right=chunk) # mutable-ok: a self-reference needs the dict to exist before it can hold itself + return chunk + + _REDACTED_HEADERS: Final = RedactedDict({"authorization": "Bearer hook-canary"}) @@ -309,6 +315,7 @@ _REDACTED_HEADERS: Final = RedactedDict({"authorization": "Bearer hook-canary"}) {"type": "response.output_text.delta", ("a", "b"): 1}, {"type": "response.output_text.delta", "score": float("nan")}, _self_referencing_chunk(), + _chunk_holding_itself_twice(), ], ids=[ "masked_string", @@ -321,6 +328,7 @@ _REDACTED_HEADERS: Final = RedactedDict({"authorization": "Bearer hook-canary"}) "tuple_key", "nan", "holds_itself", + "holds_itself_twice", ], ) def test_responses_sse_chunk_that_json_cannot_show_as_printed_keeps_the_base_frame(chunk: Mapping[object, object]) -> None: