From a0d7377afe34cb0f91cae5c5b7721c171697daed Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 22:15:26 -0700 Subject: [PATCH 01/15] test(proxy): cover per-frame model checks and session pricing on the OpenAI websocket passthrough --- ...st_auth_checks_object_access_and_lookup.py | 133 -------- .../auth/test_model_access_group_budgets.py | 63 ++++ ...i_websocket_passthrough_logging_handler.py | 290 ++++++++++++++++ .../test_pass_through_endpoints.py | 317 +++++++++++++++++- .../test_success_handler.py | 54 +++ .../test_openai_ws_passthrough_routes.py | 294 +++++++++++----- 6 files changed, 928 insertions(+), 223 deletions(-) create mode 100644 tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py diff --git a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py index effdd9725c2..5980663ecc1 100644 --- a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py +++ b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py @@ -10750,139 +10750,6 @@ async def test_key_budget_error_keeps_the_masked_key_name(key_name): assert f"Key=prod-key ({key_name}) Current cost" in message -class _UntouchedPrisma: - def __getattr__(self, name: str) -> object: - raise AssertionError(f"database reached through {name}") - - -class _MissingUserPrisma: - class db: - class litellm_usertable: - @staticmethod - async def find_unique(where: dict[str, str], include: dict[str, bool]) -> None: - return None - - -@pytest.mark.asyncio -async def test_enforced_model_allowlists_treats_a_missing_user_row_as_unrestricted(): - from litellm.proxy.auth.auth_checks import enforced_model_allowlists - from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache - from litellm.proxy.utils import ProxyLogging - - cache = UserApiKeyCache() - scopes = await enforced_model_allowlists( - valid_token=UserAPIKeyAuth(token="hashed-fake", user_id="default_user_id"), - prisma_client=_MissingUserPrisma(), - user_api_key_cache=cache, - proxy_logging_obj=ProxyLogging(user_api_key_cache=cache), - ) - - assert [list(scope) for scope in scopes] == [[], [], [], [], []] - - -class _UnreachableUserPrisma: - class db: - class litellm_usertable: - @staticmethod - async def find_unique(where: dict[str, str], include: dict[str, bool]) -> None: - raise RuntimeError("database gone") - - -@pytest.mark.asyncio -async def test_enforced_model_allowlists_surfaces_a_failed_user_lookup(): - from litellm.proxy.auth.auth_checks import enforced_model_allowlists - from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache - from litellm.proxy.utils import ProxyLogging - - cache = UserApiKeyCache() - with pytest.raises(ValueError, match="database gone"): - await enforced_model_allowlists( - valid_token=UserAPIKeyAuth(token="hashed-fake", user_id="user-fake"), - prisma_client=_UnreachableUserPrisma(), - user_api_key_cache=cache, - proxy_logging_obj=ProxyLogging(user_api_key_cache=cache), - ) - - -@pytest.mark.asyncio -async def test_enforced_model_allowlists_reads_every_level_from_cache(): - from litellm.proxy._types import ( - LiteLLM_BudgetTable, - LiteLLM_ProjectTableCachedObj, - LiteLLM_TeamMembership, - LiteLLM_TeamTableCachedObj, - LiteLLM_UserTable, - ) - from litellm.proxy.auth.auth_checks import enforced_model_allowlists - from litellm.proxy.common_utils.user_api_key_cache import ( - UserApiKeyCache, - team_membership_reservation_cache_key, - ) - from litellm.proxy.utils import ProxyLogging - - cache = UserApiKeyCache() - proxy_logging_obj = ProxyLogging(user_api_key_cache=cache) - await cache.async_set_cache( - key="team_id:team-fake", value=LiteLLM_TeamTableCachedObj(team_id="team-fake", models=["gpt-4o"]) - ) - await cache.async_set_cache( - key=team_membership_reservation_cache_key(user_id="user-fake", team_id="team-fake"), - value=LiteLLM_TeamMembership( - user_id="user-fake", - team_id="team-fake", - litellm_budget_table=LiteLLM_BudgetTable(allowed_models=["gpt-4o-mini"]), - ), - ) - await cache.async_set_cache( - key="project_id:project-fake", - value=LiteLLM_ProjectTableCachedObj(project_id="project-fake", models=["gpt-4.1"]), - ) - await cache.async_set_cache(key="user-fake", value=LiteLLM_UserTable(user_id="user-fake", models=["o3"])) - prisma_client = _UntouchedPrisma() - - team_scoped = await enforced_model_allowlists( - valid_token=UserAPIKeyAuth( - token="hashed-fake", - models=["all-team-models"], - team_models=["gpt-4o", "gpt-4o-mini"], - user_id="user-fake", - team_id="team-fake", - project_id="project-fake", - ), - prisma_client=prisma_client, - user_api_key_cache=cache, - proxy_logging_obj=proxy_logging_obj, - ) - personal = await enforced_model_allowlists( - valid_token=UserAPIKeyAuth(token="hashed-fake", user_id="user-fake"), - prisma_client=prisma_client, - user_api_key_cache=cache, - proxy_logging_obj=proxy_logging_obj, - ) - without_database = await enforced_model_allowlists( - valid_token=UserAPIKeyAuth( - token="hashed-fake", - models=["gpt-4o"], - team_models=["gpt-4o-mini"], - user_id="user-fake", - team_id="team-fake", - ), - prisma_client=None, - user_api_key_cache=cache, - proxy_logging_obj=proxy_logging_obj, - ) - - assert [list(scope) for scope in team_scoped] == [ - ["gpt-4o", "gpt-4o-mini"], - ["gpt-4o"], - ["gpt-4o-mini"], - [], - ["gpt-4.1"], - ] - assert [list(scope) for scope in personal] == [[], [], [], ["o3"], []] - assert [list(scope) for scope in without_database] == [["gpt-4o"], ["gpt-4o-mini"]] - - @pytest.mark.asyncio @pytest.mark.parametrize("channel", ["team", "key"]) async def test_access_group_model_fallback_uses_the_injected_database(channel: str) -> None: diff --git a/tests/unit/proxy/auth/test_model_access_group_budgets.py b/tests/unit/proxy/auth/test_model_access_group_budgets.py index 7506bd031d9..7b0c4d9996e 100644 --- a/tests/unit/proxy/auth/test_model_access_group_budgets.py +++ b/tests/unit/proxy/auth/test_model_access_group_budgets.py @@ -27,6 +27,7 @@ from litellm.proxy.auth.auth_checks import ( _model_access_group_max_budget_check, collect_matched_model_access_groups, common_checks, + enforce_model_access_group_budget, stamp_matched_model_access_groups, ) from litellm.proxy.common_utils.reset_budget_job import _model_access_group_counter_key @@ -551,3 +552,65 @@ async def test_common_checks_blocks_a_request_whose_group_is_over_budget(): async def test_free_model_routes_skip_the_model_access_group_budget_check(): """skip_budget_checks is how free models stay free; it has to cover this budget too.""" assert await _common_checks_with_over_budget_group(skip_budget_checks=True) is True + + +async def _enforce_for_a_route_without_a_request_body( + *rows: _MagBudgetRow, + key_models: list[str], + spend_by_counter_key: dict[str, float], +) -> UserAPIKeyAuth: + """Run the pair the websocket frame gate runs, which reaches the proxy for what it needs.""" + valid_token = UserAPIKeyAuth(api_key="hashed", models=key_models, user_id=USER_ID) + cache = await _cache() + read, _ = _spend_reader(spend_by_counter_key) + with ( + patch("litellm.proxy.proxy_server.prisma_client", _RecordingPrismaClient(*rows)), # test-quality-ok: read lazily inside the helper, which takes no client + patch("litellm.proxy.proxy_server.user_api_key_cache", cache), # test-quality-ok: same lazy import + patch("litellm.proxy.proxy_server.proxy_logging_obj", ProxyLogging(user_api_key_cache=cache)), # test-quality-ok: same lazy import + patch("litellm.proxy.proxy_server.get_current_spend", read), # test-quality-ok: get_current_spend is lazily imported inside the budget check and has no injection point + ): + await enforce_model_access_group_budget( + model="gpt-4o", + valid_token=valid_token, + llm_router=Router(model_list=MODEL_LIST), + ) + return valid_token + + +@pytest.mark.asyncio +async def test_caller_granted_the_model_only_through_an_exhausted_group_is_refused(): + """Routes with no request body, the websocket passthrough among them, run this instead of + ``common_checks``, so the group budget has to stop them the same way.""" + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _enforce_for_a_route_without_a_request_body( + _MagBudgetRow("tier-a", max_budget=1.0), + key_models=["tier-a"], + spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 1.5}, + ) + + assert exc_info.value.entity_id == "tier-a" + assert exc_info.value.current_cost == 1.5 + + +@pytest.mark.asyncio +async def test_group_that_authorized_the_call_is_stamped_for_the_spend_writer(): + """Without the stamp the call runs on the group's budget and never pays into it, so the pool + can never catch up with what it authorized.""" + valid_token = await _enforce_for_a_route_without_a_request_body( + _MagBudgetRow("tier-a", max_budget=10.0), + key_models=["tier-a"], + spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 1.0}, + ) + + assert valid_token.matched_model_access_groups == ["tier-a"] + + +@pytest.mark.asyncio +async def test_model_granted_directly_charges_no_group(): + valid_token = await _enforce_for_a_route_without_a_request_body( + _MagBudgetRow("tier-a", max_budget=1.0), + key_models=["gpt-4o"], + spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 1.5}, + ) + + assert valid_token.matched_model_access_groups is None 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 new file mode 100644 index 00000000000..91a01f02d94 --- /dev/null +++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py @@ -0,0 +1,290 @@ +from datetime import datetime +from types import MappingProxyType +from typing import Final +from unittest.mock import MagicMock + +import pytest + +import litellm +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_websocket_passthrough_logging_handler import ( + OpenAIWebsocketPassthroughLoggingHandler, +) + +REALTIME_MODEL: Final = "gpt-realtime-2.1" +RESPONSES_MODEL: Final = "gpt-5.4-mini" +REASONING_MODEL: Final = "gpt-5.4" + + +def _logging_obj() -> MagicMock: + logging_obj = MagicMock() + logging_obj.model_call_details = {} + return logging_obj + + +def _session_created(model: str) -> dict: + return {"type": "session.created", "session": {"id": "sess_fake", "model": model}} + + +def _realtime_response_done(text_in: int, audio_in: int, text_out: int, audio_out: int) -> dict: + return { + "type": "response.done", + "response": { + "id": "resp_fake", + "status": "completed", + "usage": { + "input_tokens": text_in + audio_in, + "output_tokens": text_out + audio_out, + "total_tokens": text_in + audio_in + text_out + audio_out, + "input_token_details": {"text_tokens": text_in, "audio_tokens": audio_in, "cached_tokens": 0}, + "output_token_details": {"text_tokens": text_out, "audio_tokens": audio_out}, + }, + }, + } + + +def _responses_completed(model: str, input_tokens: int, output_tokens: int) -> dict: + return { + "type": "response.completed", + "response": { + "id": "resp_fake", + "model": model, + "status": "completed", + "usage": { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "total_tokens": input_tokens + output_tokens, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + }, + } + + +def _handle(messages, kwargs=MappingProxyType({}), logging_obj=None, url_route="/openai/v1/realtime"): + return OpenAIWebsocketPassthroughLoggingHandler().openai_websocket_passthrough_handler( + websocket_messages=messages, + logging_obj=logging_obj if logging_obj is not None else _logging_obj(), + url_route=url_route, + start_time=datetime.now(), + 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] + return ( + text_in * prices["input_cost_per_token"] + + audio_in * prices["input_cost_per_audio_token"] + + text_out * prices["output_cost_per_token"] + + audio_out * prices["output_cost_per_audio_token"] + ) + + +def test_realtime_session_is_priced_on_the_model_the_session_reported(): + logging_obj = _logging_obj() + messages = [ + _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) + + expected_cost = _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" + assert handled["result"].model == REALTIME_MODEL + assert handled["result"].usage.total_tokens == 409 + assert logging_obj.model_call_details["response_cost"] == pytest.approx(expected_cost) + assert logging_obj.model_call_details["model"] == REALTIME_MODEL + + +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( + [ + _session_created(REALTIME_MODEL), + _realtime_response_done(10, 0, 5, 0), + _realtime_response_done(10, 0, 5, 0), + ] + ) + + assert two_turns["kwargs"]["response_cost"] == pytest.approx(2 * one_turn["kwargs"]["response_cost"]) + assert two_turns["result"].usage.total_tokens == 30 + + +def test_responses_turn_is_priced_on_the_model_the_response_reported(): + handled = _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"] + assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost) + assert handled["kwargs"]["model"] == RESPONSES_MODEL + assert handled["result"].usage.prompt_tokens == 13 + assert handled["result"].usage.completion_tokens == 5 + + +def test_responses_cost_covers_every_turn_of_the_connection(): + handled = _handle( + [ + _responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5), + _responses_completed(RESPONSES_MODEL, input_tokens=20, output_tokens=7), + ], + 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"] + assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost) + assert handled["result"].usage.total_tokens == 45 + + +def test_a_model_the_caller_asked_for_never_prices_a_session_the_provider_never_named(): + """ + The caller picks what to ask for, and the provider picks what to run. Only + 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() + + handled = _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, + ) + + assert handled["result"] is None + assert handled["kwargs"] == {"model": REALTIME_MODEL} + assert logging_obj.model_call_details == {} + + +def test_each_responses_turn_is_priced_on_the_model_that_turn_ran(): + """ + One socket can run several models one after another, so a connection that + 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( + [ + _responses_completed(REASONING_MODEL, input_tokens=13, output_tokens=5), + _responses_completed(RESPONSES_MODEL, input_tokens=20, output_tokens=7), + ], + url_route="/openai_passthrough/v1/responses", + ) + + first = litellm.model_cost[REASONING_MODEL] + second = litellm.model_cost[RESPONSES_MODEL] + expected_cost = ( + 13 * first["input_cost_per_token"] + + 5 * first["output_cost_per_token"] + + 20 * second["input_cost_per_token"] + + 7 * second["output_cost_per_token"] + ) + assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost) + assert handled["result"].usage.total_tokens == 45 + + +def test_a_responses_turn_that_names_no_model_is_priced_on_the_session_model(): + handled = _handle( + [ + _session_created(RESPONSES_MODEL), + { + "type": "response.completed", + "response": {"usage": {"input_tokens": 13, "output_tokens": 5, "total_tokens": 18}}, + }, + ], + 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"] + 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( + [{"type": "response.done", "response": {"usage": {"input_tokens": 10, "output_tokens": 5}}}], + kwargs=MappingProxyType({"litellm_call_id": "call-fake"}), + logging_obj=logging_obj, + ) + + assert handled["result"] is None + assert handled["kwargs"] == {"litellm_call_id": "call-fake"} + assert logging_obj.model_call_details == {} + + +def test_a_session_that_reported_no_usage_still_records_its_model(): + logging_obj = _logging_obj() + + handled = _handle([_session_created(REALTIME_MODEL)], logging_obj=logging_obj) + + assert handled["result"] is None + assert handled["kwargs"] == {"model": REALTIME_MODEL, "custom_llm_provider": "openai"} + assert logging_obj.model_call_details["model"] == REALTIME_MODEL + assert logging_obj.model_call_details["custom_llm_provider"] == "openai" + + +def test_a_responses_turn_without_readable_usage_is_not_billed(): + handled = _handle( + [{"type": "response.completed", "response": {"model": RESPONSES_MODEL, "status": "completed"}}], + url_route="/openai_passthrough/v1/responses", + ) + + assert handled["result"] is None + assert handled["kwargs"] == {"model": RESPONSES_MODEL, "custom_llm_provider": "openai"} + + +def test_the_price_of_a_multi_model_connection_reaches_the_spend_row(): + """ + The logging layer prices whatever result it is handed, and one result can + only name one model, so a connection that ran two models has to carry the + 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( + [ + _responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5), + _responses_completed(REASONING_MODEL, input_tokens=20, output_tokens=7), + ], + url_route="/openai_passthrough/v1/responses", + ) + + logged_cost = litellm.response_cost_calculator( + response_object=handled["result"], + model=handled["kwargs"]["model"], + custom_llm_provider="openai", + call_type="completion", + optional_params={}, + ) + + session_prices = litellm.model_cost[RESPONSES_MODEL] + priced_on_the_session_model = ( + 33 * session_prices["input_cost_per_token"] + 12 * session_prices["output_cost_per_token"] + ) + assert logged_cost == pytest.approx(handled["kwargs"]["response_cost"]) + assert logged_cost > priced_on_the_session_model + assert handled["kwargs"]["model"] == RESPONSES_MODEL + + +def test_a_turn_the_caller_walked_out_of_still_names_the_model_it_ran(): + """ + A client that closes before response.completed leaves no usage behind to + 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() + + handled = _handle( + [{"type": "response.created", "response": {"id": "resp_fake", "model": REASONING_MODEL}}], + logging_obj=logging_obj, + url_route="/openai_passthrough/v1/responses", + ) + + assert handled["result"] is None + assert handled["kwargs"] == {"model": REASONING_MODEL, "custom_llm_provider": "openai"} + assert logging_obj.model_call_details["model"] == REASONING_MODEL 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 db0a74d360a..6f35ebbc161 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 @@ -7,20 +7,20 @@ import os import sys import uuid import zlib -from collections.abc import AsyncIterator, Callable, Mapping +from collections.abc import AsyncIterator, Callable, Mapping, Sequence from contextlib import ExitStack, contextmanager from dataclasses import dataclass from datetime import datetime from io import BytesIO from types import MappingProxyType, ModuleType, SimpleNamespace from typing import Final, Optional, cast -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, call, patch from uuid import UUID import httpx import pytest import respx -from fastapi import APIRouter, FastAPI, HTTPException, Request, Response, UploadFile +from fastapi import APIRouter, FastAPI, HTTPException, Request, Response, UploadFile, WebSocketDisconnect from fastapi.responses import StreamingResponse from fastapi.routing import APIRoute from fastapi.testclient import TestClient @@ -33,6 +33,7 @@ from litellm._logging import verbose_proxy_logger from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._lazy_features import LazyFeature, attach_lazy_features from litellm.proxy._types import ProxyException, UserAPIKeyAuth @@ -7515,6 +7516,281 @@ async def test_websocket_passthrough_does_not_close_twice_when_success_logging_f websocket.close.assert_awaited_once_with(code=1008, reason=upstream_reason) +SUCCESS_HANDLER: Final = ( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints." + "pass_through_endpoint_logging.pass_through_async_success_handler" +) + + +class RecordingFrameModelGate: + def __init__(self, refusals: Mapping[str, str] = MappingProxyType({})): + self.refusals = refusals + self.checked: list[str] = [] + + async def __call__(self, model: str, valid_token: UserAPIKeyAuth, /) -> Optional[str]: + self.checked.append(model) + return self.refusals.get(model) + + +def _received_frame(frame): + if isinstance(frame, bytes): + return {"type": "websocket.receive", "bytes": frame} + return {"type": "websocket.receive", "text": frame} + + +async def _relay_client_frames(frames, gate, user_api_key_dict=None): + upstream_ws = RecordingUpstreamWebSocket() + websocket = _client_websocket( + AsyncMock(side_effect=[_received_frame(frame) for frame in frames] + [{"type": "websocket.disconnect"}]) + ) + + with ( + _patched_websocket_passthrough_environment(upstream_ws), + patch(SUCCESS_HANDLER, new_callable=AsyncMock) as success_handler, + ): + await websocket_passthrough_request( + websocket=websocket, + target="wss://api.openai.com/v1/realtime?model=gpt-realtime-2.1", + custom_headers={"Authorization": "Bearer sk-provider"}, + user_api_key_dict=user_api_key_dict or UserAPIKeyAuth(), + forward_headers=False, + endpoint="/openai/v1/realtime", + accept_websocket=False, + client_frame_model_gate=gate, + ) + + return SimpleNamespace(upstream=upstream_ws, client=websocket, success_handler=success_handler) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "frame, expected_model", + [ + ({"type": "response.create", "model": "gpt-5.4-mini"}, "gpt-5.4-mini"), + ({"type": "response.create", "response": {"model": "gpt-5.4"}}, "gpt-5.4"), + ({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}, "gpt-realtime-2.1"), + ( + {"type": "session.update", "session": {"input_audio_transcription": {"model": "gpt-4o-transcribe"}}}, + "gpt-4o-transcribe", + ), + ( + { + "type": "session.update", + "session": {"audio": {"input": {"transcription": {"model": "gpt-transcribe"}}}}, + }, + "gpt-transcribe", + ), + ], +) +async def test_websocket_passthrough_checks_the_model_a_client_frame_asks_for(frame, expected_model): + gate = RecordingFrameModelGate() + + relayed = 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)] + assert relayed.client.close.await_args_list == [call()] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_refuses_a_frame_naming_a_model_the_key_cannot_call(): + refusal = "key not allowed to access model. This key can only access models=['gpt-5.4-mini']." + gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": refusal})) + frame = json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}) + + relayed = await _relay_client_frames([frame], gate) + + assert relayed.upstream.send.await_args_list == [] + assert json.loads(relayed.client.send_text.await_args.args[0]) == { + "type": "error", + "error": {"type": "invalid_request_error", "message": refusal}, + } + assert relayed.client.close.await_args_list == [call(code=1008, reason=refusal)] + relayed.upstream.close.assert_awaited() + + +@pytest.mark.asyncio +async def test_websocket_passthrough_refuses_the_whole_frame_when_any_named_model_is_denied(): + refusal = "key not allowed to access model. This key can only access models=['gpt-realtime-2.1']." + gate = RecordingFrameModelGate(MappingProxyType({"gpt-4o-transcribe": refusal})) + frame = 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) + + assert gate.checked == ["gpt-realtime-2.1", "gpt-4o-transcribe"] + assert relayed.upstream.send.await_args_list == [] + assert relayed.client.close.await_args_list == [call(code=1008, reason=refusal)] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_stops_relaying_after_a_refused_frame(): + refusal = "key not allowed to access model. This key can only access models=['gpt-5.4-mini']." + gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": refusal})) + frames = [ + json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}), + json.dumps({"type": "response.create"}), + ] + + relayed = await _relay_client_frames(frames, gate) + + assert relayed.upstream.send.await_args_list == [] + assert gate.checked == ["gpt-realtime-2.1"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "frame", + [ + '{"type": "input_audio_buffer.append", "audio": "AAAA"}', + '{"type": "response.create", "response": {"modalities": ["audio"]}}', + '{"type": "session.update", "session": {"model": 42}}', + 'not json at all, but it does say "model"', + '{"type": "session.update", "session": "model"}', + ], +) +async def test_websocket_passthrough_forwards_frames_that_name_no_model(frame): + gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": "denied"})) + + relayed = await _relay_client_frames([frame], gate) + + assert gate.checked == [] + assert [call_.args[0] for call_ in relayed.upstream.send.await_args_list] == [frame] + assert relayed.client.close.await_args_list == [call()] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_never_bills_the_model_a_client_frame_named(): + """ + A frame is what the caller asked for, not what the provider ran. A caller + who names a cheap model in a frame the provider ignores must not move the + session onto that model's price, so only the provider's own frames decide + what a session is billed as. + """ + gate = RecordingFrameModelGate() + frames = [ + 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) + + assert gate.checked == ["gpt-realtime-2.1", "gpt-realtime-2.1-mini"] + assert "model" not in relayed.success_handler.call_args.kwargs + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "frame, expected_model", + [ + ('{"type": "response.create", "\\u006dodel": "gpt-5.4"}', "gpt-5.4"), + ('{"type": "session.update", "session": {"\\u006d\\u006fdel": "gpt-realtime-2.1"}}', "gpt-realtime-2.1"), + ], +) +async def test_websocket_passthrough_checks_a_model_named_with_escaped_json(frame, expected_model): + """ + JSON lets a caller spell the same key several ways, so a frame carrying + "\\u006dodel" asks for a model just as plainly as one carrying "model" and + has to face the same check. + """ + refusal = "key not allowed to access model. This key can only access models=['gpt-5.4-mini']." + gate = RecordingFrameModelGate(MappingProxyType({expected_model: refusal})) + + relayed = await _relay_client_frames([frame], gate) + + assert gate.checked == [expected_model] + assert relayed.upstream.send.await_args_list == [] + assert relayed.client.close.await_args_list == [call(code=1008, reason=refusal)] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_checks_the_model_a_binary_frame_asks_for(): + refusal = "key not allowed to access model. This key can only access models=['gpt-5.4-mini']." + gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": refusal})) + frame = json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}).encode() + + relayed = await _relay_client_frames([frame], gate) + + assert gate.checked == ["gpt-realtime-2.1"] + assert relayed.upstream.send.await_args_list == [] + assert relayed.client.close.await_args_list == [call(code=1008, reason=refusal)] + + +@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 + + relayed = await _relay_client_frames([audio], gate) + + assert gate.checked == [] + assert [call_.args[0] for call_ in relayed.upstream.send.await_args_list] == [audio] + + +@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"}}) + + relayed = 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 + + +REALTIME_SESSION_CREATED: Final = json.dumps( + {"type": "session.created", "session": {"id": "sess_fake", "model": "gpt-realtime-2.1"}} +) +REALTIME_RESPONSE_DONE: Final = json.dumps( + { + "type": "response.done", + "response": {"status": "completed", "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}}, + } +) + + +@pytest.mark.asyncio +async def test_websocket_passthrough_still_bills_a_session_the_client_walked_out_of(): + """ + A realtime client closes the moment it has the answer it asked for, so the + next frame the provider sends lands on a socket that is already gone. That + is how a session ends, not a gateway failure, and the turns the caller + already used still have to reach spend, including the last frame the client + never collected. + """ + upstream_ws = FakeUpstreamWebSocket( + REALTIME_SESSION_CREATED, + REALTIME_RESPONSE_DONE, + json.dumps({"type": "response.output_text.delta", "delta": "too late"}), + ) + websocket = _client_websocket(_pending_receive) + websocket.send_text = AsyncMock(side_effect=[None, None, WebSocketDisconnect(code=1006)]) + + with ( + _patched_websocket_passthrough_environment(upstream_ws), + patch(SUCCESS_HANDLER, new_callable=AsyncMock) as success_handler, + ): + await websocket_passthrough_request( + websocket=websocket, + target="wss://api.openai.com/v1/realtime?model=gpt-realtime-2.1", + custom_headers={"Authorization": "Bearer sk-provider"}, + user_api_key_dict=UserAPIKeyAuth(), + forward_headers=False, + endpoint="/openai_passthrough/v1/realtime", + accept_websocket=False, + ) + + relayed_frames: Final = success_handler.call_args.kwargs["response_body"] + assert [frame["type"] for frame in relayed_frames] == [ + "session.created", + "response.done", + "response.output_text.delta", + ] + + DEEPGRAM_LISTEN_TARGET = "wss://api.deepgram.com/v1/listen?model=nova-3&encoding=linear16&sample_rate=16000" DEEPGRAM_INTERIM_FRAME = json.dumps( { @@ -8442,6 +8718,41 @@ def test_passthrough_client_cannot_forge_session_id_omission(client_metadata_key ) +class StampingFrameModelGate: + def __init__(self, groups: list[str]): + self.groups = groups + + async def __call__(self, model: str, valid_token: UserAPIKeyAuth, /) -> Optional[str]: + valid_token.matched_model_access_groups = self.groups + return None + + +async def _billed_metadata(gate) -> dict: + served = await _relay_client_frames( + [json.dumps({"type": "response.create", "model": "gpt-5.4-mini"})], + gate, + UserAPIKeyAuth(api_key="hashed-key"), + ) + return served.success_handler.call_args.kwargs["litellm_params"]["metadata"] + + +@pytest.mark.asyncio +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(StampingFrameModelGate(["tier-a"])) + + assert metadata[MODEL_ACCESS_GROUP_METADATA_KEY] == ["tier-a"] + assert metadata["user_api_key"] == "hashed-key" + + +@pytest.mark.asyncio +async def test_websocket_session_that_matched_no_group_is_charged_to_none(): + metadata = await _billed_metadata(RecordingFrameModelGate()) + + assert not metadata[MODEL_ACCESS_GROUP_METADATA_KEY] + + @pytest.mark.parametrize("client_metadata_key", ["litellm_metadata", "metadata"]) def test_passthrough_logs_the_resolved_deployment_model_info_over_the_request_body(client_metadata_key: str): """A provider route that resolved a router deployment stashes its model_info on request.state. That 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 2cc03949ff1..5b3f337fb5e 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_success_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/test_success_handler.py @@ -648,3 +648,57 @@ def test_vertex_ai_live_route_without_usage_frames_yields_no_logging_response( assert normalized["standard_logging_response_object"] is None assert normalized["kwargs"] == {"model": _LIVE_MODEL} + + +REALTIME_FRAMES: Final = [ + {"type": "session.created", "session": {"id": "sess_fake", "model": "gpt-realtime-2.1"}}, + { + "type": "response.done", + "response": { + "usage": { + "input_tokens": 33, + "output_tokens": 16, + "total_tokens": 49, + "input_token_details": {"text_tokens": 33, "audio_tokens": 0, "cached_tokens": 0}, + "output_token_details": {"text_tokens": 16, "audio_tokens": 0}, + } + }, + }, +] + + +def _normalize(url_route: str, response_body): + logging_obj = MagicMock() + logging_obj.model_call_details = {} + return PassThroughEndpointLogging().normalize_llm_passthrough_logging_payload( + httpx_response=MagicMock(spec=httpx.Response), + response_body=response_body, + request_body={}, + logging_obj=logging_obj, + url_route=url_route, + result="websocket_connection_successful", + start_time=datetime.now(), + end_time=datetime.now(), + cache_hit=False, + litellm_call_id="call-fake", + ) + + +@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) + + assert normalized["kwargs"]["model"] == "gpt-realtime-2.1" + assert normalized["kwargs"]["custom_llm_provider"] == "openai" + prices = 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"] + ) + assert normalized["standard_logging_response_object"].usage.total_tokens == 49 + + +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}}) + + 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 6a9cd972dc2..bd65001111a 100644 --- a/tests/unit/proxy/test_openai_ws_passthrough_routes.py +++ b/tests/unit/proxy/test_openai_ws_passthrough_routes.py @@ -1,11 +1,12 @@ -"""OpenAI passthrough WebSocket route: registration, opt-in gating, and refusals.""" +"""OpenAI passthrough WebSocket route: registration, opt-in gating, refusals, and per-frame model checks.""" import json -from collections.abc import Mapping, Sequence +from collections.abc import Iterator, Mapping +from contextlib import contextmanager from dataclasses import dataclass from types import MappingProxyType, SimpleNamespace from typing import Final -from unittest.mock import patch +from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest @@ -14,18 +15,16 @@ from starlette.routing import WebSocketRoute import litellm from litellm.llms.openai.workload_identity import _workload_identity_auth -from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy._types import Litellm_EntityType, LiteLLM_UserTable, UserAPIKeyAuth from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( _OPENAI_WS_DISABLED_REFUSAL, - _OPENAI_WS_MODEL_RESTRICTED_REFUSAL, - _has_model_restrictions, + _OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL, _openai_websocket_refusal, - _proxy_model_allowlists, + _proxy_frame_model_gate, openai_websocket_proxy_route, router, ) - -Scopes = tuple[Sequence[str], ...] +from litellm.proxy.pass_through_endpoints.pass_through_endpoints import WebsocketFrameModelGate ENABLED: Final = MappingProxyType({"enable_openai_websocket_passthrough": True}) DISABLED_SETTINGS: Final = ( @@ -76,6 +75,7 @@ class _RelayCall: forward_headers: bool endpoint: str accept_websocket: bool + client_frame_model_gate: WebsocketFrameModelGate | None class _FakeRelay: @@ -92,6 +92,7 @@ class _FakeRelay: forward_headers: bool, endpoint: str, accept_websocket: bool, + client_frame_model_gate: WebsocketFrameModelGate | None = None, ) -> None: self.calls.append( _RelayCall( @@ -100,24 +101,25 @@ class _FakeRelay: forward_headers=forward_headers, endpoint=endpoint, accept_websocket=accept_websocket, + client_frame_model_gate=client_frame_model_gate, ) ) -class _FakeModelAllowlists: - def __init__(self, scopes: Scopes) -> None: - self.scopes = scopes - self.calls: list[UserAPIKeyAuth] = [] +class _FakeFrameModelGate: + def __init__(self, refusal: str | None = None) -> None: + self.refusal = refusal + self.checked: list[str] = [] - async def __call__(self, valid_token: UserAPIKeyAuth, /) -> Scopes: - self.calls.append(valid_token) - return self.scopes + async def __call__(self, model: str, valid_token: UserAPIKeyAuth, /) -> str | None: + self.checked.append(model) + return self.refusal @dataclass(frozen=True, slots=True) class _Served: relay: _FakeRelay - allowlists: _FakeModelAllowlists + frame_model_gate: _FakeFrameModelGate async def _serve( @@ -125,16 +127,15 @@ async def _serve( endpoint: str, user_api_key_dict: UserAPIKeyAuth, general_settings: Mapping[str, object], - scopes: Scopes = (), ) -> _Served: - served = _Served(relay=_FakeRelay(), allowlists=_FakeModelAllowlists(scopes)) + served = _Served(relay=_FakeRelay(), frame_model_gate=_FakeFrameModelGate()) await openai_websocket_proxy_route( websocket=websocket, endpoint=endpoint, user_api_key_dict=user_api_key_dict, general_settings=general_settings, relay=served.relay, - model_allowlists=served.allowlists, + frame_model_gate=served.frame_model_gate, ) return served @@ -143,18 +144,19 @@ 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-4o-realtime-preview") + websocket = _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) assert served.relay.calls == [ _RelayCall( - target="wss://api.openai.com/v1/realtime?model=gpt-4o-realtime-preview", + target="wss://api.openai.com/v1/realtime?model=gpt-realtime-2.1", custom_headers=MappingProxyType({"Authorization": "Bearer sk-provider"}), forward_headers=False, endpoint=f"/{prefix}/v1/realtime", accept_websocket=False, + client_frame_model_gate=served.frame_model_gate, ) ] assert websocket.accepts == [None] @@ -166,7 +168,7 @@ async def test_openai_websocket_forwards_query_and_keeps_provider_auth(prefix, m async def test_openai_websocket_accepts_first_client_subprotocol(): websocket = _FakeWebSocket( "/openai/v1/realtime", - "model=gpt-4o-realtime-preview", + "model=gpt-realtime-2.1", subprotocols="realtime, openai-insecure-api-key.sk-abc, openai-beta.realtime-v1", ) @@ -239,7 +241,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-4o-realtime-preview") + websocket = _FakeWebSocket("/openai/v1/realtime", "model=gpt-realtime-2.1") with patch(GET_CREDENTIALS, return_value=None): served = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), ENABLED) @@ -255,7 +257,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-4o-realtime-preview") + websocket = _FakeWebSocket(f"/{prefix}/v1/realtime", "model=gpt-realtime-2.1") served = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), general_settings) @@ -265,25 +267,21 @@ async def test_openai_websocket_refused_unless_explicitly_enabled(prefix, genera assert served.relay.calls == [] -@pytest.mark.asyncio @pytest.mark.parametrize("general_settings", DISABLED_SETTINGS) -async def test_openai_websocket_refusal_is_disabled_for_falsy_settings(general_settings): - refusal = await _openai_websocket_refusal(UserAPIKeyAuth(), general_settings, _FakeModelAllowlists(())) - assert refusal is _OPENAI_WS_DISABLED_REFUSAL +def test_openai_websocket_refusal_is_disabled_for_falsy_settings(general_settings): + assert _openai_websocket_refusal(general_settings) is _OPENAI_WS_DISABLED_REFUSAL -@pytest.mark.asyncio @pytest.mark.parametrize("value", [True, "true", "True"]) -async def test_openai_websocket_refusal_is_none_for_truthy_settings(value): - settings = MappingProxyType({"enable_openai_websocket_passthrough": value}) - assert await _openai_websocket_refusal(UserAPIKeyAuth(), settings, _FakeModelAllowlists(())) is None +def test_openai_websocket_refusal_is_none_for_truthy_settings(value): + assert _openai_websocket_refusal(MappingProxyType({"enable_openai_websocket_passthrough": value})) is None @pytest.mark.asyncio async def test_openai_websocket_refusal_echoes_requested_subprotocol(): websocket = _FakeWebSocket( "/openai_passthrough/v1/realtime", - "model=gpt-4o-realtime-preview", + "model=gpt-realtime-2.1", subprotocols="realtime, openai-beta.realtime-v1", ) @@ -294,68 +292,190 @@ async def test_openai_websocket_refusal_echoes_requested_subprotocol(): assert served.relay.calls == [] -RESTRICTED_SCOPES: Final[tuple[Scopes, ...]] = ( - (("gpt-4o",),), - ((), ("gpt-4o-realtime-preview",)), - (("all-team-models",), ("gpt-4o",)), - ((), ("all-proxy-models",), ("gpt-4o",)), - ((), (), (), ("gpt-4o",)), - (("*",), (), (), (), ("gpt-4o",)), +@pytest.mark.asyncio +@pytest.mark.parametrize( + "key_models", + [["gpt-5.4-mini"], ["gpt-realtime-2.1", "gpt-5.4"], []], ) -UNRESTRICTED_SCOPES: Final[tuple[Scopes, ...]] = ( - (), - ((),), - (("all-proxy-models",),), - (("*",),), - (("all-team-models",), ("all-proxy-models",)), - ((), (), (), (), ()), - (("*",), ("all-proxy-models",), ("all-team-models",), (), ()), -) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("scopes", RESTRICTED_SCOPES) -async def test_openai_websocket_rejects_model_restricted_identities(scopes): - websocket = _FakeWebSocket("/openai/v1/realtime", "model=gpt-4o-realtime-preview") - user_api_key_dict = UserAPIKeyAuth(token="hashed-fake", user_id="user-fake", team_id="team-fake") - - served = await _serve(websocket, "v1/realtime", user_api_key_dict, ENABLED, scopes) - - assert "model restrictions" in websocket.error_message() - assert websocket.closed == (1008, _OPENAI_WS_MODEL_RESTRICTED_REFUSAL.close_reason) - assert served.relay.calls == [] - assert served.allowlists.calls == [user_api_key_dict] - - -@pytest.mark.asyncio -@pytest.mark.parametrize("scopes", RESTRICTED_SCOPES) -async def test_openai_websocket_disabled_refusal_skips_allowlist_lookups(scopes): - allowlists = _FakeModelAllowlists(scopes) - - refusal = await _openai_websocket_refusal(UserAPIKeyAuth(), MappingProxyType({}), allowlists) - - assert refusal is _OPENAI_WS_DISABLED_REFUSAL - assert allowlists.calls == [] - - -@pytest.mark.asyncio -@pytest.mark.parametrize("scopes", UNRESTRICTED_SCOPES) -async def test_openai_websocket_allows_unrestricted_identities(scopes): +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") with patch(GET_CREDENTIALS, return_value="sk-provider"): - served = await _serve(websocket, "v1/responses", UserAPIKeyAuth(), ENABLED, scopes) + served = 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 assert websocket.sent == [] assert websocket.closed is None @pytest.mark.asyncio -async def test_proxy_model_allowlists_reads_the_token_scopes_without_a_database(): - token: Final = UserAPIKeyAuth(models=[], team_id="team-fake", team_models=["gpt-4o"]) - with patch("litellm.proxy.proxy_server.prisma_client", None): - scopes = await _proxy_model_allowlists()(token) +async def test_proxy_frame_model_gate_allows_a_model_the_key_owns(): + gate: Final = _proxy_frame_model_gate() + token: Final = UserAPIKeyAuth(models=["gpt-realtime-2.1", "gpt-5.4-mini"]) - assert tuple(tuple(scope) for scope in scopes) == ((), ("gpt-4o",)) - assert _has_model_restrictions(scopes) + with patch("litellm.proxy.proxy_server.prisma_client", None): # test-quality-ok: the gate needs no database + assert await gate("gpt-realtime-2.1", token) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("requested_model", ["gpt-5.4", "gpt-realtime-2.1-mini"]) +async def test_proxy_frame_model_gate_refuses_a_model_the_key_lacks(requested_model): + gate: Final = _proxy_frame_model_gate() + token: Final = UserAPIKeyAuth(models=["gpt-realtime-2.1", "gpt-5.4-mini"]) + + with patch("litellm.proxy.proxy_server.prisma_client", None): # test-quality-ok: the gate needs no database + refusal = await gate(requested_model, token) + + assert refusal is not None + assert requested_model in refusal + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "granted_models, refusal_names", + [(["gpt-5.4-mini"], "gpt-5.4"), ([], "no model grants")], + ids=["model_outside_the_grant", "no_grants"], +) +async def test_proxy_frame_model_gate_refuses_a_model_the_agent_was_not_granted(granted_models, refusal_names): + from litellm.types.agents import AgentResponse + + gate: Final = _proxy_frame_model_gate() + token: Final = UserAPIKeyAuth(models=[], agent_id="agent-fake") + token.managed_agent_policy = AgentResponse( + agent_id="agent-fake", + agent_name="Agent", + agent_card_params={}, + object_permission={"object_permission_id": "grant-fake", "models": granted_models}, + ) + + with patch("litellm.proxy.proxy_server.prisma_client", None): # test-quality-ok: the gate needs no database + refusal = await gate("gpt-5.4", token) + + assert refusal is not None + assert refusal_names in refusal + + +GET_USER_OBJECT: Final = "litellm.proxy.auth.auth_checks.get_user_object" + + +def _personal_user(models: list[str]) -> LiteLLM_UserTable: + return LiteLLM_UserTable(user_id="user-fake", models=models, max_budget=None, spend=0.0) + + +@contextmanager +def _personal_user_allowlist(models: list[str]) -> Iterator[AsyncMock]: + """Stand up the database-backed owner lookup the allowlist check runs, and report every call it made.""" + lookup: Final = AsyncMock(return_value=_personal_user(models)) + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), # test-quality-ok: read at call time + patch(GET_USER_OBJECT, lookup), + ): + yield lookup + + +@pytest.mark.asyncio +async def test_proxy_frame_model_gate_refuses_a_model_the_person_behind_the_key_lacks(): + """ + A key with no team of its own carries its owner's model allowlist, the same + one every HTTP route applies, so a frame naming a model that person cannot + use is refused even when the key itself names no models at all. + """ + gate: Final = _proxy_frame_model_gate() + token: Final = UserAPIKeyAuth(models=[], user_id="user-fake") + + with _personal_user_allowlist(["gpt-5.4-mini"]): + refusal = await gate("gpt-5.4", token) + + assert refusal is not None + assert "gpt-5.4" in refusal + + +@pytest.mark.asyncio +async def test_proxy_frame_model_gate_allows_a_model_the_person_behind_the_key_owns(): + gate: Final = _proxy_frame_model_gate() + token: Final = UserAPIKeyAuth(models=[], user_id="user-fake") + + with _personal_user_allowlist(["gpt-5.4-mini", "gpt-realtime-2.1"]): + assert await gate("gpt-realtime-2.1", token) is None + + +@pytest.mark.asyncio +async def test_proxy_frame_model_gate_leaves_a_team_key_to_the_teams_allowlist(): + """ + A key that belongs to a team is governed by the team's models, not by the + personal allowlist of whoever created it, so the owner's own restrictions + must not close a session the team is entitled to. + """ + gate: Final = _proxy_frame_model_gate() + token: Final = UserAPIKeyAuth(models=[], user_id="user-fake", team_id="team-fake") + + with _personal_user_allowlist(["gpt-5.4-mini"]) as owner_lookup: + assert await gate("gpt-5.4", token) is None + + assert owner_lookup.await_args_list == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("unrestricted_models", [[], ["all-proxy-models"], ["*"]]) +async def test_proxy_frame_model_gate_allows_every_model_for_unrestricted_keys(unrestricted_models): + gate: Final = _proxy_frame_model_gate() + token: Final = UserAPIKeyAuth(models=unrestricted_models) + + with patch("litellm.proxy.proxy_server.prisma_client", None): # test-quality-ok: the gate needs no database + assert await gate("gpt-realtime-2.1", token) is None + + +ENFORCE_GROUP_BUDGET: Final = ( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.enforce_model_access_group_budget" +) + + +@pytest.mark.asyncio +async def test_proxy_frame_model_gate_refuses_a_model_whose_access_group_is_out_of_budget(): + """ + A model that sits in a budgeted access group is refused on every other route once + that group has spent its budget, so a frame naming it closes the session instead + of quietly running on an empty pool. The key is allowed the model outright here, + so nothing but the group's budget can produce the refusal. + """ + exhausted: Final = litellm.BudgetExceededError( + current_cost=2.0, + max_budget=1.0, + message="Budget has been exceeded! Model access group=tier-a Current cost: 2.0, Max budget: 1.0", + entity_type=Litellm_EntityType.MODEL_ACCESS_GROUP.value, + entity_id="tier-a", + ) + gate: Final = _proxy_frame_model_gate() + token: Final = UserAPIKeyAuth(models=["gpt-5.4-mini"]) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", None), # test-quality-ok: the model check needs no database + patch(ENFORCE_GROUP_BUDGET, AsyncMock(side_effect=exhausted)), + ): + refusal = await gate("gpt-5.4-mini", token) + + assert refusal is not None + assert "Budget has been exceeded" in refusal + assert "Model access group=tier-a" in refusal + + +@pytest.mark.asyncio +async def test_proxy_frame_model_gate_refuses_out_loud_when_it_cannot_reach_the_budgets(): + """ + A gate that cannot answer has to close the session with a reason the caller can read. + Letting the failure escape drops the socket with no frame at all, which the caller + cannot tell apart from the network going away. + """ + gate: Final = _proxy_frame_model_gate() + token: Final = UserAPIKeyAuth(models=["gpt-5.4-mini"]) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", None), # test-quality-ok: the model check needs no database + patch(ENFORCE_GROUP_BUDGET, AsyncMock(side_effect=ConnectionError("connection pool is exhausted"))), + ): + refusal = await gate("gpt-5.4-mini", token) + + assert refusal == _OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL + assert "connection pool" not in refusal From 61bf7b1ae92e04b374a7ae2006ca4b996ea9e0f6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 22:15:27 -0700 Subject: [PATCH 02/15] fix(proxy): check the model in every OpenAI websocket passthrough frame and price the session --- litellm/proxy/auth/auth_checks.py | 182 +++++++++------ .../llm_passthrough_endpoints.py | 109 ++++----- ...i_websocket_passthrough_logging_handler.py | 207 ++++++++++++++++++ .../pass_through_endpoints.py | 151 ++++++++++++- .../pass_through_endpoints/success_handler.py | 70 ++++-- 5 files changed, 570 insertions(+), 149 deletions(-) create mode 100644 litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 8887d86570a..828b5b514b7 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -5291,79 +5291,6 @@ async def _granted_model_lists( ) -async def _user_object_or_none( - valid_token: UserAPIKeyAuth, - prisma_client: PrismaClient, - user_api_key_cache: UserApiKeyCache, - proxy_logging_obj: ProxyLogging, -) -> LiteLLM_UserTable | None: - try: - return await get_user_object( - user_id=valid_token.user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - user_id_upsert=False, - proxy_logging_obj=proxy_logging_obj, - ) - except UserNotFoundError: - return None - - -async def enforced_model_allowlists( - valid_token: UserAPIKeyAuth, - prisma_client: PrismaClient | None, - user_api_key_cache: UserApiKeyCache, - proxy_logging_obj: ProxyLogging, -) -> tuple[Sequence[str], ...]: - """One model allowlist per level that ``common_checks`` enforces on a request from this identity.""" - key_models: Final = _resolve_key_models_for_auth_check(valid_token=valid_token) - if prisma_client is None: - return (key_models, tuple(valid_token.team_models or ())) - team_object: Final = ( - None - if valid_token.team_id is None - else await get_team_object( - team_id=valid_token.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - ) - user_object: Final = ( - None - if team_object is not None - else await _user_object_or_none( - valid_token=valid_token, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - ) - project_object: Final = ( - None - if valid_token.project_id is None - else await get_project_object( - project_id=valid_token.project_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - ) - return ( - key_models, - team_object.models if team_object is not None else (), - await _team_member_granted_models( - valid_token=valid_token, - team_object=team_object, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ), - user_object.models if user_object is not None else (), - project_object.models if project_object is not None else (), - ) - - async def collect_matched_model_access_groups( model: str | Sequence[str] | None, valid_token: UserAPIKeyAuth | None, @@ -5677,6 +5604,115 @@ async def can_key_call_resolved_model( ) +async def can_personal_user_call_model( + model: str, + valid_token: UserAPIKeyAuth, + llm_router: Router | None, +) -> None: + """Apply the personal user's model allowlist, the check ``common_checks`` runs for a key with no team.""" + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + if valid_token.team_id is not None or valid_token.user_id is None or prisma_client is None: + return + try: + user_object: Final = await get_user_object( + user_id=valid_token.user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + parent_otel_span=valid_token.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except UserNotFoundError: + return + if user_object is None: + return + await can_user_call_model( + model=model, + llm_router=llm_router, + user_object=user_object, + key_model_aliases=key_model_aliases_for_auth_check(valid_token), + ) + + +async def _matching_team_object( + valid_token: UserAPIKeyAuth, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> LiteLLM_TeamTableCachedObj | None: + if valid_token.team_id is None: + return None + try: + return await get_team_object( + team_id=valid_token.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=valid_token.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception as e: # noqa: BLE001 # fail-safe: a team we cannot read grants no group, it never breaks auth + verbose_proxy_logger.debug("model access group attribution could not read the team: %s", e) + return None + + +async def enforce_model_access_group_budget( + model: str, + valid_token: UserAPIKeyAuth, + llm_router: Router | None, +) -> None: + """Attribute this model to the budgeted access groups that authorize it and refuse an exhausted one. + + The pair ``common_checks`` runs for every HTTP request, for callers that reach a model only + through a group: the stamp is what the spend writer charges afterwards, and the check is what + stops a group whose budget is already gone. + + Raises: + BudgetExceededError if a group that authorized this model is over its max budget. + """ + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + if prisma_client is None: + return + matched: Final = await stamp_matched_model_access_groups( + model=model, + valid_token=valid_token, + team_object=await _matching_team_object( + valid_token=valid_token, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ), + project_object=( + None + if valid_token.project_id is None + else await get_project_object( + project_id=valid_token.project_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + ), + llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + await _model_access_group_max_budget_check( + matched_model_access_groups=matched, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + + def can_org_access_model( model: str, org_object: LiteLLM_OrganizationTable | None, diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 16fbd1de376..1d1f88aaa18 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -15,7 +15,7 @@ import os import posixpath import re import sys -from collections.abc import AsyncGenerator, Callable, Mapping, Sequence +from collections.abc import AsyncGenerator, Callable, Mapping from dataclasses import dataclass from functools import partial from types import MappingProxyType @@ -50,6 +50,7 @@ from litellm.constants import ( from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException from litellm.integrations.custom_guardrail import guardrail_request_data_with_streaming, is_guardrail_intervention from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.llms.anthropic.common_utils import AnthropicModelInfo, merge_anthropic_beta_headers from litellm.llms.azure.passthrough.transformation import ( foreign_azure_deployment, @@ -72,7 +73,12 @@ from litellm.llms.oss_decision import OssDecisionProvider, oss_connection, valid from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.passthrough.main import AsyncPassthroughStreamingResponse, is_provider_error from litellm.proxy._types import * -from litellm.proxy.auth.auth_checks import enforced_model_allowlists +from litellm.proxy.auth.auth_checks import ( + can_agent_call_model, + can_key_call_resolved_model, + can_personal_user_call_model, + enforce_model_access_group_budget, +) from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy module exports @@ -110,6 +116,7 @@ from litellm.proxy.pass_through_endpoints.injected_stream_usage import ( ) from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( HttpPassThroughEndpointHelpers, + WebsocketFrameModelGate, create_pass_through_route, create_websocket_passthrough_route, websocket_passthrough_request, @@ -173,7 +180,7 @@ def is_passthrough_request_using_router_model(request_body: dict, llm_router: li """ try: model: Final = request_body.get("model") - return is_known_model(model, llm_router) + return isinstance(model, str) and is_known_model(model, llm_router) except Exception: return False @@ -203,9 +210,8 @@ def is_passthrough_request_streaming(request_body: object) -> bool: A JSON body need not be an object, so a list or scalar can reach here; it carries no streaming flag. """ - if not isinstance(request_body, dict): - return False - return bool(request_body.get("stream", False)) + body: Final = as_str_mapping(request_body) + return body is not None and bool(body.get("stream", False)) def _optional_str(value: object) -> str | None: @@ -3261,19 +3267,6 @@ def _join_url_paths(base_url: httpx.URL, path: str, custom_llm_provider: litellm return joined_path_str -_OPENAI_WS_ALL_MODEL_ACCESS: Final = frozenset( - { - SpecialModelNames.all_proxy_models.value, - SpecialModelNames.all_team_models.value, - "*", - } -) - - -def _has_model_restrictions(model_allowlists: tuple[Sequence[str], ...]) -> bool: - return any(str(model) not in _OPENAI_WS_ALL_MODEL_ACCESS for allowlist in model_allowlists for model in allowlist) - - @dataclass(frozen=True, slots=True) class _OpenAIWebsocketRefusal: close_reason: str @@ -3290,6 +3283,11 @@ class _OpenAIWebsocketErrorFrame(TypedDict): error: ReadOnly[_OpenAIWebsocketErrorDetail] +_OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL: Final = ( + "This gateway could not authorize the model named in this frame, so the session was closed. Try again, " + "and ask a proxy admin to check the gateway logs if it keeps happening." +) + _OPENAI_WS_DISABLED_REFUSAL: Final = _OpenAIWebsocketRefusal( close_reason="OpenAI websocket passthrough is disabled", message=( @@ -3298,14 +3296,6 @@ _OPENAI_WS_DISABLED_REFUSAL: Final = _OpenAIWebsocketRefusal( ), ) -_OPENAI_WS_MODEL_RESTRICTED_REFUSAL: Final = _OpenAIWebsocketRefusal( - close_reason="Keys with model restrictions cannot use OpenAI websocket passthrough", - message=( - "Keys with model restrictions cannot use OpenAI websocket passthrough, because this route " - "relays frames to the provider without reading which model they ask for." - ), -) - def _is_openai_websocket_passthrough_enabled(general_settings: Mapping[str, object]) -> bool: setting: Final = general_settings.get("enable_openai_websocket_passthrough") @@ -3314,19 +3304,9 @@ def _is_openai_websocket_passthrough_enabled(general_settings: Mapping[str, obje return setting is True -class _OpenAIWebsocketModelAllowlists(Protocol): - async def __call__(self, valid_token: UserAPIKeyAuth, /) -> tuple[Sequence[str], ...]: ... - - -async def _openai_websocket_refusal( - user_api_key_dict: UserAPIKeyAuth, - general_settings: Mapping[str, object], - model_allowlists: _OpenAIWebsocketModelAllowlists, -) -> _OpenAIWebsocketRefusal | None: +def _openai_websocket_refusal(general_settings: Mapping[str, object]) -> _OpenAIWebsocketRefusal | None: if not _is_openai_websocket_passthrough_enabled(general_settings): return _OPENAI_WS_DISABLED_REFUSAL - if _has_model_restrictions(await model_allowlists(user_api_key_dict)): - return _OPENAI_WS_MODEL_RESTRICTED_REFUSAL return None @@ -3341,6 +3321,7 @@ class _WebsocketRelay(Protocol): forward_headers: bool, endpoint: str, accept_websocket: bool, + client_frame_model_gate: WebsocketFrameModelGate | None = None, ) -> None: ... @@ -3348,18 +3329,42 @@ def _websocket_relay() -> _WebsocketRelay: return websocket_passthrough_request -def _proxy_model_allowlists() -> _OpenAIWebsocketModelAllowlists: - from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache +def _proxy_frame_model_gate() -> WebsocketFrameModelGate: + from litellm.proxy.proxy_server import ( + llm_model_list, + llm_router, + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) - async def resolve(valid_token: UserAPIKeyAuth, /) -> tuple[Sequence[str], ...]: - return await enforced_model_allowlists( - valid_token=valid_token, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) + async def gate(model: str, valid_token: UserAPIKeyAuth, /) -> str | None: + try: + await can_key_call_resolved_model( + model=model, llm_model_list=llm_model_list, valid_token=valid_token, llm_router=llm_router + ) + await can_agent_call_model( + model=model, + valid_token=valid_token, + llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + await can_personal_user_call_model(model=model, valid_token=valid_token, llm_router=llm_router) + await enforce_model_access_group_budget(model=model, valid_token=valid_token, llm_router=llm_router) + except ProxyException as denial: + return denial.message + except HTTPException as denial: + return str(denial.detail) + except litellm.BudgetExceededError as exhausted: + return exhausted.message + except Exception: # noqa: BLE001 # fail closed: a gate that cannot answer refuses, and says so in band + verbose_proxy_logger.exception("OpenAI websocket passthrough: the frame model gate failed on %s", model) + return _OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL + return None - return resolve + return gate def _negotiated_websocket_subprotocol(websocket: WebSocket) -> str | None: @@ -3379,12 +3384,12 @@ async def openai_websocket_proxy_route( user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth_websocket)], general_settings: Annotated[Mapping[str, object], Depends(_proxy_general_settings)], relay: Annotated[_WebsocketRelay, Depends(_websocket_relay)], - model_allowlists: Annotated[_OpenAIWebsocketModelAllowlists, Depends(_proxy_model_allowlists)], + frame_model_gate: Annotated[WebsocketFrameModelGate, Depends(_proxy_frame_model_gate)], ) -> None: """WebSocket passthrough for OpenAI prefixes (realtime / responses.connect).""" negotiated_subprotocol: Final = _negotiated_websocket_subprotocol(websocket) - refusal: Final = await _openai_websocket_refusal(user_api_key_dict, general_settings, model_allowlists) + refusal: Final = _openai_websocket_refusal(general_settings) if refusal is not None: await websocket.accept(subprotocol=negotiated_subprotocol) error_frame: Final[_OpenAIWebsocketErrorFrame] = { @@ -3438,6 +3443,7 @@ async def openai_websocket_proxy_route( forward_headers=False, endpoint=websocket.url.path, accept_websocket=False, + client_frame_model_gate=frame_model_gate, ) @@ -3683,7 +3689,8 @@ async def _tinyfish_json_object_field_names(request: Request) -> frozenset[str] parsed: Final = _DECODED_JSON.validate_python(json.loads(raw_body)) except (json.JSONDecodeError, UnicodeDecodeError): return None - return frozenset(parsed) if isinstance(parsed, dict) else None + fields: Final = as_str_mapping(parsed) + return None if fields is None else frozenset(fields) def _tinyfish_route_timeout() -> float | None: 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 new file mode 100644 index 00000000000..300aed08923 --- /dev/null +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py @@ -0,0 +1,207 @@ +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from datetime import datetime +from types import MappingProxyType +from typing import Final, Literal, TypeAlias, cast # noqa: TID251 # cost helpers type raw frames as a TypedDict list + +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.cost_calculator import RealtimeAPITokenUsageProcessor, handle_realtime_stream_cost_calculation +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.proxy._types import PassThroughEndpointLoggingTypedDict +from litellm.responses.utils import ResponseAPILoggingUtils +from litellm.types.llms.openai import OpenAIRealtimeStreamList, ResponseAPIUsage +from litellm.types.utils import LlmProviders, ModelResponse, Usage + +WebsocketMessages: TypeAlias = Sequence[Mapping[str, object]] + +OPENAI_WEBSOCKET_PROVIDER: Final = LlmProviders.OPENAI.value + +_SESSION_CREATED: Final = "session.created" +_RESPONSE_CREATED: Final = "response.created" +_REALTIME_RESPONSE_DONE: Final = "response.done" +_RESPONSES_COMPLETED: Final = "response.completed" +_MODEL_NAMING_EVENTS: Final = frozenset( + {_SESSION_CREATED, _RESPONSE_CREATED, _REALTIME_RESPONSE_DONE, _RESPONSES_COMPLETED} +) +_RESPONSE_COST_HEADER: Final = "llm_provider-x-litellm-response-cost" + + +class _NamedModel(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + model: str | None = None + + +class _ModelNamingEvent(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + session: _NamedModel | None = None + response: _NamedModel | None = None + + +class _CompletedResponse(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + model: str | None = None + usage: ResponseAPIUsage + + +class _ResponsesCompletedEvent(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + type: Literal["response.completed"] + response: _CompletedResponse + + +_MODEL_NAMING_EVENT: Final = TypeAdapter(_ModelNamingEvent) +_RESPONSES_COMPLETED_EVENT: Final = TypeAdapter(_ResponsesCompletedEvent) + + +@dataclass(frozen=True, slots=True) +class _Billing: + usage: Usage + cost: float + + +@dataclass(frozen=True, slots=True) +class _CompletedTurn: + model: str + usage: Usage + + +def _event_model(event: Mapping[str, object]) -> str | None: + if event.get("type") not in _MODEL_NAMING_EVENTS: + return None + try: + naming: Final = _MODEL_NAMING_EVENT.validate_python(event) + except ValidationError: + return None + return next( + ( + container.model + for container in (naming.session, naming.response) + if container is not None and container.model + ), + None, + ) + + +def observed_websocket_model(messages: WebsocketMessages) -> str | None: + return next((model for model in map(_event_model, messages) if model is not None), None) + + +def _realtime_billing(messages: WebsocketMessages, model: str, logging_obj: LiteLLMLoggingObj) -> _Billing | None: + if not any(event.get("type") == _REALTIME_RESPONSE_DONE for event in messages): + return None + results: Final = cast( # cast-ok: the realtime cost helpers read raw provider frames through their TypedDict view + OpenAIRealtimeStreamList, tuple(event for event in messages if "type" in event) + ) + usage: Final = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results(results) + cost: Final = handle_realtime_stream_cost_calculation( + results=results, + combined_usage_object=usage, + custom_llm_provider=OPENAI_WEBSOCKET_PROVIDER, + litellm_model_name=model, + litellm_logging_obj=logging_obj, + ) + return _Billing(usage=usage, cost=cost) + + +def _responses_turn(event: Mapping[str, object], session_model: str) -> _CompletedTurn | None: + if event.get("type") != _RESPONSES_COMPLETED: + return None + try: + completed: Final = _RESPONSES_COMPLETED_EVENT.validate_python(event) + except ValidationError: + verbose_proxy_logger.warning( + "OpenAI websocket passthrough: a response.completed frame carried no readable usage, so it is not billed" + ) + return None + return _CompletedTurn( + model=completed.response.model or session_model, + usage=ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( # pyright: ignore[reportPrivateUsage] # the one shared Responses usage converter + completed.response.usage + ), + ) + + +def _responses_billing(messages: WebsocketMessages, model: str) -> _Billing | None: + turns: Final = tuple(turn for turn in (_responses_turn(event, model) for event in messages) if turn is not None) + if not turns: + return None + usage: Final = RealtimeAPITokenUsageProcessor.combine_usage_objects([turn.usage for turn in turns]) + return _Billing(usage=usage, cost=sum(_responses_cost(turn.usage, turn.model) for turn in turns)) + + +def _responses_cost(usage: Usage, model: str) -> float: + try: + return litellm.completion_cost( + completion_response=ModelResponse(model=model, usage=usage), + model=model, + custom_llm_provider=OPENAI_WEBSOCKET_PROVIDER, + ) + except Exception: # noqa: BLE001 # completion_cost raises a bare Exception for an unpriced model, and the turn must still log + verbose_proxy_logger.warning( + "OpenAI websocket passthrough: could not price a Responses turn on %s, logging it at zero", model + ) + return 0.0 + + +def _priced_result(model: str, billing: _Billing, start_time: datetime) -> ModelResponse: + result: Final = ModelResponse( + id=f"openai-websocket-{start_time.timestamp()}", + object="chat.completion", + created=int(start_time.timestamp()), + model=model, + usage=billing.usage, + ) + hidden: Final = {"additional_headers": MappingProxyType({_RESPONSE_COST_HEADER: billing.cost})} + result._hidden_params = hidden # pyright: ignore[reportPrivateUsage] # the one settled-cost channel + return result + + +class OpenAIWebsocketPassthroughLoggingHandler: + def openai_websocket_passthrough_handler( + self, + websocket_messages: WebsocketMessages, + logging_obj: LiteLLMLoggingObj, + url_route: str, + start_time: datetime, + kwargs: Mapping[str, object], + ) -> PassThroughEndpointLoggingTypedDict: + model: Final = observed_websocket_model(websocket_messages) + if model is None: + verbose_proxy_logger.debug("OpenAI websocket passthrough (%s): no model named, skipping cost", url_route) + unpriced: Final[PassThroughEndpointLoggingTypedDict] = {"result": None, "kwargs": {**kwargs}} + return unpriced + + logging_obj.model = model # rebind-ok: the spend logger reads model and cost off the shared logging object + logging_obj.model_call_details.update( + MappingProxyType({"model": model, "custom_llm_provider": OPENAI_WEBSOCKET_PROVIDER}) + ) + + billing: Final = _realtime_billing(websocket_messages, model, logging_obj) or _responses_billing( + websocket_messages, model + ) + if billing is None: + verbose_proxy_logger.debug("OpenAI websocket passthrough (%s): no usage reported, skipping cost", url_route) + unbilled: Final[PassThroughEndpointLoggingTypedDict] = { + "result": None, + "kwargs": {**kwargs, "model": model, "custom_llm_provider": OPENAI_WEBSOCKET_PROVIDER}, + } + return unbilled + + logging_obj.model_call_details.update(MappingProxyType({"response_cost": billing.cost})) + priced: Final[PassThroughEndpointLoggingTypedDict] = { + "result": _priced_result(model, billing, start_time), + "kwargs": { + **kwargs, + "model": model, + "custom_llm_provider": OPENAI_WEBSOCKET_PROVIDER, + "response_cost": billing.cost, + }, + } + return priced diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index e1e8864678e..895533044c3 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -5,12 +5,13 @@ import json import posixpath import traceback from base64 import b64encode -from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterable, Mapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterable, Mapping, MutableMapping, Sequence from dataclasses import dataclass from datetime import datetime +from functools import reduce from itertools import chain, count, groupby from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, TypedDict, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict, cast from urllib.parse import urlencode, urlparse import httpx @@ -23,6 +24,7 @@ from fastapi import ( Response, UploadFile, WebSocket, + WebSocketDisconnect, status, ) from fastapi.responses import StreamingResponse @@ -30,7 +32,7 @@ from pydantic import TypeAdapter from starlette.datastructures import UploadFile as StarletteUploadFile from starlette.routing import BaseRoute, Route from starlette.websockets import WebSocketState -from websockets.asyncio.client import connect +from websockets.asyncio.client import ClientConnection, connect from websockets.exceptions import ( ConnectionClosedError, ConnectionClosedOK, @@ -55,6 +57,7 @@ from litellm.integrations.custom_guardrail import ( ) from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import ( + as_str_mapping, bind_budget_reservation_to_callbacks, get_metadata_variable_name_from_kwargs, get_or_create_metadata_bucket, @@ -64,6 +67,7 @@ from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.litellm_logging import get_masked_values from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.litellm_core_utils.realtime_errors import realtime_error_event from litellm.litellm_core_utils.redact_messages import should_redact_message_logging from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.llms.anthropic.common_utils import is_anthropic_messages_url @@ -2454,6 +2458,106 @@ def _with_trace_context(headers: Mapping[str, str], parent_span: object) -> dict return inject_trace_context(headers, parent_span=parent_span) +class WebsocketFrameModelGate(Protocol): + async def __call__(self, model: str, valid_token: UserAPIKeyAuth, /) -> str | None: ... + + +@dataclass(frozen=True, slots=True) +class _ForwardClientFrame: + pass + + +@dataclass(frozen=True, slots=True) +class _RefuseClientFrame: + message: str + + +_PLAIN_MODEL_KEY: Final = '"model"' +_JSON_STRING_ESCAPE: Final = "\\" +_CLIENT_FRAME_MODEL_PLACEMENTS: Final = ( + ("model",), + ("response", "model"), + ("session", "model"), + ("session", "input_audio_transcription", "model"), + ("session", "audio", "input", "transcription", "model"), +) + + +def _may_name_a_model(text_data: str) -> bool: + return _PLAIN_MODEL_KEY in text_data or _JSON_STRING_ESCAPE in text_data + + +def _nested_frame(frame: Mapping[str, object] | None, key: str) -> Mapping[str, object] | None: + return None if frame is None else as_str_mapping(frame.get(key)) + + +def _placed_model(frame: Mapping[str, object], placement: Sequence[str]) -> str | None: + *parents, leaf = placement + container: Final = reduce(_nested_frame, parents, frame) + if container is None: + return None + value: Final = container.get(leaf) + return value if isinstance(value, str) else None + + +def _client_frame_models(frame_data: str | bytes) -> tuple[str, ...]: + if isinstance(frame_data, str) and not _may_name_a_model(frame_data): + return () + frame: Final = _json_object_frame(frame_data) + if frame is None: + return () + placed: Final = (_placed_model(frame, placement) for placement in _CLIENT_FRAME_MODEL_PLACEMENTS) + return tuple(dict.fromkeys(model for model in placed if model is not None)) + + +async def _client_frame_verdict( + frame_data: str | bytes, gate: WebsocketFrameModelGate | None, valid_token: UserAPIKeyAuth +) -> _ForwardClientFrame | _RefuseClientFrame: + if gate is None: + return _ForwardClientFrame() + for model in _client_frame_models(frame_data): + if (refusal := await gate(model, valid_token)) is not None: + return _RefuseClientFrame(message=refusal) + return _ForwardClientFrame() + + +async def _refuse_client_frame(websocket: WebSocket, upstream_ws: ClientConnection, message: str) -> None: + await websocket.send_text(realtime_error_event(message, error_type="invalid_request_error")) + await websocket.close(code=CloseCode.POLICY_VIOLATION, reason=_truncated_close_reason(message)) + await upstream_ws.close() + + +def _stamp_billed_access_groups(kwargs: Mapping[str, object], valid_token: UserAPIKeyAuth) -> None: + """Refresh the session's logging metadata with the access groups its frames were authorized through. + + The gate only learns them once a frame names a model, which is long after the relay stamped the + key's groups onto its metadata, so the spend writer would otherwise charge no group for the session. + """ + groups: Final = valid_token.matched_model_access_groups + litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) + metadata: Final = None if litellm_params is None else litellm_params.get("metadata") + if groups and isinstance(metadata, MutableMapping): + metadata[MODEL_ACCESS_GROUP_METADATA_KEY] = groups + + +async def _relay_client_frame( + websocket: WebSocket, + upstream_ws: ClientConnection, + checked_frame: str | bytes, + forwarded_frame: str | bytes, + gate: WebsocketFrameModelGate | None, + valid_token: UserAPIKeyAuth, +) -> bool: + """Forward one client frame upstream, or refuse it and report that the session is over.""" + match await _client_frame_verdict(checked_frame, gate, valid_token): + case _RefuseClientFrame(message=refusal_message): + await _refuse_client_frame(websocket, upstream_ws, refusal_message) + return False + case _ForwardClientFrame(): + await upstream_ws.send(forwarded_frame) + return True + + async def websocket_passthrough_request( websocket: WebSocket, target: str, @@ -2464,6 +2568,7 @@ async def websocket_passthrough_request( cost_per_request: float | None = None, accept_websocket: bool = True, setup_model_rewriter: Callable[[str], str] | None = None, + client_frame_model_gate: WebsocketFrameModelGate | None = None, ): """ WebSocket passthrough request handler. @@ -2673,9 +2778,25 @@ async def websocket_passthrough_request( ) # Not a JSON message or doesn't contain setup data - await upstream_ws.send(_rewrite_vertex_live_setup_model(text_data, setup_model_rewriter)) + if not await _relay_client_frame( + websocket, + upstream_ws, + text_data, + _rewrite_vertex_live_setup_model(text_data, setup_model_rewriter), + client_frame_model_gate, + user_api_key_dict, + ): + break elif bytes_data is not None: - await upstream_ws.send(bytes_data) + if not await _relay_client_frame( + websocket, + upstream_ws, + bytes_data, + bytes_data, + client_frame_model_gate, + user_api_key_dict, + ): + break except asyncio.CancelledError: raise except Exception: @@ -2702,11 +2823,7 @@ async def websocket_passthrough_request( is_vertex_live: Final = bool(endpoint and "/vertex_ai/live" in endpoint) json_frame_ordinal: Final = count() - async def relay_upstream_frame(upstream_message: str | bytes) -> None: - if isinstance(upstream_message, bytes): - await websocket.send_bytes(upstream_message) - else: - await websocket.send_text(upstream_message) + def record_upstream_frame(upstream_message: str | bytes) -> None: message_data: Final = _json_object_frame(upstream_message) if message_data is None: return @@ -2715,6 +2832,13 @@ async def websocket_passthrough_request( return websocket_messages.append(message_data) + async def relay_upstream_frame(upstream_message: str | bytes) -> None: + record_upstream_frame(upstream_message) + if isinstance(upstream_message, bytes): + await websocket.send_bytes(upstream_message) + else: + await websocket.send_text(upstream_message) + async def forward_upstream_to_client() -> Close | None: try: while True: @@ -2722,6 +2846,11 @@ async def websocket_passthrough_request( except (ConnectionClosedOK, ConnectionClosedError) as e: verbose_proxy_logger.debug("Upstream WebSocket connection closed: %s", e) return e.rcvd + except WebSocketDisconnect: + verbose_proxy_logger.debug( + "WebSocket passthrough (%s): client left while upstream was still sending", endpoint + ) + return None except asyncio.CancelledError: verbose_proxy_logger.debug("asyncio.CancelledError in forward_upstream_to_client") raise @@ -2767,6 +2896,8 @@ async def websocket_passthrough_request( passthrough_logging_payload["response_body"] = websocket_messages passthrough_logging_payload["end_time"] = end_time + _stamp_billed_access_groups(kwargs, user_api_key_dict) + # Remove logging_obj from kwargs to avoid duplicate keyword argument success_kwargs: Final = kwargs.copy() success_kwargs.pop("logging_obj", None) diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index 9c6c3111209..433d7501958 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/litellm/proxy/pass_through_endpoints/success_handler.py @@ -1,4 +1,5 @@ import json +from collections.abc import Mapping, Sequence from datetime import datetime from types import MappingProxyType from typing import Any, Final @@ -8,7 +9,7 @@ import httpx from litellm.constants import AZURE_SPEECH_CUSTOM_LLM_PROVIDER from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.proxy._types import PassThroughEndpointLoggingResultValues +from litellm.proxy._types import PassThroughEndpointLoggingResultValues, PassThroughEndpointLoggingTypedDict from litellm.types.passthrough_endpoints.pass_through_endpoints import ( PassthroughStandardLoggingPayload, ) @@ -64,6 +65,14 @@ def _safe_response_text(httpx_response: httpx.Response) -> str: return "" +def _websocket_frames( + response_body: Mapping[str, object] | Sequence[Mapping[str, object]] | None, +) -> tuple[Mapping[str, object], ...]: + return tuple( + message for message in (response_body if isinstance(response_body, list) else ()) if isinstance(message, dict) + ) + + class PassThroughEndpointLogging: def __init__( self, @@ -114,6 +123,7 @@ class PassThroughEndpointLogging: # Vertex AI Live API WebSocket self.TRACKED_VERTEX_AI_LIVE_ROUTES = ["/vertex_ai/live"] + self.TRACKED_OPENAI_WEBSOCKET_ROUTES: Final = ("/openai/", "/openai_passthrough/") @property def _log_dispatch(self) -> PassThroughLogDispatch: @@ -385,21 +395,18 @@ class PassThroughEndpointLogging: standard_logging_response_object = vertex_ai_live_handler_result["result"] kwargs = vertex_ai_live_handler_result["kwargs"] - elif DeepgramListenPassthroughLoggingHandler.is_deepgram_listen_route(url_route): - deepgram_handler_result: Final = ( - DeepgramListenPassthroughLoggingHandler().deepgram_listen_passthrough_handler( - websocket_messages=tuple( - message - for message in (response_body if isinstance(response_body, list) else ()) - if isinstance(message, dict) - ), - logging_obj=logging_obj, - upstream_url=str(httpx_response.request.url), - kwargs=kwargs, - ) + elif ( + websocket_relay_result := self._websocket_relay_logging_result( + url_route=url_route, + websocket_frames=_websocket_frames(response_body), + logging_obj=logging_obj, + httpx_response=httpx_response, + start_time=start_time, + kwargs=kwargs, ) - standard_logging_response_object = deepgram_handler_result["result"] # rebind-ok: elif-chain - kwargs = deepgram_handler_result["kwargs"] # rebind-ok: elif-chain contract + ) is not None: + standard_logging_response_object = websocket_relay_result["result"] # rebind-ok: elif-chain + kwargs = websocket_relay_result["kwargs"] # rebind-ok: elif-chain contract elif FalAIPassthroughLoggingHandler.is_fal_ai_route(url_route, custom_llm_provider): fal_ai_handler_result: Final = FalAIPassthroughLoggingHandler().fal_ai_passthrough_handler( response_body=response_body if isinstance(response_body, dict) else MappingProxyType({}), @@ -415,6 +422,36 @@ class PassThroughEndpointLogging: return_dict["kwargs"] = kwargs return return_dict + def _websocket_relay_logging_result( + self, + url_route: str, + websocket_frames: tuple[Mapping[str, object], ...], + logging_obj: LiteLLMLoggingObj, + httpx_response: httpx.Response, + start_time: datetime, + kwargs: Mapping[str, object], + ) -> PassThroughEndpointLoggingTypedDict | None: + if self.is_openai_websocket_route(url_route): + from .llm_provider_handlers.openai_websocket_passthrough_logging_handler import ( + OpenAIWebsocketPassthroughLoggingHandler, + ) + + return OpenAIWebsocketPassthroughLoggingHandler().openai_websocket_passthrough_handler( + websocket_messages=websocket_frames, + logging_obj=logging_obj, + url_route=url_route, + start_time=start_time, + kwargs=kwargs, + ) + if DeepgramListenPassthroughLoggingHandler.is_deepgram_listen_route(url_route): + return DeepgramListenPassthroughLoggingHandler().deepgram_listen_passthrough_handler( + websocket_messages=websocket_frames, + logging_obj=logging_obj, + upstream_url=str(httpx_response.request.url), + kwargs=kwargs, + ) + return None + async def pass_through_async_success_handler( self, httpx_response: httpx.Response, @@ -584,6 +621,9 @@ class PassThroughEndpointLogging: return True return False + def is_openai_websocket_route(self, url_route: str) -> bool: + return url_route.startswith(self.TRACKED_OPENAI_WEBSOCKET_ROUTES) + def is_cursor_route(self, url_route: str, custom_llm_provider: str | None = None): """Check if the URL route is a Cursor Cloud Agents API route.""" if custom_llm_provider == "cursor": From 28ce19c4e99a85d57ce0ff4d78daf5bbbd5bfc4e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 21 Oct 2026 01:36:03 -0700 Subject: [PATCH 03/15] fix(proxy): bill each OpenAI websocket turn on its own row and run full auth per frame Each finished turn (completed, incomplete, or failed) now logs its own spend row under the model and access groups that allowed it, so group budgets see spend as turns end instead of at close. Every model a frame names goes through the key's full auth, and a stored prompt reference with no model is refused for a key held to some models. A frame still in flight when the client leaves is recorded before the relay gives up on the send --- litellm/proxy/auth/auth_checks.py | 109 ----- .../llm_passthrough_endpoints.py | 63 ++- ...i_websocket_passthrough_logging_handler.py | 41 +- .../pass_through_endpoints.py | 443 +++++++++++------- .../spend_tracking/budget_reservation.py | 2 +- .../pass_through_endpoints.py | 12 +- .../auth/test_model_access_group_budgets.py | 63 --- ...i_websocket_passthrough_logging_handler.py | 51 +- .../test_pass_through_endpoints.py | 325 +++++++++++-- .../test_openai_ws_passthrough_routes.py | 270 +++++------ 10 files changed, 843 insertions(+), 536 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 828b5b514b7..c4b934ecefc 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -5604,115 +5604,6 @@ async def can_key_call_resolved_model( ) -async def can_personal_user_call_model( - model: str, - valid_token: UserAPIKeyAuth, - llm_router: Router | None, -) -> None: - """Apply the personal user's model allowlist, the check ``common_checks`` runs for a key with no team.""" - from litellm.proxy.proxy_server import ( - prisma_client, - proxy_logging_obj, - user_api_key_cache, - ) - - if valid_token.team_id is not None or valid_token.user_id is None or prisma_client is None: - return - try: - user_object: Final = await get_user_object( - user_id=valid_token.user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - user_id_upsert=False, - parent_otel_span=valid_token.parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) - except UserNotFoundError: - return - if user_object is None: - return - await can_user_call_model( - model=model, - llm_router=llm_router, - user_object=user_object, - key_model_aliases=key_model_aliases_for_auth_check(valid_token), - ) - - -async def _matching_team_object( - valid_token: UserAPIKeyAuth, - prisma_client: PrismaClient, - user_api_key_cache: UserApiKeyCache, - proxy_logging_obj: ProxyLogging, -) -> LiteLLM_TeamTableCachedObj | None: - if valid_token.team_id is None: - return None - try: - return await get_team_object( - team_id=valid_token.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=valid_token.parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) - except Exception as e: # noqa: BLE001 # fail-safe: a team we cannot read grants no group, it never breaks auth - verbose_proxy_logger.debug("model access group attribution could not read the team: %s", e) - return None - - -async def enforce_model_access_group_budget( - model: str, - valid_token: UserAPIKeyAuth, - llm_router: Router | None, -) -> None: - """Attribute this model to the budgeted access groups that authorize it and refuse an exhausted one. - - The pair ``common_checks`` runs for every HTTP request, for callers that reach a model only - through a group: the stamp is what the spend writer charges afterwards, and the check is what - stops a group whose budget is already gone. - - Raises: - BudgetExceededError if a group that authorized this model is over its max budget. - """ - from litellm.proxy.proxy_server import ( - prisma_client, - proxy_logging_obj, - user_api_key_cache, - ) - - if prisma_client is None: - return - matched: Final = await stamp_matched_model_access_groups( - model=model, - valid_token=valid_token, - team_object=await _matching_team_object( - valid_token=valid_token, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ), - project_object=( - None - if valid_token.project_id is None - else await get_project_object( - project_id=valid_token.project_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - ), - llm_router=llm_router, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - await _model_access_group_max_budget_check( - matched_model_access_groups=matched, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - ) - - def can_org_access_model( model: str, org_object: LiteLLM_OrganizationTable | None, diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 1d1f88aaa18..8944ce10046 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -15,11 +15,11 @@ import os import posixpath import re import sys -from collections.abc import AsyncGenerator, Callable, Mapping +from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping from dataclasses import dataclass from functools import partial from types import MappingProxyType -from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, cast +from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeAlias, cast import httpx import openai @@ -73,12 +73,7 @@ from litellm.llms.oss_decision import OssDecisionProvider, oss_connection, valid from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.passthrough.main import AsyncPassthroughStreamingResponse, is_provider_error from litellm.proxy._types import * -from litellm.proxy.auth.auth_checks import ( - can_agent_call_model, - can_key_call_resolved_model, - can_personal_user_call_model, - enforce_model_access_group_budget, -) +from litellm.proxy.auth.auth_checks import can_agent_call_model, can_key_call_resolved_model from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy module exports @@ -86,6 +81,7 @@ from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy modul get_bearer_token, is_no_auth_dev_mode, user_api_key_auth, + user_api_key_auth_accepted_websocket_for_model, user_api_key_auth_websocket, user_api_key_auth_websocket_for_model, ) @@ -114,9 +110,16 @@ from litellm.proxy.pass_through_endpoints.injected_stream_usage import ( stream_usage_relay, without_injected_usage_frames, ) +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_websocket_passthrough_logging_handler import ( + openai_websocket_frame_kind, +) from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( HttpPassThroughEndpointHelpers, + WebsocketFrameModelAllowed, WebsocketFrameModelGate, + WebsocketFrameModelRefused, + WebsocketFrameModelVerdict, + WebsocketUpstreamFrameClassifier, create_pass_through_route, create_websocket_passthrough_route, websocket_passthrough_request, @@ -3322,6 +3325,7 @@ class _WebsocketRelay(Protocol): endpoint: str, accept_websocket: bool, client_frame_model_gate: WebsocketFrameModelGate | None = None, + upstream_frame_kind: WebsocketUpstreamFrameClassifier | None = None, ) -> None: ... @@ -3329,7 +3333,20 @@ def _websocket_relay() -> _WebsocketRelay: return websocket_passthrough_request -def _proxy_frame_model_gate() -> WebsocketFrameModelGate: +def _proxy_error_type(denial: ProxyException) -> str: + return denial.type.value if isinstance(denial.type, ProxyErrorTypes) else denial.type + + +FrameModelAuthorizer: TypeAlias = Callable[[WebSocket, str], Awaitable[UserAPIKeyAuth]] + + +def _frame_model_authorizer() -> FrameModelAuthorizer: + return user_api_key_auth_accepted_websocket_for_model + + +def _proxy_frame_model_gate( + authorize_frame_model: Annotated[FrameModelAuthorizer, Depends(_frame_model_authorizer)], +) -> WebsocketFrameModelGate: from litellm.proxy.proxy_server import ( llm_model_list, llm_router, @@ -3337,32 +3354,39 @@ def _proxy_frame_model_gate() -> WebsocketFrameModelGate: proxy_logging_obj, user_api_key_cache, ) + from litellm.proxy.spend_tracking.budget_reservation import release_or_invalidate_budget_reservation - async def gate(model: str, valid_token: UserAPIKeyAuth, /) -> str | None: + async def authorize(websocket: WebSocket, model: str) -> WebsocketFrameModelAllowed: + frame_auth: Final = await authorize_frame_model(websocket, model) + access_groups: Final = tuple(frame_auth.matched_model_access_groups or ()) try: await can_key_call_resolved_model( - model=model, llm_model_list=llm_model_list, valid_token=valid_token, llm_router=llm_router + model=model, llm_model_list=llm_model_list, valid_token=frame_auth, llm_router=llm_router ) await can_agent_call_model( model=model, - valid_token=valid_token, + valid_token=frame_auth, llm_router=llm_router, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) - await can_personal_user_call_model(model=model, valid_token=valid_token, llm_router=llm_router) - await enforce_model_access_group_budget(model=model, valid_token=valid_token, llm_router=llm_router) + finally: + await release_or_invalidate_budget_reservation(budget_reservation=frame_auth.budget_reservation) + return WebsocketFrameModelAllowed(access_groups=access_groups) + + async def gate(websocket: WebSocket, model: str, /) -> WebsocketFrameModelVerdict: + try: + return await authorize(websocket, model) except ProxyException as denial: - return denial.message + return WebsocketFrameModelRefused(message=denial.message, error_type=_proxy_error_type(denial)) except HTTPException as denial: - return str(denial.detail) + return WebsocketFrameModelRefused(message=str(denial.detail), error_type="invalid_request_error") except litellm.BudgetExceededError as exhausted: - return exhausted.message + return WebsocketFrameModelRefused(message=exhausted.message, error_type="budget_exceeded") except Exception: # noqa: BLE001 # fail closed: a gate that cannot answer refuses, and says so in band verbose_proxy_logger.exception("OpenAI websocket passthrough: the frame model gate failed on %s", model) - return _OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL - return None + return WebsocketFrameModelRefused(message=_OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL, error_type="server_error") return gate @@ -3444,6 +3468,7 @@ async def openai_websocket_proxy_route( endpoint=websocket.url.path, accept_websocket=False, client_frame_model_gate=frame_model_gate, + upstream_frame_kind=openai_websocket_frame_kind, ) 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 300aed08923..c2f4e85005c 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 @@ -13,6 +13,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.proxy._types import PassThroughEndpointLoggingTypedDict from litellm.responses.utils import ResponseAPILoggingUtils from litellm.types.llms.openai import OpenAIRealtimeStreamList, ResponseAPIUsage +from litellm.types.passthrough_endpoints.pass_through_endpoints import WebsocketUpstreamFrameKind from litellm.types.utils import LlmProviders, ModelResponse, Usage WebsocketMessages: TypeAlias = Sequence[Mapping[str, object]] @@ -22,10 +23,9 @@ OPENAI_WEBSOCKET_PROVIDER: Final = LlmProviders.OPENAI.value _SESSION_CREATED: Final = "session.created" _RESPONSE_CREATED: Final = "response.created" _REALTIME_RESPONSE_DONE: Final = "response.done" -_RESPONSES_COMPLETED: Final = "response.completed" -_MODEL_NAMING_EVENTS: Final = frozenset( - {_SESSION_CREATED, _RESPONSE_CREATED, _REALTIME_RESPONSE_DONE, _RESPONSES_COMPLETED} -) +_RESPONSES_TURN_ENDS: Final = frozenset({"response.completed", "response.incomplete", "response.failed"}) +_TURN_ENDS: Final = _RESPONSES_TURN_ENDS | {_REALTIME_RESPONSE_DONE} +_MODEL_NAMING_EVENTS: Final = _TURN_ENDS | {_SESSION_CREATED, _RESPONSE_CREATED} _RESPONSE_COST_HEADER: Final = "llm_provider-x-litellm-response-cost" @@ -42,22 +42,22 @@ class _ModelNamingEvent(BaseModel): response: _NamedModel | None = None -class _CompletedResponse(BaseModel): +class _EndedResponse(BaseModel): model_config = ConfigDict(extra="ignore", frozen=True) model: str | None = None - usage: ResponseAPIUsage + usage: ResponseAPIUsage | None = None -class _ResponsesCompletedEvent(BaseModel): +class _ResponsesTurnEndEvent(BaseModel): model_config = ConfigDict(extra="ignore", frozen=True) - type: Literal["response.completed"] - response: _CompletedResponse + type: Literal["response.completed", "response.incomplete", "response.failed"] + response: _EndedResponse _MODEL_NAMING_EVENT: Final = TypeAdapter(_ModelNamingEvent) -_RESPONSES_COMPLETED_EVENT: Final = TypeAdapter(_ResponsesCompletedEvent) +_RESPONSES_TURN_END_EVENT: Final = TypeAdapter(_ResponsesTurnEndEvent) @dataclass(frozen=True, slots=True) @@ -93,6 +93,15 @@ def observed_websocket_model(messages: WebsocketMessages) -> str | None: return next((model for model in map(_event_model, messages) if model is not None), None) +def openai_websocket_frame_kind(frame: Mapping[str, object]) -> WebsocketUpstreamFrameKind: + frame_type: Final = frame.get("type") + if frame_type in _TURN_ENDS: + return "ends_turn" + if frame_type == _SESSION_CREATED: + return "opens_session" + return "relayed" + + def _realtime_billing(messages: WebsocketMessages, model: str, logging_obj: LiteLLMLoggingObj) -> _Billing | None: if not any(event.get("type") == _REALTIME_RESPONSE_DONE for event in messages): return None @@ -111,19 +120,21 @@ def _realtime_billing(messages: WebsocketMessages, model: str, logging_obj: Lite def _responses_turn(event: Mapping[str, object], session_model: str) -> _CompletedTurn | None: - if event.get("type") != _RESPONSES_COMPLETED: + if event.get("type") not in _RESPONSES_TURN_ENDS: return None try: - completed: Final = _RESPONSES_COMPLETED_EVENT.validate_python(event) + ended: Final = _RESPONSES_TURN_END_EVENT.validate_python(event) except ValidationError: verbose_proxy_logger.warning( - "OpenAI websocket passthrough: a response.completed frame carried no readable usage, so it is not billed" + "OpenAI websocket passthrough: a %s frame carried no readable usage, so it is not billed", event.get("type") ) return None + if ended.response.usage is None: + return None return _CompletedTurn( - model=completed.response.model or session_model, + model=ended.response.model or session_model, usage=ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( # pyright: ignore[reportPrivateUsage] # the one shared Responses usage converter - completed.response.usage + ended.response.usage ), ) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 895533044c3..201ecee5b44 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -3,6 +3,7 @@ import asyncio import copy import json import posixpath +import re import traceback from base64 import b64encode from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterable, Mapping, MutableMapping, Sequence @@ -11,7 +12,7 @@ from datetime import datetime from functools import reduce from itertools import chain, count, groupby from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias, TypedDict, cast from urllib.parse import urlencode, urlparse import httpx @@ -142,6 +143,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, EndpointType, PassthroughStandardLoggingPayload, + WebsocketUpstreamFrameKind, ) from litellm.types.utils import TRUSTED_CALLBACK_VARS_FIELD, Usage @@ -2458,86 +2460,139 @@ def _with_trace_context(headers: Mapping[str, str], parent_span: object) -> dict return inject_trace_context(headers, parent_span=parent_span) +@dataclass(frozen=True, slots=True) +class WebsocketFrameModelAllowed: + access_groups: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class WebsocketFrameModelRefused: + message: str + error_type: str + + +WebsocketFrameModelVerdict: TypeAlias = WebsocketFrameModelAllowed | WebsocketFrameModelRefused + + class WebsocketFrameModelGate(Protocol): - async def __call__(self, model: str, valid_token: UserAPIKeyAuth, /) -> str | None: ... + async def __call__(self, websocket: WebSocket, model: str, /) -> WebsocketFrameModelVerdict: ... + + +WebsocketUpstreamFrameClassifier: TypeAlias = Callable[[Mapping[str, object]], WebsocketUpstreamFrameKind] @dataclass(frozen=True, slots=True) class _ForwardClientFrame: - pass + turn_access_groups: tuple[str, ...] | None @dataclass(frozen=True, slots=True) class _RefuseClientFrame: - message: str + refusal: WebsocketFrameModelRefused -_PLAIN_MODEL_KEY: Final = '"model"' -_JSON_STRING_ESCAPE: Final = "\\" -_CLIENT_FRAME_MODEL_PLACEMENTS: Final = ( - ("model",), - ("response", "model"), - ("session", "model"), +@dataclass(frozen=True, slots=True) +class _ClientFrameRequests: + turn_model: str | None + models: tuple[str, ...] + names_stored_prompt: bool + + +_NO_CLIENT_FRAME_REQUESTS: Final = _ClientFrameRequests(turn_model=None, models=(), names_stored_prompt=False) +_FRAME_TEXT_WORTH_PARSING: Final = ('"model"', '"prompt"', "\\") +_TURN_MODEL_PLACEMENTS: Final = (("model",), ("response", "model"), ("session", "model")) +_TRANSCRIPTION_MODEL_PLACEMENTS: Final = ( ("session", "input_audio_transcription", "model"), ("session", "audio", "input", "transcription", "model"), ) +_STORED_PROMPT_PLACEMENTS: Final = (("prompt",), ("response", "prompt"), ("session", "prompt")) +STORED_PROMPT_PROBE_MODEL: Final = "litellm-unverifiable-stored-prompt-model" +_STORED_PROMPT_REFUSAL: Final = ( + "This key is limited to specific models, and the gateway cannot check which model a stored prompt runs, " + "so frames that reference a stored prompt are refused for this key." +) +_FIRST_SENTENCE: Final = re.compile(r".*?[.!?](?=\s|$)") def _may_name_a_model(text_data: str) -> bool: - return _PLAIN_MODEL_KEY in text_data or _JSON_STRING_ESCAPE in text_data + return any(marker in text_data for marker in _FRAME_TEXT_WORTH_PARSING) def _nested_frame(frame: Mapping[str, object] | None, key: str) -> Mapping[str, object] | None: return None if frame is None else as_str_mapping(frame.get(key)) -def _placed_model(frame: Mapping[str, object], placement: Sequence[str]) -> str | None: +def _placed_value(frame: Mapping[str, object], placement: Sequence[str]) -> object: *parents, leaf = placement container: Final = reduce(_nested_frame, parents, frame) - if container is None: - return None - value: Final = container.get(leaf) - return value if isinstance(value, str) else None + return None if container is None else container.get(leaf) -def _client_frame_models(frame_data: str | bytes) -> tuple[str, ...]: +def _placed_models(frame: Mapping[str, object], placements: Iterable[Sequence[str]]) -> tuple[str, ...]: + placed: Final = (_placed_value(frame, placement) for placement in placements) + return tuple(value for value in placed if isinstance(value, str)) + + +def _client_frame_requests(frame_data: str | bytes) -> _ClientFrameRequests: if isinstance(frame_data, str) and not _may_name_a_model(frame_data): - return () + return _NO_CLIENT_FRAME_REQUESTS frame: Final = _json_object_frame(frame_data) if frame is None: - return () - placed: Final = (_placed_model(frame, placement) for placement in _CLIENT_FRAME_MODEL_PLACEMENTS) - return tuple(dict.fromkeys(model for model in placed if model is not None)) + return _NO_CLIENT_FRAME_REQUESTS + turn_models: Final = _placed_models(frame, _TURN_MODEL_PLACEMENTS) + prompts: Final = (_placed_value(frame, placement) for placement in _STORED_PROMPT_PLACEMENTS) + return _ClientFrameRequests( + turn_model=next(iter(turn_models), None), + models=tuple(dict.fromkeys((*turn_models, *_placed_models(frame, _TRANSCRIPTION_MODEL_PLACEMENTS)))), + names_stored_prompt=any(isinstance(prompt, Mapping) for prompt in prompts), + ) + + +def _stored_prompt_refusal(probe: WebsocketFrameModelRefused) -> WebsocketFrameModelRefused: + if probe.error_type.endswith("_model_access_denied"): + return WebsocketFrameModelRefused(message=_STORED_PROMPT_REFUSAL, error_type=probe.error_type) + return probe async def _client_frame_verdict( - frame_data: str | bytes, gate: WebsocketFrameModelGate | None, valid_token: UserAPIKeyAuth + websocket: WebSocket, frame_data: str | bytes, gate: WebsocketFrameModelGate | None ) -> _ForwardClientFrame | _RefuseClientFrame: if gate is None: - return _ForwardClientFrame() - for model in _client_frame_models(frame_data): - if (refusal := await gate(model, valid_token)) is not None: - return _RefuseClientFrame(message=refusal) - return _ForwardClientFrame() + return _ForwardClientFrame(turn_access_groups=None) + requests: Final = _client_frame_requests(frame_data) + verdicts: Final = tuple([await gate(websocket, model) for model in requests.models]) + refusal: Final = next((verdict for verdict in verdicts if isinstance(verdict, WebsocketFrameModelRefused)), None) + if refusal is not None: + return _RefuseClientFrame(refusal=refusal) + if requests.names_stored_prompt and isinstance( + probe := await gate(websocket, STORED_PROMPT_PROBE_MODEL), WebsocketFrameModelRefused + ): + return _RefuseClientFrame(refusal=_stored_prompt_refusal(probe)) + allowed: Final = tuple(verdict for verdict in verdicts if isinstance(verdict, WebsocketFrameModelAllowed)) + if requests.turn_model is None or not allowed: + return _ForwardClientFrame(turn_access_groups=None) + return _ForwardClientFrame(turn_access_groups=allowed[0].access_groups) -async def _refuse_client_frame(websocket: WebSocket, upstream_ws: ClientConnection, message: str) -> None: - await websocket.send_text(realtime_error_event(message, error_type="invalid_request_error")) - await websocket.close(code=CloseCode.POLICY_VIOLATION, reason=_truncated_close_reason(message)) +def _refusal_close_reason(message: str) -> str: + first_sentence: Final = _FIRST_SENTENCE.match(message) + return _truncated_close_reason(message if first_sentence is None else first_sentence.group(0)) + + +async def _refuse_client_frame( + websocket: WebSocket, upstream_ws: ClientConnection, refusal: WebsocketFrameModelRefused +) -> None: + await websocket.send_text(realtime_error_event(refusal.message, error_type=refusal.error_type)) + await websocket.close(code=CloseCode.POLICY_VIOLATION, reason=_refusal_close_reason(refusal.message)) await upstream_ws.close() -def _stamp_billed_access_groups(kwargs: Mapping[str, object], valid_token: UserAPIKeyAuth) -> None: - """Refresh the session's logging metadata with the access groups its frames were authorized through. - - The gate only learns them once a frame names a model, which is long after the relay stamped the - key's groups onto its metadata, so the spend writer would otherwise charge no group for the session. - """ - groups: Final = valid_token.matched_model_access_groups +def _stamp_billed_access_groups(kwargs: Mapping[str, object], access_groups: tuple[str, ...]) -> None: + """Charge a turn to the access groups that authorized its model, and to none when no group did.""" litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) metadata: Final = None if litellm_params is None else litellm_params.get("metadata") - if groups and isinstance(metadata, MutableMapping): - metadata[MODEL_ACCESS_GROUP_METADATA_KEY] = groups + if isinstance(metadata, MutableMapping): + metadata[MODEL_ACCESS_GROUP_METADATA_KEY] = list(access_groups) async def _relay_client_frame( @@ -2546,16 +2601,35 @@ async def _relay_client_frame( checked_frame: str | bytes, forwarded_frame: str | bytes, gate: WebsocketFrameModelGate | None, - valid_token: UserAPIKeyAuth, -) -> bool: - """Forward one client frame upstream, or refuse it and report that the session is over.""" - match await _client_frame_verdict(checked_frame, gate, valid_token): - case _RefuseClientFrame(message=refusal_message): - await _refuse_client_frame(websocket, upstream_ws, refusal_message) - return False +) -> _ForwardClientFrame | _RefuseClientFrame: + """Forward one client frame upstream, or refuse it and end the session.""" + verdict: Final = await _client_frame_verdict(websocket, checked_frame, gate) + match verdict: + case _RefuseClientFrame(refusal=refusal): + await _refuse_client_frame(websocket, upstream_ws, refusal) case _ForwardClientFrame(): await upstream_ws.send(forwarded_frame) - return True + return verdict + + +_EXPLICIT_SUCCESS_KWARGS: Final = frozenset({"logging_obj", "passthrough_logging_payload", "custom_llm_provider"}) + + +class _WebsocketTurnLedger: + """The groups that authorized the turn in flight, when it started, and how many turns the socket has logged.""" + + def __init__(self, access_groups: tuple[str, ...], started: datetime) -> None: + self.access_groups = access_groups + self.started = started + self.logged_turns = 0 + + def authorize(self, access_groups: tuple[str, ...] | None) -> None: + if access_groups is not None: + self.access_groups = access_groups + + def close_turn(self, end_time: datetime) -> None: + self.logged_turns += 1 + self.started = end_time async def websocket_passthrough_request( @@ -2569,6 +2643,7 @@ async def websocket_passthrough_request( accept_websocket: bool = True, setup_model_rewriter: Callable[[str], str] | None = None, client_frame_model_gate: WebsocketFrameModelGate | None = None, + upstream_frame_kind: WebsocketUpstreamFrameClassifier | None = None, ): """ WebSocket passthrough request handler. @@ -2582,6 +2657,9 @@ async def websocket_passthrough_request( endpoint: The endpoint path (for logging purposes) cost_per_request: Optional field - cost per request to the target endpoint setup_model_rewriter: Optional rewrite of the setup frame's model before it reaches the upstream + client_frame_model_gate: Optional per-frame check of every model a client frame names + upstream_frame_kind: Optional sort of upstream frames into the one that ends a billable turn and the ones + that open the session, so each turn logs its own spend and keeps the session's provider-named model """ from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy.proxy_server import proxy_config, proxy_logging_obj @@ -2617,26 +2695,31 @@ async def websocket_passthrough_request( proxy_config=proxy_config, route_description="websocket_passthrough", ) - logging_obj: Final = Logging( - model="unknown", - messages=[{"role": "user", "content": "WebSocket connection"}], - stream=True, # WebSockets are inherently streaming - call_type="pass_through_endpoint", - start_time=start_time, - litellm_call_id=litellm_call_id, - function_id="websocket_passthrough", - dynamic_success_callbacks=team_callbacks.success_callbacks, - dynamic_failure_callbacks=team_callbacks.failure_callbacks, - kwargs=team_callbacks.logging_kwargs, - ) - # Create passthrough logging payload - passthrough_logging_payload: Final = PassthroughStandardLoggingPayload( - url=target, - request_body={}, # WebSocket doesn't have a traditional request body - request_method="WEBSOCKET", - cost_per_request=cost_per_request, - ) + def new_call_logging(call_start: datetime, call_id: str) -> LiteLLMLoggingObj: + return Logging( + model="unknown", + messages=[{"role": "user", "content": "WebSocket connection"}], + stream=True, + call_type="pass_through_endpoint", + start_time=call_start, + litellm_call_id=call_id, + function_id="websocket_passthrough", + dynamic_success_callbacks=team_callbacks.success_callbacks, + dynamic_failure_callbacks=team_callbacks.failure_callbacks, + kwargs=team_callbacks.logging_kwargs, + ) + + def new_call_payload() -> PassthroughStandardLoggingPayload: + return PassthroughStandardLoggingPayload( + url=target, + request_body={}, + request_method="WEBSOCKET", + cost_per_request=cost_per_request, + ) + + logging_obj: Final = new_call_logging(start_time, litellm_call_id) + passthrough_logging_payload: Final = new_call_payload() # Create a dummy request object for WebSocket connections to maintain compatibility # with the existing _init_kwargs_for_pass_through_endpoint function @@ -2658,35 +2741,117 @@ async def websocket_passthrough_request( ) # Initialize kwargs for logging using the same pattern as HTTP passthrough - kwargs: Final = HttpPassThroughEndpointHelpers.init_kwargs_for_pass_through_endpoint( - user_api_key_dict=user_api_key_dict, - _parsed_body={}, # WebSocket doesn't have a traditional request body - passthrough_logging_payload=passthrough_logging_payload, - litellm_call_id=litellm_call_id, - request=dummy_request, - logging_obj=logging_obj, + def init_call_kwargs( + token: UserAPIKeyAuth, + call_payload: PassthroughStandardLoggingPayload, + call_id: str, + call_logging: LiteLLMLoggingObj, + ) -> dict[str, object]: # mutable-ok: the Vertex relay stamps the model from setup frames into these kwargs + return HttpPassThroughEndpointHelpers.init_kwargs_for_pass_through_endpoint( + user_api_key_dict=token, + _parsed_body={}, # WebSocket doesn't have a traditional request body + passthrough_logging_payload=call_payload, + litellm_call_id=call_id, + request=dummy_request, + logging_obj=call_logging, + ) + + kwargs: Final = init_call_kwargs(user_api_key_dict, passthrough_logging_payload, litellm_call_id, logging_obj) + + def start_call_logging(call_logging: LiteLLMLoggingObj, call_kwargs: Mapping[str, object]) -> None: + call_logging.update_environment_variables( + model="unknown", + user="unknown", + optional_params={}, + litellm_params=dict(as_str_mapping(call_kwargs.get("litellm_params")) or {}), + call_type="pass_through_endpoint", + ) + call_logging.pre_call( + input=[{"role": "user", "content": "WebSocket connection"}], + api_key="", + additional_args={ + "complete_input_dict": {}, + "api_base": target, + "headers": upstream_headers, + }, + ) + + start_call_logging(logging_obj, kwargs) + + class MockWebSocketRequest: + def __init__(self, target_url: str): + self.method = "WEBSOCKET" + self.url = target_url + + class MockWebSocketResponse: + def __init__(self, target_url: str): + self.status_code = 200 + self.text = "WebSocket connection successful" + self.headers: dict[str, str] = {} + self.request = MockWebSocketRequest(target_url) + + ledger: Final = _WebsocketTurnLedger( + access_groups=tuple(user_api_key_dict.matched_model_access_groups or ()), started=start_time ) - # Update logging environment variables - logging_obj.update_environment_variables( - model="unknown", - user="unknown", - optional_params={}, - litellm_params=dict(kwargs.get("litellm_params", {})), - call_type="pass_through_endpoint", - ) - logging_obj.model_call_details["litellm_call_id"] = litellm_call_id + def log_frames( + call_logging: LiteLLMLoggingObj, + call_payload: PassthroughStandardLoggingPayload, + call_kwargs: Mapping[str, object], + end_time: datetime, + ) -> None: + frames: Final = list(websocket_messages) + websocket_messages[:] = ( + () + if upstream_frame_kind is None + else (frame for frame in frames if upstream_frame_kind(frame) == "opens_session") + ) + logged_payload: Final[PassthroughStandardLoggingPayload] = { + **call_payload, + "response_body": frames, + "end_time": end_time, + } + if client_frame_model_gate is not None: + _stamp_billed_access_groups(call_kwargs, ledger.access_groups) + success_kwargs: Final = { + key: value for key, value in call_kwargs.items() if key not in _EXPLICIT_SUCCESS_KWARGS + } + custom_llm_provider: Final = call_kwargs.get("custom_llm_provider") + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( + async_coroutine=pass_through_endpoint_logging.pass_through_async_success_handler( + httpx_response=MockWebSocketResponse(target), + response_body=frames, + url_route=endpoint or "", + result="websocket_connection_successful", + start_time=ledger.started, + end_time=end_time, + logging_obj=call_logging, + cache_hit=False, + request_body={}, + passthrough_logging_payload=logged_payload, + custom_llm_provider=custom_llm_provider if isinstance(custom_llm_provider, str) else None, + **success_kwargs, + ) + ) + bind_budget_reservation_to_callbacks(call_logging.litellm_params) - # Pre-call logging - logging_obj.pre_call( - input=[{"role": "user", "content": "WebSocket connection"}], - api_key="", - additional_args={ - "complete_input_dict": {}, - "api_base": target, - "headers": upstream_headers, - }, - ) + def log_turn() -> None: + end_time: Final = datetime.now() + if ledger.logged_turns == 0: + log_frames(logging_obj, passthrough_logging_payload, kwargs, end_time) + else: + turn_call_id: Final = str(uuid.uuid4()) + turn_logging: Final = new_call_logging(ledger.started, turn_call_id) + turn_payload: Final = new_call_payload() + turn_kwargs: Final = init_call_kwargs( + user_api_key_dict.model_copy(update={"budget_reservation": None}), + turn_payload, + turn_call_id, + turn_logging, + ) + start_call_logging(turn_logging, turn_kwargs) + log_frames(turn_logging, turn_payload, turn_kwargs, end_time) + ledger.close_turn(end_time) ### CALL HOOKS ### - modify incoming data / reject request before calling the model websocket_hook_data: Final = guardrail_request_data_with_streaming(MappingProxyType({}), is_streaming=True) @@ -2778,25 +2943,27 @@ async def websocket_passthrough_request( ) # Not a JSON message or doesn't contain setup data - if not await _relay_client_frame( - websocket, - upstream_ws, - text_data, - _rewrite_vertex_live_setup_model(text_data, setup_model_rewriter), - client_frame_model_gate, - user_api_key_dict, + if isinstance( + relayed_text := await _relay_client_frame( + websocket, + upstream_ws, + text_data, + _rewrite_vertex_live_setup_model(text_data, setup_model_rewriter), + client_frame_model_gate, + ), + _RefuseClientFrame, ): break + ledger.authorize(relayed_text.turn_access_groups) elif bytes_data is not None: - if not await _relay_client_frame( - websocket, - upstream_ws, - bytes_data, - bytes_data, - client_frame_model_gate, - user_api_key_dict, + if isinstance( + relayed_bytes := await _relay_client_frame( + websocket, upstream_ws, bytes_data, bytes_data, client_frame_model_gate + ), + _RefuseClientFrame, ): break + ledger.authorize(relayed_bytes.turn_access_groups) except asyncio.CancelledError: raise except Exception: @@ -2831,9 +2998,13 @@ async def websocket_passthrough_request( _extract_vertex_live_model_from_setup_response(message_data) return websocket_messages.append(message_data) + if upstream_frame_kind is not None and upstream_frame_kind(message_data) == "ends_turn": + log_turn() async def relay_upstream_frame(upstream_message: str | bytes) -> None: record_upstream_frame(upstream_message) + if not _client_socket_is_open(websocket): + return if isinstance(upstream_message, bytes): await websocket.send_bytes(upstream_message) else: @@ -2890,64 +3061,8 @@ async def websocket_passthrough_request( reason=_truncated_close_reason(upstream_close.reason), ) - end_time: Final = datetime.now() - - # Update passthrough logging payload with response data - passthrough_logging_payload["response_body"] = websocket_messages - passthrough_logging_payload["end_time"] = end_time - - _stamp_billed_access_groups(kwargs, user_api_key_dict) - - # Remove logging_obj from kwargs to avoid duplicate keyword argument - success_kwargs: Final = kwargs.copy() - success_kwargs.pop("logging_obj", None) - - # # Add user authentication context for database logging - # if user_api_key_dict: - # success_kwargs.setdefault('litellm_params', {}) - # success_kwargs['litellm_params'].update({ - # 'proxy_server_request': { - # 'body': { - # 'user': user_api_key_dict.user_id, - # 'team_id': user_api_key_dict.team_id, - # 'end_user_id': user_api_key_dict.end_user_id, - # } - # } - # }) - # # Also add the user_api_key for direct access - # success_kwargs['user_api_key'] = user_api_key_dict.api_key - - # Create a dummy httpx.Response for WebSocket connections - class MockWebSocketResponse: - def __init__(self, target_url: str): - self.status_code = 200 - self.text = "WebSocket connection successful" - self.headers: dict[str, str] = {} - self.request = MockWebSocketRequest(target_url) - - class MockWebSocketRequest: - def __init__(self, target_url: str): - self.method = "WEBSOCKET" - self.url = target_url - - mock_response: Final = MockWebSocketResponse(target) - - # Use the same success handler as HTTP passthrough endpoints - GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( - async_coroutine=pass_through_endpoint_logging.pass_through_async_success_handler( - httpx_response=mock_response, - response_body=websocket_messages, - url_route=endpoint or "", - result="websocket_connection_successful", - start_time=start_time, - end_time=end_time, - logging_obj=logging_obj, - cache_hit=False, - request_body={}, - **success_kwargs, - ) - ) - bind_budget_reservation_to_callbacks(logging_obj.litellm_params) + if ledger.logged_turns == 0: + log_turn() # Call the proxy logging success hook if proxy_logging_obj: diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 803df9f83ca..3d3e090a013 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -422,7 +422,7 @@ async def invalidate_budget_reservation_counters( async def release_or_invalidate_budget_reservation( - budget_reservation: dict | None, # mutable-ok: stamps finalized on the caller's shared reservation dict + budget_reservation: dict[str, object] | None, # mutable-ok: stamps finalized on the caller's reservation ) -> None: """Reconcile a still-open reservation on a terminal path that settles no cost. diff --git a/litellm/types/passthrough_endpoints/pass_through_endpoints.py b/litellm/types/passthrough_endpoints/pass_through_endpoints.py index 0d99cd54a6b..b13211d8f37 100644 --- a/litellm/types/passthrough_endpoints/pass_through_endpoints.py +++ b/litellm/types/passthrough_endpoints/pass_through_endpoints.py @@ -1,8 +1,10 @@ +from collections.abc import Mapping, Sequence +from datetime import datetime from enum import Enum -from typing import Final +from typing import Final, Literal, TypeAlias from pydantic import Field -from typing_extensions import TypedDict +from typing_extensions import ReadOnly, TypedDict from litellm.types.llms.base import LiteLLMPydanticObjectBase @@ -27,6 +29,8 @@ 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"] + class BackgroundResponseSettlement(str, Enum): SETTLES = "settles" @@ -62,7 +66,7 @@ class PassthroughStandardLoggingPayload(TypedDict, total=False): """ The body of the request """ - response_body: dict | None # only tracked for non-streaming responses + response_body: dict | Sequence[Mapping[str, object]] | None # only tracked for non-streaming responses """ The body of the response """ @@ -74,6 +78,8 @@ class PassthroughStandardLoggingPayload(TypedDict, total=False): Optional field, we use this for cost tracking only if it's set. """ + end_time: ReadOnly[datetime] + class PassThroughGuardrailSettings(LiteLLMPydanticObjectBase): """ diff --git a/tests/unit/proxy/auth/test_model_access_group_budgets.py b/tests/unit/proxy/auth/test_model_access_group_budgets.py index 7b0c4d9996e..7506bd031d9 100644 --- a/tests/unit/proxy/auth/test_model_access_group_budgets.py +++ b/tests/unit/proxy/auth/test_model_access_group_budgets.py @@ -27,7 +27,6 @@ from litellm.proxy.auth.auth_checks import ( _model_access_group_max_budget_check, collect_matched_model_access_groups, common_checks, - enforce_model_access_group_budget, stamp_matched_model_access_groups, ) from litellm.proxy.common_utils.reset_budget_job import _model_access_group_counter_key @@ -552,65 +551,3 @@ async def test_common_checks_blocks_a_request_whose_group_is_over_budget(): async def test_free_model_routes_skip_the_model_access_group_budget_check(): """skip_budget_checks is how free models stay free; it has to cover this budget too.""" assert await _common_checks_with_over_budget_group(skip_budget_checks=True) is True - - -async def _enforce_for_a_route_without_a_request_body( - *rows: _MagBudgetRow, - key_models: list[str], - spend_by_counter_key: dict[str, float], -) -> UserAPIKeyAuth: - """Run the pair the websocket frame gate runs, which reaches the proxy for what it needs.""" - valid_token = UserAPIKeyAuth(api_key="hashed", models=key_models, user_id=USER_ID) - cache = await _cache() - read, _ = _spend_reader(spend_by_counter_key) - with ( - patch("litellm.proxy.proxy_server.prisma_client", _RecordingPrismaClient(*rows)), # test-quality-ok: read lazily inside the helper, which takes no client - patch("litellm.proxy.proxy_server.user_api_key_cache", cache), # test-quality-ok: same lazy import - patch("litellm.proxy.proxy_server.proxy_logging_obj", ProxyLogging(user_api_key_cache=cache)), # test-quality-ok: same lazy import - patch("litellm.proxy.proxy_server.get_current_spend", read), # test-quality-ok: get_current_spend is lazily imported inside the budget check and has no injection point - ): - await enforce_model_access_group_budget( - model="gpt-4o", - valid_token=valid_token, - llm_router=Router(model_list=MODEL_LIST), - ) - return valid_token - - -@pytest.mark.asyncio -async def test_caller_granted_the_model_only_through_an_exhausted_group_is_refused(): - """Routes with no request body, the websocket passthrough among them, run this instead of - ``common_checks``, so the group budget has to stop them the same way.""" - with pytest.raises(litellm.BudgetExceededError) as exc_info: - await _enforce_for_a_route_without_a_request_body( - _MagBudgetRow("tier-a", max_budget=1.0), - key_models=["tier-a"], - spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 1.5}, - ) - - assert exc_info.value.entity_id == "tier-a" - assert exc_info.value.current_cost == 1.5 - - -@pytest.mark.asyncio -async def test_group_that_authorized_the_call_is_stamped_for_the_spend_writer(): - """Without the stamp the call runs on the group's budget and never pays into it, so the pool - can never catch up with what it authorized.""" - valid_token = await _enforce_for_a_route_without_a_request_body( - _MagBudgetRow("tier-a", max_budget=10.0), - key_models=["tier-a"], - spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 1.0}, - ) - - assert valid_token.matched_model_access_groups == ["tier-a"] - - -@pytest.mark.asyncio -async def test_model_granted_directly_charges_no_group(): - valid_token = await _enforce_for_a_route_without_a_request_body( - _MagBudgetRow("tier-a", max_budget=1.0), - key_models=["gpt-4o"], - spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 1.5}, - ) - - assert valid_token.matched_model_access_groups is None 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 91a01f02d94..253634760fa 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 @@ -8,6 +8,7 @@ import pytest import litellm from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_websocket_passthrough_logging_handler import ( OpenAIWebsocketPassthroughLoggingHandler, + openai_websocket_frame_kind, ) REALTIME_MODEL: Final = "gpt-realtime-2.1" @@ -42,9 +43,11 @@ def _realtime_response_done(text_in: int, audio_in: int, text_out: int, audio_ou } -def _responses_completed(model: str, input_tokens: int, output_tokens: int) -> dict: +def _responses_completed( + model: str, input_tokens: int, output_tokens: int, event: str = "response.completed" +) -> dict: return { - "type": "response.completed", + "type": event, "response": { "id": "resp_fake", "model": model, @@ -288,3 +291,47 @@ def test_a_turn_the_caller_walked_out_of_still_names_the_model_it_ran(): assert handled["result"] is None assert handled["kwargs"] == {"model": REASONING_MODEL, "custom_llm_provider": "openai"} assert logging_obj.model_call_details["model"] == REASONING_MODEL + + +@pytest.mark.parametrize("event", ["response.incomplete", "response.failed"]) +def test_a_responses_turn_that_ended_early_is_billed_for_the_tokens_it_used(event): + """ + 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( + [_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"] + 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( + [{"type": "response.failed", "response": {"model": RESPONSES_MODEL, "status": "failed", "usage": None}}], + url_route="/openai_passthrough/v1/responses", + ) + + assert handled["result"] is None + assert handled["kwargs"] == {"model": RESPONSES_MODEL, "custom_llm_provider": "openai"} + + +@pytest.mark.parametrize( + ("frame_type", "kind"), + [ + ("response.completed", "ends_turn"), + ("response.incomplete", "ends_turn"), + ("response.failed", "ends_turn"), + ("response.done", "ends_turn"), + ("session.created", "opens_session"), + ("response.created", "relayed"), + ("response.output_text.delta", "relayed"), + ("error", "relayed"), + ], +) +def test_every_frame_that_settles_a_turn_closes_it_and_nothing_else_does(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 6f35ebbc161..e6a0b1be28c 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 @@ -42,9 +42,13 @@ from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import route from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( DEFAULT_PASS_THROUGH_REQUEST_TIMEOUT_SECONDS, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, + STORED_PROMPT_PROBE_MODEL, HttpPassThroughEndpointHelpers, InitPassThroughEndpointHelpers, SafeRouteAdder, + WebsocketFrameModelAllowed, + WebsocketFrameModelRefused, + WebsocketFrameModelVerdict, _registered_pass_through_routes, _truncate_upstream_error_body, _update_metadata_with_tags_in_header, @@ -58,6 +62,9 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( set_env_variables_in_header, websocket_passthrough_request, ) +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_websocket_passthrough_logging_handler import ( + openai_websocket_frame_kind, +) from litellm.proxy.pass_through_endpoints.success_handler import ( PassThroughEndpointLogging, ) @@ -7211,7 +7218,7 @@ def _patched_websocket_passthrough_environment(upstream_ws): mock_worker.ensure_initialized_and_enqueue = MagicMock( side_effect=lambda async_coroutine: async_coroutine.close() ) - yield + yield mock_proxy_logging async def _pending_receive(): @@ -7522,14 +7529,28 @@ SUCCESS_HANDLER: Final = ( ) +KEY_MODEL_REFUSAL: Final = "key not allowed to access model. This key can only access models=['gpt-5.4-mini']." +KEY_MODEL_CLOSE_REASON: Final = "key not allowed to access model." + + class RecordingFrameModelGate: - def __init__(self, refusals: Mapping[str, str] = MappingProxyType({})): + def __init__( + self, + refusals: Mapping[str, str] = MappingProxyType({}), + groups: Mapping[str, tuple[str, ...]] = MappingProxyType({}), + error_type: str = "key_model_access_denied", + ): self.refusals = refusals + self.groups = groups + self.error_type = error_type self.checked: list[str] = [] - async def __call__(self, model: str, valid_token: UserAPIKeyAuth, /) -> Optional[str]: + async def __call__(self, websocket: object, model: str, /) -> WebsocketFrameModelVerdict: self.checked.append(model) - return self.refusals.get(model) + refusal = 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, ())) def _received_frame(frame): @@ -7594,8 +7615,7 @@ 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(): - refusal = "key not allowed to access model. This key can only access models=['gpt-5.4-mini']." - gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": refusal})) + gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": KEY_MODEL_REFUSAL})) frame = json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}) relayed = await _relay_client_frames([frame], gate) @@ -7603,16 +7623,34 @@ async def test_websocket_passthrough_refuses_a_frame_naming_a_model_the_key_cann assert relayed.upstream.send.await_args_list == [] assert json.loads(relayed.client.send_text.await_args.args[0]) == { "type": "error", - "error": {"type": "invalid_request_error", "message": refusal}, + "error": {"type": "key_model_access_denied", "message": KEY_MODEL_REFUSAL}, } - assert relayed.client.close.await_args_list == [call(code=1008, reason=refusal)] + assert relayed.client.close.await_args_list == [call(code=1008, reason=KEY_MODEL_CLOSE_REASON)] relayed.upstream.close.assert_awaited() +@pytest.mark.asyncio +async def test_websocket_passthrough_closes_with_the_whole_first_sentence_of_the_refusal(): + """ + The close frame only has room for 123 bytes, so a refusal cut at that limit + 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") + + relayed = 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"] == { + "type": "budget_exceeded", + "message": refusal, + } + + @pytest.mark.asyncio async def test_websocket_passthrough_refuses_the_whole_frame_when_any_named_model_is_denied(): - refusal = "key not allowed to access model. This key can only access models=['gpt-realtime-2.1']." - gate = RecordingFrameModelGate(MappingProxyType({"gpt-4o-transcribe": refusal})) + gate = RecordingFrameModelGate(MappingProxyType({"gpt-4o-transcribe": KEY_MODEL_REFUSAL})) frame = json.dumps( { "type": "session.update", @@ -7624,13 +7662,12 @@ async def test_websocket_passthrough_refuses_the_whole_frame_when_any_named_mode assert gate.checked == ["gpt-realtime-2.1", "gpt-4o-transcribe"] assert relayed.upstream.send.await_args_list == [] - assert relayed.client.close.await_args_list == [call(code=1008, reason=refusal)] + assert relayed.client.close.await_args_list == [call(code=1008, reason=KEY_MODEL_CLOSE_REASON)] @pytest.mark.asyncio async def test_websocket_passthrough_stops_relaying_after_a_refused_frame(): - refusal = "key not allowed to access model. This key can only access models=['gpt-5.4-mini']." - gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": refusal})) + gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": KEY_MODEL_REFUSAL})) frames = [ json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}), json.dumps({"type": "response.create"}), @@ -7697,27 +7734,76 @@ 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. """ - refusal = "key not allowed to access model. This key can only access models=['gpt-5.4-mini']." - gate = RecordingFrameModelGate(MappingProxyType({expected_model: refusal})) + gate = RecordingFrameModelGate(MappingProxyType({expected_model: KEY_MODEL_REFUSAL})) relayed = await _relay_client_frames([frame], gate) assert gate.checked == [expected_model] assert relayed.upstream.send.await_args_list == [] - assert relayed.client.close.await_args_list == [call(code=1008, reason=refusal)] + assert relayed.client.close.await_args_list == [call(code=1008, reason=KEY_MODEL_CLOSE_REASON)] @pytest.mark.asyncio async def test_websocket_passthrough_checks_the_model_a_binary_frame_asks_for(): - refusal = "key not allowed to access model. This key can only access models=['gpt-5.4-mini']." - gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": refusal})) + gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": KEY_MODEL_REFUSAL})) frame = json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}).encode() relayed = await _relay_client_frames([frame], gate) assert gate.checked == ["gpt-realtime-2.1"] assert relayed.upstream.send.await_args_list == [] - assert relayed.client.close.await_args_list == [call(code=1008, reason=refusal)] + assert relayed.client.close.await_args_list == [call(code=1008, reason=KEY_MODEL_CLOSE_REASON)] + + +STORED_PROMPT_FRAMES: Final = ( + {"type": "response.create", "prompt": {"id": "pmpt_fake", "version": "2"}}, + {"type": "response.create", "response": {"prompt": {"id": "pmpt_fake"}}}, + {"type": "session.update", "session": {"prompt": {"id": "pmpt_fake"}}}, +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("frame", STORED_PROMPT_FRAMES) +async def test_websocket_passthrough_refuses_a_stored_prompt_for_a_key_limited_to_some_models(frame): + """ + A stored prompt carries its own model on the provider side, so a key held to + 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})) + + relayed = await _relay_client_frames([json.dumps(frame)], gate) + + assert gate.checked == [STORED_PROMPT_PROBE_MODEL] + assert relayed.upstream.send.await_args_list == [] + assert "stored prompt" in json.loads(relayed.client.send_text.await_args.args[0])["error"]["message"] + assert relayed.client.close.await_args.kwargs["code"] == 1008 + + +@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() + + relayed = 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()] + + +@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( + MappingProxyType({STORED_PROMPT_PROBE_MODEL: refusal}), error_type="budget_exceeded" + ) + + relayed = 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", + "message": refusal, + } @pytest.mark.asyncio @@ -7791,6 +7877,196 @@ async def test_websocket_passthrough_still_bills_a_session_the_client_walked_out ] +class ScriptedUpstreamWebSocket: + """Answers each client frame with its scripted reply, and flushes what was still in flight when it is closed""" + + def __init__(self, *opening: str, replies: Sequence[Sequence[str]] = (), in_flight: Sequence[str] = ()): + self._frames: asyncio.Queue[str | None] = asyncio.Queue() + for frame in opening: + self._frames.put_nowait(frame) + self._replies = iter(replies) + self._in_flight = in_flight + self._closed = False + self.send = AsyncMock(side_effect=self._reply) + self.close = AsyncMock(side_effect=self._flush_and_close) + + async def _reply(self, frame: str | bytes) -> None: + for reply in next(self._replies, ()): + self._frames.put_nowait(reply) + + async def _flush_and_close(self, *args: object, **kwargs: object) -> None: + if self._closed: + return + self._closed = True + for frame in self._in_flight: + self._frames.put_nowait(frame) + self._frames.put_nowait(None) + while not self._frames.empty(): + await asyncio.sleep(0) + + async def recv(self, decode: bool | None = None) -> str: + from websockets.exceptions import ConnectionClosedOK + from websockets.frames import Close + + frame = await self._frames.get() + if frame is None: + raise ConnectionClosedOK(rcvd=Close(1000, ""), sent=Close(1000, ""), rcvd_then_sent=True) + return frame + + +def _turn_taking_client(frames: Sequence[str], answered_by: str) -> MagicMock: + """A client that sends its next frame only once the answer to the last one reached it, and that Starlette + refuses to send to once it is closed""" + from starlette.websockets import WebSocketState + + delivered: list[str] = [] + script = iter(frames) + receives = iter(range(len(frames) + 1)) + + async def receive(): + answers_due = next(receives) + while sum(answered_by in text for text in delivered) < answers_due: + await asyncio.sleep(0) + frame = next(script, None) + return {"type": "websocket.disconnect"} if frame is None else _received_frame(frame) + + websocket = _client_websocket(receive) + + async def send_text(text: str) -> None: + if websocket.application_state == WebSocketState.DISCONNECTED: + raise RuntimeError('Cannot call "send" once a close message has been sent.') + delivered.append(text) + + websocket.send_text = AsyncMock(side_effect=send_text) + return websocket + + +async def _relay_turns(client, upstream_ws, gate, endpoint="/openai_passthrough/v1/responses"): + with ( + _patched_websocket_passthrough_environment(upstream_ws) as proxy_logging, + patch(SUCCESS_HANDLER, new_callable=AsyncMock) as success_handler, + ): + await asyncio.wait_for( + websocket_passthrough_request( + websocket=client, + target="wss://api.openai.com/v1/responses", + custom_headers={"Authorization": "Bearer sk-provider"}, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + forward_headers=False, + endpoint=endpoint, + accept_websocket=False, + client_frame_model_gate=gate, + upstream_frame_kind=openai_websocket_frame_kind, + ), + timeout=10, + ) + return SimpleNamespace(success_handler=success_handler, proxy_logging=proxy_logging) + + +def _responses_turn(model: str) -> list[str]: + return [ + json.dumps({"type": "response.created", "response": {"model": model}}), + json.dumps( + { + "type": "response.completed", + "response": {"model": model, "usage": {"input_tokens": 9, "output_tokens": 3, "total_tokens": 12}}, + } + ), + ] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_logs_each_turn_on_its_own_row_charged_to_the_groups_that_allowed_it(): + """ + One socket can run a model from one access group and then a model from + another. Each turn has to reach spend as it finishes, under its own call, + 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( + groups=MappingProxyType({"gpt-5.4-mini": ("tier-a",), "gpt-5.4": ("tier-b",)}), + ) + client = _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")]) + + relayed = await _relay_turns(client, upstream_ws, gate) + + rows = [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"], + ] + assert [row["litellm_params"]["metadata"][MODEL_ACCESS_GROUP_METADATA_KEY] for row in rows] == [ + ["tier-a"], + ["tier-b"], + ] + assert [row["litellm_params"]["metadata"]["user_api_key"] for row in rows] == ["hashed-key", "hashed-key"] + assert rows[0]["litellm_call_id"] != rows[1]["litellm_call_id"] + assert rows[1]["start_time"] == rows[0]["end_time"] + relayed.proxy_logging.post_call_failure_hook.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_websocket_passthrough_carries_the_session_the_provider_opened_into_every_later_turn(): + """ + A realtime provider names the session's model once, when it opens the + 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( + [json.dumps({"type": "response.create"}), json.dumps({"type": "response.create"})], + answered_by='"response.done"', + ) + upstream_ws = ScriptedUpstreamWebSocket( + REALTIME_SESSION_CREATED, replies=[[REALTIME_RESPONSE_DONE], [REALTIME_RESPONSE_DONE]] + ) + + 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"], ["session.created", "response.done"]] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_still_bills_the_answer_in_flight_when_a_later_frame_is_refused(): + """ + A refused frame closes the client while the provider may still be sending the + answer to an earlier, allowed frame. That answer cannot go to a closed client, + 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( + [ + json.dumps({"type": "response.create", "model": "gpt-5.4-mini"}), + json.dumps({"type": "response.create", "model": "gpt-5.4"}), + ], + answered_by='"response.created"', + ) + created, completed = _responses_turn("gpt-5.4-mini") + upstream_ws = ScriptedUpstreamWebSocket(replies=[[created]], in_flight=[completed]) + + relayed = 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", + "error", + ] + assert [ + [frame["type"] for frame in logged.kwargs["response_body"]] + for logged in relayed.success_handler.call_args_list + ] == [["response.created", "response.completed"]] + relayed.proxy_logging.post_call_failure_hook.assert_not_awaited() + + DEEPGRAM_LISTEN_TARGET = "wss://api.deepgram.com/v1/listen?model=nova-3&encoding=linear16&sample_rate=16000" DEEPGRAM_INTERIM_FRAME = json.dumps( { @@ -8718,15 +8994,6 @@ def test_passthrough_client_cannot_forge_session_id_omission(client_metadata_key ) -class StampingFrameModelGate: - def __init__(self, groups: list[str]): - self.groups = groups - - async def __call__(self, model: str, valid_token: UserAPIKeyAuth, /) -> Optional[str]: - valid_token.matched_model_access_groups = self.groups - return None - - async def _billed_metadata(gate) -> dict: served = await _relay_client_frames( [json.dumps({"type": "response.create", "model": "gpt-5.4-mini"})], @@ -8740,7 +9007,7 @@ async def _billed_metadata(gate) -> dict: 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(StampingFrameModelGate(["tier-a"])) + metadata = 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" diff --git a/tests/unit/proxy/test_openai_ws_passthrough_routes.py b/tests/unit/proxy/test_openai_ws_passthrough_routes.py index bd65001111a..836104e3342 100644 --- a/tests/unit/proxy/test_openai_ws_passthrough_routes.py +++ b/tests/unit/proxy/test_openai_ws_passthrough_routes.py @@ -1,21 +1,22 @@ """OpenAI passthrough WebSocket route: registration, opt-in gating, refusals, and per-frame model checks.""" import json -from collections.abc import Iterator, Mapping -from contextlib import contextmanager +from collections.abc import Mapping from dataclasses import dataclass from types import MappingProxyType, SimpleNamespace from typing import Final -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, patch import httpx import pytest import respx +from fastapi import HTTPException from starlette.routing import WebSocketRoute +from starlette.websockets import WebSocket import litellm from litellm.llms.openai.workload_identity import _workload_identity_auth -from litellm.proxy._types import Litellm_EntityType, LiteLLM_UserTable, UserAPIKeyAuth +from litellm.proxy._types import Litellm_EntityType, ProxyErrorTypes, ProxyException, UserAPIKeyAuth from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( _OPENAI_WS_DISABLED_REFUSAL, _OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL, @@ -24,7 +25,16 @@ from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( openai_websocket_proxy_route, router, ) -from litellm.proxy.pass_through_endpoints.pass_through_endpoints import WebsocketFrameModelGate +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_websocket_passthrough_logging_handler import ( + openai_websocket_frame_kind, +) +from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + WebsocketFrameModelAllowed, + WebsocketFrameModelGate, + WebsocketFrameModelRefused, + WebsocketFrameModelVerdict, + WebsocketUpstreamFrameClassifier, +) ENABLED: Final = MappingProxyType({"enable_openai_websocket_passthrough": True}) DISABLED_SETTINGS: Final = ( @@ -76,6 +86,7 @@ class _RelayCall: endpoint: str accept_websocket: bool client_frame_model_gate: WebsocketFrameModelGate | None + upstream_frame_kind: WebsocketUpstreamFrameClassifier | None class _FakeRelay: @@ -93,6 +104,7 @@ class _FakeRelay: endpoint: str, accept_websocket: bool, client_frame_model_gate: WebsocketFrameModelGate | None = None, + upstream_frame_kind: WebsocketUpstreamFrameClassifier | None = None, ) -> None: self.calls.append( _RelayCall( @@ -102,18 +114,14 @@ class _FakeRelay: endpoint=endpoint, accept_websocket=accept_websocket, client_frame_model_gate=client_frame_model_gate, + upstream_frame_kind=upstream_frame_kind, ) ) class _FakeFrameModelGate: - def __init__(self, refusal: str | None = None) -> None: - self.refusal = refusal - self.checked: list[str] = [] - - async def __call__(self, model: str, valid_token: UserAPIKeyAuth, /) -> str | None: - self.checked.append(model) - return self.refusal + async def __call__(self, websocket: object, model: str, /) -> WebsocketFrameModelVerdict: + return WebsocketFrameModelAllowed(access_groups=()) @dataclass(frozen=True, slots=True) @@ -157,6 +165,7 @@ async def test_openai_websocket_forwards_query_and_keeps_provider_auth(prefix, m endpoint=f"/{prefix}/v1/realtime", accept_websocket=False, client_frame_model_gate=served.frame_model_gate, + upstream_frame_kind=openai_websocket_frame_kind, ) ] assert websocket.accepts == [None] @@ -310,26 +319,53 @@ async def test_openai_websocket_connects_model_restricted_keys(key_models): assert websocket.closed is None -@pytest.mark.asyncio -async def test_proxy_frame_model_gate_allows_a_model_the_key_owns(): - gate: Final = _proxy_frame_model_gate() - token: Final = UserAPIKeyAuth(models=["gpt-realtime-2.1", "gpt-5.4-mini"]) +class _FrameAuthorizer: + """Stands in for the full key, team, user and budget auth the gate re-runs for every model a frame names.""" - with patch("litellm.proxy.proxy_server.prisma_client", None): # test-quality-ok: the gate needs no database - assert await gate("gpt-realtime-2.1", token) is None + def __init__(self, token: UserAPIKeyAuth | None = None, denial: Exception | None = None) -> None: + self.token = token or UserAPIKeyAuth() + self.denial = denial + self.authorized: list[str] = [] + + async def __call__(self, websocket: object, model: str) -> UserAPIKeyAuth: + self.authorized.append(model) + if self.denial is not None: + raise self.denial + return self.token + + +SOCKET: Final = WebSocket({"type": "websocket"}, receive=AsyncMock(), send=AsyncMock()) + + +async def _gate_verdict(authorizer: _FrameAuthorizer, model: str) -> WebsocketFrameModelVerdict: + with patch("litellm.proxy.proxy_server.prisma_client", None): # test-quality-ok: the model checks need no database + return await _proxy_frame_model_gate(authorizer)(SOCKET, model) + + +def _refusal(verdict: WebsocketFrameModelVerdict) -> WebsocketFrameModelRefused: + assert isinstance(verdict, WebsocketFrameModelRefused) + return verdict + + +@pytest.mark.asyncio +async def test_proxy_frame_model_gate_runs_full_auth_for_the_model_the_frame_names(): + authorizer: Final = _FrameAuthorizer(UserAPIKeyAuth(models=["gpt-realtime-2.1", "gpt-5.4-mini"])) + + verdict: Final = await _gate_verdict(authorizer, "gpt-realtime-2.1") + + assert verdict == WebsocketFrameModelAllowed(access_groups=()) + assert authorizer.authorized == ["gpt-realtime-2.1"] @pytest.mark.asyncio @pytest.mark.parametrize("requested_model", ["gpt-5.4", "gpt-realtime-2.1-mini"]) async def test_proxy_frame_model_gate_refuses_a_model_the_key_lacks(requested_model): - gate: Final = _proxy_frame_model_gate() - token: Final = UserAPIKeyAuth(models=["gpt-realtime-2.1", "gpt-5.4-mini"]) + authorizer: Final = _FrameAuthorizer(UserAPIKeyAuth(models=["gpt-realtime-2.1", "gpt-5.4-mini"])) - with patch("litellm.proxy.proxy_server.prisma_client", None): # test-quality-ok: the gate needs no database - refusal = await gate(requested_model, token) + refusal: Final = _refusal(await _gate_verdict(authorizer, requested_model)) - assert refusal is not None - assert requested_model in refusal + assert requested_model in refusal.message + assert refusal.error_type == ProxyErrorTypes.key_model_access_denied.value @pytest.mark.asyncio @@ -341,7 +377,6 @@ async def test_proxy_frame_model_gate_refuses_a_model_the_key_lacks(requested_mo async def test_proxy_frame_model_gate_refuses_a_model_the_agent_was_not_granted(granted_models, refusal_names): from litellm.types.agents import AgentResponse - gate: Final = _proxy_frame_model_gate() token: Final = UserAPIKeyAuth(models=[], agent_id="agent-fake") token.managed_agent_policy = AgentResponse( agent_id="agent-fake", @@ -350,132 +385,105 @@ async def test_proxy_frame_model_gate_refuses_a_model_the_agent_was_not_granted( object_permission={"object_permission_id": "grant-fake", "models": granted_models}, ) - with patch("litellm.proxy.proxy_server.prisma_client", None): # test-quality-ok: the gate needs no database - refusal = await gate("gpt-5.4", token) + refusal: Final = _refusal(await _gate_verdict(_FrameAuthorizer(token), "gpt-5.4")) - assert refusal is not None - assert refusal_names in refusal - - -GET_USER_OBJECT: Final = "litellm.proxy.auth.auth_checks.get_user_object" - - -def _personal_user(models: list[str]) -> LiteLLM_UserTable: - return LiteLLM_UserTable(user_id="user-fake", models=models, max_budget=None, spend=0.0) - - -@contextmanager -def _personal_user_allowlist(models: list[str]) -> Iterator[AsyncMock]: - """Stand up the database-backed owner lookup the allowlist check runs, and report every call it made.""" - lookup: Final = AsyncMock(return_value=_personal_user(models)) - with ( - patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), # test-quality-ok: read at call time - patch(GET_USER_OBJECT, lookup), - ): - yield lookup - - -@pytest.mark.asyncio -async def test_proxy_frame_model_gate_refuses_a_model_the_person_behind_the_key_lacks(): - """ - A key with no team of its own carries its owner's model allowlist, the same - one every HTTP route applies, so a frame naming a model that person cannot - use is refused even when the key itself names no models at all. - """ - gate: Final = _proxy_frame_model_gate() - token: Final = UserAPIKeyAuth(models=[], user_id="user-fake") - - with _personal_user_allowlist(["gpt-5.4-mini"]): - refusal = await gate("gpt-5.4", token) - - assert refusal is not None - assert "gpt-5.4" in refusal - - -@pytest.mark.asyncio -async def test_proxy_frame_model_gate_allows_a_model_the_person_behind_the_key_owns(): - gate: Final = _proxy_frame_model_gate() - token: Final = UserAPIKeyAuth(models=[], user_id="user-fake") - - with _personal_user_allowlist(["gpt-5.4-mini", "gpt-realtime-2.1"]): - assert await gate("gpt-realtime-2.1", token) is None - - -@pytest.mark.asyncio -async def test_proxy_frame_model_gate_leaves_a_team_key_to_the_teams_allowlist(): - """ - A key that belongs to a team is governed by the team's models, not by the - personal allowlist of whoever created it, so the owner's own restrictions - must not close a session the team is entitled to. - """ - gate: Final = _proxy_frame_model_gate() - token: Final = UserAPIKeyAuth(models=[], user_id="user-fake", team_id="team-fake") - - with _personal_user_allowlist(["gpt-5.4-mini"]) as owner_lookup: - assert await gate("gpt-5.4", token) is None - - assert owner_lookup.await_args_list == [] + assert refusal_names in refusal.message @pytest.mark.asyncio @pytest.mark.parametrize("unrestricted_models", [[], ["all-proxy-models"], ["*"]]) async def test_proxy_frame_model_gate_allows_every_model_for_unrestricted_keys(unrestricted_models): - gate: Final = _proxy_frame_model_gate() - token: Final = UserAPIKeyAuth(models=unrestricted_models) + verdict: Final = await _gate_verdict(_FrameAuthorizer(UserAPIKeyAuth(models=unrestricted_models)), "gpt-5.4") - with patch("litellm.proxy.proxy_server.prisma_client", None): # test-quality-ok: the gate needs no database - assert await gate("gpt-realtime-2.1", token) is None + assert verdict == WebsocketFrameModelAllowed(access_groups=()) -ENFORCE_GROUP_BUDGET: Final = ( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.enforce_model_access_group_budget" +@pytest.mark.asyncio +async def test_proxy_frame_model_gate_hands_back_the_access_groups_that_authorized_the_model(): + """The turn a frame starts is billed to the groups that let its model through, so the gate has to report them.""" + token: Final = UserAPIKeyAuth(models=["gpt-5.4-mini"], matched_model_access_groups=["tier-a"]) + + verdict: Final = await _gate_verdict(_FrameAuthorizer(token), "gpt-5.4-mini") + + assert verdict == WebsocketFrameModelAllowed(access_groups=("tier-a",)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "requested_model", + ["gpt-5.4-mini", "gpt-5.4"], + ids=["allowed", "refused_by_the_model_check"], ) +async def test_proxy_frame_model_gate_releases_the_budget_it_reserved_for_the_check(requested_model): + """ + Full auth reserves budget for the request it authorizes, but a frame check is + not a request, so a reservation the gate leaves open would hold the key's + spend counter above its real spend and refuse later calls on every route. + """ + reservation: Final[dict[str, object]] = {"reserved_cost": 0.5, "entries": []} + token: Final = UserAPIKeyAuth(models=["gpt-5.4-mini"], budget_reservation=reservation) + + await _gate_verdict(_FrameAuthorizer(token), requested_model) + + assert token.budget_reservation == {**reservation, "finalized": True} @pytest.mark.asyncio -async def test_proxy_frame_model_gate_refuses_a_model_whose_access_group_is_out_of_budget(): - """ - A model that sits in a budgeted access group is refused on every other route once - that group has spent its budget, so a frame naming it closes the session instead - of quietly running on an empty pool. The key is allowed the model outright here, - so nothing but the group's budget can produce the refusal. - """ - exhausted: Final = litellm.BudgetExceededError( - current_cost=2.0, - max_budget=1.0, - message="Budget has been exceeded! Model access group=tier-a Current cost: 2.0, Max budget: 1.0", - entity_type=Litellm_EntityType.MODEL_ACCESS_GROUP.value, - entity_id="tier-a", - ) - gate: Final = _proxy_frame_model_gate() - token: Final = UserAPIKeyAuth(models=["gpt-5.4-mini"]) +@pytest.mark.parametrize( + "denial, error_type, message", + [ + ( + ProxyException( + message="Budget has been exceeded! Key=hashed Spend=2.0, Budget=1.0", + type=ProxyErrorTypes.budget_exceeded, + param=None, + code=400, + ), + "budget_exceeded", + "Budget has been exceeded! Key=hashed", + ), + ( + litellm.BudgetExceededError( + current_cost=2.0, + max_budget=1.0, + message="Budget has been exceeded! Model access group=tier-a Current cost: 2.0, Max budget: 1.0", + entity_type=Litellm_EntityType.MODEL_ACCESS_GROUP.value, + entity_id="tier-a", + ), + "budget_exceeded", + "Model access group=tier-a", + ), + ( + ProxyException( + message="team not allowed to access model. This team can only access models=['gpt-5.4-mini']", + type=ProxyErrorTypes.team_model_access_denied, + param="model", + code=401, + ), + "team_model_access_denied", + "team not allowed to access model", + ), + (HTTPException(status_code=403, detail="Invalid Authorization header format"), "invalid_request_error", ""), + ], + ids=["key_budget", "group_budget", "team_allowlist", "malformed_auth"], +) +async def test_proxy_frame_model_gate_refuses_what_full_auth_denies_with_its_reason(denial, error_type, message): + """A frame is denied for the same reasons a request is, and the caller reads which reason in the error event.""" + refusal: Final = _refusal(await _gate_verdict(_FrameAuthorizer(denial=denial), "gpt-5.4-mini")) - with ( - patch("litellm.proxy.proxy_server.prisma_client", None), # test-quality-ok: the model check needs no database - patch(ENFORCE_GROUP_BUDGET, AsyncMock(side_effect=exhausted)), - ): - refusal = await gate("gpt-5.4-mini", token) - - assert refusal is not None - assert "Budget has been exceeded" in refusal - assert "Model access group=tier-a" in refusal + assert refusal.error_type == error_type + assert message in refusal.message @pytest.mark.asyncio -async def test_proxy_frame_model_gate_refuses_out_loud_when_it_cannot_reach_the_budgets(): +async def test_proxy_frame_model_gate_refuses_out_loud_when_auth_cannot_answer(): """ A gate that cannot answer has to close the session with a reason the caller can read. Letting the failure escape drops the socket with no frame at all, which the caller cannot tell apart from the network going away. """ - gate: Final = _proxy_frame_model_gate() - token: Final = UserAPIKeyAuth(models=["gpt-5.4-mini"]) + refusal: Final = _refusal( + await _gate_verdict(_FrameAuthorizer(denial=ConnectionError("connection pool is exhausted")), "gpt-5.4-mini") + ) - with ( - patch("litellm.proxy.proxy_server.prisma_client", None), # test-quality-ok: the model check needs no database - patch(ENFORCE_GROUP_BUDGET, AsyncMock(side_effect=ConnectionError("connection pool is exhausted"))), - ): - refusal = await gate("gpt-5.4-mini", token) - - assert refusal == _OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL - assert "connection pool" not in refusal + assert refusal == WebsocketFrameModelRefused(message=_OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL, error_type="server_error") From e432910379ab396a7bd8cf26a697f4fc1b2dc6b4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 21 Oct 2026 02:05:10 -0700 Subject: [PATCH 04/15] fix(proxy): keep the stored prompt refusal's first sentence inside the websocket close reason The close reason carries the refusal's first sentence, and the stored prompt sentence ran past the 123 bytes a close frame allows, so it was cut mid-word. The sentence is shorter now and the rest of the explanation stays in the error event. --- .../proxy/pass_through_endpoints/pass_through_endpoints.py | 4 ++-- .../pass_through_endpoints/test_pass_through_endpoints.py | 5 ++++- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 201ecee5b44..1c03942bce4 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -2508,8 +2508,8 @@ _TRANSCRIPTION_MODEL_PLACEMENTS: Final = ( _STORED_PROMPT_PLACEMENTS: Final = (("prompt",), ("response", "prompt"), ("session", "prompt")) STORED_PROMPT_PROBE_MODEL: Final = "litellm-unverifiable-stored-prompt-model" _STORED_PROMPT_REFUSAL: Final = ( - "This key is limited to specific models, and the gateway cannot check which model a stored prompt runs, " - "so frames that reference a stored prompt are refused for this key." + "This key is limited to specific models, so it cannot run a stored prompt. " + "The gateway cannot check which model a stored prompt runs." ) _FIRST_SENTENCE: Final = re.compile(r".*?[.!?](?=\s|$)") 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 e6a0b1be28c..8d74b29fe45 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 @@ -7774,10 +7774,13 @@ async def test_websocket_passthrough_refuses_a_stored_prompt_for_a_key_limited_t relayed = 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"] assert gate.checked == [STORED_PROMPT_PROBE_MODEL] assert relayed.upstream.send.await_args_list == [] - assert "stored prompt" in json.loads(relayed.client.send_text.await_args.args[0])["error"]["message"] assert relayed.client.close.await_args.kwargs["code"] == 1008 + assert "stored prompt" in close_reason + assert close_reason.endswith(".") and error_message.startswith(close_reason) @pytest.mark.asyncio From 2d860f05d4ae0c8b9d52b5424f6eaf3467fdaf53 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 01:17:07 -0700 Subject: [PATCH 05/15] fix(proxy): bill pipelined websocket turns to their own frame and close the review gaps --- litellm/litellm_core_utils/core_helpers.py | 6 +- litellm/proxy/common_utils/realtime_utils.py | 3 +- .../llm_passthrough_endpoints.py | 77 +++++-- ...i_websocket_passthrough_logging_handler.py | 29 ++- .../pass_through_endpoints.py | 188 +++++++++++----- .../proxy/common_utils/test_realtime_cache.py | 7 + ...i_websocket_passthrough_logging_handler.py | 45 +++- .../test_pass_through_endpoints.py | 212 ++++++++++++++++-- .../test_success_handler.py | 2 +- .../test_openai_ws_passthrough_routes.py | 55 +++++ 10 files changed, 530 insertions(+), 94 deletions(-) diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 55025508e40..58412198064 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -3,7 +3,7 @@ import copy import logging import re -from collections.abc import Collection, Iterable, Iterator, Mapping +from collections.abc import Collection, Iterable, Iterator, Mapping, Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, cast, runtime_checkable @@ -59,6 +59,10 @@ def as_str_mapping(value: object) -> Mapping[str, object] | None: return value if isinstance(value, Mapping) else None +def as_object_list(value: object) -> Sequence[object] | None: + return value if isinstance(value, list) else None + + def decoded_json_object(value: object) -> Mapping[str, object]: mapping: Final = as_str_mapping(value) if mapping is not None: diff --git a/litellm/proxy/common_utils/realtime_utils.py b/litellm/proxy/common_utils/realtime_utils.py index c641748843e..e96685bdeae 100644 --- a/litellm/proxy/common_utils/realtime_utils.py +++ b/litellm/proxy/common_utils/realtime_utils.py @@ -1,3 +1,4 @@ +import json from functools import lru_cache from typing import Final @@ -10,7 +11,7 @@ def realtime_request_body(model: str | None) -> bytes: Generate the realtime websocket request body. Cached with LRU semantics to avoid repeated string formatting work while keeping memory usage bounded. """ - return f'{{"model": "{model or ""}"}}'.encode() + return json.dumps({"model": model or ""}).encode() _realtime_request_body: Final = realtime_request_body diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 2e51bc3465b..57c0e78c462 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -15,7 +15,7 @@ import os import posixpath import re import sys -from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping +from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass from functools import partial from types import MappingProxyType @@ -50,7 +50,7 @@ from litellm.constants import ( from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException from litellm.integrations.custom_guardrail import guardrail_request_data_with_streaming, is_guardrail_intervention from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix -from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.core_helpers import as_object_list, as_str_mapping, attribute_or_none from litellm.llms.anthropic.common_utils import AnthropicModelInfo, merge_anthropic_beta_headers from litellm.llms.azure.passthrough.transformation import ( foreign_azure_deployment, @@ -73,7 +73,12 @@ from litellm.llms.oss_decision import OssDecisionProvider, oss_connection, valid from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.passthrough.main import AsyncPassthroughStreamingResponse, is_provider_error from litellm.proxy._types import * -from litellm.proxy.auth.auth_checks import can_agent_call_model, can_key_call_resolved_model +from litellm.proxy.auth.auth_checks import ( + can_agent_call_model, + can_key_call_resolved_model, + key_model_aliases_for_auth_check, + team_model_aliases_for_auth_check, +) from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy module exports @@ -102,6 +107,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg safe_get_request_headers, safe_set_request_parsed_body, ) +from litellm.proxy.common_utils.model_listing_utils import alias_map from litellm.proxy.common_utils.openai_error_payload import LITELLM_CALL_ID_HEADER from litellm.proxy.common_utils.resource_ownership import is_proxy_admin from litellm.proxy.common_utils.sse_keepalive import ( @@ -121,6 +127,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_websocket openai_websocket_frame_kind, ) from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + STORED_PROMPT_PROBE_MODEL, HttpPassThroughEndpointHelpers, WebsocketFrameModelAllowed, WebsocketFrameModelGate, @@ -132,7 +139,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( websocket_passthrough_request, ) from litellm.proxy.utils import ProxyLogging as ProxyLoggingType -from litellm.proxy.utils import is_known_model +from litellm.proxy.utils import is_known_model, normalize_route_for_root_path from litellm.proxy.vector_store_endpoints.utils import ( assert_proxy_admin_for_vector_store_index_management, assert_user_can_access_vector_store, @@ -3304,6 +3311,11 @@ _OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL: Final = ( "and ask a proxy admin to check the gateway logs if it keeps happening." ) +_OPENAI_WS_ALIASED_MODEL_REFUSAL: Final = ( + "This key cannot use a model alias over OpenAI websocket passthrough. The passthrough sends each frame to " + "OpenAI unchanged, so name the model this key may use directly." +) + _OPENAI_WS_DISABLED_REFUSAL: Final = _OpenAIWebsocketRefusal( close_reason="OpenAI websocket passthrough is disabled", message=( @@ -3357,36 +3369,65 @@ def _frame_model_authorizer() -> FrameModelAuthorizer: return user_api_key_auth_accepted_websocket_for_model +def _proxy_model_list() -> Sequence[object] | None: + from litellm.proxy import proxy_server + + return as_object_list(attribute_or_none(proxy_server, "llm_model_list")) + + +def _gateway_remaps_model(model: str, frame_auth: UserAPIKeyAuth, llm_router: Router | None) -> bool: + alias_maps: Final = ( + key_model_aliases_for_auth_check(frame_auth) or {}, + team_model_aliases_for_auth_check(frame_auth) or {}, + alias_map(litellm.model_alias_map), + ) + targets: Final = ( + *(aliases.get(model) for aliases in alias_maps), + llm_router.get_model_from_alias(model) if llm_router is not None else None, + ) + return any(target not in (None, model) for target in targets) + + +def _aliased_model_verdict( + allowed: WebsocketFrameModelAllowed, probe: WebsocketFrameModelVerdict +) -> WebsocketFrameModelVerdict: + match probe: + case WebsocketFrameModelRefused(error_type=error_type) if error_type.endswith("_model_access_denied"): + return WebsocketFrameModelRefused(message=_OPENAI_WS_ALIASED_MODEL_REFUSAL, error_type=error_type) + case WebsocketFrameModelRefused(): + return probe + case WebsocketFrameModelAllowed(): + return allowed + + def _proxy_frame_model_gate( authorize_frame_model: Annotated[FrameModelAuthorizer, Depends(_frame_model_authorizer)], ) -> WebsocketFrameModelGate: - from litellm.proxy.proxy_server import ( - llm_model_list, - llm_router, - prisma_client, - proxy_logging_obj, - user_api_key_cache, - ) + from litellm.proxy import proxy_server from litellm.proxy.spend_tracking.budget_reservation import release_or_invalidate_budget_reservation - async def authorize(websocket: WebSocket, model: str) -> WebsocketFrameModelAllowed: + async def authorize(websocket: WebSocket, model: str) -> WebsocketFrameModelVerdict: frame_auth: Final = await authorize_frame_model(websocket, model) access_groups: Final = tuple(frame_auth.matched_model_access_groups or ()) + llm_router: Final = proxy_server.llm_router try: await can_key_call_resolved_model( - model=model, llm_model_list=llm_model_list, valid_token=frame_auth, llm_router=llm_router + model=model, llm_model_list=_proxy_model_list(), valid_token=frame_auth, llm_router=llm_router ) await can_agent_call_model( model=model, valid_token=frame_auth, llm_router=llm_router, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, + prisma_client=proxy_server.prisma_client, + user_api_key_cache=proxy_server.user_api_key_cache, + proxy_logging_obj=proxy_server.proxy_logging_obj, ) finally: await release_or_invalidate_budget_reservation(budget_reservation=frame_auth.budget_reservation) - return WebsocketFrameModelAllowed(access_groups=access_groups) + allowed: Final = WebsocketFrameModelAllowed(access_groups=access_groups) + if not _gateway_remaps_model(model, frame_auth, llm_router): + return allowed + return _aliased_model_verdict(allowed, await gate(websocket, STORED_PROMPT_PROBE_MODEL)) async def gate(websocket: WebSocket, model: str, /) -> WebsocketFrameModelVerdict: try: @@ -3478,7 +3519,7 @@ async def openai_websocket_proxy_route( custom_headers=custom_headers, user_api_key_dict=user_api_key_dict, forward_headers=False, - endpoint=websocket.url.path, + endpoint=normalize_route_for_root_path(websocket.url.path) or websocket.url.path, accept_websocket=False, client_frame_model_gate=frame_model_gate, upstream_frame_kind=openai_websocket_frame_kind, 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 c2f4e85005c..951764f74b0 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 @@ -9,6 +9,7 @@ from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger from litellm.cost_calculator import RealtimeAPITokenUsageProcessor, handle_realtime_stream_cost_calculation +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import PassThroughEndpointLoggingTypedDict from litellm.responses.utils import ResponseAPILoggingUtils @@ -139,15 +140,18 @@ def _responses_turn(event: Mapping[str, object], session_model: str) -> _Complet ) -def _responses_billing(messages: WebsocketMessages, model: str) -> _Billing | None: +def _responses_billing(messages: WebsocketMessages, model: str, requested_model: str | None) -> _Billing | None: turns: Final = tuple(turn for turn in (_responses_turn(event, model) for event in messages) if turn is not None) if not turns: return None usage: Final = RealtimeAPITokenUsageProcessor.combine_usage_objects([turn.usage for turn in turns]) - return _Billing(usage=usage, cost=sum(_responses_cost(turn.usage, turn.model) for turn in turns)) + return _Billing( + usage=usage, + cost=sum(_responses_cost(turn.usage, (turn.model, requested_model)) for turn in turns), + ) -def _responses_cost(usage: Usage, model: str) -> float: +def _model_cost(usage: Usage, model: str) -> float | None: try: return litellm.completion_cost( completion_response=ModelResponse(model=model, usage=usage), @@ -155,10 +159,25 @@ def _responses_cost(usage: Usage, model: str) -> float: custom_llm_provider=OPENAI_WEBSOCKET_PROVIDER, ) except Exception: # noqa: BLE001 # completion_cost raises a bare Exception for an unpriced model, and the turn must still log + return None + + +def _responses_cost(usage: Usage, price_as: Sequence[str | None]) -> float: + costs: Final = (_model_cost(usage, model) for model in dict.fromkeys(price_as) if model is not None) + cost: Final = next((cost for cost in costs if cost), None) + if cost is None: verbose_proxy_logger.warning( - "OpenAI websocket passthrough: could not price a Responses turn on %s, logging it at zero", model + "OpenAI websocket passthrough: could not price a Responses turn on %s, logging it at zero", price_as ) return 0.0 + return cost + + +def _requested_model(kwargs: Mapping[str, object]) -> str | None: + litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) + metadata: Final = None if litellm_params is None else as_str_mapping(litellm_params.get("metadata")) + model_group: Final = None if metadata is None else metadata.get("model_group") + return model_group if isinstance(model_group, str) else None def _priced_result(model: str, billing: _Billing, start_time: datetime) -> ModelResponse: @@ -195,7 +214,7 @@ class OpenAIWebsocketPassthroughLoggingHandler: ) billing: Final = _realtime_billing(websocket_messages, model, logging_obj) or _responses_billing( - websocket_messages, model + websocket_messages, model, _requested_model(kwargs) ) if billing is None: verbose_proxy_logger.debug("OpenAI websocket passthrough (%s): no usage reported, skipping cost", url_route) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index befdc821734..af67a451c96 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -13,9 +13,10 @@ from functools import reduce from itertools import chain, count, groupby from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias, TypedDict, cast -from urllib.parse import urlencode, urlparse +from urllib.parse import parse_qs, urlencode, urlparse import httpx +import orjson from fastapi import ( APIRouter, Depends, @@ -60,6 +61,7 @@ from litellm.integrations.custom_guardrail import ( from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.otel.model.semconv import PASSTHROUGH_OPERATION_KEY, GenAIOperation from litellm.litellm_core_utils.core_helpers import ( + as_object_list, as_str_mapping, bind_budget_reservation_to_callbacks, get_metadata_variable_name_from_kwargs, @@ -2441,11 +2443,16 @@ def _resolved_vertex_live_setup( return {**setup_data, "model": setup_model_rewriter(setup_model)} +def _decoded_frame(frame: str | bytes) -> object: + try: + return _DECODED_JSON.validate_python(orjson.loads(frame)) + except orjson.JSONDecodeError: + return _DECODED_JSON.validate_python(json.loads(frame if isinstance(frame, str) else frame.decode("utf-8"))) + + def _json_object_frame(frame: str | bytes) -> dict[str, object] | None: try: - decoded: Final[object] = _DECODED_JSON.validate_python( - json.loads(frame if isinstance(frame, str) else frame.decode("utf-8")) - ) + decoded: Final = _decoded_frame(frame) except (json.JSONDecodeError, UnicodeDecodeError): return None return decoded if isinstance(decoded, dict) else None @@ -2525,9 +2532,16 @@ class WebsocketFrameModelGate(Protocol): WebsocketUpstreamFrameClassifier: TypeAlias = Callable[[Mapping[str, object]], WebsocketUpstreamFrameKind] +@dataclass(frozen=True, slots=True) +class _TurnAuthorization: + requested_model: str | None + access_groups: tuple[str, ...] + + @dataclass(frozen=True, slots=True) class _ForwardClientFrame: - turn_access_groups: tuple[str, ...] | None + starts_turn: bool + turn: _TurnAuthorization | None @dataclass(frozen=True, slots=True) @@ -2537,24 +2551,35 @@ class _RefuseClientFrame: @dataclass(frozen=True, slots=True) class _ClientFrameRequests: + starts_turn: bool turn_model: str | None models: tuple[str, ...] - names_stored_prompt: bool + unverifiable_model_refusal: str | None -_NO_CLIENT_FRAME_REQUESTS: Final = _ClientFrameRequests(turn_model=None, models=(), names_stored_prompt=False) -_FRAME_TEXT_WORTH_PARSING: Final = ('"model"', '"prompt"', "\\") +_NO_CLIENT_FRAME_REQUESTS: Final = _ClientFrameRequests( + starts_turn=False, turn_model=None, models=(), unverifiable_model_refusal=None +) +_FORWARD_WITHOUT_TURN: Final = _ForwardClientFrame(starts_turn=False, turn=None) +_TURN_STARTING_FRAME: Final = "response.create" +_FRAME_TEXT_WORTH_PARSING: Final = ('"model"', '"prompt"', '"tools"', "\\") _TURN_MODEL_PLACEMENTS: Final = (("model",), ("response", "model"), ("session", "model")) _TRANSCRIPTION_MODEL_PLACEMENTS: Final = ( ("session", "input_audio_transcription", "model"), ("session", "audio", "input", "transcription", "model"), ) _STORED_PROMPT_PLACEMENTS: Final = (("prompt",), ("response", "prompt"), ("session", "prompt")) +_TOOL_LIST_PLACEMENTS: Final = (("tools",), ("response", "tools"), ("session", "tools")) +_TOOL_THAT_RUNS_A_MODEL: Final = "image_generation" STORED_PROMPT_PROBE_MODEL: Final = "litellm-unverifiable-stored-prompt-model" _STORED_PROMPT_REFUSAL: Final = ( "This key is limited to specific models, so it cannot run a stored prompt. " "The gateway cannot check which model a stored prompt runs." ) +_UNNAMED_TOOL_MODEL_REFUSAL: Final = ( + "This key is limited to specific models, so an image_generation tool has to name its model. " + "The gateway cannot check which model an unnamed one runs." +) _FIRST_SENTENCE: Final = re.compile(r".*?[.!?](?=\s|$)") @@ -2577,6 +2602,21 @@ def _placed_models(frame: Mapping[str, object], placements: Iterable[Sequence[st return tuple(value for value in placed if isinstance(value, str)) +def _placed_tools(frame: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: + tool_lists: Final = (_placed_value(frame, placement) for placement in _TOOL_LIST_PLACEMENTS) + tools: Final = chain.from_iterable(as_object_list(tool_list) or () for tool_list in tool_lists) + return tuple(tool for tool in map(as_str_mapping, tools) if tool is not None) + + +def _unverifiable_model_refusal(frame: Mapping[str, object], tools: Sequence[Mapping[str, object]]) -> str | None: + prompts: Final = (_placed_value(frame, placement) for placement in _STORED_PROMPT_PLACEMENTS) + if any(isinstance(prompt, Mapping) for prompt in prompts): + return _STORED_PROMPT_REFUSAL + if any(tool.get("type") == _TOOL_THAT_RUNS_A_MODEL and not isinstance(tool.get("model"), str) for tool in tools): + return _UNNAMED_TOOL_MODEL_REFUSAL + return None + + def _client_frame_requests(frame_data: str | bytes) -> _ClientFrameRequests: if isinstance(frame_data, str) and not _may_name_a_model(frame_data): return _NO_CLIENT_FRAME_REQUESTS @@ -2584,17 +2624,21 @@ def _client_frame_requests(frame_data: str | bytes) -> _ClientFrameRequests: if frame is None: return _NO_CLIENT_FRAME_REQUESTS turn_models: Final = _placed_models(frame, _TURN_MODEL_PLACEMENTS) - prompts: Final = (_placed_value(frame, placement) for placement in _STORED_PROMPT_PLACEMENTS) + tools: Final = _placed_tools(frame) + tool_models: Final = (model for model in (tool.get("model") for tool in tools) if isinstance(model, str)) return _ClientFrameRequests( + starts_turn=frame.get("type") == _TURN_STARTING_FRAME, turn_model=next(iter(turn_models), None), - models=tuple(dict.fromkeys((*turn_models, *_placed_models(frame, _TRANSCRIPTION_MODEL_PLACEMENTS)))), - names_stored_prompt=any(isinstance(prompt, Mapping) for prompt in prompts), + models=tuple( + dict.fromkeys((*turn_models, *_placed_models(frame, _TRANSCRIPTION_MODEL_PLACEMENTS), *tool_models)) + ), + unverifiable_model_refusal=_unverifiable_model_refusal(frame, tools), ) -def _stored_prompt_refusal(probe: WebsocketFrameModelRefused) -> WebsocketFrameModelRefused: +def _unverifiable_model_probe_refusal(probe: WebsocketFrameModelRefused, message: str) -> WebsocketFrameModelRefused: if probe.error_type.endswith("_model_access_denied"): - return WebsocketFrameModelRefused(message=_STORED_PROMPT_REFUSAL, error_type=probe.error_type) + return WebsocketFrameModelRefused(message=message, error_type=probe.error_type) return probe @@ -2602,20 +2646,25 @@ async def _client_frame_verdict( websocket: WebSocket, frame_data: str | bytes, gate: WebsocketFrameModelGate | None ) -> _ForwardClientFrame | _RefuseClientFrame: if gate is None: - return _ForwardClientFrame(turn_access_groups=None) + return _FORWARD_WITHOUT_TURN requests: Final = _client_frame_requests(frame_data) verdicts: Final = tuple([await gate(websocket, model) for model in requests.models]) refusal: Final = next((verdict for verdict in verdicts if isinstance(verdict, WebsocketFrameModelRefused)), None) if refusal is not None: return _RefuseClientFrame(refusal=refusal) - if requests.names_stored_prompt and isinstance( + if requests.unverifiable_model_refusal is not None and isinstance( probe := await gate(websocket, STORED_PROMPT_PROBE_MODEL), WebsocketFrameModelRefused ): - return _RefuseClientFrame(refusal=_stored_prompt_refusal(probe)) + return _RefuseClientFrame(refusal=_unverifiable_model_probe_refusal(probe, requests.unverifiable_model_refusal)) allowed: Final = tuple(verdict for verdict in verdicts if isinstance(verdict, WebsocketFrameModelAllowed)) - if requests.turn_model is None or not allowed: - return _ForwardClientFrame(turn_access_groups=None) - return _ForwardClientFrame(turn_access_groups=allowed[0].access_groups) + return _ForwardClientFrame( + starts_turn=requests.starts_turn, + turn=( + _TurnAuthorization(requested_model=requests.turn_model, access_groups=allowed[0].access_groups) + if requests.turn_model is not None and allowed + else None + ), + ) def _refusal_close_reason(message: str) -> str: @@ -2631,12 +2680,54 @@ async def _refuse_client_frame( await upstream_ws.close() -def _stamp_billed_access_groups(kwargs: Mapping[str, object], access_groups: tuple[str, ...]) -> None: - """Charge a turn to the access groups that authorized its model, and to none when no group did.""" +def _stamp_billed_turn(kwargs: Mapping[str, object], turn: _TurnAuthorization) -> None: + """Charge a turn to the access groups that authorized its model, and its model budget to the model it asked for, + since the provider may answer under a dated snapshot name the budget was never set on.""" litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) metadata: Final = None if litellm_params is None else litellm_params.get("metadata") - if isinstance(metadata, MutableMapping): - metadata[MODEL_ACCESS_GROUP_METADATA_KEY] = list(access_groups) + if not isinstance(metadata, MutableMapping): + return + metadata[MODEL_ACCESS_GROUP_METADATA_KEY] = list(turn.access_groups) + if turn.requested_model is not None: + metadata["model_group"] = turn.requested_model + + +def _opening_model(target: str) -> str | None: + return next(iter(parse_qs(urlparse(target).query).get("model", ())), None) + + +class _WebsocketTurnLedger: + """The groups that authorized each turn not yet logged, when the next row starts, and how many rows were logged. + + The provider runs a socket's response.create frames one at a time in the order they were sent, so each row is + charged to the groups that allowed the oldest turn still unlogged. A realtime session keeps the model it opened on, + so its rows stay charged to the groups that opened it. + """ + + def __init__(self, opening: _TurnAuthorization, started: datetime) -> None: + self.opening = opening + self.queued_turns: tuple[_TurnAuthorization | None, ...] = () + self.session_opened = False + self.started = started + self.logged_turns = 0 + + @property + def current(self) -> _TurnAuthorization: + oldest: Final = self.queued_turns[0] if self.queued_turns else None + return self.opening if oldest is None else oldest + + def queue_turn(self, turn: _TurnAuthorization | None) -> None: + if not self.session_opened: + self.queued_turns = (*self.queued_turns, turn) + + def open_session(self) -> None: + self.session_opened = True + self.queued_turns = () + + def close_turn(self, end_time: datetime) -> None: + self.logged_turns += 1 + self.started = end_time + self.queued_turns = self.queued_turns[1:] async def _relay_client_frame( @@ -2645,12 +2736,16 @@ async def _relay_client_frame( checked_frame: str | bytes, forwarded_frame: str | bytes, gate: WebsocketFrameModelGate | None, + ledger: _WebsocketTurnLedger, ) -> _ForwardClientFrame | _RefuseClientFrame: """Forward one client frame upstream, or refuse it and end the session.""" verdict: Final = await _client_frame_verdict(websocket, checked_frame, gate) match verdict: case _RefuseClientFrame(refusal=refusal): await _refuse_client_frame(websocket, upstream_ws, refusal) + case _ForwardClientFrame(starts_turn=True, turn=turn): + ledger.queue_turn(turn) + await upstream_ws.send(forwarded_frame) case _ForwardClientFrame(): await upstream_ws.send(forwarded_frame) return verdict @@ -2659,23 +2754,6 @@ async def _relay_client_frame( _EXPLICIT_SUCCESS_KWARGS: Final = frozenset({"logging_obj", "passthrough_logging_payload", "custom_llm_provider"}) -class _WebsocketTurnLedger: - """The groups that authorized the turn in flight, when it started, and how many turns the socket has logged.""" - - def __init__(self, access_groups: tuple[str, ...], started: datetime) -> None: - self.access_groups = access_groups - self.started = started - self.logged_turns = 0 - - def authorize(self, access_groups: tuple[str, ...] | None) -> None: - if access_groups is not None: - self.access_groups = access_groups - - def close_turn(self, end_time: datetime) -> None: - self.logged_turns += 1 - self.started = end_time - - async def websocket_passthrough_request( websocket: WebSocket, target: str, @@ -2835,7 +2913,11 @@ async def websocket_passthrough_request( self.request = MockWebSocketRequest(target_url) ledger: Final = _WebsocketTurnLedger( - access_groups=tuple(user_api_key_dict.matched_model_access_groups or ()), started=start_time + opening=_TurnAuthorization( + requested_model=_opening_model(target), + access_groups=tuple(user_api_key_dict.matched_model_access_groups or ()), + ), + started=start_time, ) def log_frames( @@ -2856,7 +2938,7 @@ async def websocket_passthrough_request( "end_time": end_time, } if client_frame_model_gate is not None: - _stamp_billed_access_groups(call_kwargs, ledger.access_groups) + _stamp_billed_turn(call_kwargs, ledger.current) success_kwargs: Final = { key: value for key, value in call_kwargs.items() if key not in _EXPLICIT_SUCCESS_KWARGS } @@ -2988,26 +3070,25 @@ async def websocket_passthrough_request( # Not a JSON message or doesn't contain setup data if isinstance( - relayed_text := await _relay_client_frame( + await _relay_client_frame( websocket, upstream_ws, text_data, _rewrite_vertex_live_setup_model(text_data, setup_model_rewriter), client_frame_model_gate, + ledger, ), _RefuseClientFrame, ): break - ledger.authorize(relayed_text.turn_access_groups) elif bytes_data is not None: if isinstance( - relayed_bytes := await _relay_client_frame( - websocket, upstream_ws, bytes_data, bytes_data, client_frame_model_gate + await _relay_client_frame( + websocket, upstream_ws, bytes_data, bytes_data, client_frame_model_gate, ledger ), _RefuseClientFrame, ): break - ledger.authorize(relayed_bytes.turn_access_groups) except asyncio.CancelledError: raise except Exception: @@ -3042,8 +3123,13 @@ async def websocket_passthrough_request( _extract_vertex_live_model_from_setup_response(message_data) return websocket_messages.append(message_data) - if upstream_frame_kind is not None and upstream_frame_kind(message_data) == "ends_turn": - log_turn() + match None if upstream_frame_kind is None else upstream_frame_kind(message_data): + case "opens_session": + ledger.open_session() + case "ends_turn": + log_turn() + case _: + pass async def relay_upstream_frame(upstream_message: str | bytes) -> None: record_upstream_frame(upstream_message) @@ -3105,7 +3191,9 @@ async def websocket_passthrough_request( reason=_truncated_close_reason(upstream_close.reason), ) - if ledger.logged_turns == 0: + if upstream_frame_kind is None or any( + upstream_frame_kind(frame) != "opens_session" for frame in websocket_messages + ): log_turn() # Call the proxy logging success hook diff --git a/tests/unit/proxy/common_utils/test_realtime_cache.py b/tests/unit/proxy/common_utils/test_realtime_cache.py index 6f5751fa923..6177a00d0aa 100644 --- a/tests/unit/proxy/common_utils/test_realtime_cache.py +++ b/tests/unit/proxy/common_utils/test_realtime_cache.py @@ -1,3 +1,4 @@ +import json from typing import Any, cast import pytest @@ -40,6 +41,12 @@ def test_realtime_request_body_caches_each_model_separately(): assert gpt4o_body_first is not gpt4o_mini_body +def test_realtime_request_body_keeps_a_model_name_with_quotes_inside_its_own_field(): + body = realtime_request_body('gpt-4o", "user": "someone-else') + + assert json.loads(body) == {"model": 'gpt-4o", "user": "someone-else'} + + def test_realtime_query_params_template_caches_each_pair_separately(): params_with_intent_first = _realtime_query_params_template("gpt-4o", "intent-a") params_with_intent_second = _realtime_query_params_template("gpt-4o", "intent-a") 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 253634760fa..c259d89c07f 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 @@ -1,3 +1,4 @@ +from collections.abc import Mapping, Sequence from datetime import datetime from types import MappingProxyType from typing import Final @@ -22,11 +23,11 @@ def _logging_obj() -> MagicMock: return logging_obj -def _session_created(model: str) -> dict: +def _session_created(model: str) -> dict[str, object]: return {"type": "session.created", "session": {"id": "sess_fake", "model": model}} -def _realtime_response_done(text_in: int, audio_in: int, text_out: int, audio_out: int) -> dict: +def _realtime_response_done(text_in: int, audio_in: int, text_out: int, audio_out: int) -> dict[str, object]: return { "type": "response.done", "response": { @@ -45,7 +46,7 @@ def _realtime_response_done(text_in: int, audio_in: int, text_out: int, audio_ou def _responses_completed( model: str, input_tokens: int, output_tokens: int, event: str = "response.completed" -) -> dict: +) -> dict[str, object]: return { "type": event, "response": { @@ -63,7 +64,12 @@ def _responses_completed( } -def _handle(messages, kwargs=MappingProxyType({}), logging_obj=None, url_route="/openai/v1/realtime"): +def _handle( + messages: Sequence[Mapping[str, object]], + kwargs: Mapping[str, object] = MappingProxyType({}), + logging_obj: MagicMock | None = None, + url_route: str = "/openai/v1/realtime", +): return OpenAIWebsocketPassthroughLoggingHandler().openai_websocket_passthrough_handler( websocket_messages=messages, logging_obj=logging_obj if logging_obj is not None else _logging_obj(), @@ -320,6 +326,37 @@ def test_a_failed_responses_turn_that_reported_no_usage_is_not_billed(): assert handled["kwargs"] == {"model": RESPONSES_MODEL, "custom_llm_provider": "openai"} +@pytest.mark.parametrize( + ("turn_model", "requested_model", "priced_as"), + [ + ("gpt-unreleased-qa-snapshot", RESPONSES_MODEL, RESPONSES_MODEL), + ("gpt-5.4-mini-qa-preview", RESPONSES_MODEL, RESPONSES_MODEL), + (RESPONSES_MODEL, REASONING_MODEL, RESPONSES_MODEL), + ], + ids=["unmapped_model", "model_the_cost_map_prices_at_zero", "priced_turn_model"], +) +def test_a_responses_turn_falls_back_to_the_requested_model_only_when_its_own_has_no_price( + turn_model: str, requested_model: str, priced_as: str +): + """ + The provider can answer under a model name the cost map does not price + yet. That turn still ran the model the key was allowed to ask for, so it is + billed at that model's price instead of $0, while a turn whose own model + has a price keeps it. + """ + handled = _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"] + assert expected_cost > 0 + assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost) + assert handled["kwargs"]["model"] == turn_model + + @pytest.mark.parametrize( ("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 3d53a8d7b67..ac42491c389 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 @@ -47,6 +47,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( InitPassThroughEndpointHelpers, SafeRouteAdder, WebsocketFrameModelAllowed, + WebsocketFrameModelGate, WebsocketFrameModelRefused, WebsocketFrameModelVerdict, _registered_pass_through_routes, @@ -7651,13 +7652,17 @@ class RecordingFrameModelGate: return WebsocketFrameModelAllowed(access_groups=self.groups.get(model, ())) -def _received_frame(frame): +def _received_frame(frame: str | bytes) -> Mapping[str, str | bytes]: if isinstance(frame, bytes): return {"type": "websocket.receive", "bytes": frame} return {"type": "websocket.receive", "text": frame} -async def _relay_client_frames(frames, gate, user_api_key_dict=None): +async def _relay_client_frames( + frames: Sequence[str | bytes], + gate: WebsocketFrameModelGate | None, + user_api_key_dict: UserAPIKeyAuth | None = None, +) -> SimpleNamespace: upstream_ws = RecordingUpstreamWebSocket() websocket = _client_websocket( AsyncMock(side_effect=[_received_frame(frame) for frame in frames] + [{"type": "websocket.disconnect"}]) @@ -7907,6 +7912,53 @@ 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 = { + "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) + + assert gate.checked == ["gpt-5.4-mini", "gpt-image-2"] + assert relayed.upstream.send.await_args_list == [] + assert relayed.client.close.await_args.kwargs["code"] == 1008 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "frame", + [ + {"type": "response.create", "model": "gpt-5.4-mini", "tools": [{"type": "image_generation"}]}, + {"type": "response.create", "response": {"tools": [{"type": "image_generation", "model": None}]}}, + ], +) +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})) + + relayed = await _relay_client_frames([json.dumps(frame)], gate) + + assert gate.checked[-1] == STORED_PROMPT_PROBE_MODEL + assert relayed.upstream.send.await_args_list == [] + assert "image_generation tool has to name its model" in relayed.client.close.await_args.kwargs["reason"] + + +@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"}}' + + relayed = await _relay_client_frames([frame], gate) + + assert gate.checked == ["gpt-5.4"] + assert relayed.upstream.send.await_args_list == [] + + @pytest.mark.asyncio async def test_websocket_passthrough_forwards_binary_audio_untouched(): gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": "denied"})) @@ -8015,14 +8067,16 @@ class ScriptedUpstreamWebSocket: return frame -def _turn_taking_client(frames: Sequence[str], answered_by: str) -> MagicMock: - """A client that sends its next frame only once the answer to the last one reached it, and that Starlette - refuses to send to once it is closed""" +def _turn_taking_client( + frames: Sequence[str], answered_by: str, answers_due: Sequence[int] | None = None +) -> MagicMock: + """A client that sends each frame, and then leaves, once the answers due by then reached it (by default the + answer to every earlier frame), and that Starlette refuses to send to once it is closed""" from starlette.websockets import WebSocketState delivered: list[str] = [] script = iter(frames) - receives = iter(range(len(frames) + 1)) + receives = iter(range(len(frames) + 1) if answers_due is None else answers_due) async def receive(): answers_due = next(receives) @@ -8042,7 +8096,14 @@ def _turn_taking_client(frames: Sequence[str], answered_by: str) -> MagicMock: return websocket -async def _relay_turns(client, upstream_ws, gate, endpoint="/openai_passthrough/v1/responses"): +async def _relay_turns( + client: MagicMock, + upstream_ws: ScriptedUpstreamWebSocket, + gate: WebsocketFrameModelGate, + endpoint: str = "/openai_passthrough/v1/responses", + target: str = "wss://api.openai.com/v1/responses", + opening_access_groups: Sequence[str] = (), +) -> SimpleNamespace: with ( _patched_websocket_passthrough_environment(upstream_ws) as proxy_logging, patch(SUCCESS_HANDLER, new_callable=AsyncMock) as success_handler, @@ -8050,9 +8111,11 @@ async def _relay_turns(client, upstream_ws, gate, endpoint="/openai_passthrough/ await asyncio.wait_for( websocket_passthrough_request( websocket=client, - target="wss://api.openai.com/v1/responses", + target=target, custom_headers={"Authorization": "Bearer sk-provider"}, - user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + user_api_key_dict=UserAPIKeyAuth( + api_key="hashed-key", matched_model_access_groups=list(opening_access_groups) + ), forward_headers=False, endpoint=endpoint, accept_websocket=False, @@ -8065,15 +8128,23 @@ async def _relay_turns(client, upstream_ws, gate, endpoint="/openai_passthrough/ def _responses_turn(model: str) -> list[str]: - return [ - json.dumps({"type": "response.created", "response": {"model": model}}), + return [_responses_turn_started(model), _responses_turn_ended(model)] + + +def _responses_turn_started(model: str) -> str: + return json.dumps({"type": "response.created", "response": {"model": model}}) + + +def _responses_turn_ended(model: str) -> str: + return ( + json.dumps( { "type": "response.completed", "response": {"model": model, "usage": {"input_tokens": 9, "output_tokens": 3, "total_tokens": 12}}, } - ), - ] + ) + ) @pytest.mark.asyncio @@ -8113,6 +8184,119 @@ async def test_websocket_passthrough_logs_each_turn_on_its_own_row_charged_to_th relayed.proxy_logging.post_call_failure_hook.assert_not_awaited() +@pytest.mark.asyncio +async def test_websocket_passthrough_charges_pipelined_turns_to_the_frame_that_asked_for_each(): + """ + A client can send its next response.create before the provider finished the + last one, and the provider then runs them one after the other. Each turn + still has to reach spend under the groups that let its own frame through, + 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( + groups=MappingProxyType({"gpt-5.4-mini": ("tier-a",), "gpt-5.4": ("tier-b",)}), + ) + client = _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"', + answers_due=(0, 0, 2), + ) + upstream_ws = 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) + + rows = [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"], + ] + assert [row["litellm_params"]["metadata"][MODEL_ACCESS_GROUP_METADATA_KEY] for row in rows] == [ + ["tier-a"], + ["tier-b"], + ] + assert [row["litellm_params"]["metadata"]["model_group"] for row in rows] == ["gpt-5.4-mini", "gpt-5.4"] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_charges_realtime_turns_to_the_session_the_key_opened(): + """ + A realtime session runs the model it was opened on, so its turns stay + 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( + [ + json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1-mini"}}), + json.dumps({"type": "response.create", "response": {"model": "gpt-realtime-2.1-mini"}}), + ], + answered_by='"type"', + answers_due=(1, 2, 4), + ) + upstream_ws = ScriptedUpstreamWebSocket( + REALTIME_SESSION_CREATED, + replies=[ + [json.dumps({"type": "session.updated"})], + [json.dumps({"type": "response.created", "response": {"status": "in_progress"}}), REALTIME_RESPONSE_DONE], + ], + ) + + relayed = await _relay_turns( + client, + upstream_ws, + gate, + "/openai/v1/realtime", + target="wss://api.openai.com/v1/realtime?model=gpt-realtime-2.1", + opening_access_groups=("tier-a",), + ) + + 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] + 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"] + + +@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"') + + relayed = await _relay_turns(client, ScriptedUpstreamWebSocket(), gate) + + assert relayed.success_handler.call_args_list == [] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_logs_the_unfinished_turn_a_client_walked_out_of(): + """ + 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( + [ + 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( + replies=[_responses_turn("gpt-5.4-mini"), [_responses_turn_started("gpt-5.4")]] + ) + + relayed = 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"]] + for logged in relayed.success_handler.call_args_list + ] == [["response.created", "response.completed"], ["response.created"]] + + @pytest.mark.asyncio async def test_websocket_passthrough_carries_the_session_the_provider_opened_into_every_later_turn(): """ @@ -9095,7 +9279,7 @@ def test_passthrough_client_cannot_forge_session_id_omission(client_metadata_key ) -async def _billed_metadata(gate) -> dict: +async def _billed_metadata(gate: WebsocketFrameModelGate) -> Mapping[str, object]: served = await _relay_client_frames( [json.dumps({"type": "response.create", "model": "gpt-5.4-mini"})], gate, 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 5b3f337fb5e..28ad1795611 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_success_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/test_success_handler.py @@ -667,7 +667,7 @@ REALTIME_FRAMES: Final = [ ] -def _normalize(url_route: str, response_body): +def _normalize(url_route: str, response_body: dict[str, object] | list[dict[str, object]]) -> dict: logging_obj = MagicMock() logging_obj.model_call_details = {} return PassThroughEndpointLogging().normalize_llm_passthrough_logging_payload( diff --git a/tests/unit/proxy/test_openai_ws_passthrough_routes.py b/tests/unit/proxy/test_openai_ws_passthrough_routes.py index 836104e3342..89478a852f3 100644 --- a/tests/unit/proxy/test_openai_ws_passthrough_routes.py +++ b/tests/unit/proxy/test_openai_ws_passthrough_routes.py @@ -18,6 +18,7 @@ import litellm from litellm.llms.openai.workload_identity import _workload_identity_auth from litellm.proxy._types import Litellm_EntityType, ProxyErrorTypes, ProxyException, UserAPIKeyAuth from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + _OPENAI_WS_ALIASED_MODEL_REFUSAL, _OPENAI_WS_DISABLED_REFUSAL, _OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL, _openai_websocket_refusal, @@ -173,6 +174,18 @@ async def test_openai_websocket_forwards_query_and_keeps_provider_auth(prefix, m assert websocket.closed is None +@pytest.mark.asyncio +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", "") + + with patch(GET_CREDENTIALS, return_value="sk-provider"): + served = await _serve(websocket, "v1/responses", UserAPIKeyAuth(), ENABLED) + + assert [call.endpoint for call in served.relay.calls] == ["/openai/v1/responses"] + + @pytest.mark.asyncio async def test_openai_websocket_accepts_first_client_subprotocol(): websocket = _FakeWebSocket( @@ -390,6 +403,48 @@ async def test_proxy_frame_model_gate_refuses_a_model_the_agent_was_not_granted( assert refusal_names in refusal.message +@pytest.mark.asyncio +@pytest.mark.parametrize( + "token", + [ + UserAPIKeyAuth(models=["gpt-5.4-mini"], aliases={"gpt-5.4": "gpt-5.4-mini"}), + UserAPIKeyAuth(models=["gpt-5.4-mini"], team_model_aliases={"gpt-5.4": "gpt-5.4-mini"}), + ], + ids=["key_alias", "team_alias"], +) +async def test_proxy_frame_model_gate_refuses_an_alias_that_would_reach_openai_unchanged(token): + """ + The passthrough sends each frame to OpenAI as written, so OpenAI runs the + alias name itself. An alias pointing at an allowed model must not let a key + limited to that model run the model the alias is named after. + """ + refusal: Final = _refusal(await _gate_verdict(_FrameAuthorizer(token), "gpt-5.4")) + + assert refusal == WebsocketFrameModelRefused( + message=_OPENAI_WS_ALIASED_MODEL_REFUSAL, error_type=ProxyErrorTypes.key_model_access_denied.value + ) + + +@pytest.mark.asyncio +async def test_proxy_frame_model_gate_refuses_a_proxy_wide_alias_for_a_key_limited_to_its_target(monkeypatch): + monkeypatch.setattr(litellm, "model_alias_map", {"gpt-5.4": "gpt-5.4-mini"}) + + refusal: Final = _refusal( + await _gate_verdict(_FrameAuthorizer(UserAPIKeyAuth(models=["gpt-5.4-mini"])), "gpt-5.4") + ) + + assert refusal.message == _OPENAI_WS_ALIASED_MODEL_REFUSAL + + +@pytest.mark.asyncio +async def test_proxy_frame_model_gate_lets_a_key_that_may_use_any_model_send_an_alias(): + token: Final = UserAPIKeyAuth(models=[], aliases={"fast": "gpt-5.4-mini"}) + + verdict: Final = await _gate_verdict(_FrameAuthorizer(token), "fast") + + assert verdict == WebsocketFrameModelAllowed(access_groups=()) + + @pytest.mark.asyncio @pytest.mark.parametrize("unrestricted_models", [[], ["all-proxy-models"], ["*"]]) async def test_proxy_frame_model_gate_allows_every_model_for_unrestricted_keys(unrestricted_models): 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 06/15] 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(): """ From ad4a71c2b90e992c0433ef827e62f346acfb6b2e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 10:56:52 -0700 Subject: [PATCH 07/15] fix(proxy): price websocket Responses turns by served tier and hosted tools, and bill turns a client walked out of A websocket Responses turn is now priced through the same Responses cost path the HTTP passthrough uses, so a priority-tier answer and a hosted web search cost what they cost over HTTP. When the client leaves mid-answer, the gateway keeps reading the provider's socket until that turn ends, up to 10 minutes, so the turn the provider still bills reaches spend --- ...i_websocket_passthrough_logging_handler.py | 24 ++--- .../pass_through_endpoints.py | 42 +++++++-- ...i_websocket_passthrough_logging_handler.py | 52 ++++++++++- .../test_pass_through_endpoints.py | 93 +++++++++++++++++++ 4 files changed, 182 insertions(+), 29 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 f2b23d599f9..71c35e62ac0 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 @@ -13,7 +13,7 @@ from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import PassThroughEndpointLoggingTypedDict from litellm.responses.utils import ResponseAPILoggingUtils -from litellm.types.llms.openai import OpenAIRealtimeStreamList, ResponseAPIUsage +from litellm.types.llms.openai import OpenAIRealtimeStreamList, ResponsesAPIResponse from litellm.types.passthrough_endpoints.pass_through_endpoints import WebsocketUpstreamFrameKind from litellm.types.utils import LlmProviders, ModelResponse, Usage @@ -43,18 +43,11 @@ class _ModelNamingEvent(BaseModel): response: _NamedModel | None = None -class _EndedResponse(BaseModel): - model_config = ConfigDict(extra="ignore", frozen=True) - - model: str | None = None - usage: ResponseAPIUsage | None = None - - class _ResponsesTurnEndEvent(BaseModel): model_config = ConfigDict(extra="ignore", frozen=True) type: Literal["response.completed", "response.incomplete", "response.failed"] - response: _EndedResponse + response: ResponsesAPIResponse _MODEL_NAMING_EVENT: Final = TypeAdapter(_ModelNamingEvent) @@ -70,6 +63,7 @@ class _Billing: @dataclass(frozen=True, slots=True) class _CompletedTurn: model: str + response: ResponsesAPIResponse usage: Usage @@ -136,6 +130,7 @@ def _responses_turn(event: Mapping[str, object], session_model: str) -> _Complet return None return _CompletedTurn( model=ended.response.model or session_model, + response=ended.response, usage=ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( # pyright: ignore[reportPrivateUsage] # the one shared Responses usage converter ended.response.usage ), @@ -149,23 +144,24 @@ def _responses_billing(messages: WebsocketMessages, model: str, requested_model: usage: Final = RealtimeAPITokenUsageProcessor.combine_usage_objects([turn.usage for turn in turns]) return _Billing( usage=usage, - cost=sum(_responses_cost(turn.usage, (turn.model, requested_model)) for turn in turns), + cost=sum(_responses_cost(turn.response, (turn.model, requested_model)) for turn in turns), ) -def _model_cost(usage: Usage, model: str) -> float | None: +def _model_cost(response: ResponsesAPIResponse, model: str) -> float | None: try: return litellm.completion_cost( - completion_response=ModelResponse(model=model, usage=usage), + completion_response=response, model=model, custom_llm_provider=OPENAI_WEBSOCKET_PROVIDER, + call_type="responses", ) except Exception: # noqa: BLE001 # completion_cost raises a bare Exception for an unpriced model, and the turn must still log return None -def _responses_cost(usage: Usage, price_as: Sequence[str | None]) -> float: - costs: Final = (_model_cost(usage, model) for model in dict.fromkeys(price_as) if model is not None) +def _responses_cost(response: ResponsesAPIResponse, price_as: Sequence[str | None]) -> float: + costs: Final = (_model_cost(response, model) for model in dict.fromkeys(price_as) if model is not None) cost: Final = next((cost for cost in costs if cost), None) if cost is None: verbose_proxy_logger.warning( diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 393b193c73f..5b2789eded9 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -2743,6 +2743,14 @@ def _session_end_needs_a_row( return "starts_turn" in kinds or (ledger.logged_turns == 0 and bool(kinds - {"opens_session"})) +def _turn_in_flight( + upstream_frame_kind: WebsocketUpstreamFrameClassifier | None, unlogged_frames: Sequence[Mapping[str, object]] +) -> bool: + return upstream_frame_kind is not None and any( + upstream_frame_kind(frame) == "starts_turn" for frame in unlogged_frames + ) + + async def _relay_client_frame( websocket: WebSocket, upstream_ws: ClientConnection, @@ -2765,6 +2773,7 @@ async def _relay_client_frame( _EXPLICIT_SUCCESS_KWARGS: Final = frozenset({"logging_obj", "passthrough_logging_payload", "custom_llm_provider"}) +WALKOUT_TURN_DRAIN_SECONDS: Final = 600.0 async def websocket_passthrough_request( @@ -2779,6 +2788,7 @@ async def websocket_passthrough_request( setup_model_rewriter: Callable[[str], str] | None = None, client_frame_model_gate: WebsocketFrameModelGate | None = None, upstream_frame_kind: WebsocketUpstreamFrameClassifier | None = None, + walkout_turn_drain_seconds: float = WALKOUT_TURN_DRAIN_SECONDS, ): """ WebSocket passthrough request handler. @@ -2795,6 +2805,8 @@ async def websocket_passthrough_request( client_frame_model_gate: Optional per-frame check of every model a client frame names upstream_frame_kind: Optional sort of upstream frames into the one that ends a billable turn and the ones that open the session, so each turn logs its own spend and keeps the session's provider-named model + walkout_turn_drain_seconds: How long to keep reading a turn the provider is still running after the client + left, so the provider's usage for it still reaches spend """ from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy.proxy_server import proxy_config, proxy_logging_obj @@ -3019,7 +3031,8 @@ async def websocket_passthrough_request( message = await websocket.receive() message_type = message.get("type") if message_type == "websocket.disconnect": - await upstream_ws.close() + if not _turn_in_flight(upstream_frame_kind, websocket_messages): + await upstream_ws.close() break text_data: str | None = message.get("text") @@ -3147,11 +3160,17 @@ async def websocket_passthrough_request( async def relay_upstream_frame(upstream_message: str | bytes) -> None: record_upstream_frame(upstream_message) if not _client_socket_is_open(websocket): + if not _turn_in_flight(upstream_frame_kind, websocket_messages): + await upstream_ws.close() return - if isinstance(upstream_message, bytes): - await websocket.send_bytes(upstream_message) - else: - await websocket.send_text(upstream_message) + try: + if isinstance(upstream_message, bytes): + await websocket.send_bytes(upstream_message) + else: + await websocket.send_text(upstream_message) + except WebSocketDisconnect: + if not _turn_in_flight(upstream_frame_kind, websocket_messages): + raise async def forward_upstream_to_client() -> Close | None: try: @@ -3176,12 +3195,15 @@ async def websocket_passthrough_request( raise # Create tasks for bidirectional message forwarding - tasks: Final = [ - asyncio.create_task(forward_client_to_upstream()), - asyncio.create_task(forward_upstream_to_client()), - ] + client_task: Final = asyncio.create_task(forward_client_to_upstream()) + upstream_task: Final = asyncio.create_task(forward_upstream_to_client()) + tasks: Final = (client_task, upstream_task) - done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) + await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) + if client_task.done() and _turn_in_flight(upstream_frame_kind, websocket_messages): + await asyncio.wait((upstream_task,), timeout=walkout_turn_drain_seconds) + done: Final = tuple(task for task in tasks if task.done()) + pending: Final = tuple(task for task in tasks if not task.done()) # Cancel remaining tasks for task in pending: 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 7494c913a4d..3390b426f02 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 @@ -4,9 +4,13 @@ from types import MappingProxyType from typing import Final from unittest.mock import MagicMock +import httpx import pytest import litellm +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import ( + OpenAIPassthroughLoggingHandler, +) from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_websocket_passthrough_logging_handler import ( OpenAIWebsocketPassthroughLoggingHandler, openai_websocket_frame_kind, @@ -45,14 +49,22 @@ def _realtime_response_done(text_in: int, audio_in: int, text_out: int, audio_ou def _responses_completed( - model: str, input_tokens: int, output_tokens: int, event: str = "response.completed" + model: str, + input_tokens: int, + output_tokens: int, + event: str = "response.completed", + served: Mapping[str, object] = MappingProxyType({}), ) -> dict[str, object]: return { "type": event, "response": { "id": "resp_fake", + "object": "response", + "created_at": 1761700000, "model": model, "status": "completed", + "output": [], + **served, "usage": { "input_tokens": input_tokens, "output_tokens": output_tokens, @@ -201,10 +213,7 @@ def test_a_responses_turn_that_names_no_model_is_priced_on_the_session_model(): handled = _handle( [ _session_created(RESPONSES_MODEL), - { - "type": "response.completed", - "response": {"usage": {"input_tokens": 13, "output_tokens": 5, "total_tokens": 18}}, - }, + _responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5, served={"model": None}), ], url_route="/openai_passthrough/v1/responses", ) @@ -357,6 +366,39 @@ def test_a_responses_turn_falls_back_to_the_requested_model_only_when_its_own_ha assert handled["kwargs"]["model"] == turn_model +def _http_passthrough_cost(response: Mapping[str, object]) -> float: + http_response = httpx.Response( + 200, json=response, request=httpx.Request("POST", "https://api.openai.com/v1/responses") + ) + _, cost = OpenAIPassthroughLoggingHandler._build_responses_api_response_and_cost( + RESPONSES_MODEL, http_response, _logging_obj(), "openai" + ) + return cost + + +@pytest.mark.parametrize( + "served", + [ + {"service_tier": "priority"}, + {"output": [{"type": "web_search_call", "id": "ws_fake", "status": "completed", "action": {"type": "search"}}]}, + ], + ids=["priority_tier", "hosted_web_search"], +) +def test_a_responses_turn_costs_what_the_http_passthrough_charges_for_the_same_response(served): + """ + The provider bills a turn by the tier that served it and the hosted tools it + 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) + + handled = _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"]) + + @pytest.mark.parametrize( ("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 c4b08685d74..bae193d0da2 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 @@ -8103,6 +8103,7 @@ async def _relay_turns( endpoint: str = "/openai_passthrough/v1/responses", target: str = "wss://api.openai.com/v1/responses", opening_access_groups: Sequence[str] = (), + walkout_turn_drain_seconds: float = 0.05, ) -> SimpleNamespace: with ( _patched_websocket_passthrough_environment(upstream_ws) as proxy_logging, @@ -8121,6 +8122,7 @@ async def _relay_turns( accept_websocket=False, client_frame_model_gate=gate, upstream_frame_kind=openai_websocket_frame_kind, + walkout_turn_drain_seconds=walkout_turn_drain_seconds, ), timeout=10, ) @@ -8297,6 +8299,97 @@ async def test_websocket_passthrough_logs_the_unfinished_turn_a_client_walked_ou ] == [["response.created", "response.completed"], ["response.created"]] +class UpstreamFinishingTheTurnAfterTheClientLeft: + """Starts a turn for each client frame and sends the rest of it only after the client left, and stops once closed""" + + def __init__(self, client_left: asyncio.Event, started: str, rest_of_the_turn: Sequence[str]): + self._frames: asyncio.Queue[str | None] = asyncio.Queue() + self._client_left = client_left + self._started = started + self._rest_of_the_turn = rest_of_the_turn + self._closed = False + self._finishing: list[asyncio.Task[None]] = [] + self.send = AsyncMock(side_effect=self._start_turn) + self.close = AsyncMock(side_effect=self._close) + + def _close(self, *args: object, **kwargs: object) -> None: + self._closed = True + self._frames.put_nowait(None) + + async def _start_turn(self, frame: str | bytes) -> None: + self._frames.put_nowait(self._started) + self._finishing.append(asyncio.create_task(self._finish_once_the_client_left())) + + async def _finish_once_the_client_left(self) -> None: + await self._client_left.wait() + for frame in self._rest_of_the_turn: + await asyncio.sleep(0) + if self._closed: + return + self._frames.put_nowait(frame) + + async def recv(self, decode: bool | None = None) -> str: + from websockets.exceptions import ConnectionClosedOK + from websockets.frames import Close + + frame = await self._frames.get() + if frame is None: + raise ConnectionClosedOK(rcvd=Close(1000, ""), sent=Close(1000, ""), rcvd_then_sent=True) + return frame + + +@pytest.mark.asyncio +@pytest.mark.parametrize("hang_up_read_before_the_next_send", [True, False], ids=["read_first", "found_by_a_send"]) +async def test_websocket_passthrough_bills_the_turn_a_client_walked_out_of_once_the_provider_ends_it( + hang_up_read_before_the_next_send: bool, +): + """ + Leaving mid-answer must not make a turn free. The provider keeps running and + bills it, so the gateway keeps reading until the turn ends and logs its usage, + whether it reads the hang-up first or finds it by failing to relay a frame. + """ + from starlette.websockets import WebSocketDisconnect, WebSocketState + + client_left = asyncio.Event() + send_failed = 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 + + async def receive_until_leaving() -> Mapping[str, object]: + message = await receive() + if message["type"] == "websocket.disconnect": + client_left.set() + if not hang_up_read_before_the_next_send: + await send_failed.wait() + client.client_state = WebSocketState.DISCONNECTED + return message + + async def send_text_until_leaving(text: str) -> None: + if client_left.is_set(): + send_failed.set() + raise WebSocketDisconnect(code=1006) + await send_text(text) + + client.receive = receive_until_leaving + client.send_text = AsyncMock(side_effect=send_text_until_leaving) + upstream_ws = 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) + + assert [ + [json.loads(frame)["type"] if isinstance(frame, str) else frame["type"] for frame in logged.kwargs["response_body"]] + for logged in relayed.success_handler.call_args_list + ] == [["response.created", "response.output_text.delta", "response.completed"]] + upstream_ws.close.assert_awaited() + + @pytest.mark.asyncio async def test_websocket_passthrough_logs_no_extra_row_for_what_the_provider_sends_after_the_last_turn(): """ From 69311a55fc1913e0f762d63216e5cdd542ba9601 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 15:44:55 -0700 Subject: [PATCH 08/15] fix(proxy): bill websocket turns by lane, skip turns OpenAI refuses, and stop frame checks counting the session's own reservation An error event refuses a queued response.create, so the next turn is no longer charged to the refused frame's model. Turns on separate stream_id lanes are charged to the frame that started each one. The frame model gate settles the connect reservation before it re-runs auth, so a key with less budget left than one turn's max cost is no longer refused on its first frame. --- .../llm_passthrough_endpoints.py | 6 +- ...i_websocket_passthrough_logging_handler.py | 3 + .../pass_through_endpoints.py | 100 +++++++++++++----- .../pass_through_endpoints.py | 2 +- ...i_websocket_passthrough_logging_handler.py | 2 +- .../test_pass_through_endpoints.py | 91 ++++++++++++++++ .../test_openai_ws_passthrough_routes.py | 23 ++++ 7 files changed, 197 insertions(+), 30 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 80a9f0e8c79..a7d7e031d1b 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -3401,7 +3401,7 @@ def _gateway_remaps_model(model: str, frame_auth: UserAPIKeyAuth, llm_router: Ro ) targets: Final = ( *(aliases.get(model) for aliases in alias_maps), - llm_router.get_model_from_alias(model) if llm_router is not None else None, + llm_router.registry().resolve(model, kind="alias") if llm_router is not None else None, ) return any(target not in (None, model) for target in targets) @@ -3425,6 +3425,10 @@ def _proxy_frame_model_gate( from litellm.proxy.spend_tracking.budget_reservation import release_or_invalidate_budget_reservation async def authorize(websocket: WebSocket, model: str) -> WebsocketFrameModelVerdict: + connect_reservation: Final = attribute_or_none(websocket.state, "budget_reservation") + await release_or_invalidate_budget_reservation( + budget_reservation=connect_reservation if isinstance(connect_reservation, dict) else None + ) frame_auth: Final = await authorize_frame_model(websocket, model) access_groups: Final = tuple(frame_auth.matched_model_access_groups or ()) llm_router: Final = proxy_server.llm_router 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 71c35e62ac0..27abb333a36 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 @@ -23,6 +23,7 @@ OPENAI_WEBSOCKET_PROVIDER: Final = LlmProviders.OPENAI.value _SESSION_CREATED: Final = "session.created" _RESPONSE_CREATED: Final = "response.created" +_ERROR: Final = "error" _REALTIME_RESPONSE_DONE: Final = "response.done" _RESPONSES_TURN_ENDS: Final = frozenset({"response.completed", "response.incomplete", "response.failed"}) _TURN_ENDS: Final = _RESPONSES_TURN_ENDS | {_REALTIME_RESPONSE_DONE} @@ -96,6 +97,8 @@ def openai_websocket_frame_kind(frame: Mapping[str, object]) -> WebsocketUpstrea return "opens_session" if frame_type == _RESPONSE_CREATED: return "starts_turn" + if frame_type == _ERROR: + return "refuses_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 5fdceb3d773..d64bf3d58e0 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -2549,6 +2549,7 @@ class _TurnAuthorization: class _ForwardClientFrame: starts_turn: bool turn: _TurnAuthorization | None + lane: str | None @dataclass(frozen=True, slots=True) @@ -2559,17 +2560,18 @@ class _RefuseClientFrame: @dataclass(frozen=True, slots=True) class _ClientFrameRequests: starts_turn: bool + lane: str | None turn_model: str | None models: tuple[str, ...] unverifiable_model_refusal: str | None _NO_CLIENT_FRAME_REQUESTS: Final = _ClientFrameRequests( - starts_turn=False, turn_model=None, models=(), unverifiable_model_refusal=None + starts_turn=False, lane=None, turn_model=None, models=(), unverifiable_model_refusal=None ) -_FORWARD_WITHOUT_TURN: Final = _ForwardClientFrame(starts_turn=False, turn=None) +_FORWARD_WITHOUT_TURN: Final = _ForwardClientFrame(starts_turn=False, turn=None, lane=None) _TURN_STARTING_FRAME: Final = "response.create" -_FRAME_TEXT_WORTH_PARSING: Final = ('"model"', '"prompt"', '"tools"', "\\") +_FRAME_TEXT_WORTH_PARSING: Final = (f'"{_TURN_STARTING_FRAME}"', '"model"', '"prompt"', '"tools"', "\\") _TURN_MODEL_PLACEMENTS: Final = (("model",), ("response", "model"), ("session", "model")) _TRANSCRIPTION_MODEL_PLACEMENTS: Final = ( ("session", "input_audio_transcription", "model"), @@ -2594,6 +2596,22 @@ def _may_name_a_model(text_data: str) -> bool: return any(marker in text_data for marker in _FRAME_TEXT_WORTH_PARSING) +def _frame_lane(frame: Mapping[str, object]) -> str | None: + stream_id: Final = frame.get("stream_id") + return stream_id if isinstance(stream_id, str) else None + + +def _frames_of_lane(lane: str | None) -> Callable[[Mapping[str, object]], bool]: + def in_lane(frame: Mapping[str, object]) -> bool: + return _frame_lane(frame) == lane + + return in_lane + + +def _every_frame(_frame: Mapping[str, object]) -> bool: + return True + + def _nested_frame(frame: Mapping[str, object] | None, key: str) -> Mapping[str, object] | None: return None if frame is None else as_str_mapping(frame.get(key)) @@ -2635,6 +2653,7 @@ def _client_frame_requests(frame_data: str | bytes) -> _ClientFrameRequests: tool_models: Final = (model for model in (tool.get("model") for tool in tools) if isinstance(model, str)) return _ClientFrameRequests( starts_turn=frame.get("type") == _TURN_STARTING_FRAME, + lane=_frame_lane(frame), turn_model=next(iter(turn_models), None), models=tuple( dict.fromkeys((*turn_models, *_placed_models(frame, _TRANSCRIPTION_MODEL_PLACEMENTS), *tool_models)) @@ -2666,6 +2685,7 @@ async def _client_frame_verdict( allowed: Final = tuple(verdict for verdict in verdicts if isinstance(verdict, WebsocketFrameModelAllowed)) return _ForwardClientFrame( starts_turn=requests.starts_turn, + lane=requests.lane, turn=( _TurnAuthorization(requested_model=requests.turn_model, access_groups=allowed[0].access_groups) if requests.turn_model is not None and allowed @@ -2704,37 +2724,45 @@ def _opening_model(target: str) -> str | None: class _WebsocketTurnLedger: - """The groups that authorized each turn not yet logged, when the next row starts, and how many rows were logged. + """The groups that authorized each turn not yet logged in each lane, when the next row starts, and how many rows + were logged. - The provider runs a socket's response.create frames one at a time in the order they were sent, so each row is - charged to the groups that allowed the oldest turn still unlogged. A realtime session keeps the model it opened on, - so its rows stay charged to the groups that opened it. + The provider runs the response.create frames of one lane (one stream_id, or none) one at a time in the order they + were sent, and runs separate lanes side by side, so each row is charged to the groups that allowed the oldest turn + still unlogged in its lane. A turn the provider refuses with an error event never runs, so it leaves its lane + unlogged. A realtime session keeps the model it opened on, so its rows stay charged to the groups that opened it. """ def __init__(self, opening: _TurnAuthorization, started: datetime) -> None: self.opening = opening - self.queued_turns: tuple[_TurnAuthorization | None, ...] = () + self.queued_turns: Mapping[str | None, tuple[_TurnAuthorization | None, ...]] = MappingProxyType({}) self.session_opened = False self.started = started self.logged_turns = 0 - @property - def current(self) -> _TurnAuthorization: - oldest: Final = self.queued_turns[0] if self.queued_turns else None + def current(self, lane: str | None) -> _TurnAuthorization: + oldest: Final = next(iter(self.queued_turns.get(lane, ())), None) return self.opening if oldest is None else oldest - def queue_turn(self, turn: _TurnAuthorization | None) -> None: + def queue_turn(self, lane: str | None, turn: _TurnAuthorization | None) -> None: if not self.session_opened: - self.queued_turns = (*self.queued_turns, turn) + self.queued_turns = MappingProxyType({**self.queued_turns, lane: (*self.queued_turns.get(lane, ()), turn)}) def open_session(self) -> None: self.session_opened = True - self.queued_turns = () + self.queued_turns = MappingProxyType({}) - def close_turn(self, end_time: datetime) -> None: + def close_turn(self, lane: str | None, end_time: datetime) -> None: self.logged_turns += 1 self.started = end_time - self.queued_turns = self.queued_turns[1:] + self._forget_turn(lane, 0) + + def drop_refused_turn(self, lane: str | None, lane_is_running_a_turn: bool) -> None: + self._forget_turn(lane, 1 if lane_is_running_a_turn else 0) + + def _forget_turn(self, lane: str | None, position: int) -> None: + turns: Final = self.queued_turns.get(lane, ()) + self.queued_turns = MappingProxyType({**self.queued_turns, lane: (*turns[:position], *turns[position + 1 :])}) def _session_end_needs_a_row( @@ -2758,6 +2786,13 @@ def _turn_in_flight( ) +def _lanes_in_flight( + upstream_frame_kind: WebsocketUpstreamFrameClassifier, unlogged_frames: Sequence[Mapping[str, object]] +) -> tuple[str | None, ...]: + started: Final = (frame for frame in unlogged_frames if upstream_frame_kind(frame) == "starts_turn") + return tuple(dict.fromkeys(map(_frame_lane, started))) + + async def _relay_client_frame( websocket: WebSocket, upstream_ws: ClientConnection, @@ -2771,8 +2806,8 @@ async def _relay_client_frame( match verdict: case _RefuseClientFrame(refusal=refusal): await _refuse_client_frame(websocket, upstream_ws, refusal) - case _ForwardClientFrame(starts_turn=True, turn=turn): - ledger.queue_turn(turn) + case _ForwardClientFrame(starts_turn=True, turn=turn, lane=lane): + ledger.queue_turn(lane, turn) await upstream_ws.send(forwarded_frame) case _ForwardClientFrame(): await upstream_ws.send(forwarded_frame) @@ -2957,12 +2992,15 @@ async def websocket_passthrough_request( call_payload: PassthroughStandardLoggingPayload, call_kwargs: Mapping[str, object], end_time: datetime, + lane: str | None, + in_row: Callable[[Mapping[str, object]], bool], ) -> None: - frames: Final = list(websocket_messages) + all_frames: Final = list(websocket_messages) + frames: Final = [frame for frame in all_frames if in_row(frame)] websocket_messages[:] = ( () if upstream_frame_kind is None - else (frame for frame in frames if upstream_frame_kind(frame) == "opens_session") + else (frame for frame in all_frames if not in_row(frame) or upstream_frame_kind(frame) == "opens_session") ) logged_payload: Final[PassthroughStandardLoggingPayload] = { **call_payload, @@ -2970,7 +3008,7 @@ async def websocket_passthrough_request( "end_time": end_time, } if client_frame_model_gate is not None: - _stamp_billed_turn(call_kwargs, ledger.current) + _stamp_billed_turn(call_kwargs, ledger.current(lane)) success_kwargs: Final = { key: value for key, value in call_kwargs.items() if key not in _EXPLICIT_SUCCESS_KWARGS } @@ -2993,10 +3031,10 @@ async def websocket_passthrough_request( ) bind_budget_reservation_to_callbacks(call_logging.litellm_params) - def log_turn() -> None: + def log_turn(lane: str | None, in_row: Callable[[Mapping[str, object]], bool]) -> None: end_time: Final = datetime.now() if ledger.logged_turns == 0: - log_frames(logging_obj, passthrough_logging_payload, kwargs, end_time) + log_frames(logging_obj, passthrough_logging_payload, kwargs, end_time, lane, in_row) else: turn_call_id: Final = str(uuid.uuid4()) turn_logging: Final = new_call_logging(ledger.started, turn_call_id) @@ -3008,8 +3046,8 @@ async def websocket_passthrough_request( turn_logging, ) start_call_logging(turn_logging, turn_kwargs) - log_frames(turn_logging, turn_payload, turn_kwargs, end_time) - ledger.close_turn(end_time) + log_frames(turn_logging, turn_payload, turn_kwargs, end_time, lane, in_row) + ledger.close_turn(lane, end_time) ### CALL HOOKS ### - modify incoming data / reject request before calling the model websocket_hook_data: Final = guardrail_request_data_with_streaming(MappingProxyType({}), is_streaming=True) @@ -3160,7 +3198,12 @@ async def websocket_passthrough_request( case "opens_session": ledger.open_session() case "ends_turn": - log_turn() + log_turn(_frame_lane(message_data), _frames_of_lane(_frame_lane(message_data))) + case "refuses_turn" if upstream_frame_kind is not None: + lane: Final = _frame_lane(message_data) + ledger.drop_refused_turn( + lane, lane in _lanes_in_flight(upstream_frame_kind, websocket_messages) + ) case _: pass @@ -3234,7 +3277,10 @@ async def websocket_passthrough_request( ) if _session_end_needs_a_row(upstream_frame_kind, websocket_messages, ledger): - log_turn() + 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) # Call the proxy logging success hook if proxy_logging_obj: diff --git a/litellm/types/passthrough_endpoints/pass_through_endpoints.py b/litellm/types/passthrough_endpoints/pass_through_endpoints.py index 452f727c5ea..fe6a5e309d3 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["starts_turn", "ends_turn", "opens_session", "relayed"] +WebsocketUpstreamFrameKind: TypeAlias = Literal["starts_turn", "ends_turn", "opens_session", "refuses_turn", "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 3390b426f02..ceffbd1ba08 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 @@ -409,7 +409,7 @@ def test_a_responses_turn_costs_what_the_http_passthrough_charges_for_the_same_r ("session.created", "opens_session"), ("response.created", "starts_turn"), ("response.output_text.delta", "relayed"), - ("error", "relayed"), + ("error", "refuses_turn"), ], ) def test_each_frame_is_classified_by_what_it_does_to_a_turn(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 ba6e0154acf..bbc75972150 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 @@ -8224,6 +8224,97 @@ async def test_websocket_passthrough_charges_pipelined_turns_to_the_frame_that_a assert [row["litellm_params"]["metadata"]["model_group"] for row in rows] == ["gpt-5.4-mini", "gpt-5.4"] +def _responses_error(code: str) -> str: + return json.dumps({"type": "error", "status": 400, "error": {"type": "invalid_request_error", "code": code}}) + + +def _on_lane(frame: str, stream_id: str) -> str: + return json.dumps({**json.loads(frame), "stream_id": stream_id}) + + +@pytest.mark.asyncio +async def test_websocket_passthrough_charges_turns_after_a_refused_frame_to_the_frames_that_asked_for_them(): + """ + The provider answers a response.create it cannot run, one naming no model + or a stale previous_response_id, with an error event and keeps the socket + 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( + groups=MappingProxyType({"gpt-5.4-mini": ("tier-a",), "gpt-5.4": ("tier-b",)}), + ) + client = _turn_taking_client( + [ + json.dumps({"type": "response.create", "input": "hi"}), + json.dumps({"type": "response.create", "model": "gpt-5.4", "previous_response_id": "resp_gone"}), + json.dumps({"type": "response.create", "model": "gpt-5.4-mini"}), + json.dumps({"type": "response.create", "model": "gpt-5.4"}), + ], + answered_by='"response.completed"', + answers_due=(0, 0, 0, 0, 2), + ) + upstream_ws = ScriptedUpstreamWebSocket( + replies=[ + [_responses_error("missing_required_parameter")], + [_responses_error("previous_response_not_found")], + _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) + + rows = [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"], + ["tier-b"], + ] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_charges_turns_on_parallel_lanes_to_the_frame_that_asked_for_each(): + """ + A client can put each response.create on its own stream_id lane, and the + provider runs lanes side by side, so a later frame's turn can end first. + Each turn still has to reach spend on its own row, under the groups and + model its own frame asked for. + """ + gate = RecordingFrameModelGate( + groups=MappingProxyType({"gpt-5.4-mini": ("tier-a",), "gpt-5.4": ("tier-b",)}), + ) + client = _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.completed"', + answers_due=(0, 0, 2), + ) + upstream_ws = ScriptedUpstreamWebSocket( + replies=[ + [_on_lane(_responses_turn_started("gpt-5.4"), "slow")], + [ + *(_on_lane(frame, "fast") for frame in _responses_turn("gpt-5.4-mini")), + _on_lane(_responses_turn_ended("gpt-5.4"), "slow"), + ], + ] + ) + + relayed = await _relay_turns(client, upstream_ws, gate) + + rows = [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")], + ] + 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"], + ["tier-b"], + ] + + @pytest.mark.asyncio async def test_websocket_passthrough_charges_realtime_turns_to_the_session_the_key_opened(): """ diff --git a/tests/unit/proxy/test_openai_ws_passthrough_routes.py b/tests/unit/proxy/test_openai_ws_passthrough_routes.py index 89478a852f3..140cdb20914 100644 --- a/tests/unit/proxy/test_openai_ws_passthrough_routes.py +++ b/tests/unit/proxy/test_openai_ws_passthrough_routes.py @@ -483,6 +483,29 @@ async def test_proxy_frame_model_gate_releases_the_budget_it_reserved_for_the_ch assert token.budget_reservation == {**reservation, "finalized": True} +@pytest.mark.asyncio +async def test_proxy_frame_model_gate_does_not_count_the_sessions_own_reservation_as_key_spend(): + """ + The upgrade reserves up to the key's remaining budget for the session, and the key budget check + counts open reservations as spend, so a frame checked while that reservation is open is refused + on a key with budget left, where an HTTP call with the same key goes through. + """ + connect_reservation: Final[dict[str, object]] = {"reserved_cost": 0.0001, "entries": []} + socket: Final = WebSocket( + {"type": "websocket", "state": {"budget_reservation": connect_reservation}}, receive=AsyncMock(), send=AsyncMock() + ) + + async def key_budget_check(websocket: object, model: str) -> UserAPIKeyAuth: + if connect_reservation.get("finalized") is not True: + raise litellm.BudgetExceededError(current_cost=0.0001, max_budget=0.0001) + return UserAPIKeyAuth(models=["gpt-5.4-mini"]) + + with patch("litellm.proxy.proxy_server.prisma_client", None): # test-quality-ok: the model checks need no database + verdict: Final = await _proxy_frame_model_gate(key_budget_check)(socket, "gpt-5.4-mini") + + assert verdict == WebsocketFrameModelAllowed(access_groups=()) + + @pytest.mark.asyncio @pytest.mark.parametrize( "denial, error_type, message", 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 09/15] 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 From 06af85fbe0d28a40816d99001cdd559ca0984afc Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 08:48:27 -0700 Subject: [PATCH 10/15] test(proxy): mark the remaining new websocket test variables Final Builds the logging mocks with model_call_details in the constructor instead of setting it afterwards. --- .../test_openai_websocket_passthrough_logging_handler.py | 3 +-- .../pass_through_endpoints/test_pass_through_endpoints.py | 6 +++--- .../proxy/pass_through_endpoints/test_success_handler.py | 3 +-- 3 files changed, 5 insertions(+), 7 deletions(-) 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 1bc4e0fe4d2..6d2e22fbfcd 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 @@ -23,8 +23,7 @@ SESSION_START: Final = datetime(2026, 1, 1, 12, 0, 0) def _logging_obj() -> MagicMock: - logging_obj = MagicMock() - logging_obj.model_call_details = {} + logging_obj: Final = MagicMock(model_call_details={}) return logging_obj 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 9f97110e1af..2903b3a1ac2 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 @@ -8056,7 +8056,7 @@ async def test_websocket_passthrough_still_bills_a_session_the_client_walked_out REALTIME_RESPONSE_DONE, json.dumps({"type": "response.output_text.delta", "delta": "too late"}), ) - websocket = _client_websocket(_pending_receive) + websocket: Final = _client_websocket(_pending_receive) websocket.send_text = AsyncMock(side_effect=[None, None, WebSocketDisconnect(code=1006)]) with ( @@ -8136,7 +8136,7 @@ def _turn_taking_client( frame: Final = next(script, None) return {"type": "websocket.disconnect"} if frame is None else _received_frame(frame) - websocket = _client_websocket(receive) + websocket: Final = _client_websocket(receive) async def send_text(text: str) -> None: if websocket.application_state == WebSocketState.DISCONNECTED: @@ -8529,7 +8529,7 @@ async def test_websocket_passthrough_bills_the_turn_a_client_walked_out_of_once_ client_left: Final = asyncio.Event() send_failed: Final = asyncio.Event() - client = _turn_taking_client( + client: Final = _turn_taking_client( [json.dumps({"type": "response.create", "model": "gpt-5.4"})], answered_by='"response.created"' ) receive: Final = client.receive 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 988828f48de..2e2446b8ecf 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_success_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/test_success_handler.py @@ -668,8 +668,7 @@ REALTIME_FRAMES: Final = [ def _normalize(url_route: str, response_body: dict[str, object] | list[dict[str, object]]) -> dict: - logging_obj = MagicMock() - logging_obj.model_call_details = {} + logging_obj: Final = MagicMock(model_call_details={}) return PassThroughEndpointLogging().normalize_llm_passthrough_logging_payload( httpx_response=MagicMock(spec=httpx.Response), response_body=response_body, From 6998e0f9ab96da708ff642c52bf1df6ec2dc78e3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 17:42:17 -0700 Subject: [PATCH 11/15] fix(proxy): check every model the realtime URL names before the OpenAI websocket relay connects --- .../llm_passthrough_endpoints.py | 14 +++++ .../pass_through_endpoints.py | 8 ++- .../test_openai_ws_passthrough_routes.py | 59 ++++++++++++++++++- 3 files changed, 77 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index fd50f2ae13f..0cd9f5caf53 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -142,6 +142,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( YieldsToConfiguredPassThroughRoute, create_pass_through_route, create_websocket_passthrough_route, + refuse_websocket, websocket_passthrough_request, ) from litellm.proxy.utils import ProxyLogging as ProxyLoggingType @@ -3560,6 +3561,14 @@ def _negotiated_websocket_subprotocol(websocket: WebSocket) -> str | None: return requested_subprotocols[0] if requested_subprotocols else None +async def _opening_model_refusal( + websocket: WebSocket, query_string: str, frame_model_gate: WebsocketFrameModelGate +) -> WebsocketFrameModelRefused | None: + opening_models: Final = dict.fromkeys(httpx.QueryParams(query_string).get_list("model")) + verdicts: Final = tuple([await frame_model_gate(websocket, model) for model in opening_models]) + return next((verdict for verdict in verdicts if isinstance(verdict, WebsocketFrameModelRefused)), None) + + @router.websocket("/openai_passthrough/{endpoint:path}") @router.websocket("/openai/{endpoint:path}") async def openai_websocket_proxy_route( @@ -3619,6 +3628,11 @@ async def openai_websocket_proxy_route( await websocket.accept(subprotocol=negotiated_subprotocol) + opening_refusal: Final = await _opening_model_refusal(websocket, query_string, frame_model_gate) + if opening_refusal is not None: + await refuse_websocket(websocket, opening_refusal) + return + await relay( websocket=websocket, target=wss_target, diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 0a591f9c2f8..2ccc4fe56d2 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -2788,11 +2788,15 @@ def _refusal_close_reason(message: str) -> str: return _truncated_close_reason(message if first_sentence is None else first_sentence.group(0)) +async def refuse_websocket(websocket: WebSocket, refusal: WebsocketFrameModelRefused) -> None: + await websocket.send_text(realtime_error_event(refusal.message, error_type=refusal.error_type)) + await websocket.close(code=CloseCode.POLICY_VIOLATION, reason=_refusal_close_reason(refusal.message)) + + async def _refuse_client_frame( websocket: WebSocket, upstream_ws: ClientConnection, refusal: WebsocketFrameModelRefused ) -> None: - await websocket.send_text(realtime_error_event(refusal.message, error_type=refusal.error_type)) - await websocket.close(code=CloseCode.POLICY_VIOLATION, reason=_refusal_close_reason(refusal.message)) + await refuse_websocket(websocket, refusal) await upstream_ws.close() diff --git a/tests/unit/proxy/test_openai_ws_passthrough_routes.py b/tests/unit/proxy/test_openai_ws_passthrough_routes.py index 587faa02e35..839ab49408f 100644 --- a/tests/unit/proxy/test_openai_ws_passthrough_routes.py +++ b/tests/unit/proxy/test_openai_ws_passthrough_routes.py @@ -121,8 +121,13 @@ class _FakeRelay: class _FakeFrameModelGate: + def __init__(self, refusals: Mapping[str, WebsocketFrameModelRefused] = MappingProxyType({})) -> None: + self.refusals = refusals + self.checked: list[str] = [] + async def __call__(self, websocket: object, model: str, /) -> WebsocketFrameModelVerdict: - return WebsocketFrameModelAllowed(access_groups=()) + self.checked.append(model) + return self.refusals.get(model) or WebsocketFrameModelAllowed(access_groups=()) @dataclass(frozen=True, slots=True) @@ -136,8 +141,9 @@ async def _serve( endpoint: str, user_api_key_dict: UserAPIKeyAuth, general_settings: Mapping[str, object], + frame_model_gate: _FakeFrameModelGate | None = None, ) -> _Served: - served: Final = _Served(relay=_FakeRelay(), frame_model_gate=_FakeFrameModelGate()) + served: Final = _Served(relay=_FakeRelay(), frame_model_gate=frame_model_gate or _FakeFrameModelGate()) await openai_websocket_proxy_route( websocket=websocket, endpoint=endpoint, @@ -332,6 +338,55 @@ async def test_openai_websocket_connects_model_restricted_keys(key_models): assert websocket.closed is None +ALIAS_REFUSAL: Final = WebsocketFrameModelRefused( + message=_OPENAI_WS_ALIASED_MODEL_REFUSAL, error_type=ProxyErrorTypes.key_model_access_denied.value +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "query", + [ + "model=gpt-realtime-2.1", + "model=gpt-realtime-2.1&model=gpt-realtime-2.1-mini", + "model=gpt-realtime-2.1-mini&model=gpt-realtime-2.1", + ], + ids=["one_model", "refused_model_first", "refused_model_last"], +) +async def test_openai_websocket_refuses_an_opening_model_the_frame_gate_refuses(query): + """ + The relay sends the connection URL to OpenAI unchanged, so every model it names has to pass the + same gate as a frame, or an alias the upgrade resolves to an allowed model opens the session on the + alias name itself + """ + websocket: Final = _FakeWebSocket("/openai/v1/realtime", query) + gate: Final = _FakeFrameModelGate(MappingProxyType({"gpt-realtime-2.1": ALIAS_REFUSAL})) + + with patch(GET_CREDENTIALS, return_value="sk-provider"): + served: Final = await _serve( + websocket, "v1/realtime", UserAPIKeyAuth(models=["gpt-realtime-2.1-mini"]), ENABLED, gate + ) + + assert websocket.error_message() == _OPENAI_WS_ALIASED_MODEL_REFUSAL + assert websocket.closed == (1008, "This key cannot use a model alias over OpenAI websocket passthrough.") + assert served.relay.calls == [] + assert "gpt-realtime-2.1" in gate.checked + + +@pytest.mark.asyncio +async def test_openai_websocket_relays_a_session_once_every_opening_model_passes_the_frame_gate(): + websocket: Final = _FakeWebSocket("/openai/v1/realtime", "model=gpt-realtime-2.1-mini") + + with patch(GET_CREDENTIALS, return_value="sk-provider"): + served: Final = await _serve( + websocket, "v1/realtime", UserAPIKeyAuth(models=["gpt-realtime-2.1-mini"]), ENABLED + ) + + assert served.frame_model_gate.checked == ["gpt-realtime-2.1-mini"] + assert len(served.relay.calls) == 1 + assert websocket.closed is None + + class _FrameAuthorizer: """Stands in for the full key, team, user and budget auth the gate re-runs for every model a frame names.""" From 4ee39dac33f8ea7ab5c07c191314b8e6ae28d382 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 18:51:30 -0700 Subject: [PATCH 12/15] fix(proxy): keep reading running answers after a refused websocket frame so they are billed --- .../pass_through_endpoints.py | 17 ++++-------- .../test_pass_through_endpoints.py | 26 ++++++++++++------- 2 files changed, 21 insertions(+), 22 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 2ccc4fe56d2..7b304331f28 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -2793,13 +2793,6 @@ async def refuse_websocket(websocket: WebSocket, refusal: WebsocketFrameModelRef await websocket.close(code=CloseCode.POLICY_VIOLATION, reason=_refusal_close_reason(refusal.message)) -async def _refuse_client_frame( - websocket: WebSocket, upstream_ws: ClientConnection, refusal: WebsocketFrameModelRefused -) -> None: - await refuse_websocket(websocket, refusal) - await upstream_ws.close() - - def _stamp_billed_turn(kwargs: Mapping[str, object], turn: _TurnAuthorization) -> None: """Charge a turn to the access groups that authorized its model, and its model budget to the model it asked for, since the provider may answer under a dated snapshot name the budget was never set on.""" @@ -2894,11 +2887,11 @@ async def _relay_client_frame( gate: WebsocketFrameModelGate | None, ledger: _WebsocketTurnLedger, ) -> _ForwardClientFrame | _RefuseClientFrame: - """Forward one client frame upstream, or refuse it and end the session.""" + """Forward one client frame upstream, or refuse it and close the client.""" verdict: Final = await _client_frame_verdict(websocket, checked_frame, gate) match verdict: case _RefuseClientFrame(refusal=refusal): - await _refuse_client_frame(websocket, upstream_ws, refusal) + await refuse_websocket(websocket, refusal) case _ForwardClientFrame(starts_turn=True, turn=turn, lane=lane): ledger.queue_turn(lane, turn) await upstream_ws.send(forwarded_frame) @@ -2943,7 +2936,7 @@ async def websocket_passthrough_request( upstream_frame_kind: Optional sort of upstream frames into the one that ends a billable turn and the ones that open the session, so each turn logs its own spend and keeps the session's provider-named model walkout_turn_drain_seconds: How long to keep reading a turn the provider is still running after the client - left, so the provider's usage for it still reaches spend + left or was refused, so the provider's usage for it still reaches spend builtin_provider_route: Whether a built-in provider route serves the connection, so its URL names the provider model: The model the caller named before connecting, logged on a failure before the setup frame arrives """ @@ -3174,8 +3167,6 @@ async def websocket_passthrough_request( message = await websocket.receive() message_type = message.get("type") if message_type == "websocket.disconnect": - if not _turn_in_flight(upstream_frame_kind, websocket_messages): - await upstream_ws.close() break text_data: str | None = message.get("text") @@ -3258,6 +3249,8 @@ async def websocket_passthrough_request( _RefuseClientFrame, ): break + if not _turn_in_flight(upstream_frame_kind, websocket_messages): + await upstream_ws.close() except asyncio.CancelledError: raise except Exception: 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 95629b7f6c5..217e8781c98 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 @@ -8052,14 +8052,13 @@ async def test_websocket_passthrough_still_bills_a_session_the_client_walked_out class ScriptedUpstreamWebSocket: - """Answers each client frame with its scripted reply, and flushes what was still in flight when it is closed""" + """Answers each client frame with its scripted reply""" - def __init__(self, *opening: str, replies: Sequence[Sequence[str]] = (), in_flight: Sequence[str] = ()): + def __init__(self, *opening: str, replies: Sequence[Sequence[str]] = ()): self._frames: asyncio.Queue[str | None] = asyncio.Queue() for frame in opening: self._frames.put_nowait(frame) self._replies = iter(replies) - self._in_flight = in_flight self._closed = False self.send = AsyncMock(side_effect=self._reply) self.close = AsyncMock(side_effect=self._flush_and_close) @@ -8072,8 +8071,6 @@ class ScriptedUpstreamWebSocket: if self._closed: return self._closed = True - for frame in self._in_flight: - self._frames.put_nowait(frame) self._frames.put_nowait(None) while not self._frames.empty(): await asyncio.sleep(0) @@ -8586,8 +8583,8 @@ async def test_websocket_passthrough_still_bills_the_answer_in_flight_when_a_lat """ A refused frame closes the client while the provider may still be sending the answer to an earlier, allowed frame. That answer cannot go to a closed client, - but it was paid for, so it has to reach spend instead of turning the whole - session into a failure that bills nothing. + but the provider bills it once it ends, so the gateway keeps reading until it + ends and its usage reaches spend. """ gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-5.4": KEY_MODEL_REFUSAL})) client: Final = _turn_taking_client( @@ -8597,10 +8594,19 @@ 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: Final = ScriptedUpstreamWebSocket(replies=[[created]], in_flight=[completed]) + refused: Final = asyncio.Event() + close_client: Final = client.close.side_effect - relayed: Final = await _relay_turns(client, upstream_ws, gate) + def close_client_and_note_the_refusal(*args: object, **kwargs: object) -> None: + close_client(*args, **kwargs) + refused.set() + + client.close = AsyncMock(side_effect=close_client_and_note_the_refusal) + upstream_ws: Final = UpstreamFinishingTheTurnAfterTheClientLeft( + refused, _responses_turn_started("gpt-5.4-mini"), (_responses_turn_ended("gpt-5.4-mini"),) + ) + + relayed: Final = await _relay_turns(client, upstream_ws, gate, walkout_turn_drain_seconds=5) assert [json.loads(text.args[0])["type"] for text in client.send_text.await_args_list] == [ "response.created", From 08fd75f73cdf80102e8758cbd7063c2fc7024de8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 20:01:32 -0700 Subject: [PATCH 13/15] fix(proxy): bill realtime transcription sessions on the OpenAI websocket passthrough A transcription session names its model only in session.updated under audio.input.transcription, and reports usage on each completed transcription instead of response.done, so its row logged unknown at $0 and a budgeted key never spent. --- ...i_websocket_passthrough_logging_handler.py | 41 +++++++--- ...i_websocket_passthrough_logging_handler.py | 78 +++++++++++++++++++ 2 files changed, 108 insertions(+), 11 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 27abb333a36..eb2b4edf2fe 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 @@ -22,12 +22,15 @@ WebsocketMessages: TypeAlias = Sequence[Mapping[str, object]] OPENAI_WEBSOCKET_PROVIDER: Final = LlmProviders.OPENAI.value _SESSION_CREATED: Final = "session.created" +_SESSION_UPDATED: Final = "session.updated" _RESPONSE_CREATED: Final = "response.created" _ERROR: Final = "error" _REALTIME_RESPONSE_DONE: Final = "response.done" +_TRANSCRIPTION_COMPLETED: Final = "conversation.item.input_audio_transcription.completed" _RESPONSES_TURN_ENDS: Final = frozenset({"response.completed", "response.incomplete", "response.failed"}) _TURN_ENDS: Final = _RESPONSES_TURN_ENDS | {_REALTIME_RESPONSE_DONE} -_MODEL_NAMING_EVENTS: Final = _TURN_ENDS | {_SESSION_CREATED, _RESPONSE_CREATED} +_MODEL_NAMING_EVENTS: Final = _TURN_ENDS | {_SESSION_CREATED, _SESSION_UPDATED, _RESPONSE_CREATED} +_REALTIME_BILLED_EVENTS: Final = frozenset({_REALTIME_RESPONSE_DONE, _TRANSCRIPTION_COMPLETED}) _RESPONSE_COST_HEADER: Final = "llm_provider-x-litellm-response-cost" @@ -37,10 +40,29 @@ class _NamedModel(BaseModel): model: str | None = None +class _AudioInput(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + transcription: _NamedModel | None = None + + +class _SessionAudio(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + input: _AudioInput | None = None + + +class _NamedSession(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + model: str | None = None + audio: _SessionAudio | None = None + + class _ModelNamingEvent(BaseModel): model_config = ConfigDict(extra="ignore", frozen=True) - session: _NamedModel | None = None + session: _NamedSession | None = None response: _NamedModel | None = None @@ -75,14 +97,11 @@ def _event_model(event: Mapping[str, object]) -> str | None: naming: Final = _MODEL_NAMING_EVENT.validate_python(event) except ValidationError: return None - return next( - ( - container.model - for container in (naming.session, naming.response) - if container is not None and container.model - ), - None, - ) + session: Final = naming.session or _NamedSession() + response: Final = naming.response or _NamedModel() + audio_input: Final = (session.audio or _SessionAudio()).input or _AudioInput() + transcription: Final = audio_input.transcription or _NamedModel() + return session.model or response.model or transcription.model or None def observed_websocket_model(messages: WebsocketMessages) -> str | None: @@ -103,7 +122,7 @@ def openai_websocket_frame_kind(frame: Mapping[str, object]) -> WebsocketUpstrea def _realtime_billing(messages: WebsocketMessages, model: str, logging_obj: LiteLLMLoggingObj) -> _Billing | None: - if not any(event.get("type") == _REALTIME_RESPONSE_DONE for event in messages): + if not any(event.get("type") in _REALTIME_BILLED_EVENTS for event in messages): return None results: Final = cast( # cast-ok: the realtime cost helpers read raw provider frames through their TypedDict view OpenAIRealtimeStreamList, tuple(event for event in messages if "type" in event) 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 6d2e22fbfcd..baf52af112b 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" +TRANSCRIPTION_MODEL: Final = "gpt-4o-mini-transcribe" SESSION_START: Final = datetime(2026, 1, 1, 12, 0, 0) @@ -308,6 +309,83 @@ def test_a_turn_the_caller_walked_out_of_still_names_the_model_it_ran(): assert logging_obj.model_call_details["model"] == REASONING_MODEL +def _transcription_session(model: str) -> tuple[dict[str, object], ...]: + return ( + {"type": "session.created", "session": {"type": "transcription", "audio": {"input": {"transcription": None}}}}, + { + "type": "session.updated", + "session": {"type": "transcription", "audio": {"input": {"transcription": {"model": model}}}}, + }, + ) + + +def _transcription_completed(usage: Mapping[str, object]) -> dict[str, object]: + return { + "type": "conversation.item.input_audio_transcription.completed", + "item_id": "item_fake", + "transcript": "hello", + "usage": usage, + } + + +@pytest.mark.parametrize( + ("usage", "billed_units"), + [ + ({"type": "duration", "seconds": 3}, ((3, "input_cost_per_second"),)), + ( + { + "type": "tokens", + "input_tokens": 30, + "output_tokens": 11, + "input_token_details": {"text_tokens": 0, "audio_tokens": 30}, + }, + ((30, "input_cost_per_audio_token"), (11, "output_cost_per_token")), + ), + ], +) +def test_a_transcription_session_is_billed_for_every_clip_it_transcribed(usage, billed_units): + handled: Final = _handle( + [ + *_transcription_session(TRANSCRIPTION_MODEL), + _transcription_completed(usage), + _transcription_completed(usage), + ], + url_route="/openai_passthrough/v1/realtime", + ) + + logged_cost: Final = litellm.response_cost_calculator( + response_object=handled["result"], + model=handled["kwargs"]["model"], + custom_llm_provider="openai", + call_type="completion", + optional_params={}, + ) + + prices: Final = litellm.model_cost[TRANSCRIPTION_MODEL] + expected_cost: Final = 2 * sum(units * prices[rate] for units, rate in billed_units) + assert expected_cost > 0 + assert handled["kwargs"]["model"] == TRANSCRIPTION_MODEL + assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost) + assert logged_cost == pytest.approx(expected_cost) + + +def test_a_conversation_that_also_transcribes_the_caller_is_logged_on_its_own_model(): + handled: Final = _handle( + [ + { + "type": "session.created", + "session": { + "model": REALTIME_MODEL, + "audio": {"input": {"transcription": {"model": TRANSCRIPTION_MODEL}}}, + }, + }, + _realtime_response_done(text_in=10, audio_in=0, text_out=5, audio_out=0), + ] + ) + + assert handled["kwargs"]["model"] == REALTIME_MODEL + + @pytest.mark.parametrize("event", ["response.incomplete", "response.failed"]) def test_a_responses_turn_that_ended_early_is_billed_for_the_tokens_it_used(event): """ From 90b05c31409fcbdd20748ae68b263200786b4bf9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 20:15:26 -0700 Subject: [PATCH 14/15] test(proxy): mark the turn-taking client's delivered list Final --- .../proxy/pass_through_endpoints/test_pass_through_endpoints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 217e8781c98..28c8689518d 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 @@ -8092,7 +8092,7 @@ def _turn_taking_client( answer to every earlier frame), and that Starlette refuses to send to once it is closed""" from starlette.websockets import WebSocketState - delivered: list[str] = [] + delivered: Final[list[str]] = [] script: Final = iter(frames) receives: Final = iter(range(len(frames) + 1) if answers_due is None else answers_due) From ba239b74f491691d9202c523a58df31cc6290780 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 22:43:47 -0700 Subject: [PATCH 15/15] fix(proxy): hold a websocket session's parallel slot until the socket closes Each finished turn logged its spend row under the session's own call id, so the first turn's success callback freed the key's max_parallel_requests slot while the socket was still open, and a key limited to one session could open several at once. The session now claims its slot under its own call id, every turn row logs under a call id of its own, and the slot is released when the session ends, as before Priced turn rows also take their row's call id as their id, so two turns that start in the same instant no longer share a spend log id and drop one row --- ...i_websocket_passthrough_logging_handler.py | 6 +- .../pass_through_endpoints.py | 10 ++- ...i_websocket_passthrough_logging_handler.py | 26 +++++- .../test_pass_through_endpoints.py | 88 ++++++++++++++++++- .../test_success_handler.py | 2 +- 5 files changed, 123 insertions(+), 9 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 eb2b4edf2fe..4ca71a6fc5e 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 @@ -200,9 +200,9 @@ def _requested_model(kwargs: Mapping[str, object]) -> str | None: return model_group if isinstance(model_group, str) else None -def _priced_result(model: str, billing: _Billing, start_time: datetime) -> ModelResponse: +def _priced_result(model: str, billing: _Billing, start_time: datetime, call_id: str) -> ModelResponse: result: Final = ModelResponse( - id=f"openai-websocket-{start_time.timestamp()}", + id=call_id, object="chat.completion", created=int(start_time.timestamp()), model=model, @@ -246,7 +246,7 @@ class OpenAIWebsocketPassthroughLoggingHandler: logging_obj.model_call_details.update(MappingProxyType({"response_cost": billing.cost})) priced: Final[PassThroughEndpointLoggingTypedDict] = { - "result": _priced_result(model, billing, start_time), + "result": _priced_result(model, billing, start_time, logging_obj.litellm_call_id), "kwargs": { **kwargs, "model": model, diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 7b304331f28..61178caf35f 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -3124,14 +3124,16 @@ async def websocket_passthrough_request( def log_turn(lane: str | None, in_row: Callable[[Mapping[str, object]], bool]) -> None: end_time: Final = datetime.now() - if ledger.logged_turns == 0: + if upstream_frame_kind is None: log_frames(logging_obj, passthrough_logging_payload, kwargs, end_time, lane, in_row) else: turn_call_id: Final = str(uuid.uuid4()) turn_logging: Final = new_call_logging(ledger.started, turn_call_id) turn_payload: Final = new_call_payload() turn_kwargs: Final = init_call_kwargs( - user_api_key_dict.model_copy(update={"budget_reservation": None}), + user_api_key_dict + if ledger.logged_turns == 0 + else user_api_key_dict.model_copy(update={"budget_reservation": None}), turn_payload, turn_call_id, turn_logging, @@ -3141,7 +3143,9 @@ async def websocket_passthrough_request( ledger.close_turn(lane, end_time) ### CALL HOOKS ### - modify incoming data / reject request before calling the model - websocket_hook_data: Final = guardrail_request_data_with_streaming(MappingProxyType({}), is_streaming=True) + websocket_hook_data: Final = guardrail_request_data_with_streaming( + MappingProxyType({"litellm_call_id": litellm_call_id}), is_streaming=True + ) await proxy_logging_obj.pre_call_hook( user_api_key_dict=user_api_key_dict, data=websocket_hook_data, 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 baf52af112b..35e12800ce3 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 @@ -1,3 +1,4 @@ +import uuid from collections.abc import Mapping, Sequence from datetime import datetime from types import MappingProxyType @@ -15,6 +16,8 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_websocket OpenAIWebsocketPassthroughLoggingHandler, openai_websocket_frame_kind, ) +from litellm.proxy.spend_tracking.spend_tracking_utils import get_spend_logs_id +from litellm.types.utils import CallTypes REALTIME_MODEL: Final = "gpt-realtime-2.1" RESPONSES_MODEL: Final = "gpt-5.4-mini" @@ -24,7 +27,7 @@ SESSION_START: Final = datetime(2026, 1, 1, 12, 0, 0) def _logging_obj() -> MagicMock: - logging_obj: Final = MagicMock(model_call_details={}) + logging_obj: Final = MagicMock(model_call_details={}, litellm_call_id=str(uuid.uuid4())) return logging_obj @@ -290,6 +293,27 @@ def test_the_price_of_a_multi_model_connection_reaches_the_spend_row(): assert handled["kwargs"]["model"] == RESPONSES_MODEL +def test_turns_that_start_in_the_same_instant_each_keep_their_own_spend_row(): + """ + Spend logs key each row by the id its result carries and skip a repeated id + without a word, so two turns that start in the same instant, on one socket + or on two, each need an id of their own or one of them is never billed. + """ + turns: Final = [ + _handle( + [_responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5)], + url_route="/openai_passthrough/v1/responses", + ) + for _ in range(2) + ] + + spend_log_ids: Final = { + get_spend_logs_id(CallTypes.pass_through.value, turn["result"].model_dump(), turn["kwargs"]) for turn in turns + } + assert len(spend_log_ids) == 2 + assert None not in spend_log_ids + + def test_a_turn_the_caller_walked_out_of_still_names_the_model_it_ran(): """ A client that closes before response.completed leaves no usage behind to 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 28c8689518d..06395bf72a9 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 @@ -1,5 +1,6 @@ import asyncio import base64 +import contextvars import gzip import json import logging @@ -39,6 +40,8 @@ from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._lazy_features import LazyFeature, attach_lazy_features from litellm.proxy._types import ProxyException, UserAPIKeyAuth +from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError +from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3 from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import router as llm_passthrough_router from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler import ( @@ -77,7 +80,7 @@ from litellm.proxy.pass_through_endpoints.success_handler import ( ) from litellm.proxy.route_llm_request import ProxyModelNotFoundError from litellm.proxy.spend_tracking.spend_tracking_utils import get_spend_logs_id -from litellm.proxy.utils import ProxyLogging +from litellm.proxy.utils import InternalUsageCache, ProxyLogging, hash_token from litellm.types import utils as types_utils from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import CallTypes @@ -8204,6 +8207,89 @@ async def test_websocket_passthrough_logs_each_turn_on_its_own_row_charged_to_th relayed.proxy_logging.post_call_failure_hook.assert_not_awaited() +@pytest.mark.asyncio +async def test_websocket_passthrough_holds_the_keys_parallel_slot_until_the_session_ends(): + """ + A key allowed one parallel request holds that slot for as long as its + socket stays open. A turn that finishes and reaches spend mid-session must + not free it, or a second session gets in while the first is still running. + """ + limiter_cache: Final = DualCache() + limiter: Final = PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(limiter_cache)) + key: Final = UserAPIKeyAuth(api_key=hash_token("sk-one-socket"), max_parallel_requests=1) + logged_rows: Final[list[asyncio.Task[None]]] = [] + + async def admit(user_api_key_dict: UserAPIKeyAuth, data: dict[str, object], call_type: str) -> None: + await limiter.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, cache=limiter_cache, data=data, call_type=call_type + ) + + async def release(data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, response: object) -> None: + await limiter.async_post_call_success_hook(data=data, user_api_key_dict=user_api_key_dict, response=response) + + async def log_row(**row: object) -> None: + await limiter.async_log_success_event( + kwargs=row, response_obj=None, start_time=row["start_time"], end_time=row["end_time"] + ) + + async def opens_another_session() -> bool: + await asyncio.gather(*logged_rows) + try: + await asyncio.create_task( + admit(key, {"litellm_call_id": str(uuid.uuid4())}, "pass_through_endpoint"), + context=contextvars.Context(), + ) + except ProxyRateLimitError: + return False + return True + + client: Final = _turn_taking_client( + [json.dumps({"type": "response.create", "model": "gpt-5.4-mini"})], answered_by='"response.completed"' + ) + take_frame: Final = client.receive + while_open: Final[list[bool]] = [] + + async def receive() -> Mapping[str, str | bytes]: + message: Final = await take_frame() + if message["type"] == "websocket.disconnect": + while_open.append(await opens_another_session()) + return message + + client.receive = receive + + with ( + _patched_websocket_passthrough_environment( + ScriptedUpstreamWebSocket(replies=[_responses_turn("gpt-5.4-mini")]) + ) as proxy_logging, + patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER") as worker, + patch(SUCCESS_HANDLER, new_callable=AsyncMock, side_effect=log_row), + ): + proxy_logging.pre_call_hook = AsyncMock(side_effect=admit) + proxy_logging.post_call_success_hook = AsyncMock(side_effect=release) + worker.ensure_initialized_and_enqueue = MagicMock( + side_effect=lambda async_coroutine: logged_rows.append(asyncio.create_task(async_coroutine)) + ) + await asyncio.wait_for( + websocket_passthrough_request( + websocket=client, + target="wss://api.openai.com/v1/responses", + custom_headers={"Authorization": "Bearer sk-provider"}, + user_api_key_dict=key, + forward_headers=False, + endpoint="/openai_passthrough/v1/responses", + accept_websocket=False, + client_frame_model_gate=RecordingFrameModelGate(), + upstream_frame_kind=openai_websocket_frame_kind, + walkout_turn_drain_seconds=0.05, + ), + timeout=10, + ) + after_close: Final = await opens_another_session() + + assert len(logged_rows) == 1 + assert (while_open, after_close) == ([False], True) + + @pytest.mark.asyncio async def test_websocket_passthrough_charges_pipelined_turns_to_the_frame_that_asked_for_each(): """ 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 2e2446b8ecf..302c696682d 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_success_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/test_success_handler.py @@ -668,7 +668,7 @@ REALTIME_FRAMES: Final = [ def _normalize(url_route: str, response_body: dict[str, object] | list[dict[str, object]]) -> dict: - logging_obj: Final = MagicMock(model_call_details={}) + logging_obj: Final = MagicMock(model_call_details={}, litellm_call_id="call-fake") return PassThroughEndpointLogging().normalize_llm_passthrough_logging_payload( httpx_response=MagicMock(spec=httpx.Response), response_body=response_body,