diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index b03f1e4348c..af2d31a0f87 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/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} diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index f2273924b30..d7c02637932 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/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": ""}] + 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 "" not in json.dumps(stored) + else: + assert stored["messages"] == [{"role": "user", "content": ""}] + + assert data["messages"] == [{"role": "user", "content": ""}] + 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 # ---------------------------------------------------------------------------