fix(proxy): exclude websocket transport from spend log snapshots

litellm/proxy/litellm_pre_call_utils.py

@@ -1937,12 +1937,14 @@ def refresh_proxy_server_request_body_snapshot(
     has already stamped ``data["litellm_logging_obj"]`` with a live (non-serializable)
     ``Logging`` instance, so it must be excluded here the same way ``secret_fields``
     and ``proxy_server_request`` are.
+    The live WebSocket also references the ASGI app/router graph, not request body data.
     """
     proxy_server_request = data.get("proxy_server_request")
     if not isinstance(proxy_server_request, dict):
         return
     _body_snapshot_exclude = (
-        frozenset({"secret_fields", "proxy_server_request", "litellm_logging_obj"}) | _TRANSPORT_ONLY_CREDENTIAL_KEYS
+        frozenset({"secret_fields", "proxy_server_request", "litellm_logging_obj", "websocket"})
+        | _TRANSPORT_ONLY_CREDENTIAL_KEYS
     )
     proxy_server_request["body"] = {k: v for k, v in data.items() if k not in _body_snapshot_exclude}
 

tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py

@@ -2419,6 +2419,60 @@ def test_proxy_server_request_payload_excludes_secret_fields(mock_should_store):
     assert parsed["messages"] == [{"role": "user", "content": "hello"}]
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("redact", [False, True], ids=["stored", "redacted"])
+@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs")
+async def test_websocket_snapshot_spend_serialization_preserves_transport(mock_should_store, redact):
+    from starlette.websockets import WebSocket
+
+    from litellm.proxy.litellm_pre_call_utils import refresh_proxy_server_request_body_snapshot
+
+    class AppTraversalTrap:
+        @property
+        def __dict__(self):
+            raise AssertionError("Spend logging must not traverse the WebSocket app graph")
+
+    mock_should_store.return_value = True
+    send = AsyncMock()
+    websocket = WebSocket(
+        scope={"type": "websocket", "app": AppTraversalTrap()},
+        receive=AsyncMock(return_value={"type": "websocket.connect"}),
+        send=send,
+    )
+    await websocket.accept()
+    data = {
+        "model": "gpt-4o-realtime-preview",
+        "messages": [{"role": "user", "content": "private input"}],
+        "metadata": {"trace_id": "realtime-trace", "websocket": "user metadata"},
+        "websocket": websocket,
+        "proxy_server_request": {},
+    }
+    refresh_proxy_server_request_body_snapshot(data)
+    # Guardrails can replace messages after the initial snapshot.
+    data["messages"] = [{"role": "user", "content": "<MASKED>"}]
+    refresh_proxy_server_request_body_snapshot(data)
+
+    stored = json.loads(
+        _get_proxy_server_request_for_spend_logs_payload(
+            metadata={},
+            litellm_params={"proxy_server_request": data["proxy_server_request"]},
+            kwargs={"standard_callback_dynamic_params": {"turn_off_message_logging": redact}},
+        )
+    )
+    assert "websocket" not in stored
+    assert stored["model"] == data["model"]
+    assert stored["metadata"] == data["metadata"]
+    if redact:
+        assert "<MASKED>" not in json.dumps(stored)
+    else:
+        assert stored["messages"] == [{"role": "user", "content": "<MASKED>"}]
+
+    assert data["messages"] == [{"role": "user", "content": "<MASKED>"}]
+    assert data["websocket"] is websocket
+    await data["websocket"].send_text("still connected")
+    send.assert_awaited_with({"type": "websocket.send", "text": "still connected"})
+
+
 # ---------------------------------------------------------------------------
 # LIT-2992: error_information sanitization for spend logs
 # ---------------------------------------------------------------------------