From bbd446b476526864e6f0abeb3a6cb6728e9772ed Mon Sep 17 00:00:00 2001 From: alberto Date: Sat, 17 Oct 2026 20:52:10 +0100 Subject: [PATCH] fix(proxy): report requested model in streamed Responses response.model (#46662) * chore: Add failing test case Signed-off-by: Alberto Perdomo * fix: Restamp nested response.model in streamed Responses events Signed-off-by: Alberto Perdomo * chore: Add tests Signed-off-by: Alberto Perdomo * chore: Run format and lint, and fix Signed-off-by: Alberto Perdomo * fix: Type nested Responses restamp to satisfy basedpyright gate Signed-off-by: Alberto Perdomo * fix(proxy): restamp the nested model on streamed response.queued events Background streams send response.queued as a GenericEvent, so its nested response.model still carried the backend name. Also drop the new comments * fix(proxy): restamp a nested response kept as a dict on streamed lifecycle events * test(proxy): drive typed events through the nested restamp no-op test * fix(proxy): restamp the nested model on guardrail block frames in streamed responses * test(proxy): assert a cache-polled background response reports the requested model * fix(proxy): restamp nested model on dict events and on every SSE data line --------- Signed-off-by: Alberto Perdomo Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 105 ++++- ...est_background_response_billing_proxies.py | 4 +- .../proxy/test_response_model_sanitization.py | 415 +++++++++++++++++- 3 files changed, 500 insertions(+), 24 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1e16eb626fe..3ed3c57a96c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -954,7 +954,14 @@ from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionSystemMessage, ChatCompletionToolParam, + GenericEvent, HttpxBinaryResponseContent, + ResponseCompletedEvent, + ResponseCreatedEvent, + ResponseFailedEvent, + ResponseIncompleteEvent, + ResponseInProgressEvent, + ResponsesAPIResponse, ) from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry from litellm.types.proxy.management_endpoints.model_management_endpoints import ( @@ -9922,7 +9929,7 @@ def _restamp_streaming_chunk_model( # Note: This warning is intentionally verbose. A mismatch is a useful signal that an # internal provider/deployment identifier is leaking into the public API, and helps # maintainers/operators catch regressions while preserving OpenAI-compatible output. - if not target_model or not isinstance(chunk, (BaseModel, dict)): + if not target_model or not isinstance(chunk, (BaseModel, dict, bytes)): return chunk, model_mismatch_logged # For Azure Model Router, preserve the actual model used in each chunk @@ -9934,9 +9941,21 @@ def _restamp_streaming_chunk_model( if not fallback_was_attempted and request_data.get("fastest_response", False): return chunk, model_mismatch_logged - downstream_model: Final = chunk.get("model") if isinstance(chunk, dict) else getattr(chunk, "model", None) + if isinstance(chunk, bytes): + if request_data.get("_litellm_raw_sse_stream"): + return chunk, model_mismatch_logged + return _restamp_nested_response_model_in_sse_frame(chunk, target_model), model_mismatch_logged + + typed_chunk: Final[BaseModel | Mapping[str, object]] = chunk + restamped_chunk: Final = _restamp_nested_response_model(typed_chunk, target_model) + + downstream_model: Final = ( + restamped_chunk.get("model") + if isinstance(restamped_chunk, Mapping) + else getattr(restamped_chunk, "model", None) + ) if downstream_model == target_model: - return chunk, model_mismatch_logged + return restamped_chunk, model_mismatch_logged if not model_mismatch_logged and downstream_model != target_model: verbose_proxy_logger.debug( @@ -9954,22 +9973,92 @@ def _restamp_streaming_chunk_model( logging_obj: Final = request_data.get("litellm_logging_obj") if isinstance(logging_obj, LiteLLMLoggingObj): logging_obj.client_facing_stream_model = target_model - if isinstance(chunk, dict): - return {**chunk, "model": target_model}, model_mismatch_logged + if isinstance(restamped_chunk, Mapping): + return {**restamped_chunk, "model": target_model}, model_mismatch_logged try: - return chunk.model_copy(update={"model": target_model}), model_mismatch_logged + return restamped_chunk.model_copy(update={"model": target_model}), model_mismatch_logged except Exception as e: verbose_proxy_logger.error( "litellm_call_id=%s: failed to override chunk.model=%r on chunk_type=%s. error=%s", request_data.get("litellm_call_id"), target_model, - type(chunk), + type(restamped_chunk), str(e), exc_info=True, ) - return chunk, model_mismatch_logged + return restamped_chunk, model_mismatch_logged + + +def _nested_response_of(chunk: BaseModel) -> object: + if isinstance( + chunk, + ( + ResponseCreatedEvent, + ResponseInProgressEvent, + ResponseCompletedEvent, + ResponseFailedEvent, + ResponseIncompleteEvent, + ), + ): + return chunk.response + if isinstance(chunk, GenericEvent): + extra_fields: Final[Mapping[str, object]] = chunk.model_extra or {} + return extra_fields.get("response") + return None + + +def _restamp_nested_response_model( + chunk: BaseModel | Mapping[str, object], target_model: str +) -> BaseModel | Mapping[str, object]: + if not isinstance(chunk, BaseModel): + restamped_event: Final = _event_with_restamped_nested_model(chunk, target_model) + return chunk if restamped_event is None else restamped_event + nested_response: Final = _nested_response_of(chunk) + if isinstance(nested_response, ResponsesAPIResponse): + if nested_response.model in (None, target_model): + return chunk + return chunk.model_copy(update={"response": nested_response.model_copy(update={"model": target_model})}) + untyped_response: Final = as_str_mapping(nested_response) + if untyped_response is None or untyped_response.get("model") in (None, target_model): + return chunk + return chunk.model_copy(update={"response": {**untyped_response, "model": target_model}}) + + +def _event_with_restamped_nested_model(event: Mapping[str, object], target_model: str) -> Mapping[str, object] | None: + nested_response: Final = as_str_mapping(event.get("response")) + if nested_response is None or nested_response.get("model") in (None, target_model): + return None + return {**event, "response": {**nested_response, "model": target_model}} + + +_SSE_DATA_FIELD: Final = b"data:" +_SSE_NESTED_RESPONSE_MARKER: Final = b'"response"' +_SSE_EVENT_OBJECT: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) + + +def _restamp_nested_response_model_in_sse_frame(frame: bytes, target_model: str) -> bytes: + if _SSE_NESTED_RESPONSE_MARKER not in frame: + return frame + return b"".join( + _restamp_nested_response_model_in_sse_line(line, target_model) for line in frame.splitlines(keepends=True) + ) + + +def _restamp_nested_response_model_in_sse_line(line: bytes, target_model: str) -> bytes: + if not line.startswith(_SSE_DATA_FIELD): + return line + try: + event: Final = _SSE_EVENT_OBJECT.validate_json(line.removeprefix(_SSE_DATA_FIELD).strip()) + except ValidationError: + return line + restamped_event: Final = _event_with_restamped_nested_model(event, target_model) + if restamped_event is None: + return line + serialized_event: Final = json.dumps(restamped_event, separators=(",", ":"), ensure_ascii=False) + line_ending: Final = line[len(line.rstrip(b"\r\n")) :] + return _SSE_DATA_FIELD + b" " + serialized_event.encode() + line_ending def _fast_serialize_simple_model_response_stream( diff --git a/tests/integration/spend/test_background_response_billing_proxies.py b/tests/integration/spend/test_background_response_billing_proxies.py index e83f2629f33..4b4185f5da2 100644 --- a/tests/integration/spend/test_background_response_billing_proxies.py +++ b/tests/integration/spend/test_background_response_billing_proxies.py @@ -549,7 +549,7 @@ def _cache_polled_create(proxy: Gateway, key: str, model: str, create: Mapping[s @pytest.mark.timeout(OWNED_PROXY_SECONDS + 120) @pytest.mark.parametrize("create", CACHE_POLLED_CREATES.keys()) -def test_a_background_response_the_proxy_polls_through_its_cache_bills_its_creator_once( +def test_a_background_response_polled_through_the_cache_reports_its_alias_and_bills_its_creator_once( gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, create: str ) -> None: body, provider_sees_background = CACHE_POLLED_CREATES[create] @@ -564,6 +564,8 @@ def test_a_background_response_the_proxy_polls_through_its_cache_bills_its_creat seconds=SPEND_RETRY_SECONDS, ) assert finished.text == "pong" + polled: Final = object_value(proxy.request("GET", f"/v1/responses/{polling_id}", key=key).json()) + assert polled["model"] == handle.scenario_id assert_charged_exactly(lambda: key_spend(key), PRICE) assert [read_response(proxy, key, polling_id).text for _ in range(READS)] == ["pong"] * READS assert_charged_exactly(lambda: key_spend(key), PRICE) diff --git a/tests/unit/proxy/test_response_model_sanitization.py b/tests/unit/proxy/test_response_model_sanitization.py index ba66d27609b..f52736453c1 100644 --- a/tests/unit/proxy/test_response_model_sanitization.py +++ b/tests/unit/proxy/test_response_model_sanitization.py @@ -68,6 +68,19 @@ def _decode_sse_chunk(chunk) -> str: return chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk +def _make_response_created_event(model: str | None): + from litellm.types.llms.openai import ( + ResponseCreatedEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, + ) + + return ResponseCreatedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_CREATED, + response=ResponsesAPIResponse(id="resp_test", created_at=0, output=[], model=model), + ) + + def test_restamp_streaming_chunk_skips_matching_model(): from litellm.proxy.proxy_server import _restamp_streaming_chunk_model @@ -117,19 +130,13 @@ def test_fast_serialize_returns_none_when_model_field_is_missing(): # produce a serialized result via the slow-path fallback, and it must # not contain ``"model": null``. serialized = _serialize_streaming_chunk(chunk) - payload_str = ( - serialized.decode("utf-8") if isinstance(serialized, bytes) else serialized - ) + payload_str = serialized.decode("utf-8") if isinstance(serialized, bytes) else serialized assert '"model": null' not in payload_str assert '"model":null' not in payload_str - assert json.loads(payload_str) == json.loads( - chunk.model_dump_json(exclude_none=True, exclude_unset=True) - ) + assert json.loads(payload_str) == json.loads(chunk.model_dump_json(exclude_none=True, exclude_unset=True)) -def test_proxy_chat_completion_does_not_return_provider_prefixed_model( - tmp_path, monkeypatch -): +def test_proxy_chat_completion_does_not_return_provider_prefixed_model(tmp_path, monkeypatch): """ Regression test: @@ -159,15 +166,11 @@ def test_proxy_chat_completion_does_not_return_provider_prefixed_model( monkeypatch.setattr( proxy_server.llm_router, # type: ignore[arg-type] "acompletion", - AsyncMock( - return_value=_make_minimal_chat_completion_response(model=internal_model) - ), + AsyncMock(return_value=_make_minimal_chat_completion_response(model=internal_model)), ) # Also no-op proxy logging hooks to keep this test focused and deterministic. - monkeypatch.setattr( - proxy_server.proxy_logging_obj, "during_call_hook", AsyncMock(return_value=None) - ) + monkeypatch.setattr(proxy_server.proxy_logging_obj, "during_call_hook", AsyncMock(return_value=None)) monkeypatch.setattr( proxy_server.proxy_logging_obj, "update_request_status", @@ -482,3 +485,385 @@ async def test_proxy_streaming_fastest_response_preserves_winning_model(monkeypa payload = json.loads(first[len("data: ") :].strip()) assert payload["model"] == winning_model assert payload["model"] != comma_separated_models + + +def test_restamp_streaming_chunk_rewrites_nested_response_model(): + from litellm.proxy.proxy_server import _restamp_streaming_chunk_model + + chunk = _make_response_created_event("hosted_vllm/some-model") + result, _ = _restamp_streaming_chunk_model( + chunk=chunk, + requested_model_from_client="client-model", + request_data={"litellm_call_id": "test-call-id"}, + model_mismatch_logged=False, + ) + + assert result.response.model == "client-model" + assert chunk.response.model == "hosted_vllm/some-model" + + +def _make_response_completed_event(model: str): + from litellm.types.llms.openai import ( + ResponseCompletedEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, + ) + + return ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=ResponsesAPIResponse(id="resp_test", created_at=0, output=[], model=model), + ) + + +def test_restamp_streaming_chunk_rewrites_nested_model_on_completed_event(): + from litellm.proxy.proxy_server import _restamp_streaming_chunk_model + + chunk = _make_response_completed_event("hosted_vllm/some-model") + result, _ = _restamp_streaming_chunk_model( + chunk=chunk, + requested_model_from_client="client-model", + request_data={"litellm_call_id": "test-call-id"}, + model_mismatch_logged=False, + ) + + assert result.response.model == "client-model" + assert chunk.response.model == "hosted_vllm/some-model", "the priced original must not change" + + +@pytest.mark.parametrize("event_type", ["response.in_progress", "response.failed", "response.incomplete"]) +def test_restamp_streaming_chunk_rewrites_nested_model_on_other_lifecycle_events(event_type): + from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig + from litellm.proxy.proxy_server import _restamp_streaming_chunk_model + + payload = { + "type": event_type, + "sequence_number": 1, + "response": {"id": "resp_test", "object": "response", "created_at": 0, "output": [], "model": "gpt-backend"}, + } + chunk = OpenAIResponsesAPIConfig.get_event_model_class(event_type).model_validate(payload) + result, _ = _restamp_streaming_chunk_model( + chunk=chunk, + requested_model_from_client="client-model", + request_data={"litellm_call_id": "test-call-id"}, + model_mismatch_logged=False, + ) + + assert result.response.model == "client-model" + assert chunk.response.model == "gpt-backend", "the priced original must not change" + + +def test_restamp_streaming_chunk_leaves_nested_none_model_untouched(): + from litellm.proxy.proxy_server import _restamp_streaming_chunk_model + + chunk = _make_response_created_event(None) + result, _ = _restamp_streaming_chunk_model( + chunk=chunk, + requested_model_from_client="client-model", + request_data={"litellm_call_id": "test-call-id"}, + model_mismatch_logged=False, + ) + + assert result.response.model is None + + +@pytest.mark.parametrize("event_type", ["response.output_text.delta", "response.custom_event"]) +def test_restamp_streaming_chunk_does_not_add_nested_response_to_other_events(event_type): + from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig + from litellm.proxy.proxy_server import _restamp_streaming_chunk_model + + chunk = OpenAIResponsesAPIConfig.get_event_model_class(event_type).model_validate( + { + "type": event_type, + "sequence_number": 3, + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "delta": "hi", + } + ) + result, _ = _restamp_streaming_chunk_model( + chunk=chunk, + requested_model_from_client="client-model", + request_data={"litellm_call_id": "test-call-id"}, + model_mismatch_logged=False, + ) + + assert "response" not in result.model_dump() + assert result.delta == "hi" + + +def test_restamp_streaming_chunk_keeps_nested_model_when_raw_name_requested(): + from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY + from litellm.proxy.proxy_server import _restamp_streaming_chunk_model + + chunk = _make_response_created_event("hosted_vllm/some-model") + result, _ = _restamp_streaming_chunk_model( + chunk=chunk, + requested_model_from_client="client-model", + request_data={"metadata": {RETURN_RAW_MODEL_NAME_METADATA_KEY: True}}, + model_mismatch_logged=False, + ) + + assert result.response.model == "hosted_vllm/some-model" + + +def test_restamp_streaming_chunk_uses_fallback_model_for_nested_response(): + from litellm.proxy.proxy_server import _restamp_streaming_chunk_model + + chunk = _make_response_created_event("hosted_vllm/some-model") + result, _ = _restamp_streaming_chunk_model( + chunk=chunk, + requested_model_from_client="client-model", + request_data={"litellm_call_id": "test-call-id"}, + model_mismatch_logged=False, + fallback_was_attempted=True, + fallback_model_from_metadata="fallback-public", + ) + + assert result.response.model == "fallback-public" + + +@pytest.mark.asyncio +async def test_proxy_responses_stream_reports_requested_model_in_nested_response(monkeypatch): + """ + End-to-end regression for /v1/responses streaming: + + The nested `response.model` on lifecycle events reaching the wire must report the + client-requested model, not the provider-prefixed deployment name. + """ + client_model = "vllm-model" + internal_model = f"hosted_vllm/{client_model}" + + from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth + + async def _iterator_hook( + user_api_key_dict: UserAPIKeyAuth, + response: AsyncGenerator, + request_data: dict, + ): + yield _make_response_created_event(model=internal_model) + + monkeypatch.setattr( + proxy_server.proxy_logging_obj, + "async_post_call_streaming_iterator_hook", + _iterator_hook, + ) + monkeypatch.setattr( + proxy_server.proxy_logging_obj, + "async_post_call_streaming_hook", + AsyncMock(side_effect=lambda **kwargs: kwargs["response"]), + ) + for name in ("has_streaming_callbacks", "needs_iterator_wrap", "needs_per_chunk_streaming_hook"): + monkeypatch.setattr(proxy_server.proxy_logging_obj, name, MagicMock(return_value=True)) + + gen = proxy_server.async_data_generator( + response=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(api_key=MASTER_KEY), + request_data={"model": client_model}, + responses_stream_errors=True, + ) + + payloads = [] + async for item in gen: + text = _decode_sse_chunk(item) + for line in text.splitlines(): + if line.startswith("data: ") and line.strip() != "data: [DONE]": + payloads.append(json.loads(line[len("data: ") :].strip())) + + created = next(p for p in payloads if p.get("type") == "response.created") + assert created["response"]["model"] == client_model + assert not created["response"]["model"].startswith("hosted_vllm/") + + +def test_restamp_streaming_chunk_rewrites_nested_model_on_queued_event(): + from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig + from litellm.proxy.proxy_server import _restamp_streaming_chunk_model + + queued_payload = { + "type": "response.queued", + "sequence_number": 0, + "response": {"id": "resp_test", "object": "response", "status": "queued", "model": "gpt-backend-2026-01-01"}, + } + chunk = OpenAIResponsesAPIConfig.get_event_model_class("response.queued").model_validate(queued_payload) + result, _ = _restamp_streaming_chunk_model( + chunk=chunk, + requested_model_from_client="client-model", + request_data={"litellm_call_id": "test-call-id"}, + model_mismatch_logged=False, + ) + + assert result.model_dump()["response"] == {**queued_payload["response"], "model": "client-model"} + assert chunk.model_dump()["response"]["model"] == "gpt-backend-2026-01-01", "the priced original must not change" + + +def test_restamp_streaming_chunk_rewrites_nested_model_when_lifecycle_event_kept_raw_response(): + from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig + from litellm.proxy.proxy_server import _restamp_streaming_chunk_model + + partial_response = {"id": "resp_test", "object": "response", "status": "in_progress", "model": "gpt-backend-2026-01-01"} + chunk = OpenAIResponsesAPIConfig().transform_streaming_response( + model="gpt-backend", + parsed_chunk={"type": "response.created", "sequence_number": 0, "response": partial_response}, + logging_obj=MagicMock(), + ) + result, _ = _restamp_streaming_chunk_model( + chunk=chunk, + requested_model_from_client="client-model", + request_data={"litellm_call_id": "test-call-id"}, + model_mismatch_logged=False, + ) + + assert result.response == {**partial_response, "model": "client-model"} + assert chunk.response["model"] == "gpt-backend-2026-01-01", "the priced original must not change" + + +def test_restamp_streaming_chunk_leaves_generic_event_without_nested_response_untouched(): + from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig + from litellm.proxy.proxy_server import _restamp_streaming_chunk_model + + shell_payload = {"type": "response.shell_call.in_progress", "sequence_number": 3, "item_id": "sh_1"} + chunk = OpenAIResponsesAPIConfig.get_event_model_class(shell_payload["type"]).model_validate(shell_payload) + result, _ = _restamp_streaming_chunk_model( + chunk=chunk, + requested_model_from_client="client-model", + request_data={"litellm_call_id": "test-call-id"}, + model_mismatch_logged=False, + ) + + assert "response" not in result.model_dump() + + +@pytest.mark.parametrize("stream_started", [False, True]) +def test_restamp_streaming_chunk_rewrites_nested_model_in_guardrail_block_frames(stream_started): + from litellm.exceptions import ModifyResponseException + from litellm.llms.openai.responses.guardrail_translation.handler import OpenAIResponsesHandler + from litellm.proxy.proxy_server import _restamp_streaming_chunk_model + + sent_before_block = ( + [{"type": "response.created", "response": {"id": "resp_live", "model": "gpt-backend-2026-01-01"}}] + if stream_started + else [] + ) + block = ModifyResponseException( + message="Blocked by policy.", + model="gpt-backend-2026-01-01", + request_data={}, + original_response=sent_before_block, + ) + frames = OpenAIResponsesHandler().build_block_sse_chunks( + block, stream_started=stream_started, responses_so_far=sent_before_block + ) + client_frames = [ + _restamp_streaming_chunk_model( + chunk=frame, + requested_model_from_client="client-model", + request_data={"litellm_call_id": "test-call-id"}, + model_mismatch_logged=False, + )[0] + for frame in frames + ] + events = [json.loads(frame.decode().removeprefix("data: ")) for frame in client_frames] + + assert all(frame.endswith(b"\n\n") for frame in client_frames) + assert events[-1]["type"] == "response.completed" + assert events[-1]["response"]["output"][0]["content"][0]["text"] == "Blocked by policy." + assert {event["response"]["model"] for event in events if "response" in event} == {"client-model"} + assert [frame for frame, event in zip(client_frames, events) if "response" not in event] == [ + frame for frame in frames if b'"response"' not in frame + ] + + +def test_restamp_streaming_chunk_rewrites_nested_model_on_dict_lifecycle_events(): + from litellm.proxy.proxy_server import _restamp_streaming_chunk_model + + chunk = {"type": "response.completed", "response": {"id": "resp_test", "model": "gpt-backend-2026-01-01"}} + result, _ = _restamp_streaming_chunk_model( + chunk=chunk, + requested_model_from_client="client-model", + request_data={"litellm_call_id": "test-call-id"}, + model_mismatch_logged=False, + ) + + assert result["response"] == {"id": "resp_test", "model": "client-model"} + assert chunk["response"]["model"] == "gpt-backend-2026-01-01", "the priced original must not change" + + +_BACKEND_COMPLETED_EVENT = b'{"type":"response.completed","response":{"id":"resp_1","model":"gpt-backend-2026-01-01"}}' + + +@pytest.mark.parametrize( + "frame", + [ + pytest.param(b"event: response.completed\ndata: " + _BACKEND_COMPLETED_EVENT + b"\n\n", id="event_line_first"), + pytest.param(b"data:" + _BACKEND_COMPLETED_EVENT + b"\n\n", id="no_space_after_data"), + pytest.param( + b"data: " + _BACKEND_COMPLETED_EVENT + b"\n\ndata: " + _BACKEND_COMPLETED_EVENT + b"\r\n\r\n", + id="coalesced_frames", + ), + ], +) +def test_restamp_streaming_chunk_rewrites_nested_model_on_every_sse_data_line(frame): + from litellm.proxy.proxy_server import _restamp_streaming_chunk_model + + result, _ = _restamp_streaming_chunk_model( + chunk=frame, + requested_model_from_client="client-model", + request_data={"litellm_call_id": "test-call-id"}, + model_mismatch_logged=False, + ) + + def is_data_line(line: bytes) -> bool: + return line.startswith(b"data:") + + result_lines = result.splitlines(keepends=True) + frame_lines = frame.splitlines(keepends=True) + data_events = [json.loads(line.removeprefix(b"data:")) for line in result_lines if is_data_line(line)] + assert len(data_events) == len([line for line in frame_lines if is_data_line(line)]) + assert {event["response"]["model"] for event in data_events} == {"client-model"} + assert [line for line in result_lines if not is_data_line(line)] == [ + line for line in frame_lines if not is_data_line(line) + ] + assert [line[len(line.rstrip(b"\r\n")) :] for line in result_lines] == [ + line[len(line.rstrip(b"\r\n")) :] for line in frame_lines + ] + + +_SKIPPED_RESTAMPS = { + "azure_model_router": ("azure_ai/model_router", {}), + "fastest_response": ("openai/gpt-4o,gemini/gemini-2.5-flash", {"fastest_response": True}), +} + + +@pytest.mark.parametrize("skip", _SKIPPED_RESTAMPS.keys()) +def test_restamp_streaming_chunk_keeps_nested_model_where_the_top_level_model_is_kept(skip): + from litellm.proxy.proxy_server import _restamp_streaming_chunk_model + + requested_model, request_data = _SKIPPED_RESTAMPS[skip] + frame = b"data: " + _BACKEND_COMPLETED_EVENT + b"\n\n" + results = [ + _restamp_streaming_chunk_model( + chunk=chunk, + requested_model_from_client=requested_model, + request_data=request_data, + model_mismatch_logged=False, + )[0] + for chunk in (_make_response_created_event("gpt-backend-2026-01-01"), frame) + ] + + assert results[0].response.model == "gpt-backend-2026-01-01" + assert results[1] == frame + + +def test_restamp_streaming_chunk_leaves_raw_sse_passthrough_frames_untouched(): + from litellm.proxy.proxy_server import _restamp_streaming_chunk_model + + frame = b"data: " + _BACKEND_COMPLETED_EVENT + b"\n\n" + result, _ = _restamp_streaming_chunk_model( + chunk=frame, + requested_model_from_client="client-model", + request_data={"_litellm_raw_sse_stream": True}, + model_mismatch_logged=False, + ) + + assert result == frame