diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 448643ef59b..563de5694a3 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -5305,79 +5305,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, diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index e478096f7ad..2e51bc3465b 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, Sequence +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 @@ -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,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 enforced_model_allowlists +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 @@ -80,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, ) @@ -115,8 +117,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, @@ -180,7 +190,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 @@ -210,9 +220,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: @@ -3274,19 +3283,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 @@ -3303,6 +3299,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=( @@ -3311,14 +3312,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") @@ -3327,19 +3320,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 @@ -3354,6 +3337,8 @@ class _WebsocketRelay(Protocol): forward_headers: bool, endpoint: str, accept_websocket: bool, + client_frame_model_gate: WebsocketFrameModelGate | None = None, + upstream_frame_kind: WebsocketUpstreamFrameClassifier | None = None, ) -> None: ... @@ -3361,18 +3346,62 @@ 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_error_type(denial: ProxyException) -> str: + return denial.type.value if isinstance(denial.type, ProxyErrorTypes) else denial.type - 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, - ) - return resolve +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, + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + from litellm.proxy.spend_tracking.budget_reservation import release_or_invalidate_budget_reservation + + 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=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, + ) + 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 WebsocketFrameModelRefused(message=denial.message, error_type=_proxy_error_type(denial)) + except HTTPException as denial: + return WebsocketFrameModelRefused(message=str(denial.detail), error_type="invalid_request_error") + except litellm.BudgetExceededError as exhausted: + 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 WebsocketFrameModelRefused(message=_OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL, error_type="server_error") + + return gate def _negotiated_websocket_subprotocol(websocket: WebSocket) -> str | None: @@ -3392,12 +3421,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] = { @@ -3451,6 +3480,8 @@ async def openai_websocket_proxy_route( forward_headers=False, endpoint=websocket.url.path, accept_websocket=False, + client_frame_model_gate=frame_model_gate, + upstream_frame_kind=openai_websocket_frame_kind, ) @@ -3696,7 +3727,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..c2f4e85005c --- /dev/null +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py @@ -0,0 +1,218 @@ +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.passthrough_endpoints.pass_through_endpoints import WebsocketUpstreamFrameKind +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_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" + + +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 _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 + + +_MODEL_NAMING_EVENT: Final = TypeAdapter(_ModelNamingEvent) +_RESPONSES_TURN_END_EVENT: Final = TypeAdapter(_ResponsesTurnEndEvent) + + +@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 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 + 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") not in _RESPONSES_TURN_ENDS: + return None + try: + ended: Final = _RESPONSES_TURN_END_EVENT.validate_python(event) + except ValidationError: + verbose_proxy_logger.warning( + "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=ended.response.model or session_model, + usage=ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( # pyright: ignore[reportPrivateUsage] # the one shared Responses usage converter + ended.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 650dcae3239..33ee2f0599a 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -3,14 +3,16 @@ 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, 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, TypeAlias, TypedDict, cast from urllib.parse import urlencode, urlparse import httpx @@ -23,6 +25,7 @@ from fastapi import ( Response, UploadFile, WebSocket, + WebSocketDisconnect, status, ) from fastapi.responses import StreamingResponse @@ -30,7 +33,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, @@ -56,6 +59,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_str_mapping, bind_budget_reservation_to_callbacks, get_metadata_variable_name_from_kwargs, get_or_create_metadata_bucket, @@ -65,6 +69,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 @@ -141,6 +146,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 @@ -2486,6 +2492,178 @@ 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, websocket: WebSocket, model: str, /) -> WebsocketFrameModelVerdict: ... + + +WebsocketUpstreamFrameClassifier: TypeAlias = Callable[[Mapping[str, object]], WebsocketUpstreamFrameKind] + + +@dataclass(frozen=True, slots=True) +class _ForwardClientFrame: + turn_access_groups: tuple[str, ...] | None + + +@dataclass(frozen=True, slots=True) +class _RefuseClientFrame: + refusal: WebsocketFrameModelRefused + + +@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, so it cannot run a stored prompt. " + "The gateway cannot check which model a stored prompt runs." +) +_FIRST_SENTENCE: Final = re.compile(r".*?[.!?](?=\s|$)") + + +def _may_name_a_model(text_data: str) -> bool: + 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_value(frame: Mapping[str, object], placement: Sequence[str]) -> object: + *parents, leaf = placement + container: Final = reduce(_nested_frame, parents, frame) + return None if container is None else container.get(leaf) + + +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 _NO_CLIENT_FRAME_REQUESTS + frame: Final = _json_object_frame(frame_data) + 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) + 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( + websocket: WebSocket, frame_data: str | bytes, gate: WebsocketFrameModelGate | None +) -> _ForwardClientFrame | _RefuseClientFrame: + if gate is None: + 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) + + +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], 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 isinstance(metadata, MutableMapping): + metadata[MODEL_ACCESS_GROUP_METADATA_KEY] = list(access_groups) + + +async def _relay_client_frame( + websocket: WebSocket, + upstream_ws: ClientConnection, + checked_frame: str | bytes, + forwarded_frame: str | bytes, + gate: WebsocketFrameModelGate | None, +) -> _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 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( websocket: WebSocket, target: str, @@ -2496,6 +2674,8 @@ 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, + upstream_frame_kind: WebsocketUpstreamFrameClassifier | None = None, ): """ WebSocket passthrough request handler. @@ -2509,6 +2689,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 @@ -2544,26 +2727,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 @@ -2585,35 +2773,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) @@ -2705,9 +2975,27 @@ 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 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: - await upstream_ws.send(bytes_data) + 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: @@ -2734,11 +3022,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 @@ -2746,6 +3030,17 @@ 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: + await websocket.send_text(upstream_message) async def forward_upstream_to_client() -> Close | None: try: @@ -2754,6 +3049,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 @@ -2793,62 +3093,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 - - # 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/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": diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 0b91e57f01e..571a7c07eb1 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 88dd8c16501..3d53a8d7b67 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, ) @@ -7309,7 +7316,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(): @@ -7620,14 +7627,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): @@ -7692,8 +7713,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) @@ -7701,16 +7721,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", @@ -7722,13 +7760,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"}), @@ -7795,27 +7832,79 @@ 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) + + 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 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 +@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 @@ -7889,6 +7978,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( { @@ -8816,15 +9095,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"})], @@ -8838,7 +9108,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")