From cccfe196e5ee5cf8ee600a3b49a647a15de8429b Mon Sep 17 00:00:00 2001 From: Aasif-Simpplr Date: Tue, 6 Oct 2026 19:53:35 +0530 Subject: [PATCH] fix(proxy): keep extra chunk fields on the streaming fast path _fast_serialize_simple_model_response_stream writes a fixed set of fields, so a field set on the chunk outside its schema (model_extra), for example by an async_post_call_streaming_iterator_hook, was silently dropped from plain text chunks. Fall back to the canonical serializer when the chunk carries an extra the fast payload does not write; service_tier is still served on the fast path. --- litellm/proxy/proxy_server.py | 8 ++++++++ .../proxy/proxy_server/test_streaming_helpers.py | 15 +++++++++++++++ 2 files changed, 23 insertions(+) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 06bfb467b52..d9119388f89 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -9565,6 +9565,14 @@ def _fast_serialize_simple_model_response_stream( ): return None + # Fields set on the chunk outside its schema (e.g. by a streaming iterator + # hook) live in ``model_extra``. The payload below only writes + # ``service_tier``, so any other extra would be silently dropped here + # while the canonical serializer keeps it. + model_extra: Final = getattr(chunk, "model_extra", None) + if model_extra and any(key != "service_tier" for key in model_extra): + return None + choices: Final = getattr(chunk, "choices", None) if not isinstance(choices, list) or len(choices) != 1: return None diff --git a/tests/unit/proxy/proxy_server/test_streaming_helpers.py b/tests/unit/proxy/proxy_server/test_streaming_helpers.py index 69fa195e9d6..2a3eba51f86 100644 --- a/tests/unit/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/unit/proxy/proxy_server/test_streaming_helpers.py @@ -550,6 +550,21 @@ def test_fast_serialize_simple_model_response_stream_with_usage_returns_none_inv assert _fast_serialize_simple_model_response_stream(chunk) is None +def test_fast_serialize_simple_model_response_stream_with_extra_field_returns_none(): + """A field a callback attaches to the chunk is not in the fast payload, so + the fast path must bail rather than drop it.""" + chunk = _simple_chunk() + chunk.gateway_metadata = {"served_by": "fallback-model"} + assert _fast_serialize_simple_model_response_stream(chunk) is None + + +def test_serialize_streaming_chunk_keeps_extra_field(): + chunk = _simple_chunk() + chunk.gateway_metadata = {"served_by": "fallback-model"} + payload = json.loads(_serialize_streaming_chunk(chunk)) + assert payload["gateway_metadata"] == {"served_by": "fallback-model"} + + # --------------------------------------------------------------------------- # _serialize_streaming_chunk # ---------------------------------------------------------------------------