From 5185a536dbc45b54f33418f931feb2e411977ccc Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 05:40:32 -0700 Subject: [PATCH] fix(proxy): log no extra websocket spend row for frames after the last turn --- ...i_websocket_passthrough_logging_handler.py | 2 ++ .../pass_through_endpoints.py | 17 ++++++++++++--- .../pass_through_endpoints.py | 2 +- ...i_websocket_passthrough_logging_handler.py | 4 ++-- .../test_pass_through_endpoints.py | 21 +++++++++++++++++++ 5 files changed, 40 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py index 951764f74b0..f2b23d599f9 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py @@ -100,6 +100,8 @@ def openai_websocket_frame_kind(frame: Mapping[str, object]) -> WebsocketUpstrea return "ends_turn" if frame_type == _SESSION_CREATED: return "opens_session" + if frame_type == _RESPONSE_CREATED: + return "starts_turn" return "relayed" diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index af67a451c96..393b193c73f 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -2730,6 +2730,19 @@ class _WebsocketTurnLedger: self.queued_turns = self.queued_turns[1:] +def _session_end_needs_a_row( + upstream_frame_kind: WebsocketUpstreamFrameClassifier | None, + unlogged_frames: Sequence[Mapping[str, object]], + ledger: _WebsocketTurnLedger, +) -> bool: + """A session logs one row per turn as each ends, so at its end it only owes a row for a turn still running, or + for whatever the provider sent when no turn ever ended""" + if upstream_frame_kind is None: + return True + kinds: Final = frozenset(map(upstream_frame_kind, unlogged_frames)) + return "starts_turn" in kinds or (ledger.logged_turns == 0 and bool(kinds - {"opens_session"})) + + async def _relay_client_frame( websocket: WebSocket, upstream_ws: ClientConnection, @@ -3191,9 +3204,7 @@ async def websocket_passthrough_request( reason=_truncated_close_reason(upstream_close.reason), ) - if upstream_frame_kind is None or any( - upstream_frame_kind(frame) != "opens_session" for frame in websocket_messages - ): + if _session_end_needs_a_row(upstream_frame_kind, websocket_messages, ledger): log_turn() # Call the proxy logging success hook diff --git a/litellm/types/passthrough_endpoints/pass_through_endpoints.py b/litellm/types/passthrough_endpoints/pass_through_endpoints.py index b13211d8f37..452f727c5ea 100644 --- a/litellm/types/passthrough_endpoints/pass_through_endpoints.py +++ b/litellm/types/passthrough_endpoints/pass_through_endpoints.py @@ -29,7 +29,7 @@ LITELLM_PASS_THROUGH_ENDPOINT_MARKER: Final = "__litellm_pass_through_endpoint__ SETTLES_BACKGROUND_RESPONSE_KWARG: Final = "settles_background_response" -WebsocketUpstreamFrameKind: TypeAlias = Literal["ends_turn", "opens_session", "relayed"] +WebsocketUpstreamFrameKind: TypeAlias = Literal["starts_turn", "ends_turn", "opens_session", "relayed"] class BackgroundResponseSettlement(str, Enum): diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py index c259d89c07f..7494c913a4d 100644 --- a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py @@ -365,10 +365,10 @@ def test_a_responses_turn_falls_back_to_the_requested_model_only_when_its_own_ha ("response.failed", "ends_turn"), ("response.done", "ends_turn"), ("session.created", "opens_session"), - ("response.created", "relayed"), + ("response.created", "starts_turn"), ("response.output_text.delta", "relayed"), ("error", "relayed"), ], ) -def test_every_frame_that_settles_a_turn_closes_it_and_nothing_else_does(frame_type, kind): +def test_each_frame_is_classified_by_what_it_does_to_a_turn(frame_type, kind): assert openai_websocket_frame_kind({"type": frame_type}) == kind diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index ac42491c389..c4b08685d74 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -8297,6 +8297,27 @@ async def test_websocket_passthrough_logs_the_unfinished_turn_a_client_walked_ou ] == [["response.created", "response.completed"], ["response.created"]] +@pytest.mark.asyncio +async def test_websocket_passthrough_logs_no_extra_row_for_what_the_provider_sends_after_the_last_turn(): + """ + A realtime provider follows each finished turn with frames that start no + turn, such as rate_limits.updated. The session ends with nothing running, + so its turn is its only row and the request count stays one per turn. + """ + client = _turn_taking_client([json.dumps({"type": "response.create"})], answered_by='"rate_limits.updated"') + upstream_ws = ScriptedUpstreamWebSocket( + REALTIME_SESSION_CREATED, + replies=[[REALTIME_RESPONSE_DONE, json.dumps({"type": "rate_limits.updated", "rate_limits": []})]], + ) + + relayed = await _relay_turns(client, upstream_ws, RecordingFrameModelGate(), "/openai/v1/realtime") + + assert [ + [frame["type"] for frame in logged.kwargs["response_body"]] + for logged in relayed.success_handler.call_args_list + ] == [["session.created", "response.done"]] + + @pytest.mark.asyncio async def test_websocket_passthrough_carries_the_session_the_provider_opened_into_every_later_turn(): """