diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 828b5b514b7..c4b934ecefc 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -5604,115 +5604,6 @@ async def can_key_call_resolved_model( ) -async def can_personal_user_call_model( - model: str, - valid_token: UserAPIKeyAuth, - llm_router: Router | None, -) -> None: - """Apply the personal user's model allowlist, the check ``common_checks`` runs for a key with no team.""" - from litellm.proxy.proxy_server import ( - prisma_client, - proxy_logging_obj, - user_api_key_cache, - ) - - if valid_token.team_id is not None or valid_token.user_id is None or prisma_client is None: - return - try: - user_object: Final = await get_user_object( - user_id=valid_token.user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - user_id_upsert=False, - parent_otel_span=valid_token.parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) - except UserNotFoundError: - return - if user_object is None: - return - await can_user_call_model( - model=model, - llm_router=llm_router, - user_object=user_object, - key_model_aliases=key_model_aliases_for_auth_check(valid_token), - ) - - -async def _matching_team_object( - valid_token: UserAPIKeyAuth, - prisma_client: PrismaClient, - user_api_key_cache: UserApiKeyCache, - proxy_logging_obj: ProxyLogging, -) -> LiteLLM_TeamTableCachedObj | None: - if valid_token.team_id is None: - return None - try: - return await get_team_object( - team_id=valid_token.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=valid_token.parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) - except Exception as e: # noqa: BLE001 # fail-safe: a team we cannot read grants no group, it never breaks auth - verbose_proxy_logger.debug("model access group attribution could not read the team: %s", e) - return None - - -async def enforce_model_access_group_budget( - model: str, - valid_token: UserAPIKeyAuth, - llm_router: Router | None, -) -> None: - """Attribute this model to the budgeted access groups that authorize it and refuse an exhausted one. - - The pair ``common_checks`` runs for every HTTP request, for callers that reach a model only - through a group: the stamp is what the spend writer charges afterwards, and the check is what - stops a group whose budget is already gone. - - Raises: - BudgetExceededError if a group that authorized this model is over its max budget. - """ - from litellm.proxy.proxy_server import ( - prisma_client, - proxy_logging_obj, - user_api_key_cache, - ) - - if prisma_client is None: - return - matched: Final = await stamp_matched_model_access_groups( - model=model, - valid_token=valid_token, - team_object=await _matching_team_object( - valid_token=valid_token, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ), - project_object=( - None - if valid_token.project_id is None - else await get_project_object( - project_id=valid_token.project_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - ), - llm_router=llm_router, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - await _model_access_group_max_budget_check( - matched_model_access_groups=matched, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - ) - - def can_org_access_model( model: str, org_object: LiteLLM_OrganizationTable | None, diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 1d1f88aaa18..8944ce10046 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -15,11 +15,11 @@ import os import posixpath import re import sys -from collections.abc import AsyncGenerator, Callable, Mapping +from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping from dataclasses import dataclass from functools import partial from types import MappingProxyType -from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, cast +from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeAlias, cast import httpx import openai @@ -73,12 +73,7 @@ from litellm.llms.oss_decision import OssDecisionProvider, oss_connection, valid from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.passthrough.main import AsyncPassthroughStreamingResponse, is_provider_error from litellm.proxy._types import * -from litellm.proxy.auth.auth_checks import ( - can_agent_call_model, - can_key_call_resolved_model, - can_personal_user_call_model, - enforce_model_access_group_budget, -) +from litellm.proxy.auth.auth_checks import can_agent_call_model, can_key_call_resolved_model from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy module exports @@ -86,6 +81,7 @@ from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy modul get_bearer_token, is_no_auth_dev_mode, user_api_key_auth, + user_api_key_auth_accepted_websocket_for_model, user_api_key_auth_websocket, user_api_key_auth_websocket_for_model, ) @@ -114,9 +110,16 @@ from litellm.proxy.pass_through_endpoints.injected_stream_usage import ( stream_usage_relay, without_injected_usage_frames, ) +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_websocket_passthrough_logging_handler import ( + openai_websocket_frame_kind, +) from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( HttpPassThroughEndpointHelpers, + WebsocketFrameModelAllowed, WebsocketFrameModelGate, + WebsocketFrameModelRefused, + WebsocketFrameModelVerdict, + WebsocketUpstreamFrameClassifier, create_pass_through_route, create_websocket_passthrough_route, websocket_passthrough_request, @@ -3322,6 +3325,7 @@ class _WebsocketRelay(Protocol): endpoint: str, accept_websocket: bool, client_frame_model_gate: WebsocketFrameModelGate | None = None, + upstream_frame_kind: WebsocketUpstreamFrameClassifier | None = None, ) -> None: ... @@ -3329,7 +3333,20 @@ def _websocket_relay() -> _WebsocketRelay: return websocket_passthrough_request -def _proxy_frame_model_gate() -> WebsocketFrameModelGate: +def _proxy_error_type(denial: ProxyException) -> str: + return denial.type.value if isinstance(denial.type, ProxyErrorTypes) else denial.type + + +FrameModelAuthorizer: TypeAlias = Callable[[WebSocket, str], Awaitable[UserAPIKeyAuth]] + + +def _frame_model_authorizer() -> FrameModelAuthorizer: + return user_api_key_auth_accepted_websocket_for_model + + +def _proxy_frame_model_gate( + authorize_frame_model: Annotated[FrameModelAuthorizer, Depends(_frame_model_authorizer)], +) -> WebsocketFrameModelGate: from litellm.proxy.proxy_server import ( llm_model_list, llm_router, @@ -3337,32 +3354,39 @@ def _proxy_frame_model_gate() -> WebsocketFrameModelGate: proxy_logging_obj, user_api_key_cache, ) + from litellm.proxy.spend_tracking.budget_reservation import release_or_invalidate_budget_reservation - async def gate(model: str, valid_token: UserAPIKeyAuth, /) -> str | None: + async def authorize(websocket: WebSocket, model: str) -> WebsocketFrameModelAllowed: + frame_auth: Final = await authorize_frame_model(websocket, model) + access_groups: Final = tuple(frame_auth.matched_model_access_groups or ()) try: await can_key_call_resolved_model( - model=model, llm_model_list=llm_model_list, valid_token=valid_token, llm_router=llm_router + model=model, llm_model_list=llm_model_list, valid_token=frame_auth, llm_router=llm_router ) await can_agent_call_model( model=model, - valid_token=valid_token, + valid_token=frame_auth, llm_router=llm_router, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) - await can_personal_user_call_model(model=model, valid_token=valid_token, llm_router=llm_router) - await enforce_model_access_group_budget(model=model, valid_token=valid_token, llm_router=llm_router) + finally: + await release_or_invalidate_budget_reservation(budget_reservation=frame_auth.budget_reservation) + return WebsocketFrameModelAllowed(access_groups=access_groups) + + async def gate(websocket: WebSocket, model: str, /) -> WebsocketFrameModelVerdict: + try: + return await authorize(websocket, model) except ProxyException as denial: - return denial.message + return WebsocketFrameModelRefused(message=denial.message, error_type=_proxy_error_type(denial)) except HTTPException as denial: - return str(denial.detail) + return WebsocketFrameModelRefused(message=str(denial.detail), error_type="invalid_request_error") except litellm.BudgetExceededError as exhausted: - return exhausted.message + return WebsocketFrameModelRefused(message=exhausted.message, error_type="budget_exceeded") except Exception: # noqa: BLE001 # fail closed: a gate that cannot answer refuses, and says so in band verbose_proxy_logger.exception("OpenAI websocket passthrough: the frame model gate failed on %s", model) - return _OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL - return None + return WebsocketFrameModelRefused(message=_OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL, error_type="server_error") return gate @@ -3444,6 +3468,7 @@ async def openai_websocket_proxy_route( endpoint=websocket.url.path, accept_websocket=False, client_frame_model_gate=frame_model_gate, + upstream_frame_kind=openai_websocket_frame_kind, ) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py index 300aed08923..c2f4e85005c 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py @@ -13,6 +13,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.proxy._types import PassThroughEndpointLoggingTypedDict from litellm.responses.utils import ResponseAPILoggingUtils from litellm.types.llms.openai import OpenAIRealtimeStreamList, ResponseAPIUsage +from litellm.types.passthrough_endpoints.pass_through_endpoints import WebsocketUpstreamFrameKind from litellm.types.utils import LlmProviders, ModelResponse, Usage WebsocketMessages: TypeAlias = Sequence[Mapping[str, object]] @@ -22,10 +23,9 @@ OPENAI_WEBSOCKET_PROVIDER: Final = LlmProviders.OPENAI.value _SESSION_CREATED: Final = "session.created" _RESPONSE_CREATED: Final = "response.created" _REALTIME_RESPONSE_DONE: Final = "response.done" -_RESPONSES_COMPLETED: Final = "response.completed" -_MODEL_NAMING_EVENTS: Final = frozenset( - {_SESSION_CREATED, _RESPONSE_CREATED, _REALTIME_RESPONSE_DONE, _RESPONSES_COMPLETED} -) +_RESPONSES_TURN_ENDS: Final = frozenset({"response.completed", "response.incomplete", "response.failed"}) +_TURN_ENDS: Final = _RESPONSES_TURN_ENDS | {_REALTIME_RESPONSE_DONE} +_MODEL_NAMING_EVENTS: Final = _TURN_ENDS | {_SESSION_CREATED, _RESPONSE_CREATED} _RESPONSE_COST_HEADER: Final = "llm_provider-x-litellm-response-cost" @@ -42,22 +42,22 @@ class _ModelNamingEvent(BaseModel): response: _NamedModel | None = None -class _CompletedResponse(BaseModel): +class _EndedResponse(BaseModel): model_config = ConfigDict(extra="ignore", frozen=True) model: str | None = None - usage: ResponseAPIUsage + usage: ResponseAPIUsage | None = None -class _ResponsesCompletedEvent(BaseModel): +class _ResponsesTurnEndEvent(BaseModel): model_config = ConfigDict(extra="ignore", frozen=True) - type: Literal["response.completed"] - response: _CompletedResponse + type: Literal["response.completed", "response.incomplete", "response.failed"] + response: _EndedResponse _MODEL_NAMING_EVENT: Final = TypeAdapter(_ModelNamingEvent) -_RESPONSES_COMPLETED_EVENT: Final = TypeAdapter(_ResponsesCompletedEvent) +_RESPONSES_TURN_END_EVENT: Final = TypeAdapter(_ResponsesTurnEndEvent) @dataclass(frozen=True, slots=True) @@ -93,6 +93,15 @@ def observed_websocket_model(messages: WebsocketMessages) -> str | None: return next((model for model in map(_event_model, messages) if model is not None), None) +def openai_websocket_frame_kind(frame: Mapping[str, object]) -> WebsocketUpstreamFrameKind: + frame_type: Final = frame.get("type") + if frame_type in _TURN_ENDS: + return "ends_turn" + if frame_type == _SESSION_CREATED: + return "opens_session" + return "relayed" + + def _realtime_billing(messages: WebsocketMessages, model: str, logging_obj: LiteLLMLoggingObj) -> _Billing | None: if not any(event.get("type") == _REALTIME_RESPONSE_DONE for event in messages): return None @@ -111,19 +120,21 @@ def _realtime_billing(messages: WebsocketMessages, model: str, logging_obj: Lite def _responses_turn(event: Mapping[str, object], session_model: str) -> _CompletedTurn | None: - if event.get("type") != _RESPONSES_COMPLETED: + if event.get("type") not in _RESPONSES_TURN_ENDS: return None try: - completed: Final = _RESPONSES_COMPLETED_EVENT.validate_python(event) + ended: Final = _RESPONSES_TURN_END_EVENT.validate_python(event) except ValidationError: verbose_proxy_logger.warning( - "OpenAI websocket passthrough: a response.completed frame carried no readable usage, so it is not billed" + "OpenAI websocket passthrough: a %s frame carried no readable usage, so it is not billed", event.get("type") ) return None + if ended.response.usage is None: + return None return _CompletedTurn( - model=completed.response.model or session_model, + model=ended.response.model or session_model, usage=ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( # pyright: ignore[reportPrivateUsage] # the one shared Responses usage converter - completed.response.usage + ended.response.usage ), ) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 895533044c3..201ecee5b44 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -3,6 +3,7 @@ import asyncio import copy import json import posixpath +import re import traceback from base64 import b64encode from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterable, Mapping, MutableMapping, Sequence @@ -11,7 +12,7 @@ from datetime import datetime from functools import reduce from itertools import chain, count, groupby from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias, TypedDict, cast from urllib.parse import urlencode, urlparse import httpx @@ -142,6 +143,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, EndpointType, PassthroughStandardLoggingPayload, + WebsocketUpstreamFrameKind, ) from litellm.types.utils import TRUSTED_CALLBACK_VARS_FIELD, Usage @@ -2458,86 +2460,139 @@ def _with_trace_context(headers: Mapping[str, str], parent_span: object) -> dict return inject_trace_context(headers, parent_span=parent_span) +@dataclass(frozen=True, slots=True) +class WebsocketFrameModelAllowed: + access_groups: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class WebsocketFrameModelRefused: + message: str + error_type: str + + +WebsocketFrameModelVerdict: TypeAlias = WebsocketFrameModelAllowed | WebsocketFrameModelRefused + + class WebsocketFrameModelGate(Protocol): - async def __call__(self, model: str, valid_token: UserAPIKeyAuth, /) -> str | None: ... + async def __call__(self, websocket: WebSocket, model: str, /) -> WebsocketFrameModelVerdict: ... + + +WebsocketUpstreamFrameClassifier: TypeAlias = Callable[[Mapping[str, object]], WebsocketUpstreamFrameKind] @dataclass(frozen=True, slots=True) class _ForwardClientFrame: - pass + turn_access_groups: tuple[str, ...] | None @dataclass(frozen=True, slots=True) class _RefuseClientFrame: - message: str + refusal: WebsocketFrameModelRefused -_PLAIN_MODEL_KEY: Final = '"model"' -_JSON_STRING_ESCAPE: Final = "\\" -_CLIENT_FRAME_MODEL_PLACEMENTS: Final = ( - ("model",), - ("response", "model"), - ("session", "model"), +@dataclass(frozen=True, slots=True) +class _ClientFrameRequests: + turn_model: str | None + models: tuple[str, ...] + names_stored_prompt: bool + + +_NO_CLIENT_FRAME_REQUESTS: Final = _ClientFrameRequests(turn_model=None, models=(), names_stored_prompt=False) +_FRAME_TEXT_WORTH_PARSING: Final = ('"model"', '"prompt"', "\\") +_TURN_MODEL_PLACEMENTS: Final = (("model",), ("response", "model"), ("session", "model")) +_TRANSCRIPTION_MODEL_PLACEMENTS: Final = ( ("session", "input_audio_transcription", "model"), ("session", "audio", "input", "transcription", "model"), ) +_STORED_PROMPT_PLACEMENTS: Final = (("prompt",), ("response", "prompt"), ("session", "prompt")) +STORED_PROMPT_PROBE_MODEL: Final = "litellm-unverifiable-stored-prompt-model" +_STORED_PROMPT_REFUSAL: Final = ( + "This key is limited to specific models, and the gateway cannot check which model a stored prompt runs, " + "so frames that reference a stored prompt are refused for this key." +) +_FIRST_SENTENCE: Final = re.compile(r".*?[.!?](?=\s|$)") def _may_name_a_model(text_data: str) -> bool: - return _PLAIN_MODEL_KEY in text_data or _JSON_STRING_ESCAPE in text_data + return any(marker in text_data for marker in _FRAME_TEXT_WORTH_PARSING) def _nested_frame(frame: Mapping[str, object] | None, key: str) -> Mapping[str, object] | None: return None if frame is None else as_str_mapping(frame.get(key)) -def _placed_model(frame: Mapping[str, object], placement: Sequence[str]) -> str | None: +def _placed_value(frame: Mapping[str, object], placement: Sequence[str]) -> object: *parents, leaf = placement container: Final = reduce(_nested_frame, parents, frame) - if container is None: - return None - value: Final = container.get(leaf) - return value if isinstance(value, str) else None + return None if container is None else container.get(leaf) -def _client_frame_models(frame_data: str | bytes) -> tuple[str, ...]: +def _placed_models(frame: Mapping[str, object], placements: Iterable[Sequence[str]]) -> tuple[str, ...]: + placed: Final = (_placed_value(frame, placement) for placement in placements) + return tuple(value for value in placed if isinstance(value, str)) + + +def _client_frame_requests(frame_data: str | bytes) -> _ClientFrameRequests: if isinstance(frame_data, str) and not _may_name_a_model(frame_data): - return () + return _NO_CLIENT_FRAME_REQUESTS frame: Final = _json_object_frame(frame_data) if frame is None: - return () - placed: Final = (_placed_model(frame, placement) for placement in _CLIENT_FRAME_MODEL_PLACEMENTS) - return tuple(dict.fromkeys(model for model in placed if model is not None)) + return _NO_CLIENT_FRAME_REQUESTS + turn_models: Final = _placed_models(frame, _TURN_MODEL_PLACEMENTS) + prompts: Final = (_placed_value(frame, placement) for placement in _STORED_PROMPT_PLACEMENTS) + return _ClientFrameRequests( + turn_model=next(iter(turn_models), None), + models=tuple(dict.fromkeys((*turn_models, *_placed_models(frame, _TRANSCRIPTION_MODEL_PLACEMENTS)))), + names_stored_prompt=any(isinstance(prompt, Mapping) for prompt in prompts), + ) + + +def _stored_prompt_refusal(probe: WebsocketFrameModelRefused) -> WebsocketFrameModelRefused: + if probe.error_type.endswith("_model_access_denied"): + return WebsocketFrameModelRefused(message=_STORED_PROMPT_REFUSAL, error_type=probe.error_type) + return probe async def _client_frame_verdict( - frame_data: str | bytes, gate: WebsocketFrameModelGate | None, valid_token: UserAPIKeyAuth + websocket: WebSocket, frame_data: str | bytes, gate: WebsocketFrameModelGate | None ) -> _ForwardClientFrame | _RefuseClientFrame: if gate is None: - return _ForwardClientFrame() - for model in _client_frame_models(frame_data): - if (refusal := await gate(model, valid_token)) is not None: - return _RefuseClientFrame(message=refusal) - return _ForwardClientFrame() + return _ForwardClientFrame(turn_access_groups=None) + requests: Final = _client_frame_requests(frame_data) + verdicts: Final = tuple([await gate(websocket, model) for model in requests.models]) + refusal: Final = next((verdict for verdict in verdicts if isinstance(verdict, WebsocketFrameModelRefused)), None) + if refusal is not None: + return _RefuseClientFrame(refusal=refusal) + if requests.names_stored_prompt and isinstance( + probe := await gate(websocket, STORED_PROMPT_PROBE_MODEL), WebsocketFrameModelRefused + ): + return _RefuseClientFrame(refusal=_stored_prompt_refusal(probe)) + allowed: Final = tuple(verdict for verdict in verdicts if isinstance(verdict, WebsocketFrameModelAllowed)) + if requests.turn_model is None or not allowed: + return _ForwardClientFrame(turn_access_groups=None) + return _ForwardClientFrame(turn_access_groups=allowed[0].access_groups) -async def _refuse_client_frame(websocket: WebSocket, upstream_ws: ClientConnection, message: str) -> None: - await websocket.send_text(realtime_error_event(message, error_type="invalid_request_error")) - await websocket.close(code=CloseCode.POLICY_VIOLATION, reason=_truncated_close_reason(message)) +def _refusal_close_reason(message: str) -> str: + first_sentence: Final = _FIRST_SENTENCE.match(message) + return _truncated_close_reason(message if first_sentence is None else first_sentence.group(0)) + + +async def _refuse_client_frame( + websocket: WebSocket, upstream_ws: ClientConnection, refusal: WebsocketFrameModelRefused +) -> None: + await websocket.send_text(realtime_error_event(refusal.message, error_type=refusal.error_type)) + await websocket.close(code=CloseCode.POLICY_VIOLATION, reason=_refusal_close_reason(refusal.message)) await upstream_ws.close() -def _stamp_billed_access_groups(kwargs: Mapping[str, object], valid_token: UserAPIKeyAuth) -> None: - """Refresh the session's logging metadata with the access groups its frames were authorized through. - - The gate only learns them once a frame names a model, which is long after the relay stamped the - key's groups onto its metadata, so the spend writer would otherwise charge no group for the session. - """ - groups: Final = valid_token.matched_model_access_groups +def _stamp_billed_access_groups(kwargs: Mapping[str, object], access_groups: tuple[str, ...]) -> None: + """Charge a turn to the access groups that authorized its model, and to none when no group did.""" litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) metadata: Final = None if litellm_params is None else litellm_params.get("metadata") - if groups and isinstance(metadata, MutableMapping): - metadata[MODEL_ACCESS_GROUP_METADATA_KEY] = groups + if isinstance(metadata, MutableMapping): + metadata[MODEL_ACCESS_GROUP_METADATA_KEY] = list(access_groups) async def _relay_client_frame( @@ -2546,16 +2601,35 @@ async def _relay_client_frame( checked_frame: str | bytes, forwarded_frame: str | bytes, gate: WebsocketFrameModelGate | None, - valid_token: UserAPIKeyAuth, -) -> bool: - """Forward one client frame upstream, or refuse it and report that the session is over.""" - match await _client_frame_verdict(checked_frame, gate, valid_token): - case _RefuseClientFrame(message=refusal_message): - await _refuse_client_frame(websocket, upstream_ws, refusal_message) - return False +) -> _ForwardClientFrame | _RefuseClientFrame: + """Forward one client frame upstream, or refuse it and end the session.""" + verdict: Final = await _client_frame_verdict(websocket, checked_frame, gate) + match verdict: + case _RefuseClientFrame(refusal=refusal): + await _refuse_client_frame(websocket, upstream_ws, refusal) case _ForwardClientFrame(): await upstream_ws.send(forwarded_frame) - return True + return verdict + + +_EXPLICIT_SUCCESS_KWARGS: Final = frozenset({"logging_obj", "passthrough_logging_payload", "custom_llm_provider"}) + + +class _WebsocketTurnLedger: + """The groups that authorized the turn in flight, when it started, and how many turns the socket has logged.""" + + def __init__(self, access_groups: tuple[str, ...], started: datetime) -> None: + self.access_groups = access_groups + self.started = started + self.logged_turns = 0 + + def authorize(self, access_groups: tuple[str, ...] | None) -> None: + if access_groups is not None: + self.access_groups = access_groups + + def close_turn(self, end_time: datetime) -> None: + self.logged_turns += 1 + self.started = end_time async def websocket_passthrough_request( @@ -2569,6 +2643,7 @@ async def websocket_passthrough_request( accept_websocket: bool = True, setup_model_rewriter: Callable[[str], str] | None = None, client_frame_model_gate: WebsocketFrameModelGate | None = None, + upstream_frame_kind: WebsocketUpstreamFrameClassifier | None = None, ): """ WebSocket passthrough request handler. @@ -2582,6 +2657,9 @@ async def websocket_passthrough_request( endpoint: The endpoint path (for logging purposes) cost_per_request: Optional field - cost per request to the target endpoint setup_model_rewriter: Optional rewrite of the setup frame's model before it reaches the upstream + client_frame_model_gate: Optional per-frame check of every model a client frame names + upstream_frame_kind: Optional sort of upstream frames into the one that ends a billable turn and the ones + that open the session, so each turn logs its own spend and keeps the session's provider-named model """ from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy.proxy_server import proxy_config, proxy_logging_obj @@ -2617,26 +2695,31 @@ async def websocket_passthrough_request( proxy_config=proxy_config, route_description="websocket_passthrough", ) - logging_obj: Final = Logging( - model="unknown", - messages=[{"role": "user", "content": "WebSocket connection"}], - stream=True, # WebSockets are inherently streaming - call_type="pass_through_endpoint", - start_time=start_time, - litellm_call_id=litellm_call_id, - function_id="websocket_passthrough", - dynamic_success_callbacks=team_callbacks.success_callbacks, - dynamic_failure_callbacks=team_callbacks.failure_callbacks, - kwargs=team_callbacks.logging_kwargs, - ) - # Create passthrough logging payload - passthrough_logging_payload: Final = PassthroughStandardLoggingPayload( - url=target, - request_body={}, # WebSocket doesn't have a traditional request body - request_method="WEBSOCKET", - cost_per_request=cost_per_request, - ) + def new_call_logging(call_start: datetime, call_id: str) -> LiteLLMLoggingObj: + return Logging( + model="unknown", + messages=[{"role": "user", "content": "WebSocket connection"}], + stream=True, + call_type="pass_through_endpoint", + start_time=call_start, + litellm_call_id=call_id, + function_id="websocket_passthrough", + dynamic_success_callbacks=team_callbacks.success_callbacks, + dynamic_failure_callbacks=team_callbacks.failure_callbacks, + kwargs=team_callbacks.logging_kwargs, + ) + + def new_call_payload() -> PassthroughStandardLoggingPayload: + return PassthroughStandardLoggingPayload( + url=target, + request_body={}, + request_method="WEBSOCKET", + cost_per_request=cost_per_request, + ) + + logging_obj: Final = new_call_logging(start_time, litellm_call_id) + passthrough_logging_payload: Final = new_call_payload() # Create a dummy request object for WebSocket connections to maintain compatibility # with the existing _init_kwargs_for_pass_through_endpoint function @@ -2658,35 +2741,117 @@ async def websocket_passthrough_request( ) # Initialize kwargs for logging using the same pattern as HTTP passthrough - kwargs: Final = HttpPassThroughEndpointHelpers.init_kwargs_for_pass_through_endpoint( - user_api_key_dict=user_api_key_dict, - _parsed_body={}, # WebSocket doesn't have a traditional request body - passthrough_logging_payload=passthrough_logging_payload, - litellm_call_id=litellm_call_id, - request=dummy_request, - logging_obj=logging_obj, + def init_call_kwargs( + token: UserAPIKeyAuth, + call_payload: PassthroughStandardLoggingPayload, + call_id: str, + call_logging: LiteLLMLoggingObj, + ) -> dict[str, object]: # mutable-ok: the Vertex relay stamps the model from setup frames into these kwargs + return HttpPassThroughEndpointHelpers.init_kwargs_for_pass_through_endpoint( + user_api_key_dict=token, + _parsed_body={}, # WebSocket doesn't have a traditional request body + passthrough_logging_payload=call_payload, + litellm_call_id=call_id, + request=dummy_request, + logging_obj=call_logging, + ) + + kwargs: Final = init_call_kwargs(user_api_key_dict, passthrough_logging_payload, litellm_call_id, logging_obj) + + def start_call_logging(call_logging: LiteLLMLoggingObj, call_kwargs: Mapping[str, object]) -> None: + call_logging.update_environment_variables( + model="unknown", + user="unknown", + optional_params={}, + litellm_params=dict(as_str_mapping(call_kwargs.get("litellm_params")) or {}), + call_type="pass_through_endpoint", + ) + call_logging.pre_call( + input=[{"role": "user", "content": "WebSocket connection"}], + api_key="", + additional_args={ + "complete_input_dict": {}, + "api_base": target, + "headers": upstream_headers, + }, + ) + + start_call_logging(logging_obj, kwargs) + + class MockWebSocketRequest: + def __init__(self, target_url: str): + self.method = "WEBSOCKET" + self.url = target_url + + class MockWebSocketResponse: + def __init__(self, target_url: str): + self.status_code = 200 + self.text = "WebSocket connection successful" + self.headers: dict[str, str] = {} + self.request = MockWebSocketRequest(target_url) + + ledger: Final = _WebsocketTurnLedger( + access_groups=tuple(user_api_key_dict.matched_model_access_groups or ()), started=start_time ) - # Update logging environment variables - logging_obj.update_environment_variables( - model="unknown", - user="unknown", - optional_params={}, - litellm_params=dict(kwargs.get("litellm_params", {})), - call_type="pass_through_endpoint", - ) - logging_obj.model_call_details["litellm_call_id"] = litellm_call_id + def log_frames( + call_logging: LiteLLMLoggingObj, + call_payload: PassthroughStandardLoggingPayload, + call_kwargs: Mapping[str, object], + end_time: datetime, + ) -> None: + frames: Final = list(websocket_messages) + websocket_messages[:] = ( + () + if upstream_frame_kind is None + else (frame for frame in frames if upstream_frame_kind(frame) == "opens_session") + ) + logged_payload: Final[PassthroughStandardLoggingPayload] = { + **call_payload, + "response_body": frames, + "end_time": end_time, + } + if client_frame_model_gate is not None: + _stamp_billed_access_groups(call_kwargs, ledger.access_groups) + success_kwargs: Final = { + key: value for key, value in call_kwargs.items() if key not in _EXPLICIT_SUCCESS_KWARGS + } + custom_llm_provider: Final = call_kwargs.get("custom_llm_provider") + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( + async_coroutine=pass_through_endpoint_logging.pass_through_async_success_handler( + httpx_response=MockWebSocketResponse(target), + response_body=frames, + url_route=endpoint or "", + result="websocket_connection_successful", + start_time=ledger.started, + end_time=end_time, + logging_obj=call_logging, + cache_hit=False, + request_body={}, + passthrough_logging_payload=logged_payload, + custom_llm_provider=custom_llm_provider if isinstance(custom_llm_provider, str) else None, + **success_kwargs, + ) + ) + bind_budget_reservation_to_callbacks(call_logging.litellm_params) - # Pre-call logging - logging_obj.pre_call( - input=[{"role": "user", "content": "WebSocket connection"}], - api_key="", - additional_args={ - "complete_input_dict": {}, - "api_base": target, - "headers": upstream_headers, - }, - ) + def log_turn() -> None: + end_time: Final = datetime.now() + if ledger.logged_turns == 0: + log_frames(logging_obj, passthrough_logging_payload, kwargs, end_time) + else: + turn_call_id: Final = str(uuid.uuid4()) + turn_logging: Final = new_call_logging(ledger.started, turn_call_id) + turn_payload: Final = new_call_payload() + turn_kwargs: Final = init_call_kwargs( + user_api_key_dict.model_copy(update={"budget_reservation": None}), + turn_payload, + turn_call_id, + turn_logging, + ) + start_call_logging(turn_logging, turn_kwargs) + log_frames(turn_logging, turn_payload, turn_kwargs, end_time) + ledger.close_turn(end_time) ### CALL HOOKS ### - modify incoming data / reject request before calling the model websocket_hook_data: Final = guardrail_request_data_with_streaming(MappingProxyType({}), is_streaming=True) @@ -2778,25 +2943,27 @@ async def websocket_passthrough_request( ) # Not a JSON message or doesn't contain setup data - if not await _relay_client_frame( - websocket, - upstream_ws, - text_data, - _rewrite_vertex_live_setup_model(text_data, setup_model_rewriter), - client_frame_model_gate, - user_api_key_dict, + if isinstance( + relayed_text := await _relay_client_frame( + websocket, + upstream_ws, + text_data, + _rewrite_vertex_live_setup_model(text_data, setup_model_rewriter), + client_frame_model_gate, + ), + _RefuseClientFrame, ): break + ledger.authorize(relayed_text.turn_access_groups) elif bytes_data is not None: - if not await _relay_client_frame( - websocket, - upstream_ws, - bytes_data, - bytes_data, - client_frame_model_gate, - user_api_key_dict, + if isinstance( + relayed_bytes := await _relay_client_frame( + websocket, upstream_ws, bytes_data, bytes_data, client_frame_model_gate + ), + _RefuseClientFrame, ): break + ledger.authorize(relayed_bytes.turn_access_groups) except asyncio.CancelledError: raise except Exception: @@ -2831,9 +2998,13 @@ async def websocket_passthrough_request( _extract_vertex_live_model_from_setup_response(message_data) return websocket_messages.append(message_data) + if upstream_frame_kind is not None and upstream_frame_kind(message_data) == "ends_turn": + log_turn() async def relay_upstream_frame(upstream_message: str | bytes) -> None: record_upstream_frame(upstream_message) + if not _client_socket_is_open(websocket): + return if isinstance(upstream_message, bytes): await websocket.send_bytes(upstream_message) else: @@ -2890,64 +3061,8 @@ async def websocket_passthrough_request( reason=_truncated_close_reason(upstream_close.reason), ) - end_time: Final = datetime.now() - - # Update passthrough logging payload with response data - passthrough_logging_payload["response_body"] = websocket_messages - passthrough_logging_payload["end_time"] = end_time - - _stamp_billed_access_groups(kwargs, user_api_key_dict) - - # Remove logging_obj from kwargs to avoid duplicate keyword argument - success_kwargs: Final = kwargs.copy() - success_kwargs.pop("logging_obj", None) - - # # Add user authentication context for database logging - # if user_api_key_dict: - # success_kwargs.setdefault('litellm_params', {}) - # success_kwargs['litellm_params'].update({ - # 'proxy_server_request': { - # 'body': { - # 'user': user_api_key_dict.user_id, - # 'team_id': user_api_key_dict.team_id, - # 'end_user_id': user_api_key_dict.end_user_id, - # } - # } - # }) - # # Also add the user_api_key for direct access - # success_kwargs['user_api_key'] = user_api_key_dict.api_key - - # Create a dummy httpx.Response for WebSocket connections - class MockWebSocketResponse: - def __init__(self, target_url: str): - self.status_code = 200 - self.text = "WebSocket connection successful" - self.headers: dict[str, str] = {} - self.request = MockWebSocketRequest(target_url) - - class MockWebSocketRequest: - def __init__(self, target_url: str): - self.method = "WEBSOCKET" - self.url = target_url - - mock_response: Final = MockWebSocketResponse(target) - - # Use the same success handler as HTTP passthrough endpoints - GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( - async_coroutine=pass_through_endpoint_logging.pass_through_async_success_handler( - httpx_response=mock_response, - response_body=websocket_messages, - url_route=endpoint or "", - result="websocket_connection_successful", - start_time=start_time, - end_time=end_time, - logging_obj=logging_obj, - cache_hit=False, - request_body={}, - **success_kwargs, - ) - ) - bind_budget_reservation_to_callbacks(logging_obj.litellm_params) + if ledger.logged_turns == 0: + log_turn() # Call the proxy logging success hook if proxy_logging_obj: diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 803df9f83ca..3d3e090a013 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -422,7 +422,7 @@ async def invalidate_budget_reservation_counters( async def release_or_invalidate_budget_reservation( - budget_reservation: dict | None, # mutable-ok: stamps finalized on the caller's shared reservation dict + budget_reservation: dict[str, object] | None, # mutable-ok: stamps finalized on the caller's reservation ) -> None: """Reconcile a still-open reservation on a terminal path that settles no cost. diff --git a/litellm/types/passthrough_endpoints/pass_through_endpoints.py b/litellm/types/passthrough_endpoints/pass_through_endpoints.py index 0d99cd54a6b..b13211d8f37 100644 --- a/litellm/types/passthrough_endpoints/pass_through_endpoints.py +++ b/litellm/types/passthrough_endpoints/pass_through_endpoints.py @@ -1,8 +1,10 @@ +from collections.abc import Mapping, Sequence +from datetime import datetime from enum import Enum -from typing import Final +from typing import Final, Literal, TypeAlias from pydantic import Field -from typing_extensions import TypedDict +from typing_extensions import ReadOnly, TypedDict from litellm.types.llms.base import LiteLLMPydanticObjectBase @@ -27,6 +29,8 @@ LITELLM_PASS_THROUGH_ENDPOINT_MARKER: Final = "__litellm_pass_through_endpoint__ SETTLES_BACKGROUND_RESPONSE_KWARG: Final = "settles_background_response" +WebsocketUpstreamFrameKind: TypeAlias = Literal["ends_turn", "opens_session", "relayed"] + class BackgroundResponseSettlement(str, Enum): SETTLES = "settles" @@ -62,7 +66,7 @@ class PassthroughStandardLoggingPayload(TypedDict, total=False): """ The body of the request """ - response_body: dict | None # only tracked for non-streaming responses + response_body: dict | Sequence[Mapping[str, object]] | None # only tracked for non-streaming responses """ The body of the response """ @@ -74,6 +78,8 @@ class PassthroughStandardLoggingPayload(TypedDict, total=False): Optional field, we use this for cost tracking only if it's set. """ + end_time: ReadOnly[datetime] + class PassThroughGuardrailSettings(LiteLLMPydanticObjectBase): """ diff --git a/tests/unit/proxy/auth/test_model_access_group_budgets.py b/tests/unit/proxy/auth/test_model_access_group_budgets.py index 7b0c4d9996e..7506bd031d9 100644 --- a/tests/unit/proxy/auth/test_model_access_group_budgets.py +++ b/tests/unit/proxy/auth/test_model_access_group_budgets.py @@ -27,7 +27,6 @@ from litellm.proxy.auth.auth_checks import ( _model_access_group_max_budget_check, collect_matched_model_access_groups, common_checks, - enforce_model_access_group_budget, stamp_matched_model_access_groups, ) from litellm.proxy.common_utils.reset_budget_job import _model_access_group_counter_key @@ -552,65 +551,3 @@ async def test_common_checks_blocks_a_request_whose_group_is_over_budget(): async def test_free_model_routes_skip_the_model_access_group_budget_check(): """skip_budget_checks is how free models stay free; it has to cover this budget too.""" assert await _common_checks_with_over_budget_group(skip_budget_checks=True) is True - - -async def _enforce_for_a_route_without_a_request_body( - *rows: _MagBudgetRow, - key_models: list[str], - spend_by_counter_key: dict[str, float], -) -> UserAPIKeyAuth: - """Run the pair the websocket frame gate runs, which reaches the proxy for what it needs.""" - valid_token = UserAPIKeyAuth(api_key="hashed", models=key_models, user_id=USER_ID) - cache = await _cache() - read, _ = _spend_reader(spend_by_counter_key) - with ( - patch("litellm.proxy.proxy_server.prisma_client", _RecordingPrismaClient(*rows)), # test-quality-ok: read lazily inside the helper, which takes no client - patch("litellm.proxy.proxy_server.user_api_key_cache", cache), # test-quality-ok: same lazy import - patch("litellm.proxy.proxy_server.proxy_logging_obj", ProxyLogging(user_api_key_cache=cache)), # test-quality-ok: same lazy import - patch("litellm.proxy.proxy_server.get_current_spend", read), # test-quality-ok: get_current_spend is lazily imported inside the budget check and has no injection point - ): - await enforce_model_access_group_budget( - model="gpt-4o", - valid_token=valid_token, - llm_router=Router(model_list=MODEL_LIST), - ) - return valid_token - - -@pytest.mark.asyncio -async def test_caller_granted_the_model_only_through_an_exhausted_group_is_refused(): - """Routes with no request body, the websocket passthrough among them, run this instead of - ``common_checks``, so the group budget has to stop them the same way.""" - with pytest.raises(litellm.BudgetExceededError) as exc_info: - await _enforce_for_a_route_without_a_request_body( - _MagBudgetRow("tier-a", max_budget=1.0), - key_models=["tier-a"], - spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 1.5}, - ) - - assert exc_info.value.entity_id == "tier-a" - assert exc_info.value.current_cost == 1.5 - - -@pytest.mark.asyncio -async def test_group_that_authorized_the_call_is_stamped_for_the_spend_writer(): - """Without the stamp the call runs on the group's budget and never pays into it, so the pool - can never catch up with what it authorized.""" - valid_token = await _enforce_for_a_route_without_a_request_body( - _MagBudgetRow("tier-a", max_budget=10.0), - key_models=["tier-a"], - spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 1.0}, - ) - - assert valid_token.matched_model_access_groups == ["tier-a"] - - -@pytest.mark.asyncio -async def test_model_granted_directly_charges_no_group(): - valid_token = await _enforce_for_a_route_without_a_request_body( - _MagBudgetRow("tier-a", max_budget=1.0), - key_models=["gpt-4o"], - spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 1.5}, - ) - - assert valid_token.matched_model_access_groups is None diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py index 91a01f02d94..253634760fa 100644 --- a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py @@ -8,6 +8,7 @@ import pytest import litellm from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_websocket_passthrough_logging_handler import ( OpenAIWebsocketPassthroughLoggingHandler, + openai_websocket_frame_kind, ) REALTIME_MODEL: Final = "gpt-realtime-2.1" @@ -42,9 +43,11 @@ def _realtime_response_done(text_in: int, audio_in: int, text_out: int, audio_ou } -def _responses_completed(model: str, input_tokens: int, output_tokens: int) -> dict: +def _responses_completed( + model: str, input_tokens: int, output_tokens: int, event: str = "response.completed" +) -> dict: return { - "type": "response.completed", + "type": event, "response": { "id": "resp_fake", "model": model, @@ -288,3 +291,47 @@ def test_a_turn_the_caller_walked_out_of_still_names_the_model_it_ran(): assert handled["result"] is None assert handled["kwargs"] == {"model": REASONING_MODEL, "custom_llm_provider": "openai"} assert logging_obj.model_call_details["model"] == REASONING_MODEL + + +@pytest.mark.parametrize("event", ["response.incomplete", "response.failed"]) +def test_a_responses_turn_that_ended_early_is_billed_for_the_tokens_it_used(event): + """ + A turn cut short by max_output_tokens or a provider failure still reports + the tokens it burned, and OpenAI bills them, so the row has to as well. + """ + handled = _handle( + [_responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5, event=event)], + url_route="/openai_passthrough/v1/responses", + ) + + prices = litellm.model_cost[RESPONSES_MODEL] + expected_cost = 13 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"] + assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost) + assert handled["result"].usage.total_tokens == 18 + + +def test_a_failed_responses_turn_that_reported_no_usage_is_not_billed(): + handled = _handle( + [{"type": "response.failed", "response": {"model": RESPONSES_MODEL, "status": "failed", "usage": None}}], + url_route="/openai_passthrough/v1/responses", + ) + + assert handled["result"] is None + assert handled["kwargs"] == {"model": RESPONSES_MODEL, "custom_llm_provider": "openai"} + + +@pytest.mark.parametrize( + ("frame_type", "kind"), + [ + ("response.completed", "ends_turn"), + ("response.incomplete", "ends_turn"), + ("response.failed", "ends_turn"), + ("response.done", "ends_turn"), + ("session.created", "opens_session"), + ("response.created", "relayed"), + ("response.output_text.delta", "relayed"), + ("error", "relayed"), + ], +) +def test_every_frame_that_settles_a_turn_closes_it_and_nothing_else_does(frame_type, kind): + assert openai_websocket_frame_kind({"type": frame_type}) == kind diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 6f35ebbc161..e6a0b1be28c 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -42,9 +42,13 @@ from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import route from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( DEFAULT_PASS_THROUGH_REQUEST_TIMEOUT_SECONDS, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, + STORED_PROMPT_PROBE_MODEL, HttpPassThroughEndpointHelpers, InitPassThroughEndpointHelpers, SafeRouteAdder, + WebsocketFrameModelAllowed, + WebsocketFrameModelRefused, + WebsocketFrameModelVerdict, _registered_pass_through_routes, _truncate_upstream_error_body, _update_metadata_with_tags_in_header, @@ -58,6 +62,9 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( set_env_variables_in_header, websocket_passthrough_request, ) +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_websocket_passthrough_logging_handler import ( + openai_websocket_frame_kind, +) from litellm.proxy.pass_through_endpoints.success_handler import ( PassThroughEndpointLogging, ) @@ -7211,7 +7218,7 @@ def _patched_websocket_passthrough_environment(upstream_ws): mock_worker.ensure_initialized_and_enqueue = MagicMock( side_effect=lambda async_coroutine: async_coroutine.close() ) - yield + yield mock_proxy_logging async def _pending_receive(): @@ -7522,14 +7529,28 @@ SUCCESS_HANDLER: Final = ( ) +KEY_MODEL_REFUSAL: Final = "key not allowed to access model. This key can only access models=['gpt-5.4-mini']." +KEY_MODEL_CLOSE_REASON: Final = "key not allowed to access model." + + class RecordingFrameModelGate: - def __init__(self, refusals: Mapping[str, str] = MappingProxyType({})): + def __init__( + self, + refusals: Mapping[str, str] = MappingProxyType({}), + groups: Mapping[str, tuple[str, ...]] = MappingProxyType({}), + error_type: str = "key_model_access_denied", + ): self.refusals = refusals + self.groups = groups + self.error_type = error_type self.checked: list[str] = [] - async def __call__(self, model: str, valid_token: UserAPIKeyAuth, /) -> Optional[str]: + async def __call__(self, websocket: object, model: str, /) -> WebsocketFrameModelVerdict: self.checked.append(model) - return self.refusals.get(model) + refusal = self.refusals.get(model) + if refusal is not None: + return WebsocketFrameModelRefused(message=refusal, error_type=self.error_type) + return WebsocketFrameModelAllowed(access_groups=self.groups.get(model, ())) def _received_frame(frame): @@ -7594,8 +7615,7 @@ async def test_websocket_passthrough_checks_the_model_a_client_frame_asks_for(fr @pytest.mark.asyncio async def test_websocket_passthrough_refuses_a_frame_naming_a_model_the_key_cannot_call(): - refusal = "key not allowed to access model. This key can only access models=['gpt-5.4-mini']." - gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": refusal})) + gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": KEY_MODEL_REFUSAL})) frame = json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}) relayed = await _relay_client_frames([frame], gate) @@ -7603,16 +7623,34 @@ async def test_websocket_passthrough_refuses_a_frame_naming_a_model_the_key_cann assert relayed.upstream.send.await_args_list == [] assert json.loads(relayed.client.send_text.await_args.args[0]) == { "type": "error", - "error": {"type": "invalid_request_error", "message": refusal}, + "error": {"type": "key_model_access_denied", "message": KEY_MODEL_REFUSAL}, } - assert relayed.client.close.await_args_list == [call(code=1008, reason=refusal)] + assert relayed.client.close.await_args_list == [call(code=1008, reason=KEY_MODEL_CLOSE_REASON)] relayed.upstream.close.assert_awaited() +@pytest.mark.asyncio +async def test_websocket_passthrough_closes_with_the_whole_first_sentence_of_the_refusal(): + """ + The close frame only has room for 123 bytes, so a refusal cut at that limit + reads as a truncated number. The first sentence names the reason whole, and + the error event before the close still carries every word with its type. + """ + refusal = "Budget has been exceeded! Key=" + "f" * 64 + " Spend=1.23456789, Budget=1.0" + gate = RecordingFrameModelGate(MappingProxyType({"gpt-5.4": refusal}), error_type="budget_exceeded") + + relayed = await _relay_client_frames([json.dumps({"type": "response.create", "model": "gpt-5.4"})], gate) + + assert relayed.client.close.await_args_list == [call(code=1008, reason="Budget has been exceeded!")] + assert json.loads(relayed.client.send_text.await_args.args[0])["error"] == { + "type": "budget_exceeded", + "message": refusal, + } + + @pytest.mark.asyncio async def test_websocket_passthrough_refuses_the_whole_frame_when_any_named_model_is_denied(): - refusal = "key not allowed to access model. This key can only access models=['gpt-realtime-2.1']." - gate = RecordingFrameModelGate(MappingProxyType({"gpt-4o-transcribe": refusal})) + gate = RecordingFrameModelGate(MappingProxyType({"gpt-4o-transcribe": KEY_MODEL_REFUSAL})) frame = json.dumps( { "type": "session.update", @@ -7624,13 +7662,12 @@ async def test_websocket_passthrough_refuses_the_whole_frame_when_any_named_mode assert gate.checked == ["gpt-realtime-2.1", "gpt-4o-transcribe"] assert relayed.upstream.send.await_args_list == [] - assert relayed.client.close.await_args_list == [call(code=1008, reason=refusal)] + assert relayed.client.close.await_args_list == [call(code=1008, reason=KEY_MODEL_CLOSE_REASON)] @pytest.mark.asyncio async def test_websocket_passthrough_stops_relaying_after_a_refused_frame(): - refusal = "key not allowed to access model. This key can only access models=['gpt-5.4-mini']." - gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": refusal})) + gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": KEY_MODEL_REFUSAL})) frames = [ json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}), json.dumps({"type": "response.create"}), @@ -7697,27 +7734,76 @@ async def test_websocket_passthrough_checks_a_model_named_with_escaped_json(fram "\\u006dodel" asks for a model just as plainly as one carrying "model" and has to face the same check. """ - refusal = "key not allowed to access model. This key can only access models=['gpt-5.4-mini']." - gate = RecordingFrameModelGate(MappingProxyType({expected_model: refusal})) + gate = RecordingFrameModelGate(MappingProxyType({expected_model: KEY_MODEL_REFUSAL})) relayed = await _relay_client_frames([frame], gate) assert gate.checked == [expected_model] assert relayed.upstream.send.await_args_list == [] - assert relayed.client.close.await_args_list == [call(code=1008, reason=refusal)] + assert relayed.client.close.await_args_list == [call(code=1008, reason=KEY_MODEL_CLOSE_REASON)] @pytest.mark.asyncio async def test_websocket_passthrough_checks_the_model_a_binary_frame_asks_for(): - refusal = "key not allowed to access model. This key can only access models=['gpt-5.4-mini']." - gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": refusal})) + gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": KEY_MODEL_REFUSAL})) frame = json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}).encode() relayed = await _relay_client_frames([frame], gate) assert gate.checked == ["gpt-realtime-2.1"] assert relayed.upstream.send.await_args_list == [] - assert relayed.client.close.await_args_list == [call(code=1008, reason=refusal)] + assert relayed.client.close.await_args_list == [call(code=1008, reason=KEY_MODEL_CLOSE_REASON)] + + +STORED_PROMPT_FRAMES: Final = ( + {"type": "response.create", "prompt": {"id": "pmpt_fake", "version": "2"}}, + {"type": "response.create", "response": {"prompt": {"id": "pmpt_fake"}}}, + {"type": "session.update", "session": {"prompt": {"id": "pmpt_fake"}}}, +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("frame", STORED_PROMPT_FRAMES) +async def test_websocket_passthrough_refuses_a_stored_prompt_for_a_key_limited_to_some_models(frame): + """ + A stored prompt carries its own model on the provider side, so a key held to + a model list could reach any model through one. The gateway cannot see that + model, so a key that is not allowed every model cannot run a stored prompt. + """ + gate = RecordingFrameModelGate(MappingProxyType({STORED_PROMPT_PROBE_MODEL: KEY_MODEL_REFUSAL})) + + relayed = await _relay_client_frames([json.dumps(frame)], gate) + + assert gate.checked == [STORED_PROMPT_PROBE_MODEL] + assert relayed.upstream.send.await_args_list == [] + assert "stored prompt" in json.loads(relayed.client.send_text.await_args.args[0])["error"]["message"] + assert relayed.client.close.await_args.kwargs["code"] == 1008 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("frame", STORED_PROMPT_FRAMES) +async def test_websocket_passthrough_forwards_a_stored_prompt_for_a_key_allowed_every_model(frame): + gate = RecordingFrameModelGate() + + relayed = await _relay_client_frames([json.dumps(frame)], gate) + + assert [call_.args[0] for call_ in relayed.upstream.send.await_args_list] == [json.dumps(frame)] + assert relayed.client.close.await_args_list == [call()] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_keeps_the_budget_reason_when_a_stored_prompt_is_refused_for_spend(): + refusal = "Budget has been exceeded! Key=hashed Spend=2.0, Budget=1.0" + gate = RecordingFrameModelGate( + MappingProxyType({STORED_PROMPT_PROBE_MODEL: refusal}), error_type="budget_exceeded" + ) + + relayed = await _relay_client_frames([json.dumps(STORED_PROMPT_FRAMES[0])], gate) + + assert json.loads(relayed.client.send_text.await_args.args[0])["error"] == { + "type": "budget_exceeded", + "message": refusal, + } @pytest.mark.asyncio @@ -7791,6 +7877,196 @@ async def test_websocket_passthrough_still_bills_a_session_the_client_walked_out ] +class ScriptedUpstreamWebSocket: + """Answers each client frame with its scripted reply, and flushes what was still in flight when it is closed""" + + def __init__(self, *opening: str, replies: Sequence[Sequence[str]] = (), in_flight: Sequence[str] = ()): + self._frames: asyncio.Queue[str | None] = asyncio.Queue() + for frame in opening: + self._frames.put_nowait(frame) + self._replies = iter(replies) + self._in_flight = in_flight + self._closed = False + self.send = AsyncMock(side_effect=self._reply) + self.close = AsyncMock(side_effect=self._flush_and_close) + + async def _reply(self, frame: str | bytes) -> None: + for reply in next(self._replies, ()): + self._frames.put_nowait(reply) + + async def _flush_and_close(self, *args: object, **kwargs: object) -> None: + if self._closed: + return + self._closed = True + for frame in self._in_flight: + self._frames.put_nowait(frame) + self._frames.put_nowait(None) + while not self._frames.empty(): + await asyncio.sleep(0) + + async def recv(self, decode: bool | None = None) -> str: + from websockets.exceptions import ConnectionClosedOK + from websockets.frames import Close + + frame = await self._frames.get() + if frame is None: + raise ConnectionClosedOK(rcvd=Close(1000, ""), sent=Close(1000, ""), rcvd_then_sent=True) + return frame + + +def _turn_taking_client(frames: Sequence[str], answered_by: str) -> MagicMock: + """A client that sends its next frame only once the answer to the last one reached it, and that Starlette + refuses to send to once it is closed""" + from starlette.websockets import WebSocketState + + delivered: list[str] = [] + script = iter(frames) + receives = iter(range(len(frames) + 1)) + + async def receive(): + answers_due = next(receives) + while sum(answered_by in text for text in delivered) < answers_due: + await asyncio.sleep(0) + frame = next(script, None) + return {"type": "websocket.disconnect"} if frame is None else _received_frame(frame) + + websocket = _client_websocket(receive) + + async def send_text(text: str) -> None: + if websocket.application_state == WebSocketState.DISCONNECTED: + raise RuntimeError('Cannot call "send" once a close message has been sent.') + delivered.append(text) + + websocket.send_text = AsyncMock(side_effect=send_text) + return websocket + + +async def _relay_turns(client, upstream_ws, gate, endpoint="/openai_passthrough/v1/responses"): + with ( + _patched_websocket_passthrough_environment(upstream_ws) as proxy_logging, + patch(SUCCESS_HANDLER, new_callable=AsyncMock) as success_handler, + ): + await asyncio.wait_for( + websocket_passthrough_request( + websocket=client, + target="wss://api.openai.com/v1/responses", + custom_headers={"Authorization": "Bearer sk-provider"}, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + forward_headers=False, + endpoint=endpoint, + accept_websocket=False, + client_frame_model_gate=gate, + upstream_frame_kind=openai_websocket_frame_kind, + ), + timeout=10, + ) + return SimpleNamespace(success_handler=success_handler, proxy_logging=proxy_logging) + + +def _responses_turn(model: str) -> list[str]: + return [ + json.dumps({"type": "response.created", "response": {"model": model}}), + json.dumps( + { + "type": "response.completed", + "response": {"model": model, "usage": {"input_tokens": 9, "output_tokens": 3, "total_tokens": 12}}, + } + ), + ] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_logs_each_turn_on_its_own_row_charged_to_the_groups_that_allowed_it(): + """ + One socket can run a model from one access group and then a model from + another. Each turn has to reach spend as it finishes, under its own call, + and be charged to the group that let its model through, so a group budget + sees its share during the session and never pays for the other group's turn. + """ + gate = RecordingFrameModelGate( + groups=MappingProxyType({"gpt-5.4-mini": ("tier-a",), "gpt-5.4": ("tier-b",)}), + ) + client = _turn_taking_client( + [ + json.dumps({"type": "response.create", "model": "gpt-5.4-mini"}), + json.dumps({"type": "response.create", "model": "gpt-5.4"}), + ], + answered_by='"response.completed"', + ) + upstream_ws = ScriptedUpstreamWebSocket(replies=[_responses_turn("gpt-5.4-mini"), _responses_turn("gpt-5.4")]) + + relayed = await _relay_turns(client, upstream_ws, gate) + + rows = [logged.kwargs for logged in relayed.success_handler.call_args_list] + assert [[frame["response"]["model"] for frame in row["response_body"]] for row in rows] == [ + ["gpt-5.4-mini", "gpt-5.4-mini"], + ["gpt-5.4", "gpt-5.4"], + ] + assert [row["litellm_params"]["metadata"][MODEL_ACCESS_GROUP_METADATA_KEY] for row in rows] == [ + ["tier-a"], + ["tier-b"], + ] + assert [row["litellm_params"]["metadata"]["user_api_key"] for row in rows] == ["hashed-key", "hashed-key"] + assert rows[0]["litellm_call_id"] != rows[1]["litellm_call_id"] + assert rows[1]["start_time"] == rows[0]["end_time"] + relayed.proxy_logging.post_call_failure_hook.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_websocket_passthrough_carries_the_session_the_provider_opened_into_every_later_turn(): + """ + A realtime provider names the session's model once, when it opens the + session, and never again. Each later turn logs on its own, so it has to + carry that opening frame or it would be priced as a model nobody named. + """ + client = _turn_taking_client( + [json.dumps({"type": "response.create"}), json.dumps({"type": "response.create"})], + answered_by='"response.done"', + ) + upstream_ws = ScriptedUpstreamWebSocket( + REALTIME_SESSION_CREATED, replies=[[REALTIME_RESPONSE_DONE], [REALTIME_RESPONSE_DONE]] + ) + + relayed = await _relay_turns(client, upstream_ws, RecordingFrameModelGate(), "/openai/v1/realtime") + + assert [ + [frame["type"] for frame in logged.kwargs["response_body"]] + for logged in relayed.success_handler.call_args_list + ] == [["session.created", "response.done"], ["session.created", "response.done"]] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_still_bills_the_answer_in_flight_when_a_later_frame_is_refused(): + """ + A refused frame closes the client while the provider may still be sending the + answer to an earlier, allowed frame. That answer cannot go to a closed client, + but it was paid for, so it has to reach spend instead of turning the whole + session into a failure that bills nothing. + """ + gate = RecordingFrameModelGate(MappingProxyType({"gpt-5.4": KEY_MODEL_REFUSAL})) + client = _turn_taking_client( + [ + json.dumps({"type": "response.create", "model": "gpt-5.4-mini"}), + json.dumps({"type": "response.create", "model": "gpt-5.4"}), + ], + answered_by='"response.created"', + ) + created, completed = _responses_turn("gpt-5.4-mini") + upstream_ws = ScriptedUpstreamWebSocket(replies=[[created]], in_flight=[completed]) + + relayed = await _relay_turns(client, upstream_ws, gate) + + assert [json.loads(text.args[0])["type"] for text in client.send_text.await_args_list] == [ + "response.created", + "error", + ] + assert [ + [frame["type"] for frame in logged.kwargs["response_body"]] + for logged in relayed.success_handler.call_args_list + ] == [["response.created", "response.completed"]] + relayed.proxy_logging.post_call_failure_hook.assert_not_awaited() + + DEEPGRAM_LISTEN_TARGET = "wss://api.deepgram.com/v1/listen?model=nova-3&encoding=linear16&sample_rate=16000" DEEPGRAM_INTERIM_FRAME = json.dumps( { @@ -8718,15 +8994,6 @@ def test_passthrough_client_cannot_forge_session_id_omission(client_metadata_key ) -class StampingFrameModelGate: - def __init__(self, groups: list[str]): - self.groups = groups - - async def __call__(self, model: str, valid_token: UserAPIKeyAuth, /) -> Optional[str]: - valid_token.matched_model_access_groups = self.groups - return None - - async def _billed_metadata(gate) -> dict: served = await _relay_client_frames( [json.dumps({"type": "response.create", "model": "gpt-5.4-mini"})], @@ -8740,7 +9007,7 @@ async def _billed_metadata(gate) -> dict: async def test_websocket_session_is_charged_to_the_group_its_frames_were_authorized_through(): """A websocket names its model in a frame, long after the relay built its logging metadata, so the groups the gate matched have to be picked up at the end or the group pays for nothing it allowed.""" - metadata = await _billed_metadata(StampingFrameModelGate(["tier-a"])) + metadata = await _billed_metadata(RecordingFrameModelGate(groups=MappingProxyType({"gpt-5.4-mini": ("tier-a",)}))) assert metadata[MODEL_ACCESS_GROUP_METADATA_KEY] == ["tier-a"] assert metadata["user_api_key"] == "hashed-key" diff --git a/tests/unit/proxy/test_openai_ws_passthrough_routes.py b/tests/unit/proxy/test_openai_ws_passthrough_routes.py index bd65001111a..836104e3342 100644 --- a/tests/unit/proxy/test_openai_ws_passthrough_routes.py +++ b/tests/unit/proxy/test_openai_ws_passthrough_routes.py @@ -1,21 +1,22 @@ """OpenAI passthrough WebSocket route: registration, opt-in gating, refusals, and per-frame model checks.""" import json -from collections.abc import Iterator, Mapping -from contextlib import contextmanager +from collections.abc import Mapping from dataclasses import dataclass from types import MappingProxyType, SimpleNamespace from typing import Final -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, patch import httpx import pytest import respx +from fastapi import HTTPException from starlette.routing import WebSocketRoute +from starlette.websockets import WebSocket import litellm from litellm.llms.openai.workload_identity import _workload_identity_auth -from litellm.proxy._types import Litellm_EntityType, LiteLLM_UserTable, UserAPIKeyAuth +from litellm.proxy._types import Litellm_EntityType, ProxyErrorTypes, ProxyException, UserAPIKeyAuth from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( _OPENAI_WS_DISABLED_REFUSAL, _OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL, @@ -24,7 +25,16 @@ from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( openai_websocket_proxy_route, router, ) -from litellm.proxy.pass_through_endpoints.pass_through_endpoints import WebsocketFrameModelGate +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_websocket_passthrough_logging_handler import ( + openai_websocket_frame_kind, +) +from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + WebsocketFrameModelAllowed, + WebsocketFrameModelGate, + WebsocketFrameModelRefused, + WebsocketFrameModelVerdict, + WebsocketUpstreamFrameClassifier, +) ENABLED: Final = MappingProxyType({"enable_openai_websocket_passthrough": True}) DISABLED_SETTINGS: Final = ( @@ -76,6 +86,7 @@ class _RelayCall: endpoint: str accept_websocket: bool client_frame_model_gate: WebsocketFrameModelGate | None + upstream_frame_kind: WebsocketUpstreamFrameClassifier | None class _FakeRelay: @@ -93,6 +104,7 @@ class _FakeRelay: endpoint: str, accept_websocket: bool, client_frame_model_gate: WebsocketFrameModelGate | None = None, + upstream_frame_kind: WebsocketUpstreamFrameClassifier | None = None, ) -> None: self.calls.append( _RelayCall( @@ -102,18 +114,14 @@ class _FakeRelay: endpoint=endpoint, accept_websocket=accept_websocket, client_frame_model_gate=client_frame_model_gate, + upstream_frame_kind=upstream_frame_kind, ) ) class _FakeFrameModelGate: - def __init__(self, refusal: str | None = None) -> None: - self.refusal = refusal - self.checked: list[str] = [] - - async def __call__(self, model: str, valid_token: UserAPIKeyAuth, /) -> str | None: - self.checked.append(model) - return self.refusal + async def __call__(self, websocket: object, model: str, /) -> WebsocketFrameModelVerdict: + return WebsocketFrameModelAllowed(access_groups=()) @dataclass(frozen=True, slots=True) @@ -157,6 +165,7 @@ async def test_openai_websocket_forwards_query_and_keeps_provider_auth(prefix, m endpoint=f"/{prefix}/v1/realtime", accept_websocket=False, client_frame_model_gate=served.frame_model_gate, + upstream_frame_kind=openai_websocket_frame_kind, ) ] assert websocket.accepts == [None] @@ -310,26 +319,53 @@ async def test_openai_websocket_connects_model_restricted_keys(key_models): assert websocket.closed is None -@pytest.mark.asyncio -async def test_proxy_frame_model_gate_allows_a_model_the_key_owns(): - gate: Final = _proxy_frame_model_gate() - token: Final = UserAPIKeyAuth(models=["gpt-realtime-2.1", "gpt-5.4-mini"]) +class _FrameAuthorizer: + """Stands in for the full key, team, user and budget auth the gate re-runs for every model a frame names.""" - with patch("litellm.proxy.proxy_server.prisma_client", None): # test-quality-ok: the gate needs no database - assert await gate("gpt-realtime-2.1", token) is None + def __init__(self, token: UserAPIKeyAuth | None = None, denial: Exception | None = None) -> None: + self.token = token or UserAPIKeyAuth() + self.denial = denial + self.authorized: list[str] = [] + + async def __call__(self, websocket: object, model: str) -> UserAPIKeyAuth: + self.authorized.append(model) + if self.denial is not None: + raise self.denial + return self.token + + +SOCKET: Final = WebSocket({"type": "websocket"}, receive=AsyncMock(), send=AsyncMock()) + + +async def _gate_verdict(authorizer: _FrameAuthorizer, model: str) -> WebsocketFrameModelVerdict: + with patch("litellm.proxy.proxy_server.prisma_client", None): # test-quality-ok: the model checks need no database + return await _proxy_frame_model_gate(authorizer)(SOCKET, model) + + +def _refusal(verdict: WebsocketFrameModelVerdict) -> WebsocketFrameModelRefused: + assert isinstance(verdict, WebsocketFrameModelRefused) + return verdict + + +@pytest.mark.asyncio +async def test_proxy_frame_model_gate_runs_full_auth_for_the_model_the_frame_names(): + authorizer: Final = _FrameAuthorizer(UserAPIKeyAuth(models=["gpt-realtime-2.1", "gpt-5.4-mini"])) + + verdict: Final = await _gate_verdict(authorizer, "gpt-realtime-2.1") + + assert verdict == WebsocketFrameModelAllowed(access_groups=()) + assert authorizer.authorized == ["gpt-realtime-2.1"] @pytest.mark.asyncio @pytest.mark.parametrize("requested_model", ["gpt-5.4", "gpt-realtime-2.1-mini"]) async def test_proxy_frame_model_gate_refuses_a_model_the_key_lacks(requested_model): - gate: Final = _proxy_frame_model_gate() - token: Final = UserAPIKeyAuth(models=["gpt-realtime-2.1", "gpt-5.4-mini"]) + authorizer: Final = _FrameAuthorizer(UserAPIKeyAuth(models=["gpt-realtime-2.1", "gpt-5.4-mini"])) - with patch("litellm.proxy.proxy_server.prisma_client", None): # test-quality-ok: the gate needs no database - refusal = await gate(requested_model, token) + refusal: Final = _refusal(await _gate_verdict(authorizer, requested_model)) - assert refusal is not None - assert requested_model in refusal + assert requested_model in refusal.message + assert refusal.error_type == ProxyErrorTypes.key_model_access_denied.value @pytest.mark.asyncio @@ -341,7 +377,6 @@ async def test_proxy_frame_model_gate_refuses_a_model_the_key_lacks(requested_mo async def test_proxy_frame_model_gate_refuses_a_model_the_agent_was_not_granted(granted_models, refusal_names): from litellm.types.agents import AgentResponse - gate: Final = _proxy_frame_model_gate() token: Final = UserAPIKeyAuth(models=[], agent_id="agent-fake") token.managed_agent_policy = AgentResponse( agent_id="agent-fake", @@ -350,132 +385,105 @@ async def test_proxy_frame_model_gate_refuses_a_model_the_agent_was_not_granted( object_permission={"object_permission_id": "grant-fake", "models": granted_models}, ) - with patch("litellm.proxy.proxy_server.prisma_client", None): # test-quality-ok: the gate needs no database - refusal = await gate("gpt-5.4", token) + refusal: Final = _refusal(await _gate_verdict(_FrameAuthorizer(token), "gpt-5.4")) - assert refusal is not None - assert refusal_names in refusal - - -GET_USER_OBJECT: Final = "litellm.proxy.auth.auth_checks.get_user_object" - - -def _personal_user(models: list[str]) -> LiteLLM_UserTable: - return LiteLLM_UserTable(user_id="user-fake", models=models, max_budget=None, spend=0.0) - - -@contextmanager -def _personal_user_allowlist(models: list[str]) -> Iterator[AsyncMock]: - """Stand up the database-backed owner lookup the allowlist check runs, and report every call it made.""" - lookup: Final = AsyncMock(return_value=_personal_user(models)) - with ( - patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), # test-quality-ok: read at call time - patch(GET_USER_OBJECT, lookup), - ): - yield lookup - - -@pytest.mark.asyncio -async def test_proxy_frame_model_gate_refuses_a_model_the_person_behind_the_key_lacks(): - """ - A key with no team of its own carries its owner's model allowlist, the same - one every HTTP route applies, so a frame naming a model that person cannot - use is refused even when the key itself names no models at all. - """ - gate: Final = _proxy_frame_model_gate() - token: Final = UserAPIKeyAuth(models=[], user_id="user-fake") - - with _personal_user_allowlist(["gpt-5.4-mini"]): - refusal = await gate("gpt-5.4", token) - - assert refusal is not None - assert "gpt-5.4" in refusal - - -@pytest.mark.asyncio -async def test_proxy_frame_model_gate_allows_a_model_the_person_behind_the_key_owns(): - gate: Final = _proxy_frame_model_gate() - token: Final = UserAPIKeyAuth(models=[], user_id="user-fake") - - with _personal_user_allowlist(["gpt-5.4-mini", "gpt-realtime-2.1"]): - assert await gate("gpt-realtime-2.1", token) is None - - -@pytest.mark.asyncio -async def test_proxy_frame_model_gate_leaves_a_team_key_to_the_teams_allowlist(): - """ - A key that belongs to a team is governed by the team's models, not by the - personal allowlist of whoever created it, so the owner's own restrictions - must not close a session the team is entitled to. - """ - gate: Final = _proxy_frame_model_gate() - token: Final = UserAPIKeyAuth(models=[], user_id="user-fake", team_id="team-fake") - - with _personal_user_allowlist(["gpt-5.4-mini"]) as owner_lookup: - assert await gate("gpt-5.4", token) is None - - assert owner_lookup.await_args_list == [] + assert refusal_names in refusal.message @pytest.mark.asyncio @pytest.mark.parametrize("unrestricted_models", [[], ["all-proxy-models"], ["*"]]) async def test_proxy_frame_model_gate_allows_every_model_for_unrestricted_keys(unrestricted_models): - gate: Final = _proxy_frame_model_gate() - token: Final = UserAPIKeyAuth(models=unrestricted_models) + verdict: Final = await _gate_verdict(_FrameAuthorizer(UserAPIKeyAuth(models=unrestricted_models)), "gpt-5.4") - with patch("litellm.proxy.proxy_server.prisma_client", None): # test-quality-ok: the gate needs no database - assert await gate("gpt-realtime-2.1", token) is None + assert verdict == WebsocketFrameModelAllowed(access_groups=()) -ENFORCE_GROUP_BUDGET: Final = ( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.enforce_model_access_group_budget" +@pytest.mark.asyncio +async def test_proxy_frame_model_gate_hands_back_the_access_groups_that_authorized_the_model(): + """The turn a frame starts is billed to the groups that let its model through, so the gate has to report them.""" + token: Final = UserAPIKeyAuth(models=["gpt-5.4-mini"], matched_model_access_groups=["tier-a"]) + + verdict: Final = await _gate_verdict(_FrameAuthorizer(token), "gpt-5.4-mini") + + assert verdict == WebsocketFrameModelAllowed(access_groups=("tier-a",)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "requested_model", + ["gpt-5.4-mini", "gpt-5.4"], + ids=["allowed", "refused_by_the_model_check"], ) +async def test_proxy_frame_model_gate_releases_the_budget_it_reserved_for_the_check(requested_model): + """ + Full auth reserves budget for the request it authorizes, but a frame check is + not a request, so a reservation the gate leaves open would hold the key's + spend counter above its real spend and refuse later calls on every route. + """ + reservation: Final[dict[str, object]] = {"reserved_cost": 0.5, "entries": []} + token: Final = UserAPIKeyAuth(models=["gpt-5.4-mini"], budget_reservation=reservation) + + await _gate_verdict(_FrameAuthorizer(token), requested_model) + + assert token.budget_reservation == {**reservation, "finalized": True} @pytest.mark.asyncio -async def test_proxy_frame_model_gate_refuses_a_model_whose_access_group_is_out_of_budget(): - """ - A model that sits in a budgeted access group is refused on every other route once - that group has spent its budget, so a frame naming it closes the session instead - of quietly running on an empty pool. The key is allowed the model outright here, - so nothing but the group's budget can produce the refusal. - """ - exhausted: Final = litellm.BudgetExceededError( - current_cost=2.0, - max_budget=1.0, - message="Budget has been exceeded! Model access group=tier-a Current cost: 2.0, Max budget: 1.0", - entity_type=Litellm_EntityType.MODEL_ACCESS_GROUP.value, - entity_id="tier-a", - ) - gate: Final = _proxy_frame_model_gate() - token: Final = UserAPIKeyAuth(models=["gpt-5.4-mini"]) +@pytest.mark.parametrize( + "denial, error_type, message", + [ + ( + ProxyException( + message="Budget has been exceeded! Key=hashed Spend=2.0, Budget=1.0", + type=ProxyErrorTypes.budget_exceeded, + param=None, + code=400, + ), + "budget_exceeded", + "Budget has been exceeded! Key=hashed", + ), + ( + litellm.BudgetExceededError( + current_cost=2.0, + max_budget=1.0, + message="Budget has been exceeded! Model access group=tier-a Current cost: 2.0, Max budget: 1.0", + entity_type=Litellm_EntityType.MODEL_ACCESS_GROUP.value, + entity_id="tier-a", + ), + "budget_exceeded", + "Model access group=tier-a", + ), + ( + ProxyException( + message="team not allowed to access model. This team can only access models=['gpt-5.4-mini']", + type=ProxyErrorTypes.team_model_access_denied, + param="model", + code=401, + ), + "team_model_access_denied", + "team not allowed to access model", + ), + (HTTPException(status_code=403, detail="Invalid Authorization header format"), "invalid_request_error", ""), + ], + ids=["key_budget", "group_budget", "team_allowlist", "malformed_auth"], +) +async def test_proxy_frame_model_gate_refuses_what_full_auth_denies_with_its_reason(denial, error_type, message): + """A frame is denied for the same reasons a request is, and the caller reads which reason in the error event.""" + refusal: Final = _refusal(await _gate_verdict(_FrameAuthorizer(denial=denial), "gpt-5.4-mini")) - with ( - patch("litellm.proxy.proxy_server.prisma_client", None), # test-quality-ok: the model check needs no database - patch(ENFORCE_GROUP_BUDGET, AsyncMock(side_effect=exhausted)), - ): - refusal = await gate("gpt-5.4-mini", token) - - assert refusal is not None - assert "Budget has been exceeded" in refusal - assert "Model access group=tier-a" in refusal + assert refusal.error_type == error_type + assert message in refusal.message @pytest.mark.asyncio -async def test_proxy_frame_model_gate_refuses_out_loud_when_it_cannot_reach_the_budgets(): +async def test_proxy_frame_model_gate_refuses_out_loud_when_auth_cannot_answer(): """ A gate that cannot answer has to close the session with a reason the caller can read. Letting the failure escape drops the socket with no frame at all, which the caller cannot tell apart from the network going away. """ - gate: Final = _proxy_frame_model_gate() - token: Final = UserAPIKeyAuth(models=["gpt-5.4-mini"]) + refusal: Final = _refusal( + await _gate_verdict(_FrameAuthorizer(denial=ConnectionError("connection pool is exhausted")), "gpt-5.4-mini") + ) - with ( - patch("litellm.proxy.proxy_server.prisma_client", None), # test-quality-ok: the model check needs no database - patch(ENFORCE_GROUP_BUDGET, AsyncMock(side_effect=ConnectionError("connection pool is exhausted"))), - ): - refusal = await gate("gpt-5.4-mini", token) - - assert refusal == _OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL - assert "connection pool" not in refusal + assert refusal == WebsocketFrameModelRefused(message=_OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL, error_type="server_error")