diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 249c6370b86..f9216c405ce 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -3107,6 +3107,9 @@ _MANAGED_WS_SKIP_KWARGS: Final[frozenset[str]] = frozenset( "aresponses", "_aresponses_websocket", "user_api_key_dict", + # Kept on the handler and copied into every turn, where the turn's body is + # overlaid; as a pass-through kwarg it would overwrite that copy (#49878). + "proxy_server_request", } ) @@ -3169,6 +3172,12 @@ class ManagedResponsesWebSocketHandler: self.quota_callbacks: tuple[ProjectQuotaCallback, ...] = tuple(quota_callbacks) if quota_callbacks else () self.turn_admission: ResponsesWsTurnAdmission | None = turn_admission self._failure: Exception | None = None + # The handshake's request record (headers, url, arrival time); each turn gets a + # copy with its own body so spend logs see the turn, not the handshake. + handshake_request: Final = kwargs.get("proxy_server_request") + self.proxy_server_request: Mapping[str, object] = ( + handshake_request if _is_json_object(handshake_request) else {} + ) # Carry through safe pass-through kwargs (e.g. extra_headers) self.extra_kwargs: dict[str, object] = {k: v for k, v in kwargs.items() if k not in _MANAGED_WS_SKIP_KWARGS} # In-memory session history: response_id → full accumulated message list. @@ -3460,6 +3469,8 @@ class ManagedResponsesWebSocketHandler: call_kwargs["custom_llm_provider"] = self.custom_llm_provider if self.litellm_metadata: call_kwargs["litellm_metadata"] = dict(self.litellm_metadata) + if self.proxy_server_request: + call_kwargs["proxy_server_request"] = dict(self.proxy_server_request) async def _guard_frame(self, raw_message: str) -> GuardedResponsesWsFrame | None: """Same per-frame pre_call guardrail walk as the native socket: a rejected frame errors and is not sent.""" @@ -3484,8 +3495,13 @@ class ManagedResponsesWebSocketHandler: @staticmethod def _update_proxy_request(call_kwargs: dict[str, Any], model: str) -> None: - """Update proxy_server_request body so spend logs record the full request.""" - proxy_server_request = (call_kwargs.get("litellm_metadata") or {}).get("proxy_server_request") or {} + """Give ``proxy_server_request`` the turn's body so spend logs record the turn. + + Only the top-level kwarg is written. That is the one the spend-log writer + redacts and sanitises into the ``proxy_server_request`` column; a copy under + ``litellm_metadata`` would be serialised raw into the ``metadata`` column (#49878). + """ + proxy_server_request = call_kwargs.get("proxy_server_request") or {} if not isinstance(proxy_server_request, dict): return body: Final = dict(proxy_server_request.get("body") or {}) @@ -3495,11 +3511,7 @@ class ManagedResponsesWebSocketHandler: for k in ("tools", "tool_choice", "instructions", "metadata"): if k in call_kwargs and call_kwargs[k] is not None: body[k] = call_kwargs[k] - proxy_server_request = {**proxy_server_request, "body": body} - if "litellm_metadata" not in call_kwargs: - call_kwargs["litellm_metadata"] = {} - call_kwargs["litellm_metadata"]["proxy_server_request"] = proxy_server_request - call_kwargs["proxy_server_request"] = proxy_server_request + call_kwargs["proxy_server_request"] = {**proxy_server_request, "body": body} async def _stream_and_forward(self, model: str, call_kwargs: dict[str, object]) -> _MutableJsonObject | None: """ diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index 7e044b55f5a..0c1859087f9 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -935,9 +935,7 @@ class TestUpdateProxyRequest: call_kwargs = { "input": "hello", "store": True, - "litellm_metadata": { - "proxy_server_request": {"headers": {}, "body": {}}, - }, + "proxy_server_request": {"headers": {}, "body": {}}, } ManagedResponsesWebSocketHandler._update_proxy_request( @@ -950,23 +948,66 @@ class TestUpdateProxyRequest: ) assert call_kwargs["proxy_server_request"]["body"]["input"] == "hello" - def test_proxy_server_request_matches_metadata(self): + def test_body_stays_out_of_litellm_metadata(self): + """The body must only reach the spend-log writer through the top-level kwarg, + which it redacts and sanitises; a copy in litellm_metadata lands raw in the + metadata column (#49878).""" from litellm.responses.streaming_iterator import ( ManagedResponsesWebSocketHandler, ) call_kwargs = { "input": "hi", - "litellm_metadata": {"proxy_server_request": {"body": {}}}, + "litellm_metadata": {"model_group": "gpt-4o"}, + "proxy_server_request": {"headers": {"x-id": "1"}, "url": "wss://p/v1/responses", "body": {}}, } ManagedResponsesWebSocketHandler._update_proxy_request(call_kwargs, "gpt-4o") - assert ( - call_kwargs["proxy_server_request"] - == call_kwargs["litellm_metadata"]["proxy_server_request"] + assert call_kwargs["litellm_metadata"] == {"model_group": "gpt-4o"} + assert call_kwargs["proxy_server_request"]["headers"] == {"x-id": "1"} + assert call_kwargs["proxy_server_request"]["url"] == "wss://p/v1/responses" + assert call_kwargs["proxy_server_request"]["body"]["input"] == "hi" + + def test_handshake_request_is_copied_per_turn_and_not_passed_through(self): + """The handshake's proxy_server_request seeds each turn's record and does not + come back through extra_kwargs, which would overwrite the turn's body.""" + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.responses.streaming_iterator import ( + ManagedResponsesWebSocketHandler, ) + handshake = {"headers": {"x-id": "1"}, "body": {"first": True}} + handler = ManagedResponsesWebSocketHandler( + websocket=MagicMock(), + model="test-model", + logging_obj=Logging( + model="test-model", + messages=[], + stream=True, + call_type="aresponses", + start_time=0, + litellm_call_id="test-id", + function_id="test-func", + ), + litellm_metadata={}, + proxy_server_request=handshake, + extra_headers={"x-extra": "1"}, + ) + + assert "proxy_server_request" not in handler.extra_kwargs + assert handler.extra_kwargs["extra_headers"] == {"x-extra": "1"} + + call_kwargs: dict = {"input": "turn one"} + handler._inject_credentials(call_kwargs, model="test-model") + handler._update_proxy_request(call_kwargs, "test-model") + call_kwargs.update(handler.extra_kwargs) + + assert call_kwargs["proxy_server_request"]["headers"] == {"x-id": "1"} + assert call_kwargs["proxy_server_request"]["body"]["input"] == "turn one" + # The handshake record itself is untouched for the next turn. + assert handshake == {"headers": {"x-id": "1"}, "body": {"first": True}} + class TestWebSocketEventTypes: """Test that all WebSocket event types are properly handled with dict-based chunks"""