From 1d70354c0697b7741ca04aa624cb8da64f8db014 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 05:31:08 -0700 Subject: [PATCH] fix(proxy): bill each websocket lane cut off at session end on its own row When a client leaves while turns on several stream_id lanes are still running, each unfinished turn now gets its own spend row under the groups and model its own frame asked for, instead of one merged row. Also pins the PR's new single-assigned test variables as Final and swaps datetime.now() in the logging-handler and success-handler fixtures for fixed timestamps. --- .../pass_through_endpoints.py | 5 +- .../proxy/common_utils/test_realtime_cache.py | 4 +- ...i_websocket_passthrough_logging_handler.py | 91 +++---- .../test_pass_through_endpoints.py | 245 ++++++++++-------- .../test_success_handler.py | 10 +- .../test_openai_ws_passthrough_routes.py | 18 +- 6 files changed, 206 insertions(+), 167 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 06359fa7da3..eaaf8aa03d0 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -3313,7 +3313,10 @@ async def websocket_passthrough_request( lanes_in_flight: Final = ( () if upstream_frame_kind is None else _lanes_in_flight(upstream_frame_kind, websocket_messages) ) - log_turn(next(iter(lanes_in_flight), None), _every_frame) + *earlier_lanes, last_lane = lanes_in_flight or (None,) + for lane in earlier_lanes: + log_turn(lane, _frames_of_lane(lane)) + log_turn(last_lane, _every_frame) # Call the proxy logging success hook if proxy_logging_obj: diff --git a/tests/unit/proxy/common_utils/test_realtime_cache.py b/tests/unit/proxy/common_utils/test_realtime_cache.py index 6177a00d0aa..9e76caff054 100644 --- a/tests/unit/proxy/common_utils/test_realtime_cache.py +++ b/tests/unit/proxy/common_utils/test_realtime_cache.py @@ -1,5 +1,5 @@ import json -from typing import Any, cast +from typing import Any, Final, cast import pytest @@ -42,7 +42,7 @@ def test_realtime_request_body_caches_each_model_separately(): def test_realtime_request_body_keeps_a_model_name_with_quotes_inside_its_own_field(): - body = realtime_request_body('gpt-4o", "user": "someone-else') + body: Final = realtime_request_body('gpt-4o", "user": "someone-else') assert json.loads(body) == {"model": 'gpt-4o", "user": "someone-else'} 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 ceffbd1ba08..1bc4e0fe4d2 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 @@ -19,6 +19,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_websocket REALTIME_MODEL: Final = "gpt-realtime-2.1" RESPONSES_MODEL: Final = "gpt-5.4-mini" REASONING_MODEL: Final = "gpt-5.4" +SESSION_START: Final = datetime(2026, 1, 1, 12, 0, 0) def _logging_obj() -> MagicMock: @@ -86,13 +87,13 @@ def _handle( websocket_messages=messages, logging_obj=logging_obj if logging_obj is not None else _logging_obj(), url_route=url_route, - start_time=datetime.now(), + start_time=SESSION_START, kwargs=kwargs, ) def _realtime_cost(model: str, text_in: int, audio_in: int, text_out: int, audio_out: int) -> float: - prices = litellm.model_cost[model] + prices: Final = litellm.model_cost[model] return ( text_in * prices["input_cost_per_token"] + audio_in * prices["input_cost_per_audio_token"] @@ -102,16 +103,16 @@ def _realtime_cost(model: str, text_in: int, audio_in: int, text_out: int, audio def test_realtime_session_is_priced_on_the_model_the_session_reported(): - logging_obj = _logging_obj() - messages = [ + logging_obj: Final = _logging_obj() + messages: Final = [ _session_created(REALTIME_MODEL), {"type": "response.created", "response": {"id": "resp_fake"}}, _realtime_response_done(text_in=33, audio_in=120, text_out=16, audio_out=240), ] - handled = _handle(messages, logging_obj=logging_obj) + handled: Final = _handle(messages, logging_obj=logging_obj) - expected_cost = _realtime_cost(REALTIME_MODEL, 33, 120, 16, 240) + expected_cost: Final = _realtime_cost(REALTIME_MODEL, 33, 120, 16, 240) assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost) assert handled["kwargs"]["model"] == REALTIME_MODEL assert handled["kwargs"]["custom_llm_provider"] == "openai" @@ -122,8 +123,8 @@ def test_realtime_session_is_priced_on_the_model_the_session_reported(): def test_realtime_cost_covers_every_turn_of_the_session(): - one_turn = _handle([_session_created(REALTIME_MODEL), _realtime_response_done(10, 0, 5, 0)]) - two_turns = _handle( + one_turn: Final = _handle([_session_created(REALTIME_MODEL), _realtime_response_done(10, 0, 5, 0)]) + two_turns: Final = _handle( [ _session_created(REALTIME_MODEL), _realtime_response_done(10, 0, 5, 0), @@ -136,13 +137,13 @@ def test_realtime_cost_covers_every_turn_of_the_session(): def test_responses_turn_is_priced_on_the_model_the_response_reported(): - handled = _handle( + handled: Final = _handle( [_responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5)], url_route="/openai_passthrough/v1/responses", ) - prices = litellm.model_cost[RESPONSES_MODEL] - expected_cost = 13 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"] + prices: Final = litellm.model_cost[RESPONSES_MODEL] + expected_cost: Final = 13 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"] assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost) assert handled["kwargs"]["model"] == RESPONSES_MODEL assert handled["result"].usage.prompt_tokens == 13 @@ -150,7 +151,7 @@ def test_responses_turn_is_priced_on_the_model_the_response_reported(): def test_responses_cost_covers_every_turn_of_the_connection(): - handled = _handle( + handled: Final = _handle( [ _responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5), _responses_completed(RESPONSES_MODEL, input_tokens=20, output_tokens=7), @@ -158,8 +159,8 @@ def test_responses_cost_covers_every_turn_of_the_connection(): url_route="/openai_passthrough/v1/responses", ) - prices = litellm.model_cost[RESPONSES_MODEL] - expected_cost = 33 * prices["input_cost_per_token"] + 12 * prices["output_cost_per_token"] + prices: Final = litellm.model_cost[RESPONSES_MODEL] + expected_cost: Final = 33 * prices["input_cost_per_token"] + 12 * prices["output_cost_per_token"] assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost) assert handled["result"].usage.total_tokens == 45 @@ -170,9 +171,9 @@ def test_a_model_the_caller_asked_for_never_prices_a_session_the_provider_never_ the provider's own frames may set the price of a session, so a model the caller put in the request cannot stand in for one the provider never named. """ - logging_obj = _logging_obj() + logging_obj: Final = _logging_obj() - handled = _handle( + handled: Final = _handle( [_realtime_response_done(text_in=10, audio_in=0, text_out=5, audio_out=0)], kwargs=MappingProxyType({"model": REALTIME_MODEL}), logging_obj=logging_obj, @@ -189,7 +190,7 @@ def test_each_responses_turn_is_priced_on_the_model_that_turn_ran(): starts on a big model and finishes on a small one has to pay the big model's price for the turn that used it. """ - handled = _handle( + handled: Final = _handle( [ _responses_completed(REASONING_MODEL, input_tokens=13, output_tokens=5), _responses_completed(RESPONSES_MODEL, input_tokens=20, output_tokens=7), @@ -197,9 +198,9 @@ def test_each_responses_turn_is_priced_on_the_model_that_turn_ran(): url_route="/openai_passthrough/v1/responses", ) - first = litellm.model_cost[REASONING_MODEL] - second = litellm.model_cost[RESPONSES_MODEL] - expected_cost = ( + first: Final = litellm.model_cost[REASONING_MODEL] + second: Final = litellm.model_cost[RESPONSES_MODEL] + expected_cost: Final = ( 13 * first["input_cost_per_token"] + 5 * first["output_cost_per_token"] + 20 * second["input_cost_per_token"] @@ -210,7 +211,7 @@ def test_each_responses_turn_is_priced_on_the_model_that_turn_ran(): def test_a_responses_turn_that_names_no_model_is_priced_on_the_session_model(): - handled = _handle( + handled: Final = _handle( [ _session_created(RESPONSES_MODEL), _responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5, served={"model": None}), @@ -218,14 +219,14 @@ def test_a_responses_turn_that_names_no_model_is_priced_on_the_session_model(): url_route="/openai_passthrough/v1/responses", ) - prices = litellm.model_cost[RESPONSES_MODEL] - expected_cost = 13 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"] + prices: Final = litellm.model_cost[RESPONSES_MODEL] + expected_cost: Final = 13 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"] assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost) def test_a_session_nobody_named_a_model_for_is_left_unpriced(): - logging_obj = _logging_obj() - handled = _handle( + logging_obj: Final = _logging_obj() + handled: Final = _handle( [{"type": "response.done", "response": {"usage": {"input_tokens": 10, "output_tokens": 5}}}], kwargs=MappingProxyType({"litellm_call_id": "call-fake"}), logging_obj=logging_obj, @@ -237,9 +238,9 @@ def test_a_session_nobody_named_a_model_for_is_left_unpriced(): def test_a_session_that_reported_no_usage_still_records_its_model(): - logging_obj = _logging_obj() + logging_obj: Final = _logging_obj() - handled = _handle([_session_created(REALTIME_MODEL)], logging_obj=logging_obj) + handled: Final = _handle([_session_created(REALTIME_MODEL)], logging_obj=logging_obj) assert handled["result"] is None assert handled["kwargs"] == {"model": REALTIME_MODEL, "custom_llm_provider": "openai"} @@ -248,7 +249,7 @@ def test_a_session_that_reported_no_usage_still_records_its_model(): def test_a_responses_turn_without_readable_usage_is_not_billed(): - handled = _handle( + handled: Final = _handle( [{"type": "response.completed", "response": {"model": RESPONSES_MODEL, "status": "completed"}}], url_route="/openai_passthrough/v1/responses", ) @@ -264,7 +265,7 @@ def test_the_price_of_a_multi_model_connection_reaches_the_spend_row(): price it already worked out or the cheaper model pays for both turns. The row it lands in names the model the connection opened on. """ - handled = _handle( + handled: Final = _handle( [ _responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5), _responses_completed(REASONING_MODEL, input_tokens=20, output_tokens=7), @@ -272,7 +273,7 @@ def test_the_price_of_a_multi_model_connection_reaches_the_spend_row(): url_route="/openai_passthrough/v1/responses", ) - logged_cost = litellm.response_cost_calculator( + logged_cost: Final = litellm.response_cost_calculator( response_object=handled["result"], model=handled["kwargs"]["model"], custom_llm_provider="openai", @@ -280,8 +281,8 @@ def test_the_price_of_a_multi_model_connection_reaches_the_spend_row(): optional_params={}, ) - session_prices = litellm.model_cost[RESPONSES_MODEL] - priced_on_the_session_model = ( + session_prices: Final = litellm.model_cost[RESPONSES_MODEL] + priced_on_the_session_model: Final = ( 33 * session_prices["input_cost_per_token"] + 12 * session_prices["output_cost_per_token"] ) assert logged_cost == pytest.approx(handled["kwargs"]["response_cost"]) @@ -295,9 +296,9 @@ def test_a_turn_the_caller_walked_out_of_still_names_the_model_it_ran(): bill, but response.created already named the model, so the row says which model answered instead of falling back to unknown. """ - logging_obj = _logging_obj() + logging_obj: Final = _logging_obj() - handled = _handle( + handled: Final = _handle( [{"type": "response.created", "response": {"id": "resp_fake", "model": REASONING_MODEL}}], logging_obj=logging_obj, url_route="/openai_passthrough/v1/responses", @@ -314,19 +315,19 @@ def test_a_responses_turn_that_ended_early_is_billed_for_the_tokens_it_used(even A turn cut short by max_output_tokens or a provider failure still reports the tokens it burned, and OpenAI bills them, so the row has to as well. """ - handled = _handle( + handled: Final = _handle( [_responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5, event=event)], url_route="/openai_passthrough/v1/responses", ) - prices = litellm.model_cost[RESPONSES_MODEL] - expected_cost = 13 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"] + prices: Final = litellm.model_cost[RESPONSES_MODEL] + expected_cost: Final = 13 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"] assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost) assert handled["result"].usage.total_tokens == 18 def test_a_failed_responses_turn_that_reported_no_usage_is_not_billed(): - handled = _handle( + handled: Final = _handle( [{"type": "response.failed", "response": {"model": RESPONSES_MODEL, "status": "failed", "usage": None}}], url_route="/openai_passthrough/v1/responses", ) @@ -353,21 +354,21 @@ def test_a_responses_turn_falls_back_to_the_requested_model_only_when_its_own_ha billed at that model's price instead of $0, while a turn whose own model has a price keeps it. """ - handled = _handle( + handled: Final = _handle( [_responses_completed(turn_model, input_tokens=13, output_tokens=5)], kwargs=MappingProxyType({"litellm_params": {"metadata": {"model_group": requested_model}}}), url_route="/openai_passthrough/v1/responses", ) - prices = litellm.model_cost[priced_as] - expected_cost = 13 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"] + prices: Final = litellm.model_cost[priced_as] + expected_cost: Final = 13 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"] assert expected_cost > 0 assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost) assert handled["kwargs"]["model"] == turn_model def _http_passthrough_cost(response: Mapping[str, object]) -> float: - http_response = httpx.Response( + http_response: Final = httpx.Response( 200, json=response, request=httpx.Request("POST", "https://api.openai.com/v1/responses") ) _, cost = OpenAIPassthroughLoggingHandler._build_responses_api_response_and_cost( @@ -390,10 +391,10 @@ def test_a_responses_turn_costs_what_the_http_passthrough_charges_for_the_same_r called, not only its tokens, so a websocket turn has to cost exactly what the same response costs over HTTP, never the cheaper token-only price. """ - turn = _responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5, served=served) - plain_turn = _responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5) + turn: Final = _responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5, served=served) + plain_turn: Final = _responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5) - handled = _handle([turn], url_route="/openai_passthrough/v1/responses") + handled: Final = _handle([turn], url_route="/openai_passthrough/v1/responses") assert handled["kwargs"]["response_cost"] == pytest.approx(_http_passthrough_cost(turn["response"])) assert handled["kwargs"]["response_cost"] > _http_passthrough_cost(plain_turn["response"]) 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 7b2894e9aa7..9f97110e1af 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 @@ -7697,7 +7697,7 @@ class RecordingFrameModelGate: async def __call__(self, websocket: object, model: str, /) -> WebsocketFrameModelVerdict: self.checked.append(model) - refusal = self.refusals.get(model) + refusal: Final = self.refusals.get(model) if refusal is not None: return WebsocketFrameModelRefused(message=refusal, error_type=self.error_type) return WebsocketFrameModelAllowed(access_groups=self.groups.get(model, ())) @@ -7714,8 +7714,8 @@ async def _relay_client_frames( gate: WebsocketFrameModelGate | None, user_api_key_dict: UserAPIKeyAuth | None = None, ) -> SimpleNamespace: - upstream_ws = RecordingUpstreamWebSocket() - websocket = _client_websocket( + upstream_ws: Final = RecordingUpstreamWebSocket() + websocket: Final = _client_websocket( AsyncMock(side_effect=[_received_frame(frame) for frame in frames] + [{"type": "websocket.disconnect"}]) ) @@ -7758,9 +7758,9 @@ async def _relay_client_frames( ], ) async def test_websocket_passthrough_checks_the_model_a_client_frame_asks_for(frame, expected_model): - gate = RecordingFrameModelGate() + gate: Final = RecordingFrameModelGate() - relayed = await _relay_client_frames([json.dumps(frame)], gate) + relayed: Final = await _relay_client_frames([json.dumps(frame)], gate) assert gate.checked == [expected_model] assert [call.args[0] for call in relayed.upstream.send.await_args_list] == [json.dumps(frame)] @@ -7769,10 +7769,10 @@ async def test_websocket_passthrough_checks_the_model_a_client_frame_asks_for(fr @pytest.mark.asyncio async def test_websocket_passthrough_refuses_a_frame_naming_a_model_the_key_cannot_call(): - gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": KEY_MODEL_REFUSAL})) - frame = json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}) + gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": KEY_MODEL_REFUSAL})) + frame: Final = json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}) - relayed = await _relay_client_frames([frame], gate) + relayed: Final = await _relay_client_frames([frame], gate) assert relayed.upstream.send.await_args_list == [] assert json.loads(relayed.client.send_text.await_args.args[0]) == { @@ -7790,10 +7790,10 @@ async def test_websocket_passthrough_closes_with_the_whole_first_sentence_of_the reads as a truncated number. The first sentence names the reason whole, and the error event before the close still carries every word with its type. """ - refusal = "Budget has been exceeded! Key=" + "f" * 64 + " Spend=1.23456789, Budget=1.0" - gate = RecordingFrameModelGate(MappingProxyType({"gpt-5.4": refusal}), error_type="budget_exceeded") + refusal: Final = "Budget has been exceeded! Key=" + "f" * 64 + " Spend=1.23456789, Budget=1.0" + gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-5.4": refusal}), error_type="budget_exceeded") - relayed = await _relay_client_frames([json.dumps({"type": "response.create", "model": "gpt-5.4"})], gate) + relayed: Final = await _relay_client_frames([json.dumps({"type": "response.create", "model": "gpt-5.4"})], gate) assert relayed.client.close.await_args_list == [call(code=1008, reason="Budget has been exceeded!")] assert json.loads(relayed.client.send_text.await_args.args[0])["error"] == { @@ -7804,15 +7804,15 @@ async def test_websocket_passthrough_closes_with_the_whole_first_sentence_of_the @pytest.mark.asyncio async def test_websocket_passthrough_refuses_the_whole_frame_when_any_named_model_is_denied(): - gate = RecordingFrameModelGate(MappingProxyType({"gpt-4o-transcribe": KEY_MODEL_REFUSAL})) - frame = json.dumps( + gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-4o-transcribe": KEY_MODEL_REFUSAL})) + frame: Final = json.dumps( { "type": "session.update", "session": {"model": "gpt-realtime-2.1", "input_audio_transcription": {"model": "gpt-4o-transcribe"}}, } ) - relayed = await _relay_client_frames([frame], gate) + relayed: Final = await _relay_client_frames([frame], gate) assert gate.checked == ["gpt-realtime-2.1", "gpt-4o-transcribe"] assert relayed.upstream.send.await_args_list == [] @@ -7821,13 +7821,13 @@ async def test_websocket_passthrough_refuses_the_whole_frame_when_any_named_mode @pytest.mark.asyncio async def test_websocket_passthrough_stops_relaying_after_a_refused_frame(): - gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": KEY_MODEL_REFUSAL})) - frames = [ + gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": KEY_MODEL_REFUSAL})) + frames: Final = [ json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}), json.dumps({"type": "response.create"}), ] - relayed = await _relay_client_frames(frames, gate) + relayed: Final = await _relay_client_frames(frames, gate) assert relayed.upstream.send.await_args_list == [] assert gate.checked == ["gpt-realtime-2.1"] @@ -7845,9 +7845,9 @@ async def test_websocket_passthrough_stops_relaying_after_a_refused_frame(): ], ) async def test_websocket_passthrough_forwards_frames_that_name_no_model(frame): - gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": "denied"})) + gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": "denied"})) - relayed = await _relay_client_frames([frame], gate) + relayed: Final = await _relay_client_frames([frame], gate) assert gate.checked == [] assert [call_.args[0] for call_ in relayed.upstream.send.await_args_list] == [frame] @@ -7862,13 +7862,13 @@ async def test_websocket_passthrough_never_bills_the_model_a_client_frame_named( session onto that model's price, so only the provider's own frames decide what a session is billed as. """ - gate = RecordingFrameModelGate() - frames = [ + gate: Final = RecordingFrameModelGate() + frames: Final = [ json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}), json.dumps({"type": "response.create", "response": {"model": "gpt-realtime-2.1-mini"}}), ] - relayed = await _relay_client_frames(frames, gate) + relayed: Final = await _relay_client_frames(frames, gate) assert gate.checked == ["gpt-realtime-2.1", "gpt-realtime-2.1-mini"] assert "model" not in relayed.success_handler.call_args.kwargs @@ -7888,9 +7888,9 @@ async def test_websocket_passthrough_checks_a_model_named_with_escaped_json(fram "\\u006dodel" asks for a model just as plainly as one carrying "model" and has to face the same check. """ - gate = RecordingFrameModelGate(MappingProxyType({expected_model: KEY_MODEL_REFUSAL})) + gate: Final = RecordingFrameModelGate(MappingProxyType({expected_model: KEY_MODEL_REFUSAL})) - relayed = await _relay_client_frames([frame], gate) + relayed: Final = await _relay_client_frames([frame], gate) assert gate.checked == [expected_model] assert relayed.upstream.send.await_args_list == [] @@ -7899,10 +7899,10 @@ async def test_websocket_passthrough_checks_a_model_named_with_escaped_json(fram @pytest.mark.asyncio async def test_websocket_passthrough_checks_the_model_a_binary_frame_asks_for(): - gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": KEY_MODEL_REFUSAL})) - frame = json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}).encode() + gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": KEY_MODEL_REFUSAL})) + frame: Final = json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}).encode() - relayed = await _relay_client_frames([frame], gate) + relayed: Final = await _relay_client_frames([frame], gate) assert gate.checked == ["gpt-realtime-2.1"] assert relayed.upstream.send.await_args_list == [] @@ -7924,9 +7924,9 @@ async def test_websocket_passthrough_refuses_a_stored_prompt_for_a_key_limited_t a model list could reach any model through one. The gateway cannot see that model, so a key that is not allowed every model cannot run a stored prompt. """ - gate = RecordingFrameModelGate(MappingProxyType({STORED_PROMPT_PROBE_MODEL: KEY_MODEL_REFUSAL})) + gate: Final = RecordingFrameModelGate(MappingProxyType({STORED_PROMPT_PROBE_MODEL: KEY_MODEL_REFUSAL})) - relayed = await _relay_client_frames([json.dumps(frame)], gate) + relayed: Final = await _relay_client_frames([json.dumps(frame)], gate) error_message: Final = json.loads(relayed.client.send_text.await_args.args[0])["error"]["message"] close_reason: Final = relayed.client.close.await_args.kwargs["reason"] @@ -7940,9 +7940,9 @@ async def test_websocket_passthrough_refuses_a_stored_prompt_for_a_key_limited_t @pytest.mark.asyncio @pytest.mark.parametrize("frame", STORED_PROMPT_FRAMES) async def test_websocket_passthrough_forwards_a_stored_prompt_for_a_key_allowed_every_model(frame): - gate = RecordingFrameModelGate() + gate: Final = RecordingFrameModelGate() - relayed = await _relay_client_frames([json.dumps(frame)], gate) + relayed: Final = await _relay_client_frames([json.dumps(frame)], gate) assert [call_.args[0] for call_ in relayed.upstream.send.await_args_list] == [json.dumps(frame)] assert relayed.client.close.await_args_list == [call()] @@ -7950,12 +7950,12 @@ async def test_websocket_passthrough_forwards_a_stored_prompt_for_a_key_allowed_ @pytest.mark.asyncio async def test_websocket_passthrough_keeps_the_budget_reason_when_a_stored_prompt_is_refused_for_spend(): - refusal = "Budget has been exceeded! Key=hashed Spend=2.0, Budget=1.0" - gate = RecordingFrameModelGate( + refusal: Final = "Budget has been exceeded! Key=hashed Spend=2.0, Budget=1.0" + gate: Final = RecordingFrameModelGate( MappingProxyType({STORED_PROMPT_PROBE_MODEL: refusal}), error_type="budget_exceeded" ) - relayed = await _relay_client_frames([json.dumps(STORED_PROMPT_FRAMES[0])], gate) + relayed: Final = await _relay_client_frames([json.dumps(STORED_PROMPT_FRAMES[0])], gate) assert json.loads(relayed.client.send_text.await_args.args[0])["error"] == { "type": "budget_exceeded", @@ -7966,14 +7966,14 @@ async def test_websocket_passthrough_keeps_the_budget_reason_when_a_stored_promp @pytest.mark.asyncio async def test_websocket_passthrough_checks_the_model_a_hosted_tool_names(): """A hosted tool such as image generation runs its own model, so that model gets the same check as the turn's""" - gate = RecordingFrameModelGate(MappingProxyType({"gpt-image-2": KEY_MODEL_REFUSAL})) - frame = { + gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-image-2": KEY_MODEL_REFUSAL})) + frame: Final = { "type": "response.create", "model": "gpt-5.4-mini", "tools": [{"type": "web_search"}, {"type": "image_generation", "model": "gpt-image-2"}], } - relayed = await _relay_client_frames([json.dumps(frame)], gate) + relayed: Final = await _relay_client_frames([json.dumps(frame)], gate) assert gate.checked == ["gpt-5.4-mini", "gpt-image-2"] assert relayed.upstream.send.await_args_list == [] @@ -7989,9 +7989,9 @@ async def test_websocket_passthrough_checks_the_model_a_hosted_tool_names(): ], ) async def test_websocket_passthrough_refuses_an_image_tool_without_a_model_for_a_key_limited_to_some_models(frame): - gate = RecordingFrameModelGate(MappingProxyType({STORED_PROMPT_PROBE_MODEL: KEY_MODEL_REFUSAL})) + gate: Final = RecordingFrameModelGate(MappingProxyType({STORED_PROMPT_PROBE_MODEL: KEY_MODEL_REFUSAL})) - relayed = await _relay_client_frames([json.dumps(frame)], gate) + relayed: Final = await _relay_client_frames([json.dumps(frame)], gate) assert gate.checked[-1] == STORED_PROMPT_PROBE_MODEL assert relayed.upstream.send.await_args_list == [] @@ -8001,10 +8001,10 @@ async def test_websocket_passthrough_refuses_an_image_tool_without_a_model_for_a @pytest.mark.asyncio async def test_websocket_passthrough_checks_a_frame_only_the_standard_json_parser_reads(): """A lone surrogate escape is valid JSON syntax the fast parser refuses, so the frame is still checked""" - gate = RecordingFrameModelGate(MappingProxyType({"gpt-5.4": KEY_MODEL_REFUSAL})) - frame = '{"type": "response.create", "model": "gpt-5.4", "metadata": {"note": "\\ud800"}}' + gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-5.4": KEY_MODEL_REFUSAL})) + frame: Final = '{"type": "response.create", "model": "gpt-5.4", "metadata": {"note": "\\ud800"}}' - relayed = await _relay_client_frames([frame], gate) + relayed: Final = await _relay_client_frames([frame], gate) assert gate.checked == ["gpt-5.4"] assert relayed.upstream.send.await_args_list == [] @@ -8012,10 +8012,10 @@ async def test_websocket_passthrough_checks_a_frame_only_the_standard_json_parse @pytest.mark.asyncio async def test_websocket_passthrough_forwards_binary_audio_untouched(): - gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": "denied"})) - audio = bytes(range(256)) * 4 + gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": "denied"})) + audio: Final = bytes(range(256)) * 4 - relayed = await _relay_client_frames([audio], gate) + relayed: Final = await _relay_client_frames([audio], gate) assert gate.checked == [] assert [call_.args[0] for call_ in relayed.upstream.send.await_args_list] == [audio] @@ -8023,9 +8023,9 @@ async def test_websocket_passthrough_forwards_binary_audio_untouched(): @pytest.mark.asyncio async def test_websocket_passthrough_without_a_frame_gate_forwards_every_frame(): - frame = json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}) + frame: Final = json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}) - relayed = await _relay_client_frames([frame], None) + relayed: Final = await _relay_client_frames([frame], None) assert [call_.args[0] for call_ in relayed.upstream.send.await_args_list] == [frame] assert "model" not in relayed.success_handler.call_args.kwargs @@ -8051,7 +8051,7 @@ async def test_websocket_passthrough_still_bills_a_session_the_client_walked_out already used still have to reach spend, including the last frame the client never collected. """ - upstream_ws = FakeUpstreamWebSocket( + upstream_ws: Final = FakeUpstreamWebSocket( REALTIME_SESSION_CREATED, REALTIME_RESPONSE_DONE, json.dumps({"type": "response.output_text.delta", "delta": "too late"}), @@ -8112,7 +8112,7 @@ class ScriptedUpstreamWebSocket: from websockets.exceptions import ConnectionClosedOK from websockets.frames import Close - frame = await self._frames.get() + frame: Final = await self._frames.get() if frame is None: raise ConnectionClosedOK(rcvd=Close(1000, ""), sent=Close(1000, ""), rcvd_then_sent=True) return frame @@ -8126,14 +8126,14 @@ def _turn_taking_client( from starlette.websockets import WebSocketState delivered: list[str] = [] - script = iter(frames) - receives = iter(range(len(frames) + 1) if answers_due is None else answers_due) + script: Final = iter(frames) + receives: Final = iter(range(len(frames) + 1) if answers_due is None else answers_due) async def receive(): - answers_due = next(receives) + answers_due: Final = next(receives) while sum(answered_by in text for text in delivered) < answers_due: await asyncio.sleep(0) - frame = next(script, None) + frame: Final = next(script, None) return {"type": "websocket.disconnect"} if frame is None else _received_frame(frame) websocket = _client_websocket(receive) @@ -8208,21 +8208,21 @@ async def test_websocket_passthrough_logs_each_turn_on_its_own_row_charged_to_th and be charged to the group that let its model through, so a group budget sees its share during the session and never pays for the other group's turn. """ - gate = RecordingFrameModelGate( + gate: Final = RecordingFrameModelGate( groups=MappingProxyType({"gpt-5.4-mini": ("tier-a",), "gpt-5.4": ("tier-b",)}), ) - client = _turn_taking_client( + client: Final = _turn_taking_client( [ json.dumps({"type": "response.create", "model": "gpt-5.4-mini"}), json.dumps({"type": "response.create", "model": "gpt-5.4"}), ], answered_by='"response.completed"', ) - upstream_ws = ScriptedUpstreamWebSocket(replies=[_responses_turn("gpt-5.4-mini"), _responses_turn("gpt-5.4")]) + upstream_ws: Final = ScriptedUpstreamWebSocket(replies=[_responses_turn("gpt-5.4-mini"), _responses_turn("gpt-5.4")]) - relayed = await _relay_turns(client, upstream_ws, gate) + relayed: Final = await _relay_turns(client, upstream_ws, gate) - rows = [logged.kwargs for logged in relayed.success_handler.call_args_list] + rows: Final = [logged.kwargs for logged in relayed.success_handler.call_args_list] assert [[frame["response"]["model"] for frame in row["response_body"]] for row in rows] == [ ["gpt-5.4-mini", "gpt-5.4-mini"], ["gpt-5.4", "gpt-5.4"], @@ -8246,10 +8246,10 @@ async def test_websocket_passthrough_charges_pipelined_turns_to_the_frame_that_a and under the model that frame asked for, even when the provider answers under a dated snapshot name a per-model budget was never set on. """ - gate = RecordingFrameModelGate( + gate: Final = RecordingFrameModelGate( groups=MappingProxyType({"gpt-5.4-mini": ("tier-a",), "gpt-5.4": ("tier-b",)}), ) - client = _turn_taking_client( + client: Final = _turn_taking_client( [ json.dumps({"type": "response.create", "model": "gpt-5.4-mini"}), json.dumps({"type": "response.create", "model": "gpt-5.4"}), @@ -8257,13 +8257,13 @@ async def test_websocket_passthrough_charges_pipelined_turns_to_the_frame_that_a answered_by='"response.completed"', answers_due=(0, 0, 2), ) - upstream_ws = ScriptedUpstreamWebSocket( + upstream_ws: Final = ScriptedUpstreamWebSocket( replies=[_responses_turn("gpt-5.4-mini-2099-01-01"), _responses_turn("gpt-5.4-2099-01-01")] ) - relayed = await _relay_turns(client, upstream_ws, gate) + relayed: Final = await _relay_turns(client, upstream_ws, gate) - rows = [logged.kwargs for logged in relayed.success_handler.call_args_list] + rows: Final = [logged.kwargs for logged in relayed.success_handler.call_args_list] assert [[frame["response"]["model"] for frame in row["response_body"]] for row in rows] == [ ["gpt-5.4-mini-2099-01-01", "gpt-5.4-mini-2099-01-01"], ["gpt-5.4-2099-01-01", "gpt-5.4-2099-01-01"], @@ -8291,10 +8291,10 @@ async def test_websocket_passthrough_charges_turns_after_a_refused_frame_to_the_ open. Those frames never run, so each later turn still has to reach spend under the groups and model its own frame asked for. """ - gate = RecordingFrameModelGate( + gate: Final = RecordingFrameModelGate( groups=MappingProxyType({"gpt-5.4-mini": ("tier-a",), "gpt-5.4": ("tier-b",)}), ) - client = _turn_taking_client( + client: Final = _turn_taking_client( [ json.dumps({"type": "response.create", "input": "hi"}), json.dumps({"type": "response.create", "model": "gpt-5.4", "previous_response_id": "resp_gone"}), @@ -8304,7 +8304,7 @@ async def test_websocket_passthrough_charges_turns_after_a_refused_frame_to_the_ answered_by='"response.completed"', answers_due=(0, 0, 0, 0, 2), ) - upstream_ws = ScriptedUpstreamWebSocket( + upstream_ws: Final = ScriptedUpstreamWebSocket( replies=[ [_responses_error("missing_required_parameter")], [_responses_error("previous_response_not_found")], @@ -8313,9 +8313,9 @@ async def test_websocket_passthrough_charges_turns_after_a_refused_frame_to_the_ ] ) - relayed = await _relay_turns(client, upstream_ws, gate) + relayed: Final = await _relay_turns(client, upstream_ws, gate) - rows = [logged.kwargs for logged in relayed.success_handler.call_args_list] + rows: Final = [logged.kwargs for logged in relayed.success_handler.call_args_list] assert [row["litellm_params"]["metadata"]["model_group"] for row in rows] == ["gpt-5.4-mini", "gpt-5.4"] assert [row["litellm_params"]["metadata"][MODEL_ACCESS_GROUP_METADATA_KEY] for row in rows] == [ ["tier-a"], @@ -8331,10 +8331,10 @@ async def test_websocket_passthrough_charges_turns_on_parallel_lanes_to_the_fram Each turn still has to reach spend on its own row, under the groups and model its own frame asked for. """ - gate = RecordingFrameModelGate( + gate: Final = RecordingFrameModelGate( groups=MappingProxyType({"gpt-5.4-mini": ("tier-a",), "gpt-5.4": ("tier-b",)}), ) - client = _turn_taking_client( + client: Final = _turn_taking_client( [ json.dumps({"type": "response.create", "stream_id": "slow", "model": "gpt-5.4"}), json.dumps({"type": "response.create", "stream_id": "fast", "model": "gpt-5.4-mini"}), @@ -8342,7 +8342,7 @@ async def test_websocket_passthrough_charges_turns_on_parallel_lanes_to_the_fram answered_by='"response.completed"', answers_due=(0, 0, 2), ) - upstream_ws = ScriptedUpstreamWebSocket( + upstream_ws: Final = ScriptedUpstreamWebSocket( replies=[ [_on_lane(_responses_turn_started("gpt-5.4"), "slow")], [ @@ -8352,9 +8352,9 @@ async def test_websocket_passthrough_charges_turns_on_parallel_lanes_to_the_fram ] ) - relayed = await _relay_turns(client, upstream_ws, gate) + relayed: Final = await _relay_turns(client, upstream_ws, gate) - rows = [logged.kwargs for logged in relayed.success_handler.call_args_list] + rows: Final = [logged.kwargs for logged in relayed.success_handler.call_args_list] assert [[(frame["stream_id"], frame["response"]["model"]) for frame in row["response_body"]] for row in rows] == [ [("fast", "gpt-5.4-mini"), ("fast", "gpt-5.4-mini")], [("slow", "gpt-5.4"), ("slow", "gpt-5.4")], @@ -8373,8 +8373,8 @@ async def test_websocket_passthrough_charges_realtime_turns_to_the_session_the_k charged to the groups and the model the key opened it with, whatever model a later session.update or response.create names. """ - gate = RecordingFrameModelGate(groups=MappingProxyType({"gpt-realtime-2.1-mini": ("tier-b",)})) - client = _turn_taking_client( + gate: Final = RecordingFrameModelGate(groups=MappingProxyType({"gpt-realtime-2.1-mini": ("tier-b",)})) + client: Final = _turn_taking_client( [ json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1-mini"}}), json.dumps({"type": "response.create", "response": {"model": "gpt-realtime-2.1-mini"}}), @@ -8382,7 +8382,7 @@ async def test_websocket_passthrough_charges_realtime_turns_to_the_session_the_k answered_by='"type"', answers_due=(1, 2, 4), ) - upstream_ws = ScriptedUpstreamWebSocket( + upstream_ws: Final = ScriptedUpstreamWebSocket( REALTIME_SESSION_CREATED, replies=[ [json.dumps({"type": "session.updated"})], @@ -8390,7 +8390,7 @@ async def test_websocket_passthrough_charges_realtime_turns_to_the_session_the_k ], ) - relayed = await _relay_turns( + relayed: Final = await _relay_turns( client, upstream_ws, gate, @@ -8400,7 +8400,7 @@ async def test_websocket_passthrough_charges_realtime_turns_to_the_session_the_k ) assert gate.checked == ["gpt-realtime-2.1-mini", "gpt-realtime-2.1-mini"] - rows = [logged.kwargs for logged in relayed.success_handler.call_args_list] + rows: Final = [logged.kwargs for logged in relayed.success_handler.call_args_list] assert [row["litellm_params"]["metadata"][MODEL_ACCESS_GROUP_METADATA_KEY] for row in rows] == [["tier-a"]] assert [row["litellm_params"]["metadata"]["model_group"] for row in rows] == ["gpt-realtime-2.1"] @@ -8408,10 +8408,10 @@ async def test_websocket_passthrough_charges_realtime_turns_to_the_session_the_k @pytest.mark.asyncio async def test_websocket_passthrough_logs_no_spend_row_for_a_session_refused_before_any_answer(): """A session refused at its first frame ran nothing, so its only record is the refusal, not an extra $0 row""" - gate = RecordingFrameModelGate(MappingProxyType({"gpt-5.4": KEY_MODEL_REFUSAL})) - client = _turn_taking_client([json.dumps({"type": "response.create", "model": "gpt-5.4"})], answered_by='"error"') + gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-5.4": KEY_MODEL_REFUSAL})) + client: Final = _turn_taking_client([json.dumps({"type": "response.create", "model": "gpt-5.4"})], answered_by='"error"') - relayed = await _relay_turns(client, ScriptedUpstreamWebSocket(), gate) + relayed: Final = await _relay_turns(client, ScriptedUpstreamWebSocket(), gate) assert relayed.success_handler.call_args_list == [] @@ -8422,18 +8422,18 @@ async def test_websocket_passthrough_logs_the_unfinished_turn_a_client_walked_ou A client can leave while a later turn is still running. That turn never ends, but it was started and may be billed, so it still gets its own row. """ - client = _turn_taking_client( + client: Final = _turn_taking_client( [ json.dumps({"type": "response.create", "model": "gpt-5.4-mini"}), json.dumps({"type": "response.create", "model": "gpt-5.4"}), ], answered_by='"response.created"', ) - upstream_ws = ScriptedUpstreamWebSocket( + upstream_ws: Final = ScriptedUpstreamWebSocket( replies=[_responses_turn("gpt-5.4-mini"), [_responses_turn_started("gpt-5.4")]] ) - relayed = await _relay_turns(client, upstream_ws, RecordingFrameModelGate()) + relayed: Final = await _relay_turns(client, upstream_ws, RecordingFrameModelGate()) assert [ [json.loads(frame)["type"] if isinstance(frame, str) else frame["type"] for frame in logged.kwargs["response_body"]] @@ -8441,6 +8441,41 @@ 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_charges_each_lane_a_client_walked_out_of_to_the_frame_that_started_it(): + """ + A client can leave while turns on several stream_id lanes are still + running. Each unfinished turn still gets its own row, under the groups and + model its own frame asked for. + """ + gate: Final = RecordingFrameModelGate( + groups=MappingProxyType({"gpt-5.4-mini": ("tier-a",), "gpt-5.4": ("tier-b",)}), + ) + client: Final = _turn_taking_client( + [ + json.dumps({"type": "response.create", "stream_id": "slow", "model": "gpt-5.4"}), + json.dumps({"type": "response.create", "stream_id": "fast", "model": "gpt-5.4-mini"}), + ], + answered_by='"response.created"', + ) + upstream_ws: Final = ScriptedUpstreamWebSocket( + replies=[ + [_on_lane(_responses_turn_started("gpt-5.4"), "slow")], + [_on_lane(_responses_turn_started("gpt-5.4-mini"), "fast")], + ] + ) + + relayed: Final = await _relay_turns(client, upstream_ws, gate) + + rows: Final = [logged.kwargs for logged in relayed.success_handler.call_args_list] + assert [[frame["stream_id"] for frame in row["response_body"]] for row in rows] == [["slow"], ["fast"]] + assert [row["litellm_params"]["metadata"]["model_group"] for row in rows] == ["gpt-5.4", "gpt-5.4-mini"] + assert [row["litellm_params"]["metadata"][MODEL_ACCESS_GROUP_METADATA_KEY] for row in rows] == [ + ["tier-b"], + ["tier-a"], + ] + + class UpstreamFinishingTheTurnAfterTheClientLeft: """Starts a turn for each client frame and sends the rest of it only after the client left, and stops once closed""" @@ -8474,7 +8509,7 @@ class UpstreamFinishingTheTurnAfterTheClientLeft: from websockets.exceptions import ConnectionClosedOK from websockets.frames import Close - frame = await self._frames.get() + frame: Final = await self._frames.get() if frame is None: raise ConnectionClosedOK(rcvd=Close(1000, ""), sent=Close(1000, ""), rcvd_then_sent=True) return frame @@ -8492,16 +8527,16 @@ async def test_websocket_passthrough_bills_the_turn_a_client_walked_out_of_once_ """ from starlette.websockets import WebSocketDisconnect, WebSocketState - client_left = asyncio.Event() - send_failed = asyncio.Event() + client_left: Final = asyncio.Event() + send_failed: Final = asyncio.Event() client = _turn_taking_client( [json.dumps({"type": "response.create", "model": "gpt-5.4"})], answered_by='"response.created"' ) - receive = client.receive - send_text = client.send_text + receive: Final = client.receive + send_text: Final = client.send_text async def receive_until_leaving() -> Mapping[str, object]: - message = await receive() + message: Final = await receive() if message["type"] == "websocket.disconnect": client_left.set() if not hang_up_read_before_the_next_send: @@ -8517,13 +8552,13 @@ async def test_websocket_passthrough_bills_the_turn_a_client_walked_out_of_once_ client.receive = receive_until_leaving client.send_text = AsyncMock(side_effect=send_text_until_leaving) - upstream_ws = UpstreamFinishingTheTurnAfterTheClientLeft( + upstream_ws: Final = UpstreamFinishingTheTurnAfterTheClientLeft( client_left, _responses_turn_started("gpt-5.4"), (json.dumps({"type": "response.output_text.delta", "delta": "Hi"}), _responses_turn_ended("gpt-5.4")), ) - relayed = await _relay_turns(client, upstream_ws, RecordingFrameModelGate(), walkout_turn_drain_seconds=5) + relayed: Final = await _relay_turns(client, upstream_ws, RecordingFrameModelGate(), walkout_turn_drain_seconds=5) assert [ [json.loads(frame)["type"] if isinstance(frame, str) else frame["type"] for frame in logged.kwargs["response_body"]] @@ -8539,13 +8574,13 @@ async def test_websocket_passthrough_logs_no_extra_row_for_what_the_provider_sen 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( + client: Final = _turn_taking_client([json.dumps({"type": "response.create"})], answered_by='"rate_limits.updated"') + upstream_ws: Final = 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") + relayed: Final = await _relay_turns(client, upstream_ws, RecordingFrameModelGate(), "/openai/v1/realtime") assert [ [frame["type"] for frame in logged.kwargs["response_body"]] @@ -8560,15 +8595,15 @@ async def test_websocket_passthrough_carries_the_session_the_provider_opened_int session, and never again. Each later turn logs on its own, so it has to carry that opening frame or it would be priced as a model nobody named. """ - client = _turn_taking_client( + client: Final = _turn_taking_client( [json.dumps({"type": "response.create"}), json.dumps({"type": "response.create"})], answered_by='"response.done"', ) - upstream_ws = ScriptedUpstreamWebSocket( + upstream_ws: Final = ScriptedUpstreamWebSocket( REALTIME_SESSION_CREATED, replies=[[REALTIME_RESPONSE_DONE], [REALTIME_RESPONSE_DONE]] ) - relayed = await _relay_turns(client, upstream_ws, RecordingFrameModelGate(), "/openai/v1/realtime") + relayed: Final = await _relay_turns(client, upstream_ws, RecordingFrameModelGate(), "/openai/v1/realtime") assert [ [frame["type"] for frame in logged.kwargs["response_body"]] @@ -8584,8 +8619,8 @@ async def test_websocket_passthrough_still_bills_the_answer_in_flight_when_a_lat but it was paid for, so it has to reach spend instead of turning the whole session into a failure that bills nothing. """ - gate = RecordingFrameModelGate(MappingProxyType({"gpt-5.4": KEY_MODEL_REFUSAL})) - client = _turn_taking_client( + gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-5.4": KEY_MODEL_REFUSAL})) + client: Final = _turn_taking_client( [ json.dumps({"type": "response.create", "model": "gpt-5.4-mini"}), json.dumps({"type": "response.create", "model": "gpt-5.4"}), @@ -8593,9 +8628,9 @@ async def test_websocket_passthrough_still_bills_the_answer_in_flight_when_a_lat answered_by='"response.created"', ) created, completed = _responses_turn("gpt-5.4-mini") - upstream_ws = ScriptedUpstreamWebSocket(replies=[[created]], in_flight=[completed]) + upstream_ws: Final = ScriptedUpstreamWebSocket(replies=[[created]], in_flight=[completed]) - relayed = await _relay_turns(client, upstream_ws, gate) + relayed: Final = await _relay_turns(client, upstream_ws, gate) assert [json.loads(text.args[0])["type"] for text in client.send_text.await_args_list] == [ "response.created", @@ -9536,7 +9571,7 @@ def test_passthrough_client_cannot_forge_session_id_omission(client_metadata_key async def _billed_metadata(gate: WebsocketFrameModelGate) -> Mapping[str, object]: - served = await _relay_client_frames( + served: Final = await _relay_client_frames( [json.dumps({"type": "response.create", "model": "gpt-5.4-mini"})], gate, UserAPIKeyAuth(api_key="hashed-key"), @@ -9548,7 +9583,7 @@ async def _billed_metadata(gate: WebsocketFrameModelGate) -> Mapping[str, object async def test_websocket_session_is_charged_to_the_group_its_frames_were_authorized_through(): """A websocket names its model in a frame, long after the relay built its logging metadata, so the groups the gate matched have to be picked up at the end or the group pays for nothing it allowed.""" - metadata = await _billed_metadata(RecordingFrameModelGate(groups=MappingProxyType({"gpt-5.4-mini": ("tier-a",)}))) + metadata: Final = await _billed_metadata(RecordingFrameModelGate(groups=MappingProxyType({"gpt-5.4-mini": ("tier-a",)}))) assert metadata[MODEL_ACCESS_GROUP_METADATA_KEY] == ["tier-a"] assert metadata["user_api_key"] == "hashed-key" @@ -9556,7 +9591,7 @@ async def test_websocket_session_is_charged_to_the_group_its_frames_were_authori @pytest.mark.asyncio async def test_websocket_session_that_matched_no_group_is_charged_to_none(): - metadata = await _billed_metadata(RecordingFrameModelGate()) + metadata: Final = await _billed_metadata(RecordingFrameModelGate()) assert not metadata[MODEL_ACCESS_GROUP_METADATA_KEY] diff --git a/tests/unit/proxy/pass_through_endpoints/test_success_handler.py b/tests/unit/proxy/pass_through_endpoints/test_success_handler.py index d71cf2b5249..988828f48de 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_success_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/test_success_handler.py @@ -677,8 +677,8 @@ def _normalize(url_route: str, response_body: dict[str, object] | list[dict[str, logging_obj=logging_obj, url_route=url_route, result="websocket_connection_successful", - start_time=datetime.now(), - end_time=datetime.now(), + start_time=_START, + end_time=_END, cache_hit=False, litellm_call_id="call-fake", ) @@ -686,11 +686,11 @@ def _normalize(url_route: str, response_body: dict[str, object] | list[dict[str, @pytest.mark.parametrize("url_route", ["/openai/v1/realtime", "/openai_passthrough/v1/realtime"]) def test_openai_websocket_frames_reach_the_websocket_cost_handler(url_route): - normalized = _normalize(url_route, REALTIME_FRAMES) + normalized: Final = _normalize(url_route, REALTIME_FRAMES) assert normalized["kwargs"]["model"] == "gpt-realtime-2.1" assert normalized["kwargs"]["custom_llm_provider"] == "openai" - prices = litellm.model_cost["gpt-realtime-2.1"] + prices: Final = litellm.model_cost["gpt-realtime-2.1"] assert normalized["kwargs"]["response_cost"] == pytest.approx( 33 * prices["input_cost_per_token"] + 16 * prices["output_cost_per_token"] ) @@ -698,7 +698,7 @@ def test_openai_websocket_frames_reach_the_websocket_cost_handler(url_route): def test_openai_http_passthrough_bodies_are_left_to_the_http_handlers(): - normalized = _normalize("/openai/v1/responses", {"model": "gpt-5.4-mini", "usage": {"input_tokens": 13}}) + normalized: Final = _normalize("/openai/v1/responses", {"model": "gpt-5.4-mini", "usage": {"input_tokens": 13}}) assert normalized["standard_logging_response_object"] is None assert normalized["kwargs"] == {"litellm_call_id": "call-fake"} diff --git a/tests/unit/proxy/test_openai_ws_passthrough_routes.py b/tests/unit/proxy/test_openai_ws_passthrough_routes.py index 140cdb20914..587faa02e35 100644 --- a/tests/unit/proxy/test_openai_ws_passthrough_routes.py +++ b/tests/unit/proxy/test_openai_ws_passthrough_routes.py @@ -137,7 +137,7 @@ async def _serve( user_api_key_dict: UserAPIKeyAuth, general_settings: Mapping[str, object], ) -> _Served: - served = _Served(relay=_FakeRelay(), frame_model_gate=_FakeFrameModelGate()) + served: Final = _Served(relay=_FakeRelay(), frame_model_gate=_FakeFrameModelGate()) await openai_websocket_proxy_route( websocket=websocket, endpoint=endpoint, @@ -153,7 +153,7 @@ async def _serve( @pytest.mark.parametrize("prefix", ["openai", "openai_passthrough"]) async def test_openai_websocket_forwards_query_and_keeps_provider_auth(prefix, monkeypatch): monkeypatch.delenv("OPENAI_API_BASE", raising=False) - websocket = _FakeWebSocket(f"/{prefix}/v1/realtime", "model=gpt-realtime-2.1") + websocket: Final = _FakeWebSocket(f"/{prefix}/v1/realtime", "model=gpt-realtime-2.1") with patch(GET_CREDENTIALS, return_value="sk-provider"): served = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), ENABLED) @@ -178,10 +178,10 @@ async def test_openai_websocket_forwards_query_and_keeps_provider_auth(prefix, m async def test_openai_websocket_reports_its_route_without_the_proxy_root_path(monkeypatch): """Route allowlists and spend rows key on the route the proxy serves, which a root path prefix is not part of""" monkeypatch.setenv("SERVER_ROOT_PATH", "/litellm") - websocket = _FakeWebSocket("/litellm/openai/v1/responses", "") + websocket: Final = _FakeWebSocket("/litellm/openai/v1/responses", "") with patch(GET_CREDENTIALS, return_value="sk-provider"): - served = await _serve(websocket, "v1/responses", UserAPIKeyAuth(), ENABLED) + served: Final = await _serve(websocket, "v1/responses", UserAPIKeyAuth(), ENABLED) assert [call.endpoint for call in served.relay.calls] == ["/openai/v1/responses"] @@ -263,7 +263,7 @@ async def test_openai_websocket_closes_cleanly_when_workload_identity_exchange_f @pytest.mark.asyncio async def test_openai_websocket_closes_cleanly_when_provider_credentials_missing(): - websocket = _FakeWebSocket("/openai/v1/realtime", "model=gpt-realtime-2.1") + websocket: Final = _FakeWebSocket("/openai/v1/realtime", "model=gpt-realtime-2.1") with patch(GET_CREDENTIALS, return_value=None): served = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), ENABLED) @@ -279,7 +279,7 @@ async def test_openai_websocket_closes_cleanly_when_provider_credentials_missing @pytest.mark.parametrize("prefix", ["openai", "openai_passthrough"]) @pytest.mark.parametrize("general_settings", DISABLED_SETTINGS) async def test_openai_websocket_refused_unless_explicitly_enabled(prefix, general_settings): - websocket = _FakeWebSocket(f"/{prefix}/v1/realtime", "model=gpt-realtime-2.1") + websocket: Final = _FakeWebSocket(f"/{prefix}/v1/realtime", "model=gpt-realtime-2.1") served = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), general_settings) @@ -320,11 +320,11 @@ async def test_openai_websocket_refusal_echoes_requested_subprotocol(): [["gpt-5.4-mini"], ["gpt-realtime-2.1", "gpt-5.4"], []], ) async def test_openai_websocket_connects_model_restricted_keys(key_models): - websocket = _FakeWebSocket("/openai/v1/responses", "") - user_api_key_dict = UserAPIKeyAuth(models=key_models, token="hashed-fake", user_id="user-fake") + websocket: Final = _FakeWebSocket("/openai/v1/responses", "") + user_api_key_dict: Final = UserAPIKeyAuth(models=key_models, token="hashed-fake", user_id="user-fake") with patch(GET_CREDENTIALS, return_value="sk-provider"): - served = await _serve(websocket, "v1/responses", user_api_key_dict, ENABLED) + served: Final = await _serve(websocket, "v1/responses", user_api_key_dict, ENABLED) assert len(served.relay.calls) == 1 assert served.relay.calls[0].client_frame_model_gate is served.frame_model_gate