Merge 5a3ddf2b19 into 9752869d84

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:
         """

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"""