scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-23 00:41:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 81

Merge remote-tracking branch 'origin/main' into litellm_settings_store_immutable

Browse source
This commit is contained in:
Yuneng Jiang 2026-09-18 22:39:29 -07:00
commit 3a480a5d6c
No known key found for this signature in database
30 changed files with 2613 additions and 188 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

33
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -2223,6 +2223,7 @@ class BaseLLMHTTPHandler:
# Prepare headers
kwargs = kwargs or {}
kwargs_for_agentic: Final = self._agentic_hook_kwargs(kwargs=kwargs, api_key=api_key, api_base=api_base)
provider_specific_header: Final = cast(
litellm.types.utils.ProviderSpecificHeader | Sequence[litellm.types.utils.ProviderSpecificHeader] | None,
kwargs.get("provider_specific_header", None),
@ -2410,7 +2411,7 @@ class BaseLLMHTTPHandler:
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
kwargs={**kwargs, "api_key": api_key} if api_key else kwargs,
kwargs=kwargs_for_agentic,
hold_back=bool(held_back_tool_names),
server_fulfilled_tool_names=held_back_tool_names,
)
@ -2433,8 +2434,7 @@ class BaseLLMHTTPHandler:
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
api_key=api_key,
kwargs=kwargs,
kwargs=kwargs_for_agentic,
)
async def _finalize_anthropic_messages_response(
@ -2447,14 +2447,8 @@ class BaseLLMHTTPHandler:
anthropic_messages_optional_request_params: dict,
logging_obj: LiteLLMLoggingObj,
custom_llm_provider: str,
api_key: str | None,
kwargs: dict,
kwargs: dict[str, object],
) -> AnthropicMessagesResponse | AsyncIterator:
# Inject api_key into kwargs so follow-up calls in agentic hooks can
# authenticate. api_key is a named param here (not in kwargs), so
# _prepare_followup_kwargs would miss it otherwise.
kwargs_for_agentic: Final = {**kwargs, "api_key": api_key} if api_key else kwargs
# Call agentic completion hooks (non-streaming path only)
final_response: Final = await self._call_agentic_completion_hooks(
response=initial_response,
model=model,
@ -2464,7 +2458,7 @@ class BaseLLMHTTPHandler:
logging_obj=logging_obj,
stream=False,
custom_llm_provider=custom_llm_provider,
kwargs=kwargs_for_agentic,
kwargs=kwargs,
)
return self._maybe_wrap_in_fake_stream(
@ -5312,6 +5306,15 @@ class BaseLLMHTTPHandler:
fingerprints: Final = list(kwargs.get("_agentic_loop_fingerprints", []) or [])
return depth, max_loops, fingerprints
@staticmethod
def _agentic_hook_kwargs(
kwargs: Mapping[str, object], api_key: str | None, api_base: str | None
) -> dict[str, object]:
"""``api_key`` and ``api_base`` are named parameters of ``anthropic_messages`` rather than kwargs, so the
follow-up call an agentic hook makes only reaches the same deployment if they are re-added here."""
deployment_params: Final = {"api_key": api_key, "api_base": api_base}
return {**kwargs, **{key: value for key, value in deployment_params.items() if value}}
@staticmethod
def _has_agentic_completion_hook(logging_obj: LiteLLMLoggingObj) -> bool:
"""
@ -6604,7 +6607,7 @@ class BaseLLMHTTPHandler:
first_message: str | None = None,
request_defaults: ResponsesWebSocketRequestDefaults | None = None,
**kwargs: Any,
):
) -> Exception | None:
"""
Handles Responses API WebSocket mode.
@ -6638,7 +6641,7 @@ class BaseLLMHTTPHandler:
**kwargs,
)
await handler.run()
return
return None
import websockets
from websockets.asyncio.client import ClientConnection
@ -6757,9 +6760,10 @@ class BaseLLMHTTPHandler:
output_guardrail_callbacks=_ws_output_guardrail_callbacks,
quota_callbacks=_ws_quota_callbacks,
authorized_model=model,
custom_llm_provider=custom_llm_provider,
request_defaults=request_defaults,
)
await streaming.bidirectional_forward()
return await streaming.bidirectional_forward()
except websockets.exceptions.InvalidStatusCode as e:
verbose_logger.exception("Error connecting to responses WS backend: %s", e)
@ -6773,6 +6777,7 @@ class BaseLLMHTTPHandler:
pass
else:
raise Exception(f"Unexpected error while closing WebSocket: {close_error}")
return None
def image_edit_handler(
self,

14
litellm/proxy/_experimental/mcp_server/bridge_token_flow.py
View file

@ -262,7 +262,12 @@ async def _reload_active_user_by_id(user_id: str) -> "_KeyResolutionFailure | No
return loaded if isinstance(loaded, str) else None
async def load_active_user_by_id(user_id: str) -> "LiteLLM_UserTable | _KeyResolutionFailure":
UserRowSource = Literal["cache", "database"]
async def load_active_user_by_id(
user_id: str, source: UserRowSource = "cache"
) -> "LiteLLM_UserTable | _KeyResolutionFailure":
"""Load a live litellm user by id, returning the record when the user is active or a precise
failure otherwise. The interactive DCR client authenticates via SSO, so its refresh envelope seals a
user subject; renewing it must re-check the user is still live (present and not SCIM-deactivated) so a
@ -273,7 +278,11 @@ async def load_active_user_by_id(user_id: str) -> "LiteLLM_UserTable | _KeyResol
``HTTPException``, a SCIM-deactivated user, and, unlike the key path, a missing user. ``get_user_object``
catches every DB failure and re-raises a bare ``ValueError`` (a deleted user and a real outage look
identical, the original error surviving only as ``__context__``), so the outage check walks the cause
chain, and a missing user falls through to ``no_active_key`` rather than an opaque gateway fault."""
chain, and a missing user falls through to ``no_active_key`` rather than an opaque gateway fault.
``source="database"`` reads the row from the database, never the cache, so the credential mint refuses
a user that a writer deactivated or deleted without evicting the cached row, and it leaves the fresh
row in the cache for the requests the credential makes next. Every other caller keeps the cache read,
so introspection, which a resource server may call per request, stays off the database."""
from litellm.proxy._types import (
ProxyException, # noqa: PLC0415 # inline import avoids a module-load circular import
)
@ -296,6 +305,7 @@ async def load_active_user_by_id(user_id: str) -> "LiteLLM_UserTable | _KeyResol
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
user_id_upsert=False,
check_db_only=source == "database",
)
except (ProxyException, HTTPException):
return "no_active_key"

26
litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py
View file

@ -59,6 +59,11 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import (
register_aggregate_client,
relative_request_url,
revoke_refresh_token,
supported_grant_types,
)
from litellm.proxy._experimental.mcp_server.idp_token_exchange import (
exchange_idp_subject_token,
token_exchange_available,
)
from litellm.proxy._experimental.mcp_server.oauth_identity_binding import (
RefreshOwnershipProven,
@ -1980,6 +1985,9 @@ async def token_endpoint(
refresh_token: str | None = Form(None),
scope: str | None = Form(None),
resource: str | None = Form(None),
subject_token: str | None = Form(None),
subject_token_type: str | None = Form(None),
requested_token_type: str | None = Form(None),
mcp_server_name: str | None = None,
):
"""
@ -2010,6 +2018,10 @@ async def token_endpoint(
cache=user_api_key_cache,
resource=resource,
mint_proxy_credential=mint_proxy_credential,
subject_token=subject_token,
subject_token_type=subject_token_type,
requested_token_type=requested_token_type,
exchange_subject_token=exchange_idp_subject_token,
)
lookup_name: Final = mcp_server_name or client_id
@ -2131,7 +2143,9 @@ async def introspect_endpoint(token: str = Form(...)) -> Response:
async def native_client_auth_discovery(request: Request) -> JSONResponse:
"""The versioned contract a native client (``lite login --pkce``, or a CLI in any other
language) reads to sign a user in through the browser and obtain a proxy credential."""
return JSONResponse(native_client_auth_contract(request), headers=TOKEN_NO_CACHE_HEADERS)
return JSONResponse(
native_client_auth_contract(request, token_exchange_available()), headers=TOKEN_NO_CACHE_HEADERS
)
# Per RFC 6749 §4.1.2.1, an IdP that rejects an OAuth authorization request
@ -2619,7 +2633,7 @@ def _build_aggregate_protected_resource_response(request: Request) -> dict:
}
def _build_aggregate_authorization_server_response(request: Request) -> dict:
def _build_aggregate_authorization_server_response(request: Request, token_exchange_available: bool) -> dict:
"""RFC 8414 metadata for the gateway as the aggregate authorization server.
The issuer is ``{base}/mcp`` and must stay equal to the value the
@ -2638,7 +2652,7 @@ def _build_aggregate_authorization_server_response(request: Request) -> dict:
"registration_endpoint": f"{request_base_url}/register",
"response_types_supported": ["code"],
"scopes_supported": [],
"grant_types_supported": ["authorization_code", "refresh_token"],
"grant_types_supported": supported_grant_types(token_exchange_available),
"code_challenge_methods_supported": ["S256"],
"token_endpoint_auth_methods_supported": ["none", "client_secret_post"],
}
@ -2676,7 +2690,7 @@ async def oauth_authorization_server_aggregate(request: Request):
per-server row win here instead would serve an issuer of {base} against a resource that
advertised {base}/mcp, which fails the RFC 8414 issuer check and breaks the front door.
"""
return _build_aggregate_authorization_server_response(request)
return _build_aggregate_authorization_server_response(request, token_exchange_available())
# Standard MCP pattern: /.well-known/oauth-protected-resource/mcp/{server_name}
@ -2902,7 +2916,9 @@ async def register_client(request: Request, mcp_server_name: str | None = None):
# advertises that), so this does not affect it. A request without redirect_uris is not
# a DCR request, so the legacy single-server-or-dummy fallback is kept for it.
if data.get("redirect_uris"):
return await register_aggregate_client(request=request, request_body=data)
return await register_aggregate_client(
request=request, request_body=data, token_exchange_available=token_exchange_available()
)
resolved: Final = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip)
if resolved:
return await register_client_with_server(

175
litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py
View file

@ -51,7 +51,7 @@ from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse
from fastapi import HTTPException, Request
from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Response
from pydantic import BaseModel, ConfigDict, Field, ValidationError
from typing_extensions import ReadOnly, TypedDict, assert_never
from typing_extensions import NotRequired, ReadOnly, TypedDict, assert_never
from litellm._logging import verbose_logger
from litellm.caching.caching import DualCache
@ -187,6 +187,52 @@ class MintProxyCredential(Protocol):
) -> Awaitable[MintedProxyCredential | ProxyCredentialMintFailure]: ...
TOKEN_EXCHANGE_GRANT_TYPE: Final = "urn:ietf:params:oauth:grant-type:token-exchange"
def supported_grant_types(token_exchange_available: bool) -> tuple[str, ...]:
"""The grants ``/token`` can serve on this deployment. The RFC 8693 exchange is listed
only where the JWT auth that proves a subject token is on, backed by a database, and
licensed, so a client never selects a grant the gateway would then refuse."""
if token_exchange_available:
return ("authorization_code", "refresh_token", TOKEN_EXCHANGE_GRANT_TYPE)
return ("authorization_code", "refresh_token")
"""RFC 8693: a native client that already holds a token from the customer's identity
provider trades it for the proxy-API credential without a browser round trip."""
_IssuedTokenType = Literal["urn:ietf:params:oauth:token-type:access_token"]
ACCESS_TOKEN_TOKEN_TYPE: Final[_IssuedTokenType] = "urn:ietf:params:oauth:token-type:access_token"
SUBJECT_TOKEN_TYPES: Final = frozenset(
{
"urn:ietf:params:oauth:token-type:jwt",
"urn:ietf:params:oauth:token-type:id_token",
ACCESS_TOKEN_TOKEN_TYPE,
}
)
class SubjectIdentity(BaseModel):
model_config = ConfigDict(frozen=True)
user_id: str = Field(min_length=1)
team_id: str | None = None
class SubjectTokenRefusal(BaseModel):
model_config = ConfigDict(frozen=True)
error: Literal["unsupported_grant_type", "invalid_request", "temporarily_unavailable"]
description: str = Field(min_length=1)
class ExchangeSubjectToken(Protocol):
"""Injected RFC 8693 subject-token verifier ``(subject_token, request)``: proves the
IdP token the way the proxy's own JWT auth does and names the litellm user and team it
stands for, or says why this gateway will not take it."""
def __call__(self, subject_token: str, request: Request, /) -> Awaitable[SubjectIdentity | SubjectTokenRefusal]: ...
class ConsentTeam(BaseModel):
model_config = ConfigDict(frozen=True)
team_id: str = Field(min_length=1)
@ -213,6 +259,12 @@ async def _refuse_proxy_credential(user_id: str, team_id: str | None) -> ProxyCr
return "unresolvable"
async def _refuse_subject_token(subject_token: str, request: Request) -> SubjectTokenRefusal:
return SubjectTokenRefusal(
error="unsupported_grant_type", description="this gateway is not configured to exchange IdP tokens"
)
async def _unavailable_vendor_credential(user_id: str, server_id: str) -> VendorCredentialState:
return "unavailable"
@ -318,7 +370,9 @@ def open_gateway_dcr_client(client_id: str) -> GatewayDcrClient | None:
return _open_sealed(client_id, GATEWAY_DCR_CLIENT_ID_PREFIX, GatewayDcrClient, _CLIENT_RECORD_DEBUG_KEY)
async def register_aggregate_client(request: Request, request_body: Mapping[str, object]) -> Response:
async def register_aggregate_client(
request: Request, request_body: Mapping[str, object], token_exchange_available: bool
) -> Response:
"""RFC 7591 dynamic registration against the gateway itself, statelessly.
Only ``redirect_uris`` is authoritative; every client is registered as a public
@ -382,7 +436,7 @@ async def register_aggregate_client(request: Request, request_body: Mapping[str,
"client_id_issued_at": int(now.timestamp()),
"redirect_uris": list(raw_uris),
"token_endpoint_auth_method": "none",
"grant_types": ["authorization_code", "refresh_token"],
"grant_types": list(supported_grant_types(token_exchange_available)),
"response_types": ["code"],
},
)
@ -580,7 +634,7 @@ class NativeClientAuthContract(TypedDict):
revocation_endpoint_auth_methods_supported: ReadOnly[tuple[str, ...]]
def native_client_auth_contract(request: Request) -> NativeClientAuthContract:
def native_client_auth_contract(request: Request, token_exchange_available: bool) -> NativeClientAuthContract:
"""The versioned discovery document at ``/.well-known/litellm-cli-auth``: everything a
native client (in any language) needs to run the sign-in without reading LiteLLM
source. ``resource`` is the exact value to send as the RFC 8707 ``resource`` parameter
@ -595,7 +649,7 @@ def native_client_auth_contract(request: Request) -> NativeClientAuthContract:
"revocation_endpoint": f"{base_url}/revoke",
"resource": base_url,
"response_types_supported": ("code",),
"grant_types_supported": ("authorization_code", "refresh_token"),
"grant_types_supported": supported_grant_types(token_exchange_available),
"code_challenge_methods_supported": ("S256",),
"token_endpoint_auth_methods_supported": ("none",),
"revocation_endpoint_auth_methods_supported": ("none",),
@ -1033,20 +1087,26 @@ class _ProxyCredentialTokenResponse(TypedDict):
refresh_token: ReadOnly[str]
user_id: ReadOnly[str]
team_id: ReadOnly[str | None]
issued_token_type: NotRequired[ReadOnly[_IssuedTokenType]]
def _proxy_credential_response(
minted: MintedProxyCredential, principal: SessionPrincipal, keys: SessionSigningKeys, now: datetime
minted: MintedProxyCredential,
principal: SessionPrincipal,
keys: SessionSigningKeys,
now: datetime,
issued_token_type: _IssuedTokenType | None = None,
) -> Response:
"""The proxy-API token response: the access token is the very credential ``lite
login`` stores (accepted on every proxy route with user and team attribution), and
the refresh token is a gateway-sealed rotating token bound to the team the credential
was minted for, so a renewal keeps the team the user consented to."""
was minted for, so a renewal keeps the team the user consented to. A token exchange
also states ``issued_token_type``, which RFC 8693 section 2.2.1 requires."""
bound_principal: Final = principal.model_copy(update=MappingProxyType({"team_id": minted.team_id}))
refresh: Final = mint_session_refresh_token(bound_principal, keys, now)
if not isinstance(refresh, MintedSessionToken):
return _oauth_error(500, "server_error", "failed to mint the session credential")
body: Final[_ProxyCredentialTokenResponse] = {
credential: Final[_ProxyCredentialTokenResponse] = {
"access_token": minted.key,
"token_type": "Bearer",
"expires_in": minted.expires_in,
@ -1054,7 +1114,10 @@ def _proxy_credential_response(
"user_id": minted.user_id,
"team_id": minted.team_id,
}
return JSONResponse(status_code=200, content=body, headers=TOKEN_NO_CACHE_HEADERS)
if issued_token_type is None:
return JSONResponse(status_code=200, content=credential, headers=TOKEN_NO_CACHE_HEADERS)
exchanged: Final[_ProxyCredentialTokenResponse] = {**credential, "issued_token_type": issued_token_type}
return JSONResponse(status_code=200, content=exchanged, headers=TOKEN_NO_CACHE_HEADERS)
def _reload_failure_response(failure: ReloadUserFailure) -> Response:
@ -1073,6 +1136,16 @@ def _reload_failure_response(failure: ReloadUserFailure) -> Response:
assert_never(failure)
def _subject_token_refusal_response(refusal: SubjectTokenRefusal) -> Response:
match refusal.error:
case "temporarily_unavailable":
return _oauth_error(503, refusal.error, refusal.description)
case "unsupported_grant_type" | "invalid_request":
return _oauth_error(400, refusal.error, refusal.description)
case _:
assert_never(refusal.error)
def _mint_failure_response(failure: ProxyCredentialMintFailure) -> Response:
match failure:
case "not_a_member":
@ -1116,11 +1189,16 @@ async def aggregate_token(
cache: DualCache,
resource: str | None = None,
mint_proxy_credential: MintProxyCredential = _refuse_proxy_credential,
subject_token: str | None = None,
subject_token_type: str | None = None,
requested_token_type: str | None = None,
exchange_subject_token: ExchangeSubjectToken = _refuse_subject_token,
) -> Response:
"""The aggregate token verb: authorization_code and refresh_token grants for the
identity-only session pair, or for the proxy-API credential when the grant was issued
with that audience. Every path re-validates the litellm user live before minting, so a
deactivated user cannot obtain or renew a session."""
with that audience, and the RFC 8693 token exchange that turns an IdP token straight
into the proxy-API credential. Every path re-validates the litellm user live before
minting, so a deactivated user cannot obtain or renew a session."""
if master_key is None:
verbose_logger.error("mcp_gateway_dcr token grant rejected: no master_key configured")
return _oauth_error(500, "server_error", "the gateway has no master key configured")
@ -1159,7 +1237,20 @@ async def aggregate_token(
now=now,
issue=issue,
)
return _oauth_error(400, "unsupported_grant_type", "grant_type must be authorization_code or refresh_token")
if grant_type == TOKEN_EXCHANGE_GRANT_TYPE:
return await _token_exchange_grant(
subject_token=subject_token,
subject_token_type=subject_token_type,
requested_token_type=requested_token_type,
client_id=client_id,
exchange_subject_token=exchange_subject_token,
issue=issue,
)
return _oauth_error(
400,
"unsupported_grant_type",
f"grant_type must be authorization_code, refresh_token, or {TOKEN_EXCHANGE_GRANT_TYPE}",
)
class _GrantIssuer:
@ -1211,10 +1302,9 @@ class _GrantIssuer:
async def _issue_proxy_credential(
self, principal: SessionPrincipal, claim_key: str, claim_ttl_seconds: int, replayed: str
) -> Response:
if self._resource is not None and not is_proxy_api_resource(self._request, self._resource):
return _oauth_error(
400, "invalid_target", "resource does not match the proxy API this grant was issued for"
)
target_refusal: Final = self._proxy_api_target_refusal()
if target_refusal is not None:
return target_refusal
minted: Final = await self._mint_proxy_credential(principal.user_id, principal.team_id)
if not isinstance(minted, MintedProxyCredential):
return _mint_failure_response(minted)
@ -1223,6 +1313,33 @@ class _GrantIssuer:
return refusal
return _proxy_credential_response(minted, principal, self._keys, self._now)
async def exchange(
self, subject_token: str, client_id: str, exchange_subject_token: ExchangeSubjectToken
) -> Response:
"""The RFC 8693 tail: prove the IdP token, then mint. No single-use marker, because
the subject token stays a valid proof for as long as the IdP says it is and every
exchange mints a fresh credential and refresh token of its own."""
target_refusal: Final = self._proxy_api_target_refusal()
if target_refusal is not None:
return target_refusal
identity: Final = await exchange_subject_token(subject_token, self._request)
if isinstance(identity, SubjectTokenRefusal):
return _subject_token_refusal_response(identity)
principal: Final = SessionPrincipal(
user_id=identity.user_id, client_id=client_id, audience=PROXY_API_AUDIENCE, team_id=identity.team_id
)
minted: Final = await self._mint_proxy_credential(principal.user_id, principal.team_id)
if not isinstance(minted, MintedProxyCredential):
return _mint_failure_response(minted)
return _proxy_credential_response(
minted, principal, self._keys, self._now, issued_token_type=ACCESS_TOKEN_TOKEN_TYPE
)
def _proxy_api_target_refusal(self) -> Response | None:
if self._resource is None or is_proxy_api_resource(self._request, self._resource):
return None
return _oauth_error(400, "invalid_target", "resource does not match the proxy API this grant was issued for")
async def _claim_refusal(self, claim_key: str, claim_ttl_seconds: int, replayed: str) -> Response | None:
return _claim_refusal(
await self._guard.claim(claim_key, claim_ttl_seconds), replayed=_oauth_error(400, "invalid_grant", replayed)
@ -1297,6 +1414,32 @@ async def _refresh_token_grant(
)
async def _token_exchange_grant(
subject_token: str | None,
subject_token_type: str | None,
requested_token_type: str | None,
client_id: str,
exchange_subject_token: ExchangeSubjectToken,
issue: _GrantIssuer,
) -> Response:
"""RFC 8693 token exchange for a registered native client that already holds an IdP
token: the gateway proves the token the way its JWT auth does and answers with the
proxy-API credential, so a fresh laptop with only an IdP login gets a gateway key
without a browser round trip. The client must be registered because the refresh token
in the answer is bound to it."""
if not is_gateway_dcr_client_id(client_id) or open_gateway_dcr_client(client_id) is None:
return _oauth_error(401, "invalid_client", "unknown or malformed client_id")
if not subject_token or not subject_token_type:
return _oauth_error(400, "invalid_request", "subject_token and subject_token_type are required")
if subject_token_type not in SUBJECT_TOKEN_TYPES:
return _oauth_error(
400, "invalid_request", f"subject_token_type must be one of {', '.join(sorted(SUBJECT_TOKEN_TYPES))}"
)
if requested_token_type is not None and requested_token_type != ACCESS_TOKEN_TOKEN_TYPE:
return _oauth_error(400, "invalid_request", f"requested_token_type must be {ACCESS_TOKEN_TOKEN_TYPE}")
return await issue.exchange(subject_token, client_id, exchange_subject_token)
async def revoke_refresh_token(token: str, client_id: str, master_key: str | None, cache: DualCache) -> Response:
"""RFC 7009 revocation for the gateway's refresh tokens: burn the presented token's
``jti`` so neither the holder nor a thief can rotate it again. Access tokens are

217
litellm/proxy/_experimental/mcp_server/idp_token_exchange.py Normal file
View file

@ -0,0 +1,217 @@
"""The identity-provider side of the RFC 8693 token exchange on ``POST /token``: a native
client that already holds a JWT from the customer's IdP trades it for the same proxy-API
credential ``lite login`` stores, proven by the proxy's own JWT auth (signature, claims,
and the user and team sync it performs), so no browser round trip is needed."""
from __future__ import annotations
from collections.abc import Awaitable, Callable, Mapping
from dataclasses import dataclass
from typing import Final, Literal, Protocol
from fastapi import HTTPException, Request
from typing_extensions import assert_never
from litellm._logging import verbose_proxy_logger
from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import SubjectIdentity, SubjectTokenRefusal
from litellm.proxy._types import JWTAuthBuilderResult, ProxyException
from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
EXCHANGE_ROUTE: Final = "/token"
REJECTED_SUBJECT_TOKEN: Final = "subject_token was rejected by the gateway's JWT auth"
SUBJECT_TOKEN_CHECK_UNAVAILABLE: Final = (
"the gateway could not verify subject_token because its identity provider or database is unavailable; retry"
)
SUBJECT_TOKEN_CHECK_FAULTED: Final = (
"the gateway could not verify subject_token because its database reported a fault that is not a transient "
"outage; retrying will not help until the gateway deployment is repaired"
)
GatewayOutage = Literal["retryable", "faulted"]
@dataclass(frozen=True, slots=True)
class TokenExchangePrerequisites:
"""The deployment-level gates ``user_api_key_auth`` applies before it verifies any JWT
bearer, plus the JWT-to-virtual-key mapping it consults first: a gateway that maps
tokens authenticates a JWT as its mapped key, with that key's models and budget, or
refuses an unmapped one, and the exchange proves the token through ``auth_builder``
alone, so it would mint the user's own credential past that policy. Discovery and
registration advertise the exchange grant only when every gate holds, and an exchange
attempt is refused naming the first one that does not."""
jwt_auth_enabled: bool
has_database: bool
licensed: bool
maps_jwts_to_virtual_keys: bool
@property
def available(self) -> bool:
return self.jwt_auth_enabled and self.has_database and self.licensed and not self.maps_jwts_to_virtual_keys
def refusal(self) -> SubjectTokenRefusal | None:
if not self.jwt_auth_enabled:
return SubjectTokenRefusal(
error="unsupported_grant_type",
description="JWT auth is not enabled on this gateway, so it cannot exchange IdP tokens",
)
if not self.has_database:
return SubjectTokenRefusal(
error="unsupported_grant_type",
description="this gateway has no database, so it cannot exchange IdP tokens",
)
if not self.licensed:
return SubjectTokenRefusal(
error="unsupported_grant_type",
description="JWT auth is an enterprise only feature; no license is set",
)
if self.maps_jwts_to_virtual_keys:
return SubjectTokenRefusal(
error="unsupported_grant_type",
description="this gateway maps IdP tokens to virtual keys, which the exchange does not serve",
)
return None
def read_token_exchange_prerequisites() -> TokenExchangePrerequisites:
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # rebound after startup, so read them per call
general_settings,
jwt_handler,
premium_user,
prisma_client,
)
return TokenExchangePrerequisites(
jwt_auth_enabled=general_settings.get("enable_jwt_auth", False) is True,
has_database=prisma_client is not None,
licensed=premium_user is True,
maps_jwts_to_virtual_keys=_maps_jwts_to_virtual_keys(jwt_handler),
)
def _maps_jwts_to_virtual_keys(jwt_handler: JWTHandler) -> bool:
if not hasattr(jwt_handler, "litellm_jwtauth"):
return False
return jwt_handler.litellm_jwtauth.is_virtual_key_mapping_configured()
def token_exchange_available() -> bool:
return read_token_exchange_prerequisites().available
class AuthorizeSubjectToken(Protocol):
"""Injected JWT authorization ``(subject_token, request_headers)``: the proxy's
``JWTAuthManager.auth_builder`` in production, which raises when the token is not
acceptable and otherwise names the user and team it resolved."""
def __call__(
self, subject_token: str, request_headers: Mapping[str, str], /
) -> Awaitable[JWTAuthBuilderResult]: ...
async def exchange_idp_subject_token(subject_token: str, request: Request) -> SubjectIdentity | SubjectTokenRefusal:
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # rebound after startup, so read them per call
general_settings,
jwt_handler,
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)
async def authorize(token: str, request_headers: Mapping[str, str]) -> JWTAuthBuilderResult:
return await JWTAuthManager.auth_builder(
api_key=token,
jwt_handler=jwt_handler,
request_data={},
general_settings=general_settings,
route=EXCHANGE_ROUTE,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=proxy_logging_obj,
request_headers=request_headers,
request_method="POST",
)
return await identity_from_subject_token(
subject_token,
request_headers=request.headers,
prerequisites=read_token_exchange_prerequisites(),
is_jwt=jwt_handler.is_jwt,
authorize=authorize,
)
async def identity_from_subject_token(
subject_token: str,
request_headers: Mapping[str, str],
prerequisites: TokenExchangePrerequisites,
is_jwt: Callable[[str], bool],
authorize: AuthorizeSubjectToken,
) -> SubjectIdentity | SubjectTokenRefusal:
"""Apply the same gates ``user_api_key_auth`` applies to a JWT bearer, then let the
proxy's JWT auth prove the token. A rejection comes back as ``invalid_request``, which
RFC 8693 section 2.2.2 prescribes for an invalid or unacceptable subject token, and a
check the gateway could not complete (the IdP's JWKS unreachable with no cached copy,
the auth database down) as ``temporarily_unavailable``, so the client retries instead
of treating a valid token as bad, worded by whether retrying can help. The reason stays
in the proxy log: this endpoint is public and JWT auth's own wording can name the JWKS
URL it fetched or quote the IdP's response."""
unmet: Final = prerequisites.refusal()
if unmet is not None:
return unmet
if not is_jwt(subject_token):
return SubjectTokenRefusal(error="invalid_request", description="subject_token is not a JWT")
try:
result: Final = await authorize(subject_token, request_headers)
except HTTPException as denied:
return _refusal_for(denied, denied.detail)
except ProxyException as denied:
return _refusal_for(denied, denied.message)
except Exception as denied: # noqa: BLE001 # auth_jwt raises a plain Exception on signature and claim failures
return _refusal_for(denied, denied)
user_id: Final = result["user_id"]
if user_id is None:
return SubjectTokenRefusal(error="invalid_request", description="subject_token names no user the gateway knows")
return SubjectIdentity(user_id=user_id, team_id=result["team_id"])
def _refusal_for(denied: Exception, reason: object) -> SubjectTokenRefusal:
outage: Final = _gateway_could_not_verify(denied)
if outage is None:
verbose_proxy_logger.warning("token exchange refused a subject_token: %s", reason)
return SubjectTokenRefusal(error="invalid_request", description=REJECTED_SUBJECT_TOKEN)
verbose_proxy_logger.error("token exchange could not verify a subject_token, %s: %s", outage, reason)
return SubjectTokenRefusal(error="temporarily_unavailable", description=_check_unavailable_description(outage))
def _check_unavailable_description(outage: GatewayOutage) -> str:
match outage:
case "retryable":
return SUBJECT_TOKEN_CHECK_UNAVAILABLE
case "faulted":
return SUBJECT_TOKEN_CHECK_FAULTED
case _:
assert_never(outage)
def _gateway_could_not_verify(denied: Exception) -> GatewayOutage | None:
"""A database fault anywhere in the chain (``get_user_object`` wraps prisma failures in a
bare ``ValueError``) or a 5xx from JWT auth (the IdP's JWKS unreachable with no cached
copy) is the gateway failing, not the token. A fault retrying cannot clear (a missing or
version-skewed query engine) is named as such, the way the mint path words it, so the
client is not told to wait on a deployment that needs repair."""
fault: Final = PrismaDBExceptionHandler.find_database_service_unavailable_error_in_chain(denied)
if fault is not None:
return "faulted" if PrismaDBExceptionHandler.is_permanent_database_fault(fault) else "retryable"
return "retryable" if _is_server_error(denied) else None
def _is_server_error(denied: Exception) -> bool:
match denied:
case HTTPException(status_code=status_code):
return status_code >= 500
case ProxyException(code=code):
return code.isdigit() and int(code) >= 500
case _:
return False

16
litellm/proxy/_experimental/mcp_server/proxy_api_credentials.py
View file

@ -16,7 +16,7 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import (
ReloadUserFailure,
)
from litellm.proxy._types import LiteLLM_UserTable
from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken
from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken, effective_user_role
from litellm.proxy.management_endpoints.ui_sso import (
CliSsoTeamDetail,
fetch_cli_sso_team_details,
@ -42,7 +42,7 @@ async def mint_proxy_credential(
user_id: str, team_id: str | None
) -> MintedProxyCredential | ProxyCredentialMintFailure:
"""Mint the ``lite login`` credential for a consented grant. Membership is checked
live, so a team the user left between consent and redemption (or between refreshes)
live against the database row, so a team the user left between consent and redemption (or between refreshes)
refuses the grant instead of minting a credential attributed to a team they are no
longer on. The team is exactly the one the consent page sealed into the grant; nothing
is picked on the user's behalf here, so a refresh can never move the credential, and a
@ -51,12 +51,12 @@ async def mint_proxy_credential(
posting the consent form without one. Memberships whose team rows are gone count as no
team at all, the way ``lite login`` treats them, so they can never lock a user out. The
user row handed to the minter carries no team list, exactly like ``lite login``'s, so
the minter's own first-team fallback stays inert."""
user: Final = await load_active_user_by_id(user_id)
the minter's own first-team fallback stays inert. The credential carries the role the
proxy already enforces for the user on every request, so a row with no role (JWT auth's
upsert writes none) mints as an internal user instead of being refused."""
user: Final = await load_active_user_by_id(user_id, source="database")
if isinstance(user, str):
return user
if user.user_role is None:
return "no_active_key"
if team_id is not None and team_id not in user.teams:
return "not_a_member"
details: Final = await _team_details(user.teams) if user.teams else ()
@ -68,7 +68,9 @@ async def mint_proxy_credential(
if selected is None:
return "not_a_member"
key: Final = ExperimentalUIJWTToken.get_cli_jwt_auth_token(
user_info=LiteLLM_UserTable(user_id=user.user_id, user_role=user.user_role, models=user.models),
user_info=LiteLLM_UserTable(
user_id=user.user_id, user_role=effective_user_role(user.user_role).value, models=user.models
),
team_id=team_id,
team_alias=selected.team_alias,
team_models=selected.team_models,

66
litellm/proxy/_lazy_openapi_snapshot.json
View file

@ -23684,6 +23684,17 @@
],
"title": "Refresh Token"
},
"requested_token_type": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Requested Token Type"
},
"resource": {
"anyOf": [
{
@ -23705,6 +23716,28 @@
}
],
"title": "Scope"
},
"subject_token": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Subject Token"
},
"subject_token_type": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Subject Token Type"
}
},
"required": [
@ -23758,6 +23791,17 @@
],
"title": "Refresh Token"
},
"requested_token_type": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Requested Token Type"
},
"resource": {
"anyOf": [
{
@ -23779,6 +23823,28 @@
}
],
"title": "Scope"
},
"subject_token": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Subject Token"
},
"subject_token_type": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Subject Token Type"
}
},
"required": [

1
litellm/proxy/_types.py
View file

@ -534,6 +534,7 @@ class LiteLLMRoutes(enum.Enum):
"/mcp-rest/tools/call",
"/v1/mcp/tools",
"/introspect",
"/token",
]
# MCP server CRUD routes — control-plane. Gated by DISABLE_ADMIN_ENDPOINTS.

20
litellm/proxy/auth/auth_checks.py
View file

@ -1216,21 +1216,19 @@ async def common_checks(
return True
def effective_user_role(user_role: str | None) -> LitellmUserRoles:
try:
return LitellmUserRoles(user_role)
except ValueError:
return LitellmUserRoles.INTERNAL_USER
def _get_user_role(
user_obj: LiteLLM_UserTable | None,
) -> LitellmUserRoles | None:
if user_obj is None:
return None
_user: Final = user_obj
_user_role: Final = _user.user_role
try:
role: Final = LitellmUserRoles(_user_role)
except ValueError:
return LitellmUserRoles.INTERNAL_USER
return role
return effective_user_role(user_obj.user_role)
def _is_api_route_allowed(
@ -2577,7 +2575,7 @@ async def get_user_object(
raise Exception("No db connected")
try:
db_access_time_key: Final = f"user_id:{user_id}"
should_check_db: Final = _should_check_db(
should_check_db: Final = bool(check_db_only) or _should_check_db(
key=db_access_time_key,
last_db_access_time=last_db_access_time,
db_cache_expiry=db_cache_expiry,

8
litellm/proxy/auth/handle_jwt.py
View file

@ -1867,7 +1867,7 @@ class JWTAuthManager:
@staticmethod
def get_team_id_from_header(
request_headers: dict | None,
request_headers: Mapping[str, str] | None,
allowed_team_ids: set[str],
fallback_to_db_teams: bool = False,
) -> str | None:
@ -2037,7 +2037,7 @@ class JWTAuthManager:
async def _attach_team_from_header_for_admin(
admin_result: JWTAuthBuilderResult,
route: str,
request_headers: dict | None,
request_headers: Mapping[str, str] | None,
jwt_handler: JWTHandler,
prisma_client: PrismaClient | None,
user_api_key_cache: UserApiKeyCache,
@ -2293,7 +2293,7 @@ class JWTAuthManager:
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Span | None,
proxy_logging_obj: ProxyLogging,
request_headers: dict | None = None,
request_headers: Mapping[str, str] | None = None,
request_method: str | None = None,
) -> JWTAuthBuilderResult:
return await JWTAuthManager.authorize_jwt(
@ -2390,7 +2390,7 @@ class JWTAuthManager:
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Span | None,
proxy_logging_obj: ProxyLogging,
request_headers: dict[str, str] | None = None,
request_headers: Mapping[str, str] | None = None,
request_method: str | None = None,
provisioning: _JWTProvisioning | None = None,
) -> JWTAuthBuilderResult:

5
litellm/proxy/management_endpoints/team_endpoints.py
View file

@ -3335,6 +3335,7 @@ async def team_member_add(
```
"""
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast
from litellm.proxy.proxy_server import (
litellm_proxy_admin_name,
premium_user,
@ -3429,6 +3430,10 @@ async def team_member_add(
litellm_proxy_admin_name=litellm_proxy_admin_name,
)
await evict_and_broadcast(
cache_keys=tuple(sorted(user.user_id for user in updated_users)),
user_api_key_cache=user_api_key_cache,
)
await _evict_created_membership_caches(
user_ids=(tm.user_id for tm in updated_team_memberships),
team_id=data.team_id,

88
litellm/proxy/response_api_endpoints/endpoints.py
View file

@ -1,7 +1,8 @@
import asyncio
import contextlib
import json
import time
from collections.abc import AsyncIterator, Awaitable, Mapping
from collections.abc import AsyncIterator, Awaitable, Mapping, Sequence
from enum import Enum
from functools import partial
from types import MappingProxyType
@ -12,10 +13,12 @@ import fastapi
from fastapi import APIRouter, Depends, HTTPException, Request, Response
from fastapi.responses import JSONResponse
from openai.types.responses.response_create_params import ResponseInputParam
from pydantic import BaseModel, ConfigDict, ValidationError
from starlette.websockets import WebSocket, WebSocketDisconnect
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.constants import EMPTY_MAPPING
from litellm.integrations.custom_guardrail import ModifyResponseException
from litellm.llms.base_llm.guardrail_translation.utils import (
blocked_responses_api_usage as _blocked_responses_api_usage,
@ -1291,7 +1294,8 @@ async def cancel_response(
async def _read_ws_model_from_first_frame(
websocket: WebSocket,
) -> tuple | None:
query_model: str | None = None,
) -> tuple[str, str] | None:
"""Read the first WS frame and return (model, raw_message), or None on error.
Sends an appropriate error frame and closes the socket before returning None.
@ -1340,7 +1344,7 @@ async def _read_ws_model_from_first_frame(
await websocket.close(code=1008, reason="Invalid first message")
return None
model: Final = _extract_model_from_first_ws_event(first_event)
model: Final = query_model or _extract_model_from_first_ws_event(first_event)
if not model:
await websocket.send_text(
json.dumps(
@ -1371,6 +1375,38 @@ def _extract_model_from_first_ws_event(first_event: Any) -> str | None:
return (nested.get("model") if isinstance(nested, dict) else None) or first_event.get("model")
class _ResponseCreateRoutingHints(BaseModel):
model_config = ConfigDict(extra="ignore", frozen=True)
input: str | Sequence[object] | None = None
previous_response_id: str | None = None
response: "_ResponseCreateRoutingHints | None" = None
def _routing_hints_from_first_ws_frame(first_message: str) -> Mapping[str, object]:
try:
frame: Final = _ResponseCreateRoutingHints.model_validate_json(first_message)
except ValidationError:
return EMPTY_MAPPING
nested: Final = frame.response or frame
hints: Final = {
"input": frame.input if nested.input is None else nested.input,
"previous_response_id": (
frame.previous_response_id if nested.previous_response_id is None else nested.previous_response_id
),
}
return MappingProxyType({key: value for key, value in hints.items() if value is not None})
def _responses_ws_failure_frame(failure: Exception) -> str:
raw_status: Final = getattr(failure, "status_code", None)
status: Final = raw_status if isinstance(raw_status, int) and not isinstance(raw_status, bool) else 500
error_type: Final = (
"rate_limit_exceeded" if status == 429 else "invalid_request_error" if 400 <= status < 500 else "server_error"
)
return json.dumps({"type": "error", "status": status, "error": {"type": error_type, "message": str(failure)}})
async def _enforce_responses_ws_first_frame_model_auth(
request: Request,
model: str,
@ -1457,19 +1493,16 @@ async def responses_websocket_endpoint(
accept_kwargs["subprotocol"] = requested_protocols[0]
await websocket.accept(**accept_kwargs)
first_message: str | None = None
if not model:
result: Final = await _read_ws_model_from_first_frame(websocket)
if result is None:
return
model, first_message = result
result: Final = await _read_ws_model_from_first_frame(websocket, query_model=model)
if result is None:
return
resolved_model, first_message = result
data: dict[str, object] = {
"model": model,
"model": resolved_model,
"websocket": websocket,
"first_message": first_message,
}
if first_message is not None:
data["first_message"] = first_message
# Construct a synthetic Request for pre-call processing
headers_list: Final = list(websocket.scope.get("headers") or [])
@ -1482,7 +1515,7 @@ async def responses_websocket_endpoint(
request: Final = Request(scope=scope)
request._url = websocket.url
_body_bytes: Final = json.dumps({"model": model}).encode()
_body_bytes: Final = json.dumps({"model": resolved_model}).encode()
async def return_body():
return _body_bytes
@ -1492,10 +1525,10 @@ async def responses_websocket_endpoint(
# Phase 1: pre-call processing (auth, guardrails, rate limits)
base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data)
try:
if first_message is not None:
if not model:
await _enforce_responses_ws_first_frame_model_auth(
request=request,
model=model,
model=resolved_model,
user_api_key_dict=user_api_key_dict,
llm_router=llm_router,
)
@ -1514,7 +1547,7 @@ async def responses_websocket_endpoint(
user_request_timeout=user_request_timeout,
user_max_tokens=user_max_tokens,
user_api_base=user_api_base,
model=model,
model=resolved_model,
route_type="_aresponses_websocket",
)
except Exception as e:
@ -1536,16 +1569,31 @@ async def responses_websocket_endpoint(
await websocket.close(code=1008, reason="Pre-call error")
return
routed_data: Final = dict(
data, user_api_key_dict=user_api_key_dict, **_routing_hints_from_first_ws_frame(first_message)
)
# Phase 2: route to upstream provider
try:
data["user_api_key_dict"] = user_api_key_dict
llm_call: Final = await route_request(
data=data,
data=routed_data,
route_type="_aresponses_websocket",
llm_router=llm_router,
user_model=user_model,
)
await llm_call
except Exception:
failure: Final = await llm_call
if isinstance(failure, Exception):
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict,
original_exception=failure,
request_data=routed_data,
)
except Exception as e:
verbose_proxy_logger.exception("Responses WebSocket error")
with contextlib.suppress(Exception):
await websocket.send_text(_responses_ws_failure_frame(e))
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict,
original_exception=e,
request_data=routed_data,
)
await websocket.close(code=1011, reason="Internal server error")

118
litellm/responses/litellm_completion_transformation/streaming_iterator.py
View file

@ -22,7 +22,6 @@ from litellm.types.llms.openai import (
ContentPartAddedEvent,
ContentPartDoneEvent,
ContentPartDonePartOutputText,
ContentPartDonePartReasoningText,
FunctionCallArgumentsDeltaEvent,
FunctionCallArgumentsDoneEvent,
OutputItemAddedEvent,
@ -102,6 +101,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
self.sent_response_created_event: bool = False
self.sent_response_in_progress_event: bool = False
self.sent_output_item_added_event: bool = False
self.sent_message_item_added_event: bool = False
self.sent_content_part_added_event: bool = False
self.sent_output_text_done_event: bool = False
self.sent_output_content_part_done_event: bool = False
@ -111,6 +111,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
self.completed_response = None
self.final_text: str = ""
self._cached_item_id: str | None = None
self._message_output_index: int = 0
self._cached_response_id: str | None = None
self._buffered_chunk: ModelResponseStream | None = None
self._upstream_exhausted: bool = False
@ -563,7 +564,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
self._sequence_number += 1
event: Final = OutputItemAddedEvent(
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
output_index=0,
output_index=self._message_output_index,
item=BaseLiteLLMOpenAIResponseObject(
**{
"id": self._cached_item_id,
@ -585,13 +586,41 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
event: Final = ContentPartAddedEvent(
type=ResponsesAPIStreamEvents.CONTENT_PART_ADDED,
item_id=self._cached_item_id,
output_index=0,
output_index=self._message_output_index,
content_index=0,
part=BaseLiteLLMOpenAIResponseObject(**{"type": "output_text", "text": "", "annotations": []}),
)
event.__dict__["sequence_number"] = self._sequence_number
return event
def _queue_message_item_added_events(self) -> None:
if self._cached_item_id is None:
self._cached_item_id = f"msg_{uuid.uuid4()}"
self.sent_message_item_added_event = True
self.sent_content_part_added_event = True
if self._cached_reasoning_item_id is not None:
self._message_output_index = self._next_tool_output_index
self._next_tool_output_index += 1
else:
self._message_output_index = 0
self._sequence_number += 1
event: Final = OutputItemAddedEvent(
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
output_index=self._message_output_index,
item=BaseLiteLLMOpenAIResponseObject(
**{
"id": self._cached_item_id,
"type": "message",
"role": "assistant",
"status": "in_progress",
"content": [],
}
),
)
event.__dict__["sequence_number"] = self._sequence_number
self._pending_response_events.append(event)
self._pending_response_events.append(self.create_content_part_added_event())
def _merge_provider_specific_fields(self, src: dict) -> None:
"""Merge provider_specific_fields using last-value-wins for lists.
@ -711,7 +740,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
return OutputTextDoneEvent(
type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE,
item_id=self._cached_item_id,
output_index=0,
output_index=self._message_output_index,
content_index=0,
text=getattr(litellm_complete_object.choices[0].message, "content", "") or "",
)
@ -721,33 +750,24 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
self._cached_item_id = f"msg_{uuid.uuid4()}"
text: Final = getattr(litellm_complete_object.choices[0].message, "content", "") or ""
reasoning_content = getattr(litellm_complete_object.choices[0].message, "reasoning_content", "") or ""
annotations: Final = getattr(litellm_complete_object.choices[0].message, "annotations", None)
part: PART_UNION_TYPES | None = None
if reasoning_content:
part = ContentPartDonePartReasoningText(
type="reasoning_text",
reasoning=reasoning_content,
)
else:
response_annotations: Final = (
LiteLLMCompletionResponsesConfig._transform_chat_completion_annotations_to_response_output_annotations(
annotations=annotations
)
)
part = ContentPartDonePartOutputText(
type="output_text",
text=text,
annotations=response_annotations,
logprobs=None,
response_annotations: Final = (
LiteLLMCompletionResponsesConfig._transform_chat_completion_annotations_to_response_output_annotations(
annotations=annotations
)
)
part: Final[PART_UNION_TYPES] = ContentPartDonePartOutputText(
type="output_text",
text=text,
annotations=response_annotations,
logprobs=None,
)
return ContentPartDoneEvent(
type=ResponsesAPIStreamEvents.CONTENT_PART_DONE,
item_id=self._cached_item_id,
output_index=0,
output_index=self._message_output_index,
content_index=0,
part=part,
)
@ -766,7 +786,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
)
return OutputItemDoneEvent(
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
output_index=0,
output_index=self._message_output_index,
sequence_number=1,
item=BaseLiteLLMOpenAIResponseObject(
**{
@ -832,6 +852,15 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
def return_default_done_events(
self, litellm_complete_object: ModelResponse
) -> BaseLiteLLMOpenAIResponseObject | None:
if self.sent_message_item_added_event is False:
final_content: Final = litellm_complete_object.choices[0].message.content or ""
if not final_content:
self.sent_output_text_done_event = True
self.sent_output_content_part_done_event = True
self.sent_output_item_done_event = True
return None
self._queue_message_item_added_events()
return self._pending_response_events.pop(0)
if self.sent_output_text_done_event is False:
self.sent_output_text_done_event = True
return self.create_output_text_done_event(litellm_complete_object)
@ -936,31 +965,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
return
# Default: message
self._cached_item_id = self._cached_item_id or f"msg_{uuid.uuid4()}"
event = OutputItemAddedEvent(
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
output_index=0,
item=BaseLiteLLMOpenAIResponseObject(
**{
"id": self._cached_item_id,
"type": "message",
"role": "assistant",
"status": "in_progress",
"content": [],
}
),
)
event.__dict__["sequence_number"] = self._sequence_number
self._pending_response_events.append(event)
# Emit content_part.added immediately after output_item.added for message
# items. The OpenAI Responses spec requires this event before any
# output_text.delta events so downstream parsers can initialize the
# text part structure.
if not self.sent_content_part_added_event:
self.sent_content_part_added_event = True
content_part_event: Final = self.create_content_part_added_event()
self._pending_response_events.append(content_part_event)
self._queue_message_item_added_events()
return
async def __anext__(
@ -1115,12 +1120,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
self.collected_chat_completion_chunks.append(
self._snapshot_chunk_for_stream_chunk_builder(cast(ModelResponseStream, chunk))
)
# Emit any just-queued output_item event
if self._pending_response_events:
return self._pending_response_events.pop(0)
response_api_chunk = self._transform_chat_completion_chunk_to_response_api_chunk(chunk)
if response_api_chunk:
return response_api_chunk
self._pending_response_events.append(response_api_chunk)
if self._pending_response_events:
return self._pending_response_events.pop(0)
# Otherwise, loop to next chunk
except StopIteration:
return self.common_done_event_logic(sync_mode=True)
@ -1162,7 +1166,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
event = OutputTextAnnotationAddedEvent(
type=ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED,
item_id=item_id,
output_index=0,
output_index=self._message_output_index,
content_index=0,
annotation_index=idx,
annotation=annotation_dict,
@ -1189,11 +1193,13 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
# Priority 2: Handle text deltas
delta_content: Final = self._get_delta_string_from_streaming_choices(chunk.choices)
if delta_content:
if not self.sent_message_item_added_event:
self._queue_message_item_added_events()
self._sequence_number += 1
text_delta_event: Final = OutputTextDeltaEvent(
type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA,
item_id=item_id,
output_index=0,
output_index=self._message_output_index,
content_index=0,
delta=delta_content,
)

41
litellm/responses/main.py
View file

@ -1,5 +1,6 @@
import asyncio
import contextvars
import json
from collections.abc import Coroutine, Generator, Iterable, Mapping, Sequence
from contextlib import contextmanager
from dataclasses import dataclass
@ -8,7 +9,7 @@ from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, TypeAlias, cast
import httpx
from pydantic import BaseModel, TypeAdapter
from pydantic import BaseModel, TypeAdapter, ValidationError
from typing_extensions import assert_never
import litellm
@ -2274,6 +2275,27 @@ def _deployment_reasoning_default(kwargs: Mapping[str, object]) -> Reasoning | d
return _JSON_OBJECT_ADAPTER.validate_python(reasoning_effort) if isinstance(reasoning_effort, Mapping) else None
_RESPONSES_WS_ROUTING_HINT_KEYS: Final = frozenset({"input", "previous_response_id"})
def _first_ws_frame_with_routed_input(first_message: str, routed_input: object) -> str:
try:
frame: Final = _JSON_OBJECT_ADAPTER.validate_json(first_message)
except ValidationError:
return first_message
if frame is None or routed_input is None:
return first_message
raw_nested: Final = frame.get("response")
nested: Final = _JSON_OBJECT_ADAPTER.validate_python(raw_nested) if isinstance(raw_nested, Mapping) else None
if nested is not None and nested.get("input") is not None:
if nested["input"] == routed_input:
return first_message
return json.dumps({**frame, "response": {**nested, "input": routed_input}})
if frame.get("input") == routed_input:
return first_message
return json.dumps({**frame, "input": routed_input})
def _build_responses_websocket_request_defaults(kwargs: Mapping[str, object]) -> ResponsesWebSocketRequestDefaults:
default_reasoning: Final = _deployment_reasoning_default(kwargs)
candidate_params: Final[dict[str, object]] = {
@ -2295,11 +2317,11 @@ async def _aresponses_websocket(
api_key: str | None = None,
timeout: float | None = None,
**kwargs,
):
) -> Exception | None:
"""
Private function to handle the Responses API WebSocket mode.
For PROXY use only.
For PROXY use only. Returns the provider failure that ended the connection, if any.
Resolves the LLM provider from ``model``, looks up the matching
``BaseResponsesAPIConfig``, and hands off to
@ -2364,10 +2386,14 @@ async def _aresponses_websocket(
"api_base",
"api_key",
"timeout",
"first_message",
*_RESPONSES_WS_ROUTING_HINT_KEYS,
}
remaining_kwargs: Final = {k: v for k, v in kwargs.items() if k not in _explicit_keys}
deployment_kwargs: Final = {k: v for k, v in kwargs.items() if k not in _RESPONSES_WS_ROUTING_HINT_KEYS}
first_message: Final = kwargs.get("first_message")
await base_llm_http_handler.async_responses_websocket(
return await base_llm_http_handler.async_responses_websocket(
model=resolved_model,
websocket=websocket,
logging_obj=litellm_logging_obj,
@ -2375,9 +2401,14 @@ async def _aresponses_websocket(
api_base=resolved_api_base,
api_key=resolved_api_key,
timeout=timeout,
first_message=(
_first_ws_frame_with_routed_input(first_message, kwargs.get("input"))
if isinstance(first_message, str)
else None
),
user_api_key_dict=kwargs.get("user_api_key_dict"),
litellm_metadata=_build_litellm_metadata_for_ws(kwargs),
custom_llm_provider=_custom_llm_provider,
request_defaults=_build_responses_websocket_request_defaults(kwargs),
request_defaults=_build_responses_websocket_request_defaults(deployment_kwargs),
**remaining_kwargs,
)

186
litellm/responses/streaming_iterator.py
View file

@ -1,6 +1,7 @@
from __future__ import annotations
import asyncio
import copy
import json
import time
import traceback
@ -154,7 +155,7 @@ def _load_json_value(payload: str | bytes) -> object:
return json.loads(payload)
def _model_id_from_metadata(litellm_metadata: dict[str, object] | None) -> str | None:
def _model_id_from_metadata(litellm_metadata: Mapping[str, object] | None) -> str | None:
model_info: Final = litellm_metadata.get("model_info") if litellm_metadata else None
model_id: Final = model_info.get("id") if _is_json_object(model_info) else None
return model_id if isinstance(model_id, str) else None
@ -229,6 +230,29 @@ def _status_code_for_error_fields(error_type: str | None, error_code: str | None
return next((status for status in map(_status_code_for_error_field, fields) if status is not None), 500)
def _map_stream_error_to_exception(error_obj: object, model: str, custom_llm_provider: str) -> Exception:
from litellm.llms.base_llm.chat.transformation import BaseLLMException
error_message, error_type, error_code = _error_event_fields(error_obj)
status_code: Final = _status_code_for_error_fields(error_type, error_code)
error_body: Final = {"message": error_message, "type": error_type, "code": error_code}
provider_exception: Final = BaseLLMException(
status_code=status_code,
message=f"Error code: {status_code} - {{'error': {error_body}}}",
body=error_body,
)
try:
return litellm.exception_type(
model=model,
custom_llm_provider=custom_llm_provider,
original_exception=provider_exception,
completion_kwargs={},
extra_kwargs={},
)
except Exception as mapped_exception:
return mapped_exception
def _mid_stream_fallback_eligible(mapped_exception: Exception) -> bool:
if isinstance(mapped_exception, litellm.ContentPolicyViolationError):
return True
@ -592,26 +616,7 @@ class BaseResponsesAPIStreamingIterator:
)
def _map_error_event_exception(self, error_obj: object) -> Exception:
from litellm.llms.base_llm.chat.transformation import BaseLLMException
error_message, error_type, error_code = _error_event_fields(error_obj)
status_code: Final = _status_code_for_error_fields(error_type, error_code)
error_body: Final = {"message": error_message, "type": error_type, "code": error_code}
provider_exception: Final = BaseLLMException(
status_code=status_code,
message=f"Error code: {status_code} - {{'error': {error_body}}}",
body=error_body,
)
try:
return litellm.exception_type(
model=self.model or "",
custom_llm_provider=self.custom_llm_provider or "",
original_exception=provider_exception,
completion_kwargs={},
extra_kwargs={},
)
except Exception as mapped_exception:
return mapped_exception
return _map_stream_error_to_exception(error_obj, self.model or "", self.custom_llm_provider or "")
def _maybe_raise_for_error_event(self, result: object) -> None:
chunk_type: Final = getattr(result, "type", None)
@ -1695,6 +1700,65 @@ RESPONSES_WS_LOGGED_EVENT_TYPES: Final = [
RESPONSES_WS_MASKABLE_TEXT_BLOCK_TYPES: Final = frozenset({"input_text", "output_text", "text"})
_RESPONSES_WS_FAILURE_EVENT_TYPES: Final = frozenset({"error", "response.failed"})
_RESPONSES_WS_OUTPUT_ITEM_EVENT_TYPES: Final = frozenset({"response.output_item.added", "response.output_item.done"})
def _ws_event_error(event: Mapping[str, object]) -> object:
if event.get("type") == "error":
return event.get("error")
response: Final = event.get("response")
return response.get("error") if _is_json_object(response) else None
def _restore_input_item_ids(items: Sequence[object]) -> Sequence[object]:
return ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input(copy.deepcopy(list(items))) # pyright: ignore[reportPrivateUsage] # same restore the HTTP responses path runs
def _restored_container_fields(container: Mapping[str, object]) -> Mapping[str, object]:
input_items: Final = container.get("input")
previous_response_id: Final = container.get("previous_response_id")
restored: Final = {
"input": _restore_input_item_ids(input_items) if _is_json_array(input_items) else input_items,
"previous_response_id": (
ResponsesAPIRequestUtils.decode_previous_response_id_to_original_previous_response_id(previous_response_id)
if isinstance(previous_response_id, str)
else previous_response_id
),
}
return MappingProxyType({key: value for key, value in restored.items() if value != container.get(key)})
def _restore_wrapped_ids_in_response_create(msg_obj: Mapping[str, object]) -> dict[str, object] | None:
nested: Final = msg_obj.get("response")
nested_fields: Final = _restored_container_fields(nested) if _is_json_object(nested) else EMPTY_MAPPING
top_fields: Final = _restored_container_fields(msg_obj)
if not nested_fields and not top_fields:
return None
restored_nested: Final = (
{"response": {**nested, **nested_fields}} if _is_json_object(nested) and nested_fields else EMPTY_MAPPING
)
return {**msg_obj, **top_fields, **restored_nested}
def _wrap_output_item_encrypted_content(
event_obj: Mapping[str, object], litellm_metadata: Mapping[str, object]
) -> dict[str, object] | None:
if not litellm_metadata.get("encrypted_content_affinity_enabled"):
return None
model_id: Final = _model_id_from_metadata(litellm_metadata)
item: Final = event_obj.get("item")
if model_id is None or not _is_json_object(item):
return None
encrypted_content: Final = item.get("encrypted_content")
if not isinstance(encrypted_content, str) or not encrypted_content:
return None
wrapped_content: Final = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( # pyright: ignore[reportPrivateUsage] # same wrap the HTTP streaming path applies
encrypted_content=encrypted_content, model_id=model_id
)
return {**event_obj, "item": {**item, "encrypted_content": wrapped_content}}
class ResponsesWebSocketStreaming:
"""
@ -1721,6 +1785,7 @@ class ResponsesWebSocketStreaming:
output_guardrail_callbacks: list[PresidioGuardrailCallback] | None = None,
quota_callbacks: Sequence[ProjectQuotaCallback] | None = None,
authorized_model: str | None = None,
custom_llm_provider: str | None = None,
request_defaults: ResponsesWebSocketRequestDefaults | None = None,
):
self.websocket = websocket
@ -1728,6 +1793,9 @@ class ResponsesWebSocketStreaming:
self.logging_obj = logging_obj
self.user_api_key_dict = user_api_key_dict
self.request_data: dict[str, object] = request_data or {}
litellm_metadata: Final = self.request_data.get("litellm_metadata")
self.litellm_metadata: dict[str, object] = litellm_metadata if _is_json_object(litellm_metadata) else {}
self.custom_llm_provider: str | None = custom_llm_provider
self.messages: list[_MutableJsonObject] = []
self.input_messages: list[dict[str, object]] = []
self.first_message = first_message
@ -1796,13 +1864,65 @@ class ResponsesWebSocketStreaming:
if self.logging_obj:
self.logging_obj.pre_call(input=message, api_key="")
def _failure_exception(self) -> Exception | None:
failed_event: Final = next(
(event for event in self.messages if event.get("type") in _RESPONSES_WS_FAILURE_EVENT_TYPES), None
)
if failed_event is None:
return None
return _map_stream_error_to_exception(
_ws_event_error(failed_event), self.authorized_model or "", self.custom_llm_provider or ""
)
async def _log_messages(self) -> None:
if not self.logging_obj:
return
if self.input_messages:
self.logging_obj.model_call_details["messages"] = self.input_messages
if self.messages:
if not self.messages:
return
exception: Final = self._failure_exception()
if exception is None:
asyncio.create_task(self.logging_obj.dispatch_success_handlers(self.messages, prefer_async_handlers=True))
return
self._record_usage_for_failure()
traceback_exception: Final = "".join(traceback.format_exception(exception))
asyncio.create_task(
self.logging_obj.dispatch_failure_handlers(exception, traceback_exception, prefer_async_handlers=True)
)
def _record_usage_for_failure(self) -> None:
from litellm.cost_calculator import ResponsesWebSocketTokenUsageProcessor
from litellm.types.utils import LiteLLMRealtimeStreamLoggingObject
usage: Final = ResponsesWebSocketTokenUsageProcessor.collect_and_combine_usage_from_responses_ws_results(
self.messages
)
tier_partition: Final = ResponsesWebSocketTokenUsageProcessor.partition_results_by_service_tier(self.messages)
service_tier: Final = next(iter(tier_partition)) if len(tier_partition) == 1 else None
logging_result: Final = LiteLLMRealtimeStreamLoggingObject(
usage=usage, results=self.messages, service_tier=service_tier
)
response_cost: Final = self.logging_obj._response_cost_calculator(result=logging_result) or 0.0 # pyright: ignore[reportPrivateUsage] # as the HTTP streaming iterator does
self.logging_obj.record_partial_usage_for_failure(usage, response_cost)
def _wrap_response_event(self, response_str: str) -> str:
try:
event_obj: Final = _load_json_object(response_str)
except (json.JSONDecodeError, TypeError):
return response_str
response: Final = event_obj.get("response")
if _is_json_object(response):
wrapped_response: Final = ResponsesAPIRequestUtils._update_responses_api_response_id_with_model_id( # pyright: ignore[reportPrivateUsage] # same wrap the HTTP streaming path applies
responses_api_response=response,
custom_llm_provider=self.custom_llm_provider,
litellm_metadata=self.litellm_metadata,
)
return json.dumps({**event_obj, "response": wrapped_response})
if event_obj.get("type") not in _RESPONSES_WS_OUTPUT_ITEM_EVENT_TYPES:
return response_str
wrapped_event: Final = _wrap_output_item_encrypted_content(event_obj, self.litellm_metadata)
return response_str if wrapped_event is None else json.dumps(wrapped_event)
async def backend_to_client(self) -> None:
"""Forward events from backend WebSocket to the client."""
@ -1839,12 +1959,13 @@ class ResponsesWebSocketStreaming:
unmasked_str = self._unmask_response_event(response_str)
output_masked_str = await self._mask_response_completed(unmasked_str)
wrapped_str = self._wrap_response_event(output_masked_str)
# Log the output-masked form so PII redacted by apply_to_output
# guardrails does not appear in success logs.
self._store_event(output_masked_str)
self._store_event(wrapped_str)
await self.websocket.send_text(output_masked_str)
await self.websocket.send_text(wrapped_str)
except websockets.exceptions.ConnectionClosed as e:
verbose_logger.debug("Responses WS backend connection closed: %s", e)
@ -1913,19 +2034,22 @@ class ResponsesWebSocketStreaming:
if parsed.get("type") != "response.create":
return message
msg_obj: Final = self._with_request_defaults(parsed)
defaults_applied: Final = msg_obj != parsed
authorized_obj: Final = self._with_request_defaults(parsed)
defaults_applied: Final = authorized_obj != parsed
# Always enforce the authorized model, even when PII masking is off.
model_modified: Final = self._enforce_authorized_model(msg_obj)
model_modified: Final = self._enforce_authorized_model(authorized_obj)
restored_obj: Final = _restore_wrapped_ids_in_response_create(authorized_obj)
msg_obj: Final = authorized_obj if restored_obj is None else restored_obj
frame_modified: Final = model_modified or restored_obj is not None or defaults_applied
if not self.guardrail_callbacks:
return json.dumps(msg_obj) if model_modified or defaults_applied else message
return json.dumps(msg_obj) if frame_modified else message
if "metadata" not in self.request_data:
self.request_data["metadata"] = {}
modified = model_modified or defaults_applied
modified = frame_modified
guardrail_cbs: Final[tuple[PresidioGuardrailCallback, ...]] = tuple(self.guardrail_callbacks)
for cb in guardrail_cbs:
presidio_config = cb.get_presidio_settings_from_request_data(self.request_data)
@ -2209,8 +2333,7 @@ class ResponsesWebSocketStreaming:
except Exception as e:
verbose_logger.debug("Responses WS client_to_backend ended: %s", e)
async def bidirectional_forward(self) -> None:
"""Run both forwarding directions concurrently."""
async def bidirectional_forward(self) -> Exception | None:
forward_task: Final = asyncio.create_task(self.backend_to_client())
try:
await self.client_to_backend()
@ -2227,6 +2350,7 @@ class ResponsesWebSocketStreaming:
await self.backend_ws.close()
except Exception:
pass
return self._failure_exception()
# ---------------------------------------------------------------------------

2
litellm/utils.py
View file

@ -2008,7 +2008,7 @@ def client(original_function):
result=result,
call_type=call_type,
)
elif call_type == CallTypes.arealtime.value:
elif call_type in (CallTypes.arealtime.value, CallTypes.aresponses_websocket.value):
return result
### POST-CALL RULES ###
post_call_processing(

29
tests/test_litellm/litellm_core_utils/test_litellm_logging.py
View file

@ -1068,6 +1068,35 @@ async def test_arealtime_marks_litellm_params_async(monkeypatch):
assert LitellmLogging._is_sync_litellm_request(captured["litellm_params"]) is False
@pytest.mark.asyncio
async def test_aresponses_websocket_hands_back_the_provider_failure_without_a_success_log(monkeypatch: pytest.MonkeyPatch):
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.responses.main import base_llm_http_handler
success_events = []
class CaptureLogger(CustomLogger):
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
success_events.append(response_obj)
monkeypatch.setattr(litellm, "callbacks", [CaptureLogger()])
monkeypatch.setattr(litellm, "failure_callback", [])
monkeypatch.setattr(litellm, "_async_failure_callback", [])
monkeypatch.setattr(litellm, "success_callback", [])
monkeypatch.setattr(litellm, "_async_success_callback", [])
failure = litellm.BadRequestError(message="invalid_encrypted_content", model="gpt-4o", llm_provider="openai")
with patch.object( # test-quality-ok: the provider socket is the seam; how the wrapper treats the relay's outcome is under test
base_llm_http_handler, "async_responses_websocket", AsyncMock(return_value=failure)
):
outcome = await litellm._aresponses_websocket(model="openai/gpt-4o", websocket=MagicMock(), api_key="sk-test")
await asyncio.sleep(0)
with contextlib.suppress(asyncio.TimeoutError):
await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10.0)
assert outcome is failure
assert success_events == []
@pytest.mark.asyncio
async def test_agenerate_content_marks_litellm_params_async():
"""LIT-4475: the async ``agenerate_content`` entrypoint must plant

90
tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py
View file

@ -1957,6 +1957,96 @@ async def test_async_anthropic_messages_handler_passes_api_key_to_agentic_hooks(
)
_FOUNDRY_API_BASE: Final = "https://lit5418.services.ai.azure.com/anthropic"
_FOUNDRY_SSE_BODY: Final = (
b'event: message_start\ndata: {"type": "message_start", "message": {"id": "msg_1", "type": "message", '
b'"role": "assistant", "model": "claude-fable-5-1", "content": [], "stop_reason": null, '
b'"usage": {"input_tokens": 1, "output_tokens": 0}}}\n\n'
b'event: content_block_start\ndata: {"type": "content_block_start", "index": 0, '
b'"content_block": {"type": "text", "text": ""}}\n\n'
b'event: content_block_delta\ndata: {"type": "content_block_delta", "index": 0, '
b'"delta": {"type": "text_delta", "text": "ready"}}\n\n'
b'event: content_block_stop\ndata: {"type": "content_block_stop", "index": 0}\n\n'
b'event: message_delta\ndata: {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, '
b'"usage": {"output_tokens": 1}}\n\n'
b'event: message_stop\ndata: {"type": "message_stop"}\n\n'
)
@pytest.mark.parametrize("stream", [False, True])
@pytest.mark.asyncio
async def test_async_anthropic_messages_handler_passes_deployment_api_base_to_agentic_hooks(stream, monkeypatch):
"""
Regression for LIT-5418: an azure_ai deployment carries its Foundry endpoint as
``api_base``, a named parameter that never lands in kwargs. The agentic hooks
(websearch interception's follow-up call after the search) must receive it on
both the non-streaming and the streaming path, or the follow-up fails with
"Missing Azure API Base" and the client gets the dangling tool_use back.
"""
from litellm.integrations.custom_logger import CustomLogger
from litellm.llms.azure_ai.anthropic.messages_transformation import AzureAnthropicMessagesConfig
monkeypatch.delenv("AZURE_API_BASE", raising=False)
class CapturingAgenticCallback(CustomLogger):
def __init__(self):
super().__init__()
self.hook_kwargs: dict | None = None
async def async_should_run_agentic_loop(self, response, model, messages, tools, stream, custom_llm_provider, kwargs):
self.hook_kwargs = dict(kwargs)
return False, {}
callback = CapturingAgenticCallback()
handler = BaseLLMHTTPHandler()
upstream_request = httpx.Request("POST", f"{_FOUNDRY_API_BASE}/v1/messages")
upstream_response = (
httpx.Response(200, content=_FOUNDRY_SSE_BODY, request=upstream_request)
if stream
else httpx.Response(
200,
json={
"id": "msg_1",
"type": "message",
"role": "assistant",
"model": "claude-fable-5-1",
"content": [{"type": "text", "text": "ready"}],
"stop_reason": "end_turn",
"usage": {"input_tokens": 1, "output_tokens": 1},
},
request=upstream_request,
)
)
mock_client = AsyncMock(spec=AsyncHTTPHandler)
mock_client.post = AsyncMock(return_value=upstream_response)
mock_logging_obj = Mock()
mock_logging_obj.model_call_details = {}
mock_logging_obj.dynamic_success_callbacks = [callback]
result = await handler.async_anthropic_messages_handler(
model="claude-fable-5-1",
messages=[{"role": "user", "content": "Say ready"}],
anthropic_messages_provider_config=AzureAnthropicMessagesConfig(),
anthropic_messages_optional_request_params={"max_tokens": 32},
custom_llm_provider="azure_ai",
litellm_params=GenericLiteLLMParams(api_key="foundry-key", api_base=_FOUNDRY_API_BASE),
logging_obj=mock_logging_obj,
client=mock_client,
api_key="foundry-key",
api_base=_FOUNDRY_API_BASE,
stream=stream,
kwargs={},
)
if stream:
_ = [chunk async for chunk in result]
assert mock_client.post.call_args.kwargs["url"] == f"{_FOUNDRY_API_BASE}/v1/messages"
assert callback.hook_kwargs is not None, "agentic hook never ran"
assert callback.hook_kwargs.get("api_base") == _FOUNDRY_API_BASE
assert callback.hook_kwargs.get("api_key") == "foundry-key"
class _FakeWSExceptions:
class WebSocketException(Exception):
pass

112
tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
View file

@ -7568,6 +7568,69 @@ async def test_reload_active_user_by_id_permanent_engine_fault_is_faulted(proxy_
assert await _reload_active_user_by_id("sso-user-7") == "faulted"
@pytest.mark.asyncio
async def test_load_active_user_by_id_reads_the_row_from_the_database_not_the_cache(proxy_globals):
"""JWT auth caches the user it creates before it adds that user to the JWT's team, and adding a
member never evicts the cached row, so a credential minted off the cached row refused the very first
token exchange as not a member. The database source has to read the row from the database and leave
the fresh row in the cache for the requests the credential makes next."""
from litellm.proxy._experimental.mcp_server.bridge_token_flow import load_active_user_by_id
from litellm.proxy._types import LiteLLM_UserTable
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
cache = UserApiKeyCache()
await cache.async_set_cache(
key="fresh-jwt-user", value=LiteLLM_UserTable(user_id="fresh-jwt-user", teams=[]), model_type=LiteLLM_UserTable
)
prisma = MagicMock()
prisma.db.litellm_usertable.find_unique = AsyncMock(
return_value=LiteLLM_UserTable(user_id="fresh-jwt-user", teams=["team-a"])
)
proxy_globals.user_api_key_cache = cache
proxy_globals.prisma_client = prisma
loaded = await load_active_user_by_id("fresh-jwt-user", source="database")
assert not isinstance(loaded, str)
assert loaded.teams == ["team-a"]
cached = await cache.async_get_cache(key="fresh-jwt-user", model_type=LiteLLM_UserTable)
assert cached is not None
assert cached.teams == ["team-a"]
@pytest.mark.asyncio
async def test_load_active_user_by_id_serves_a_cached_row_without_a_database_read(proxy_globals):
"""Introspection and refresh revalidation run per call, so the loader's default source is the cache: a
cached row answers without a database read, and only a caller that asks for the database row pays for
one."""
from litellm.proxy._experimental.mcp_server.bridge_token_flow import (
_reload_active_user_by_id,
load_active_user_by_id,
)
from litellm.proxy._types import LiteLLM_UserTable
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
cache = UserApiKeyCache()
await cache.async_set_cache(
key="cached-jwt-user",
value=LiteLLM_UserTable(user_id="cached-jwt-user", teams=["team-a"]),
model_type=LiteLLM_UserTable,
)
prisma = MagicMock()
prisma.db.litellm_usertable.find_unique = AsyncMock(
return_value=LiteLLM_UserTable(user_id="cached-jwt-user", teams=[])
)
proxy_globals.user_api_key_cache = cache
proxy_globals.prisma_client = prisma
loaded = await load_active_user_by_id("cached-jwt-user")
assert not isinstance(loaded, str)
assert loaded.teams == ["team-a"]
assert await _reload_active_user_by_id("cached-jwt-user") is None
prisma.db.litellm_usertable.find_unique.assert_not_awaited()
@pytest.mark.asyncio
async def test_token_endpoint_uses_client_secret_basic_when_configured():
"""LIT-4091: a server with token_endpoint_auth_method=client_secret_basic must send the
@ -11048,6 +11111,43 @@ def test_native_client_login_walks_discovery_consent_token_refresh_and_revoke(mo
assert stranger.json()["error"] == "invalid_client"
@pytest.mark.parametrize(
"jwt_auth_enabled, virtual_key_claim_field, exchange_servable",
[(True, None, True), (False, None, False), (True, "client_id", False)],
ids=["jwt auth on", "jwt auth off", "jwts mapped to virtual keys"],
)
def test_discovery_advertises_the_exchange_grant_only_where_the_gateway_can_serve_it(
monkeypatch, jwt_auth_enabled, virtual_key_claim_field, exchange_servable
):
"""Every document a native client reads before it picks a grant (the versioned contract, the
aggregate authorization-server metadata, and the registration response) lists the RFC 8693
exchange exactly when the running proxy can serve it: JWT auth on, a database, a license, and
no JWT-to-virtual-key mapping, since the exchange would mint past the mapped key's policy."""
from litellm.caching.caching import DualCache
from litellm.proxy._types import LiteLLM_JWTAuth
from litellm.proxy.auth.handle_jwt import JWTHandler
client, _session_cookie, _minted = _native_client_app(monkeypatch)
handler: Final = JWTHandler()
handler.update_environment(
prisma_client=None,
user_api_key_cache=DualCache(),
litellm_jwtauth=LiteLLM_JWTAuth(virtual_key_claim_field=virtual_key_claim_field),
)
monkeypatch.setattr("litellm.proxy.proxy_server.jwt_handler", handler)
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"enable_jwt_auth": jwt_auth_enabled})
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", object())
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
exchange_grant = ["urn:ietf:params:oauth:grant-type:token-exchange"] if exchange_servable else []
expected = ["authorization_code", "refresh_token", *exchange_grant]
assert client.get("/.well-known/litellm-cli-auth").json()["grant_types_supported"] == expected
assert client.get("/.well-known/oauth-authorization-server/mcp").json()["grant_types_supported"] == expected
registered = client.post("/register", json={"redirect_uris": ["http://127.0.0.1:51234/callback"]})
assert registered.status_code == 201
assert registered.json()["grant_types"] == expected
def test_native_client_authorize_without_the_proxy_resource_keeps_the_mcp_flow(monkeypatch):
"""A registered client asking for the MCP resource (or no resource) never sees the consent
page, so existing MCP clients are untouched by the native-client arm."""
@ -11847,13 +11947,17 @@ async def test_oauth_refresh_revalidates_the_same_active_user_rule(
from litellm.proxy._experimental.mcp_server.bridge_token_flow import _reload_active_user_by_id
handler, _ = jwt_oauth_identity
handler.user_api_key_cache.set_cache(
"jwt-owner", LiteLLM_UserTable(user_id="jwt-owner", metadata={"scim_active": state != "inactive"})
)
user_id: Final = f"jwt-owner-{state}"
row: Final = LiteLLM_UserTable(user_id=user_id, metadata={"scim_active": state != "inactive"})
proxy_server.prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=row)
if state == "missing_database":
monkeypatch.setattr(proxy_server, "prisma_client", None)
expected: Final = None if state == "active" else "no_active_key" if state == "inactive" else "unresolvable"
assert await _reload_active_user_by_id("jwt-owner") == expected
assert await _reload_active_user_by_id(user_id) == expected
if state != "missing_database":
cached: Final = handler.user_api_key_cache.get_cache(user_id, model_type=LiteLLM_UserTable)
assert cached is not None
assert cached.metadata == row.metadata
@pytest.mark.asyncio

216
tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py
View file

@ -15,13 +15,18 @@ from starlette.requests import Request
from litellm.caching.caching import DualCache
from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import (
_AUTH_CODE_DEBUG_KEY,
ACCESS_TOKEN_TOKEN_TYPE,
CONNECT_FLOW_COOKIE_PREFIX,
GATEWAY_AUTH_CODE_PREFIX,
GATEWAY_AUTH_CODE_TTL_SECONDS,
MANUAL_DELIVERY_AUTH_CODE_TTL_SECONDS,
MAX_CLIENT_ID_LENGTH,
SUBJECT_TOKEN_TYPES,
TOKEN_EXCHANGE_GRANT_TYPE,
ConsentTeam,
MintedProxyCredential,
SubjectIdentity,
SubjectTokenRefusal,
_GatewayAuthCode,
_open_sealed,
_seal,
@ -90,9 +95,11 @@ def _request(path="/authorize", query="", cookies=None, method="GET"):
)
async def _register(redirect_uris) -> dict:
async def _register(redirect_uris, token_exchange_available=True) -> dict:
response = await register_aggregate_client(
request=_request(path="/register", method="POST"), request_body={"redirect_uris": redirect_uris}
request=_request(path="/register", method="POST"),
request_body={"redirect_uris": redirect_uris},
token_exchange_available=token_exchange_available,
)
return json.loads(response.body)
@ -105,6 +112,7 @@ async def _reload_user_active(user_id: str):
async def test_register_mints_stateless_public_client():
body = await _register([REDIRECT_URI])
assert body["token_endpoint_auth_method"] == "none"
assert body["grant_types"] == ["authorization_code", "refresh_token", TOKEN_EXCHANGE_GRANT_TYPE]
assert "client_secret" not in body
assert body["redirect_uris"] == [REDIRECT_URI]
assert is_gateway_dcr_client_id(body["client_id"])
@ -113,11 +121,18 @@ async def test_register_mints_stateless_public_client():
assert record.redirect_uris == (REDIRECT_URI,)
@pytest.mark.asyncio
async def test_register_omits_the_exchange_grant_where_the_gateway_cannot_serve_it():
body = await _register([REDIRECT_URI], token_exchange_available=False)
assert body["grant_types"] == ["authorization_code", "refresh_token"]
@pytest.mark.asyncio
@pytest.mark.parametrize("redirect_uris", [VSCODE_REDIRECT_URIS, MAX_LENGTH_REDIRECT_URIS])
async def test_register_four_callbacks_preserves_metadata(redirect_uris: tuple[str, ...]) -> None:
response: Final = await register_aggregate_client(
request=_request(path="/register", method="POST"),
token_exchange_available=True,
request_body={
"client_name": "Visual Studio Code",
"client_uri": "https://code.visualstudio.com",
@ -143,6 +158,7 @@ async def test_register_four_callbacks_preserves_metadata(redirect_uris: tuple[s
async def test_register_rejects_five_valid_callbacks() -> None:
response: Final = await register_aggregate_client(
request=_request(path="/register", method="POST"),
token_exchange_available=True,
request_body={"redirect_uris": [*VSCODE_REDIRECT_URIS, "http://127.0.0.1:33419/"]},
)
assert response.status_code == 400
@ -156,6 +172,7 @@ async def test_register_rejects_five_valid_callbacks() -> None:
async def test_register_four_callbacks_preserves_encoded_size_guard() -> None:
response: Final = await register_aggregate_client(
request=_request(path="/register", method="POST"),
token_exchange_available=True,
request_body={"redirect_uris": [f"https://client.example/{index}/".ljust(256, "é") for index in range(4)]},
)
assert response.status_code == 400
@ -208,6 +225,7 @@ async def test_register_rejects_userinfo_spoofed_origin():
response = await register_aggregate_client(
request=_request(path="/register", method="POST"),
request_body={"redirect_uris": ["https://claude.ai@attacker.example/callback"]},
token_exchange_available=True,
)
assert response.status_code == 400
assert json.loads(response.body)["error"] == "invalid_redirect_uri"
@ -228,7 +246,9 @@ async def test_register_rejects_userinfo_spoofed_origin():
)
async def test_register_rejects_bad_redirect_uris(redirect_uris):
response = await register_aggregate_client(
request=_request(path="/register", method="POST"), request_body={"redirect_uris": redirect_uris}
request=_request(path="/register", method="POST"),
request_body={"redirect_uris": redirect_uris},
token_exchange_available=True,
)
assert response.status_code == 400
assert json.loads(response.body)["error"] in ("invalid_redirect_uri", "invalid_client_metadata")
@ -1948,7 +1968,7 @@ async def test_revoke_refuses_unknown_clients_and_a_missing_master_key():
def test_native_client_auth_contract_points_every_endpoint_at_this_proxy():
assert json.loads(json.dumps(native_client_auth_contract(_request("/.well-known/litellm-cli-auth")))) == {
assert json.loads(json.dumps(native_client_auth_contract(_request("/.well-known/litellm-cli-auth"), True))) == {
"contract_version": 1,
"issuer": "https://llm.example.com",
"authorization_endpoint": "https://llm.example.com/authorize",
@ -1957,13 +1977,22 @@ def test_native_client_auth_contract_points_every_endpoint_at_this_proxy():
"revocation_endpoint": "https://llm.example.com/revoke",
"resource": "https://llm.example.com",
"response_types_supported": ["code"],
"grant_types_supported": ["authorization_code", "refresh_token"],
"grant_types_supported": [
"authorization_code",
"refresh_token",
"urn:ietf:params:oauth:grant-type:token-exchange",
],
"code_challenge_methods_supported": ["S256"],
"token_endpoint_auth_methods_supported": ["none"],
"revocation_endpoint_auth_methods_supported": ["none"],
}
def test_native_client_auth_contract_omits_the_exchange_grant_where_the_gateway_cannot_serve_it():
contract = native_client_auth_contract(_request("/.well-known/litellm-cli-auth"), False)
assert list(contract["grant_types_supported"]) == ["authorization_code", "refresh_token"]
@pytest.mark.parametrize(
"resource, expected",
[
@ -2148,3 +2177,180 @@ async def test_gateway_owned_resource_stays_scoped_through_consent_and_refresh(a
)
assert renewed.status_code == 200
assert _opened_principal(json.loads(renewed.body)).resource_server_id == "github-id"
JWT_SUBJECT_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:jwt"
IDP_TOKEN = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1MSJ9.idp-signature"
class _Exchanger:
def __init__(self, result=None):
self.calls = []
self.result = result
async def __call__(self, subject_token, request):
self.calls.append((subject_token, request.url.path))
if self.result is not None:
return self.result
return SubjectIdentity(user_id="u1", team_id="team-b")
async def _exchange_native(client_id, minter, exchanger, cache=None, **overrides):
arguments = {
"grant_type": TOKEN_EXCHANGE_GRANT_TYPE,
"subject_token": IDP_TOKEN,
"subject_token_type": JWT_SUBJECT_TOKEN_TYPE,
"exchange_subject_token": exchanger,
}
return await _redeem_native(None, client_id, minter, cache=cache, **{**arguments, **overrides})
@pytest.mark.asyncio
async def test_token_exchange_mints_the_proxy_credential_for_the_idp_subject():
"""RFC 8693: a registered native client trades the IdP token it already holds for the
same credential the consent flow mints, attributed to the user and team the gateway's
JWT auth resolved, with a rotating refresh token bound to that team and the client.
The exchange can be repeated while the IdP token lives; nothing is burned."""
client_id = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"]
minter, exchanger, cache = _Minter(), _Exchanger(), DualCache()
response = await _exchange_native(client_id, minter, exchanger, cache=cache)
assert response.status_code == 200
assert response.headers["cache-control"] == "no-store"
body = json.loads(response.body)
assert exchanger.calls == [(IDP_TOKEN, "/token")]
assert minter.calls == [("u1", "team-b")]
assert body["issued_token_type"] == ACCESS_TOKEN_TOKEN_TYPE
assert body["access_token"] == "sk-cli-u1"
assert body["token_type"] == "Bearer"
assert body["expires_in"] == 3600
assert (body["user_id"], body["team_id"]) == ("u1", "team-b")
principal = _opened_refresh(body["refresh_token"], client_id)
assert (principal.user_id, principal.client_id, principal.audience, principal.team_id) == (
"u1",
client_id,
"proxy_api",
"team-b",
)
again = await _exchange_native(client_id, minter, exchanger, cache=cache)
assert again.status_code == 200
assert json.loads(again.body)["refresh_token"] != body["refresh_token"]
assert minter.calls == [("u1", "team-b"), ("u1", "team-b")]
@pytest.mark.asyncio
async def test_exchanged_credential_refreshes_and_rotates_like_a_consented_one():
client_id = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"]
minter, cache = _Minter(), DualCache()
exchanged = json.loads((await _exchange_native(client_id, minter, _Exchanger(), cache=cache)).body)
refreshed = await _refresh_native(exchanged["refresh_token"], client_id, minter, cache)
assert refreshed.status_code == 200
body = json.loads(refreshed.body)
assert "issued_token_type" not in body
assert (body["access_token"], body["user_id"], body["team_id"]) == ("sk-cli-u1", "u1", "team-b")
assert body["refresh_token"] != exchanged["refresh_token"]
assert minter.calls == [("u1", "team-b"), ("u1", "team-b")]
replay = await _refresh_native(exchanged["refresh_token"], client_id, minter, cache)
assert replay.status_code == 400
assert json.loads(replay.body)["error"] == "invalid_grant"
@pytest.mark.asyncio
async def test_token_exchange_for_a_teamless_subject_mints_a_teamless_credential():
client_id = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"]
minter = _Minter()
response = await _exchange_native(client_id, minter, _Exchanger(SubjectIdentity(user_id="u2")))
assert response.status_code == 200
body = json.loads(response.body)
assert minter.calls == [("u2", None)]
assert (body["user_id"], body["team_id"]) == ("u2", None)
assert _opened_refresh(body["refresh_token"], client_id).team_id is None
@pytest.mark.asyncio
@pytest.mark.parametrize("subject_token_type", sorted(SUBJECT_TOKEN_TYPES))
async def test_token_exchange_accepts_every_advertised_subject_token_type(subject_token_type):
client_id = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"]
response = await _exchange_native(client_id, _Minter(), _Exchanger(), subject_token_type=subject_token_type)
assert response.status_code == 200
@pytest.mark.asyncio
async def test_token_exchange_without_an_idp_exchanger_is_unsupported():
"""A gateway that wires no IdP verifier into the endpoint answers the way it always
answered an unknown grant, and never reaches the minter."""
client_id = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"]
minter = _Minter()
response = await _redeem_native(
None,
client_id,
minter,
grant_type=TOKEN_EXCHANGE_GRANT_TYPE,
subject_token=IDP_TOKEN,
subject_token_type=JWT_SUBJECT_TOKEN_TYPE,
)
assert response.status_code == 400
assert json.loads(response.body)["error"] == "unsupported_grant_type"
assert minter.calls == []
@pytest.mark.asyncio
@pytest.mark.parametrize(
"overrides, status, error",
[
({"subject_token": None}, 400, "invalid_request"),
({"subject_token": ""}, 400, "invalid_request"),
({"subject_token_type": None}, 400, "invalid_request"),
({"subject_token_type": "urn:ietf:params:oauth:token-type:saml2"}, 400, "invalid_request"),
({"requested_token_type": "urn:ietf:params:oauth:token-type:refresh_token"}, 400, "invalid_request"),
({"resource": "https://other.example.com"}, 400, "invalid_target"),
({"resource": "https://llm.example.com/mcp"}, 400, "invalid_target"),
({"client_id": "llm_dcrc_forged"}, 401, "invalid_client"),
({"client_id": "not-a-gateway-client"}, 401, "invalid_client"),
],
)
async def test_token_exchange_refuses_a_malformed_request_before_touching_the_idp_token(overrides, status, error):
registered = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"]
minter, exchanger = _Minter(), _Exchanger()
response = await _exchange_native(
overrides.get("client_id", registered),
minter,
exchanger,
**{name: value for name, value in overrides.items() if name != "client_id"},
)
assert response.status_code == status
assert json.loads(response.body)["error"] == error
assert exchanger.calls == []
assert minter.calls == []
@pytest.mark.asyncio
@pytest.mark.parametrize(
"error, status",
[("unsupported_grant_type", 400), ("invalid_request", 400), ("temporarily_unavailable", 503)],
)
async def test_token_exchange_relays_the_idp_refusal_and_never_mints(error, status):
client_id = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"]
minter = _Minter()
exchanger = _Exchanger(SubjectTokenRefusal(error=error, description="subject_token was rejected: bad signature"))
response = await _exchange_native(client_id, minter, exchanger)
assert response.status_code == status
body = json.loads(response.body)
assert (body["error"], body["error_description"]) == (error, "subject_token was rejected: bad signature")
assert minter.calls == []
@pytest.mark.asyncio
@pytest.mark.parametrize(
"failure, status, error",
[
("not_a_member", 400, "invalid_grant"),
("team_required", 400, "invalid_grant"),
("no_active_key", 400, "invalid_grant"),
("unavailable", 503, "temporarily_unavailable"),
],
)
async def test_token_exchange_relays_a_mint_refusal(failure, status, error):
client_id = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"]
response = await _exchange_native(client_id, _Minter(failure), _Exchanger())
assert response.status_code == status
assert json.loads(response.body)["error"] == error

237
tests/test_litellm/proxy/_experimental/mcp_server/test_idp_token_exchange.py Normal file
View file

@ -0,0 +1,237 @@
import logging
import pytest
from fastapi import HTTPException
from prisma.engine.errors import BinaryNotFoundError
from prisma.errors import DataError
from litellm.caching.caching import DualCache
from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import SubjectIdentity, SubjectTokenRefusal
from litellm.proxy._experimental.mcp_server.idp_token_exchange import (
REJECTED_SUBJECT_TOKEN,
SUBJECT_TOKEN_CHECK_FAULTED,
SUBJECT_TOKEN_CHECK_UNAVAILABLE,
TokenExchangePrerequisites,
identity_from_subject_token,
token_exchange_available,
)
from litellm.proxy._types import JWTIssuerConfig, LiteLLM_JWTAuth, ProxyException
from litellm.proxy.auth.handle_jwt import JWKSUnreachableError, JWTHandler, jwks_unavailable_exception
IDP_JWT = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1MSJ9.idp-signature"
REQUEST_HEADERS = {"x-litellm-team-id": "team-b", "user-agent": "lite/0.1"}
EVERY_GATE_HOLDS = {
"jwt_auth_enabled": True,
"has_database": True,
"licensed": True,
"maps_jwts_to_virtual_keys": False,
}
JWKS_URL = "https://idp.example.com/.well-known/jwks.json"
JWKS_DOWN = jwks_unavailable_exception(JWKSUnreachableError(f"ConnectError fetching {JWKS_URL} after 3 attempts"))
def _authorized(user_id="u1", team_id="team-b"):
return {
"is_proxy_admin": False,
"team_object": None,
"user_object": None,
"end_user_object": None,
"org_object": None,
"token": IDP_JWT,
"team_id": team_id,
"user_id": user_id,
"user_email": None,
"end_user_id": None,
"org_id": None,
"team_membership": None,
"jwt_claims": {"sub": user_id},
"agent_id": None,
}
class _Authorizer:
def __init__(self, result=None, raises=None):
self.calls = []
self.result = result if result is not None else _authorized()
self.raises = raises
async def __call__(self, subject_token, request_headers):
self.calls.append((subject_token, dict(request_headers)))
if self.raises is not None:
raise self.raises
return self.result
async def _identity(authorizer, subject_token=IDP_JWT, **unmet):
return await identity_from_subject_token(
subject_token,
request_headers=REQUEST_HEADERS,
prerequisites=TokenExchangePrerequisites(**{**EVERY_GATE_HOLDS, **unmet}),
is_jwt=JWTHandler.is_jwt,
authorize=authorizer,
)
@pytest.mark.asyncio
async def test_a_jwt_the_proxy_accepts_names_its_user_and_team():
"""The subject token goes to the proxy's own JWT auth with the caller's headers (that is
where the team header is read), and the identity it resolved is what gets minted."""
authorizer = _Authorizer()
assert await _identity(authorizer) == SubjectIdentity(user_id="u1", team_id="team-b")
assert authorizer.calls == [(IDP_JWT, REQUEST_HEADERS)]
@pytest.mark.asyncio
async def test_a_jwt_that_resolves_no_team_names_a_teamless_identity():
assert await _identity(_Authorizer(_authorized(team_id=None))) == SubjectIdentity(user_id="u1", team_id=None)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"unmet, subject_token, error, mentions",
[
({"jwt_auth_enabled": False}, IDP_JWT, "unsupported_grant_type", "JWT auth is not enabled"),
({"has_database": False}, IDP_JWT, "unsupported_grant_type", "no database"),
({"licensed": False}, IDP_JWT, "unsupported_grant_type", "enterprise"),
({"maps_jwts_to_virtual_keys": True}, IDP_JWT, "unsupported_grant_type", "virtual keys"),
({}, "sk-litellm-virtual-key", "invalid_request", "not a JWT"),
],
)
async def test_the_gates_user_api_key_auth_applies_refuse_before_any_verification(
unmet, subject_token, error, mentions
):
authorizer = _Authorizer()
refusal = await _identity(authorizer, subject_token=subject_token, **unmet)
assert isinstance(refusal, SubjectTokenRefusal)
assert refusal.error == error
assert mentions in refusal.description
assert authorizer.calls == []
@pytest.mark.parametrize(
"unmet",
[
{},
{"jwt_auth_enabled": False},
{"has_database": False},
{"licensed": False},
{"maps_jwts_to_virtual_keys": True},
],
)
def test_the_grant_is_available_exactly_when_every_gate_holds(unmet):
prerequisites = TokenExchangePrerequisites(**{**EVERY_GATE_HOLDS, **unmet})
assert prerequisites.available is (unmet == {})
assert (prerequisites.refusal() is None) is prerequisites.available
MAPPED_ISSUER = JWTIssuerConfig(
issuer="https://idp.example.test", audience="litellm-gateway", virtual_key_claim_field="client_id"
)
def _running_jwt_handler(litellm_jwtauth):
handler = JWTHandler()
if litellm_jwtauth is not None:
handler.update_environment(prisma_client=None, user_api_key_cache=DualCache(), litellm_jwtauth=litellm_jwtauth)
return handler
@pytest.mark.parametrize(
"general_settings, prisma_client, premium_user, litellm_jwtauth, expected",
[
({"enable_jwt_auth": True}, object(), True, LiteLLM_JWTAuth(), True),
({"enable_jwt_auth": True}, object(), True, None, True),
({}, object(), True, LiteLLM_JWTAuth(), False),
({"enable_jwt_auth": True}, None, True, LiteLLM_JWTAuth(), False),
({"enable_jwt_auth": True}, object(), False, LiteLLM_JWTAuth(), False),
({"enable_jwt_auth": True}, object(), True, LiteLLM_JWTAuth(virtual_key_claim_field="client_id"), False),
({"enable_jwt_auth": True}, object(), True, LiteLLM_JWTAuth(issuers=[MAPPED_ISSUER]), False),
],
)
def test_availability_is_read_from_the_running_proxy(
monkeypatch, general_settings, prisma_client, premium_user, litellm_jwtauth, expected
):
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client)
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", premium_user)
monkeypatch.setattr("litellm.proxy.proxy_server.jwt_handler", _running_jwt_handler(litellm_jwtauth))
assert token_exchange_available() is expected
@pytest.mark.asyncio
@pytest.mark.parametrize(
"raised, reason",
[
(HTTPException(status_code=403, detail="User not allowed to access this route"), "not allowed"),
(ProxyException(message="Token expired", type="auth_error", param="token", code=401), "Token expired"),
(Exception("Validation fails: signature verification failed"), "signature verification failed"),
(Exception("Invalid JWT Submitted"), "Invalid JWT"),
(Exception(f"Failed to fetch keys from {JWKS_URL}: 502 Bad Gateway from the IdP"), JWKS_URL),
(ValueError("User doesn't exist in db. 'user_id'=u1. Got error - not found"), "not found"),
],
)
async def test_a_jwt_the_proxy_rejects_is_refused_with_the_reason_kept_in_the_log(raised, reason, caplog):
"""The endpoint is public, so the response never quotes JWT auth's wording (it can name
the JWKS URL or relay the IdP's reply); the operator reads the reason in the proxy log."""
caplog.set_level(logging.WARNING, logger="LiteLLM Proxy")
refusal = await _identity(_Authorizer(raises=raised))
assert refusal == SubjectTokenRefusal(error="invalid_request", description=REJECTED_SUBJECT_TOKEN)
assert reason in caplog.text
@pytest.mark.asyncio
async def test_a_jwt_that_resolves_no_user_cannot_be_exchanged():
refusal = await _identity(_Authorizer(_authorized(user_id=None)))
assert refusal == SubjectTokenRefusal(
error="invalid_request", description="subject_token names no user the gateway knows"
)
def _user_lookup_wrapping_a_database_outage():
p1001 = DataError(
data={"user_facing_error": {"message": "Can't reach database server at `127.0.0.1`:`5432`", "meta": {}}}
)
try:
raise p1001
except DataError as outage:
try:
raise ValueError(f"User doesn't exist in db. 'user_id'=u1. Got error - {outage}")
except ValueError as wrapped:
return wrapped
@pytest.mark.asyncio
@pytest.mark.parametrize(
"raised, reason",
[
(JWKS_DOWN, JWKS_URL),
(HTTPException(status_code=503, detail="the auth database is not reachable"), "not reachable"),
(_user_lookup_wrapping_a_database_outage(), "Can't reach database server"),
],
)
async def test_an_idp_or_gateway_outage_is_reported_as_retryable_not_as_a_bad_token(raised, reason, caplog):
caplog.set_level(logging.ERROR, logger="LiteLLM Proxy")
refusal = await _identity(_Authorizer(raises=raised))
assert refusal == SubjectTokenRefusal(error="temporarily_unavailable", description=SUBJECT_TOKEN_CHECK_UNAVAILABLE)
assert reason in caplog.text
def _user_lookup_wrapping_a_fault_retrying_cannot_clear():
try:
raise BinaryNotFoundError("query engine binary not found")
except BinaryNotFoundError as fault:
try:
raise ValueError(f"User doesn't exist in db. 'user_id'=u1. Got error - {fault}")
except ValueError as wrapped:
return wrapped
@pytest.mark.asyncio
async def test_a_database_fault_retrying_cannot_clear_is_not_reported_as_a_transient_outage(caplog):
"""The status stays 503 (the only OAuth error a client reads as the server's fault, and what
the mint path answers to the same fault) but the wording must not tell the client to wait."""
caplog.set_level(logging.ERROR, logger="LiteLLM Proxy")
refusal = await _identity(_Authorizer(raises=_user_lookup_wrapping_a_fault_retrying_cannot_clear()))
assert refusal == SubjectTokenRefusal(error="temporarily_unavailable", description=SUBJECT_TOKEN_CHECK_FAULTED)
assert "retrying will not help" in refusal.description
assert "faulted: " in caplog.text and "query engine binary not found" in caplog.text

72
tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py
View file

@ -1,6 +1,6 @@
"""Tests for minting the ``lite login`` credential from a consented native-client grant."""
from unittest.mock import ANY, AsyncMock
from unittest.mock import ANY, AsyncMock, MagicMock
import pytest
@ -8,7 +8,9 @@ from litellm.constants import CLI_JWT_EXPIRATION_HOURS
from litellm.models.user import LiteLLM_UserTable
from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ConsentTeam, MintedProxyCredential
from litellm.proxy._experimental.mcp_server.proxy_api_credentials import lookup_consent_teams, mint_proxy_credential
from litellm.proxy._types import LitellmUserRoles
from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.management_endpoints.ui_sso import CliSsoTeamDetail
_LOAD_USER = "litellm.proxy._experimental.mcp_server.proxy_api_credentials.load_active_user_by_id"
@ -67,10 +69,21 @@ async def test_mint_passes_user_lookup_failures_through(failure, load_user, fetc
@pytest.mark.asyncio
async def test_mint_refuses_a_user_without_a_role(load_user, fetch_teams):
load_user.return_value = _user(user_role=None)
assert await mint_proxy_credential("u1", None) == "no_active_key"
fetch_teams.assert_not_awaited()
@pytest.mark.parametrize(
"stored_role, minted_role",
[
(None, LitellmUserRoles.INTERNAL_USER),
("made_up_role", LitellmUserRoles.INTERNAL_USER),
("proxy_admin", LitellmUserRoles.PROXY_ADMIN),
],
)
async def test_mint_carries_the_role_the_proxy_enforces_for_the_user(load_user, fetch_teams, stored_role, minted_role):
"""A user JWT auth upserted has no role in the database, and the proxy already treats
such a user as an internal user on every request, so the credential says the same."""
load_user.return_value = _user(user_role=stored_role)
minted = await mint_proxy_credential("u1", "team-a")
assert isinstance(minted, MintedProxyCredential)
assert _decoded(minted).user_role == minted_role
@pytest.mark.asyncio
@ -79,7 +92,7 @@ async def test_mint_refuses_a_teamless_grant_for_a_team_member(load_user, fetch_
is refused for a user with teams instead of minting an unscoped credential or drifting
onto the first team, on redemption and on every refresh alike."""
assert await mint_proxy_credential("u1", None) == "team_required"
load_user.assert_awaited_once_with("u1")
load_user.assert_awaited_once_with("u1", source="database")
fetch_teams.assert_awaited_once_with(ANY, ["team-a", "team-b"])
@ -114,6 +127,53 @@ async def test_mint_honors_the_consented_team(load_user, fetch_teams):
assert decoded.team_model_aliases == {"fast": "gpt-5.4-mini"}
@pytest.mark.asyncio
async def test_mint_reads_the_users_teams_from_the_database_not_a_stale_cached_row(fetch_teams, monkeypatch):
"""JWT auth caches the user it creates before it adds that user to the JWT's team, and adding a member
never evicts the cached row, so a mint off the cached row refused the very first token exchange as not
a member. The mint has to read the database row, whatever the cache holds."""
from litellm.proxy import proxy_server
cache = UserApiKeyCache()
await cache.async_set_cache(
key="stale-cache-user", value=_user(user_id="stale-cache-user", teams=[]), model_type=LiteLLM_UserTable
)
prisma = MagicMock()
prisma.db.litellm_usertable.find_unique = AsyncMock(
return_value=_user(user_id="stale-cache-user", teams=["team-a"])
)
monkeypatch.setattr(proxy_server, "user_api_key_cache", cache)
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
minted = await mint_proxy_credential("stale-cache-user", "team-a")
assert isinstance(minted, MintedProxyCredential)
assert minted.team_id == "team-a"
assert _decoded(minted).team_id == "team-a"
@pytest.mark.asyncio
async def test_mint_refuses_a_user_scim_deactivated_after_the_cache_last_saw_them_active(fetch_teams, monkeypatch):
"""SCIM deactivation writes the user row without evicting the cached copy, so a mint off the cache would
keep issuing credentials for the management-object TTL. The mint reads the database row, so the
deactivated user is refused on the first refresh after the deactivation."""
from litellm.proxy import proxy_server
cache = UserApiKeyCache()
await cache.async_set_cache(
key="deactivated-user", value=_user(user_id="deactivated-user", teams=["team-a"]), model_type=LiteLLM_UserTable
)
prisma = MagicMock()
prisma.db.litellm_usertable.find_unique = AsyncMock(
return_value=_user(user_id="deactivated-user", teams=["team-a"], metadata={"scim_active": False})
)
monkeypatch.setattr(proxy_server, "user_api_key_cache", cache)
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
assert await mint_proxy_credential("deactivated-user", "team-a") == "no_active_key"
fetch_teams.assert_not_awaited()
@pytest.mark.asyncio
async def test_mint_refuses_a_team_the_user_is_not_on(load_user, fetch_teams):
assert await mint_proxy_credential("u1", "team-c") == "not_a_member"

44
tests/test_litellm/proxy/auth/test_auth_checks.py
View file

@ -1,5 +1,6 @@
import asyncio
import json
import time
from collections.abc import Mapping
from types import SimpleNamespace
from typing import TYPE_CHECKING, Final, Literal, Optional
@ -916,6 +917,32 @@ async def test_get_user_object_wraps_db_outage_as_valueerror_preserving_context(
assert isinstance(exc_info.value.__context__, ConnectionError)
@pytest.mark.asyncio
async def test_get_user_object_check_db_only_ignores_recent_miss(monkeypatch):
"""A database-only read is never answered by the per-worker negative memo: a row created after a miss on
this worker is returned within db_cache_expiry seconds instead of raising UserNotFoundError, so the token
exchange mints for a user JWT auth just accepted."""
from litellm.proxy.auth import auth_checks
user_id = "memo-probe-user"
monkeypatch.setitem(auth_checks.last_db_access_time, f"user_id:{user_id}", (None, time.time()))
db_row = LiteLLM_UserTable(user_id=user_id, user_email=None, user_role="internal_user")
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=db_row)
result = await get_user_object(
user_id=user_id,
prisma_client=mock_prisma_client,
user_api_key_cache=UserApiKeyCache(),
user_id_upsert=False,
check_db_only=True,
)
assert result is not None
assert result.user_id == user_id
mock_prisma_client.db.litellm_usertable.find_unique.assert_awaited_once()
@pytest.mark.asyncio
async def test_get_user_object_upsert_includes_user_email():
"""Test that user_email is included when creating a new user via get_user_object upsert"""
@ -8740,6 +8767,23 @@ async def test_access_group_model_fallback_uses_the_injected_database(channel: s
reader.assert_awaited_once_with(where={"access_group_id": "group-a"})
def test_jwt_team_role_reaches_the_gateway_token_endpoint_by_default():
"""The RFC 8693 token exchange authorizes the IdP JWT against ``POST /token`` itself, and JWT
auth only binds a team from a multi-team claim when that team may call the route, so the
default team allowlist has to cover the gateway's token endpoint or the exchange would mint
teamless credentials for every ``team_ids_jwt_field`` deployment."""
from litellm.proxy._types import LiteLLM_JWTAuth
from litellm.proxy.auth.auth_checks import allowed_routes_check
assert allowed_routes_check(
user_role=LitellmUserRoles.TEAM, user_route="/token", litellm_proxy_roles=LiteLLM_JWTAuth()
)
assert not allowed_routes_check(
user_role=LitellmUserRoles.TEAM,
user_route="/token",
litellm_proxy_roles=LiteLLM_JWTAuth(team_allowed_routes=[]),
)
def test_route_skips_budget_checks_marks_only_spend_free_routes() -> None:
assert route_skips_budget_checks(route="/v1/models") is True
assert route_skips_budget_checks(route="/spend/logs") is True

1
tests/test_litellm/proxy/auth/test_route_checks.py
View file

@ -627,6 +627,7 @@ def test_virtual_key_llm_api_routes_denies_spend_logs_v2():
"/mcp/tools/call",
"/mcp-rest/tools/call",
"/mcp/tools/list",
"/token",
],
)
def test_mcp_inference_routes_classified_as_llm_api(route):

71
tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
View file

@ -13316,6 +13316,77 @@ async def test_team_member_add_audits_a_user_created_from_a_list_payload(monkeyp
assert created_user_id not in mock_audit.call_args.kwargs["existing_user_ids"]
@pytest.mark.asyncio
async def test_team_member_add_evicts_the_new_members_cached_user_row_on_every_worker(monkeypatch):
"""Auth admits a team-bound credential off the teams list of the cached user row. The add wrote the
new team to the database row only, so a worker still holding the old row refused the member's
credential with 403 until the management-object TTL expired. The add now evicts the row here and
broadcasts the eviction to the other workers, the way /team/member_delete already does."""
from litellm.proxy._types import TeamMemberAddRequest
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.management_endpoints.team_endpoints import team_member_add
team_id = "team-b"
user_id = "dev-1"
cache = UserApiKeyCache()
await cache.async_set_cache(
key=user_id, value=LiteLLM_UserTable(user_id=user_id, teams=["team-a"]), model_type=LiteLLM_UserTable
)
broadcast = AsyncMock()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", AsyncMock())
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id")
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache)
monkeypatch.setattr(
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", broadcast
)
updated_team = MagicMock()
updated_team.model_dump.return_value = {
"team_id": team_id,
"members_with_roles": [{"user_id": user_id, "role": "user"}],
}
async def fake_add_team_members_to_team(**kwargs):
return updated_team, [LiteLLM_UserTable(user_id=user_id, teams=["team-a", team_id])], []
with (
patch( # test-quality-ok: team_member_add has no injection seam for its prisma-backed helpers
"litellm.proxy.management_endpoints.team_endpoints.get_team_object",
new_callable=AsyncMock,
return_value=LiteLLM_TeamTable(team_id=team_id, members_with_roles=[]),
),
patch( # test-quality-ok: team_member_add has no injection seam for its prisma-backed helpers
"litellm.proxy.management_endpoints.team_endpoints._validate_team_member_add_permissions",
new_callable=AsyncMock,
),
patch( # test-quality-ok: team_member_add has no injection seam for its prisma-backed helpers
"litellm.proxy.management_endpoints.team_endpoints._validate_and_populate_member_user_info",
new_callable=AsyncMock,
),
patch( # test-quality-ok: team_member_add has no injection seam for its prisma-backed helpers
"litellm.proxy.management_endpoints.team_endpoints._resolve_existing_member_user_ids",
new_callable=AsyncMock,
return_value=frozenset({user_id}),
),
patch( # test-quality-ok: team_member_add has no injection seam for its prisma-backed helpers
"litellm.proxy.management_endpoints.team_endpoints._add_team_members_to_team",
side_effect=fake_add_team_members_to_team,
),
patch( # test-quality-ok: team_member_add has no injection seam for its prisma-backed helpers
"litellm.proxy.management_endpoints.team_endpoints._create_team_member_add_audit_logs",
new_callable=AsyncMock,
),
):
await team_member_add(
data=TeamMemberAddRequest(team_id=team_id, member=Member(user_id=user_id, role="user")),
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-1"),
)
assert await cache.async_get_cache(key=user_id, model_type=LiteLLM_UserTable) is None
broadcast.assert_awaited_once_with(cache_key=user_id)
def test_validate_member_user_id_provisioning_caps_the_ids_it_echoes_back():
"""A large member list must not echo every id back in the error body."""
from litellm.proxy.management_endpoints.team_endpoints import (

230
tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py
View file

@ -617,6 +617,201 @@ class TestResponsesWSFirstFrameModelAuth:
mock_model_auth.assert_awaited_once()
@pytest.mark.asyncio
@pytest.mark.parametrize("nested", [False, True])
@pytest.mark.parametrize("query_model", [None, "gpt-4o-mini"])
async def test_endpoint_routes_on_first_frame_input_and_previous_response_id(
self, nested: bool, query_model: str | None
):
from litellm.proxy.response_api_endpoints.endpoints import (
responses_websocket_endpoint,
)
replayed_input = [{"type": "reasoning", "id": "encitem_abc", "encrypted_content": "litellm_enc:abc;blob"}]
payload = {"model": "gpt-4o-mini", "input": replayed_input, "previous_response_id": "resp_prev"}
first_frame = {"type": "response.create", "response": payload} if nested else {"type": "response.create", **payload}
raw_first_frame = json.dumps(first_frame)
ws = MagicMock()
ws.headers = {}
ws.query_params = {}
ws.scope = {"headers": []}
ws.url = "ws://testserver/v1/responses"
ws.accept = AsyncMock()
ws.receive_text = AsyncMock(return_value=raw_first_frame)
ws.close = AsyncMock()
processor = MagicMock()
processor.common_processing_pre_call_logic = AsyncMock(
return_value=({"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock())
)
async def fake_llm_call():
return None
with (
patch( # test-quality-ok: first-frame model auth needs a live router and key table and has its own tests below
"litellm.proxy.response_api_endpoints.endpoints._enforce_responses_ws_first_frame_model_auth",
new_callable=AsyncMock,
),
patch( # test-quality-ok: the pre-call processor needs a live proxy; the payload it hands to routing is what is under test
"litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing",
return_value=processor,
),
patch( # test-quality-ok: routing is the seam where the first frame's input and previous_response_id become observable
"litellm.proxy.route_llm_request.route_request",
new_callable=AsyncMock,
return_value=fake_llm_call(),
) as mock_route_request,
):
await responses_websocket_endpoint(
websocket=ws,
model=query_model,
user_api_key_dict=MagicMock(),
)
ws.receive_text.assert_awaited_once()
routed = mock_route_request.await_args.kwargs["data"]
assert routed["model"] == "gpt-4o-mini"
assert routed["input"] == replayed_input
assert routed["previous_response_id"] == "resp_prev"
assert processor.common_processing_pre_call_logic.await_args.kwargs["model"] == "gpt-4o-mini"
assert mock_route_request.await_args.kwargs["route_type"] == "_aresponses_websocket"
ws.close.assert_not_awaited()
@pytest.mark.asyncio
@pytest.mark.parametrize("provider_rejected", [True, False])
async def test_endpoint_books_a_provider_rejected_connection_as_a_failed_request(self, provider_rejected: bool):
from litellm.proxy.response_api_endpoints.endpoints import (
responses_websocket_endpoint,
)
ws = MagicMock()
ws.headers = {}
ws.query_params = {}
ws.scope = {"headers": []}
ws.url = "ws://testserver/v1/responses"
ws.accept = AsyncMock()
ws.receive_text = AsyncMock(
return_value=json.dumps({"type": "response.create", "model": "gpt-4o-mini", "input": []})
)
ws.close = AsyncMock()
processor = MagicMock()
processor.common_processing_pre_call_logic = AsyncMock(
return_value=({"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock())
)
failure = litellm.BadRequestError(
message="invalid_encrypted_content", model="gpt-4o-mini", llm_provider="openai"
)
async def fake_llm_call():
return failure if provider_rejected else None
proxy_logging_obj = MagicMock()
proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None)
user_api_key_dict = MagicMock()
with (
patch( # test-quality-ok: first-frame model auth needs a live router and key table and has its own tests above
"litellm.proxy.response_api_endpoints.endpoints._enforce_responses_ws_first_frame_model_auth",
new_callable=AsyncMock,
),
patch( # test-quality-ok: the pre-call processor needs a live proxy; what the endpoint does with the relay's outcome is under test
"litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing",
return_value=processor,
),
patch( # test-quality-ok: routing is the seam that hands back the relay's outcome
"litellm.proxy.route_llm_request.route_request",
new_callable=AsyncMock,
return_value=fake_llm_call(),
),
patch( # test-quality-ok: the failure hook is the proxy's only path to a failed spend log row
"litellm.proxy.proxy_server.proxy_logging_obj",
proxy_logging_obj,
),
):
await responses_websocket_endpoint(
websocket=ws,
model=None,
user_api_key_dict=user_api_key_dict,
)
ws.close.assert_not_awaited()
if not provider_rejected:
proxy_logging_obj.post_call_failure_hook.assert_not_awaited()
return
proxy_logging_obj.post_call_failure_hook.assert_awaited_once()
booked = proxy_logging_obj.post_call_failure_hook.await_args.kwargs
assert booked["original_exception"] is failure
assert booked["user_api_key_dict"] is user_api_key_dict
assert booked["request_data"]["model"] == "gpt-4o-mini"
@pytest.mark.asyncio
async def test_endpoint_sends_an_error_frame_when_routing_rejects_the_connection(self):
from litellm.proxy.response_api_endpoints.endpoints import (
responses_websocket_endpoint,
)
ws = MagicMock()
ws.headers = {}
ws.query_params = {}
ws.scope = {"headers": []}
ws.url = "ws://testserver/v1/responses"
ws.accept = AsyncMock()
ws.receive_text = AsyncMock(
return_value=json.dumps({"type": "response.create", "model": "gpt-4o-mini", "input": []})
)
ws.send_text = AsyncMock()
ws.close = AsyncMock()
processor = MagicMock()
processor.common_processing_pre_call_logic = AsyncMock(
return_value=({"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock())
)
rejection = litellm.RateLimitError(
message="origin deployment is cooling down", model="gpt-4o-mini", llm_provider="openai"
)
proxy_logging_obj = MagicMock()
proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None)
user_api_key_dict = MagicMock()
with (
patch( # test-quality-ok: first-frame model auth needs a live router and key table and has its own tests above
"litellm.proxy.response_api_endpoints.endpoints._enforce_responses_ws_first_frame_model_auth",
new_callable=AsyncMock,
),
patch( # test-quality-ok: the pre-call processor needs a live proxy; what the endpoint tells the client is under test
"litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing",
return_value=processor,
),
patch( # test-quality-ok: routing is the seam that raises the affinity rejection
"litellm.proxy.route_llm_request.route_request",
new_callable=AsyncMock,
side_effect=rejection,
),
patch( # test-quality-ok: the failure hook is the proxy's only path to a failed spend log row
"litellm.proxy.proxy_server.proxy_logging_obj",
proxy_logging_obj,
),
):
await responses_websocket_endpoint(
websocket=ws,
model=None,
user_api_key_dict=user_api_key_dict,
)
frame = json.loads(ws.send_text.await_args.args[0])
assert frame["type"] == "error"
assert frame["status"] == 429
assert frame["error"]["type"] == "rate_limit_exceeded"
assert "cooling down" in frame["error"]["message"]
ws.close.assert_awaited_once_with(code=1011, reason="Internal server error")
booked = proxy_logging_obj.post_call_failure_hook.await_args.kwargs
assert booked["original_exception"] is rejection
assert booked["user_api_key_dict"] is user_api_key_dict
assert booked["request_data"]["model"] == "gpt-4o-mini"
@pytest.mark.asyncio
async def test_reruns_model_auth_for_first_frame_model(self):
from starlette.requests import Request
@ -743,6 +938,41 @@ class TestReadWSModelFromFirstFrameErrors:
ws.send_text.assert_not_awaited()
ws.close.assert_not_awaited()
@pytest.mark.asyncio
async def test_query_model_wins_over_first_frame_model(self):
from litellm.proxy.response_api_endpoints.endpoints import (
_read_ws_model_from_first_frame,
)
raw = json.dumps({"type": "response.create", "model": "gpt-4o", "input": []})
ws = MagicMock()
ws.receive_text = AsyncMock(return_value=raw)
ws.send_text = AsyncMock()
ws.close = AsyncMock()
result = await _read_ws_model_from_first_frame(ws, query_model="reasoning-group")
assert result == ("reasoning-group", raw)
ws.close.assert_not_awaited()
@pytest.mark.asyncio
async def test_query_model_satisfies_a_first_frame_without_model(self):
from litellm.proxy.response_api_endpoints.endpoints import (
_read_ws_model_from_first_frame,
)
raw = json.dumps({"type": "response.create", "input": []})
ws = MagicMock()
ws.receive_text = AsyncMock(return_value=raw)
ws.send_text = AsyncMock()
ws.close = AsyncMock()
result = await _read_ws_model_from_first_frame(ws, query_model="reasoning-group")
assert result == ("reasoning-group", raw)
ws.send_text.assert_not_awaited()
ws.close.assert_not_awaited()
class TestManagedResponsesSameProvider:
def _handler(self, model, custom_llm_provider=None):

176
tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py
View file

@ -20,7 +20,10 @@ from litellm.responses.litellm_completion_transformation.streaming_iterator impo
LiteLLMCompletionStreamingIterator,
)
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.llms.openai import ResponsesAPIStreamEvents
from litellm.types.llms.openai import (
BaseLiteLLMOpenAIResponseObject,
ResponsesAPIStreamEvents,
)
from litellm.types.responses.main import build_web_search_call
from litellm.types.utils import (
Delta,
@ -957,3 +960,174 @@ def test_streamed_unrecognized_tool_choice_is_echoed_as_auto() -> None:
]
assert [event.response.tool_choice for event in response_events] == ["auto", "auto", "auto"]
def _reasoning_chunk(reasoning: str, finish_reason: str | None = None) -> ModelResponseStream:
return ModelResponseStream(
id=CHAT_COMPLETION_ID,
created=1748575031,
model="claude-haiku-4-5",
object="chat.completion.chunk",
choices=[
StreamingChoices(
index=0,
delta=Delta(role="assistant", reasoning_content=reasoning),
finish_reason=finish_reason,
)
],
)
async def _collect_events(
iterator: LiteLLMCompletionStreamingIterator, sync_mode: bool
) -> list[BaseLiteLLMOpenAIResponseObject]:
if sync_mode:
return list(iterator)
return [event async for event in iterator]
def _is_message_item(event: BaseLiteLLMOpenAIResponseObject) -> bool:
return getattr(getattr(event, "item", None), "type", None) == "message"
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_tool_only_stream_emits_no_message_item_events(sync_mode: bool):
iterator: Final = _build_iterator([_tool_call_chunk(), _chunk("", finish_reason="tool_calls")])
events: Final = await _collect_events(iterator, sync_mode)
message_item_events = [
event
for event in events
if getattr(event, "type", None)
in (ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE)
and _is_message_item(event)
]
assert message_item_events == []
assert [
event
for event in events
if str(getattr(event, "type", "")).startswith("response.output_text")
or getattr(event, "type", None)
in (ResponsesAPIStreamEvents.CONTENT_PART_ADDED, ResponsesAPIStreamEvents.CONTENT_PART_DONE)
] == []
assert any(getattr(event, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED for event in events)
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_reasoning_then_text_announces_message_item_before_text_events(sync_mode: bool):
iterator: Final = _build_iterator(
[
_reasoning_chunk("let me think"),
_chunk("Hello"),
_chunk("!", finish_reason="stop"),
]
)
events: Final = await _collect_events(iterator, sync_mode)
announced_message_ids: set[str] = set()
announced_indexes_by_item_type: dict[str, int] = {}
content_part_added_seen = False
saw_text_delta = False
for event in events:
event_type = getattr(event, "type", None)
if event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED:
announced_indexes_by_item_type[event.item.type] = event.output_index
if _is_message_item(event):
announced_message_ids.add(event.item.id)
elif event_type == ResponsesAPIStreamEvents.CONTENT_PART_ADDED:
content_part_added_seen = True
elif event_type in (
ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA,
ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE,
ResponsesAPIStreamEvents.CONTENT_PART_DONE,
):
assert event.item_id in announced_message_ids
if event_type == ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA:
assert content_part_added_seen
saw_text_delta = True
elif event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE and _is_message_item(event):
assert event.item.id in announced_message_ids
assert saw_text_delta
assert "".join(
event.delta for event in events if getattr(event, "type", None) == ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA
) == "Hello!"
assert announced_indexes_by_item_type["message"] != announced_indexes_by_item_type["reasoning"]
@pytest.mark.asyncio
async def test_reasoning_item_closes_before_message_item_opens():
iterator: Final = _build_iterator(
[
_reasoning_chunk("let me think"),
_chunk("Hello"),
_chunk("!", finish_reason="stop"),
]
)
events: Final = await _collect_events(iterator, sync_mode=False)
item_lifecycle: Final = [
(event.type, event.item.type)
for event in events
if getattr(event, "type", None)
in (ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE)
]
assert item_lifecycle == [
(ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, "reasoning"),
(ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, "reasoning"),
(ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, "message"),
(ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, "message"),
]
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_tool_then_reasoning_then_text_gives_message_its_own_output_index(sync_mode: bool):
iterator: Final = _build_iterator(
[
_tool_call_chunk(),
_reasoning_chunk("thinking"),
_chunk("Hello"),
_chunk("!", finish_reason="stop"),
]
)
events: Final = await _collect_events(iterator, sync_mode)
output_item_added_events: Final = [
event for event in events if getattr(event, "type", None) == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED
]
message_item_adds: Final = [event for event in output_item_added_events if _is_message_item(event)]
function_call_adds: Final = [
event for event in output_item_added_events if getattr(event.item, "type", None) == "function_call"
]
assert len(message_item_adds) == 1
assert all(message_item_adds[0].output_index != event.output_index for event in function_call_adds)
output_indexes_by_item_id: Final = {event.item.id: event.output_index for event in output_item_added_events}
assert len(output_indexes_by_item_id) == len(set(output_indexes_by_item_id.values()))
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_plain_text_stream_announces_exactly_one_message_item(sync_mode: bool):
iterator: Final = _build_iterator([_chunk("Hel"), _chunk("lo", finish_reason="stop")])
events: Final = await _collect_events(iterator, sync_mode)
message_item_adds = [
event
for event in events
if getattr(event, "type", None) == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED and _is_message_item(event)
]
assert len(message_item_adds) == 1
for event in events:
if getattr(event, "type", None) in (
ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA,
ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE,
):
assert event.item_id == message_item_adds[0].item.id

88
tests/test_litellm/responses/test_responses_api_request_body.py
View file

@ -424,6 +424,94 @@ async def test_aresponses_websocket_strips_responses_routing_prefix_from_openai_
assert mock_ws.call_args.kwargs["custom_llm_provider"] == "openai"
@pytest.mark.asyncio
async def test_aresponses_websocket_keeps_routing_hints_out_of_the_relay_kwargs(): # test-quality-ok: the relay kwargs are the only place a dropped key is observable; the provider socket behind them is the boundary
from unittest.mock import MagicMock
from litellm.responses.main import _aresponses_websocket
with patch.object(
import_module("litellm.responses.main").base_llm_http_handler, "async_responses_websocket",
new_callable=AsyncMock,
) as mock_ws:
await _aresponses_websocket(
model="openai/gpt-5.6",
websocket=MagicMock(),
api_key="sk-test",
litellm_logging_obj=MagicMock(),
input=[{"type": "message", "role": "user", "content": "hi"}],
previous_response_id="resp_prev",
)
mock_ws.assert_awaited_once()
assert "input" not in mock_ws.call_args.kwargs
assert "previous_response_id" not in mock_ws.call_args.kwargs
_STRIPPED_WS_INPUT = [{"role": "user", "content": "hi"}]
_ORIGINAL_WS_INPUT = [
{"type": "reasoning", "id": "rs_1", "encrypted_content": "blob-from-a-removed-deployment", "summary": []},
*_STRIPPED_WS_INPUT,
]
@pytest.mark.asyncio
@pytest.mark.parametrize("nested", [False, True])
async def test_aresponses_websocket_forwards_the_routed_input_in_the_first_frame(nested: bool): # test-quality-ok: the first frame handed to the relay is the only place the routed input is observable before the provider socket
from unittest.mock import MagicMock
from litellm.responses.main import _aresponses_websocket
body = {"model": "gpt-5.6", "input": _ORIGINAL_WS_INPUT, "store": False}
first_message = json.dumps(
{"type": "response.create", "response": body} if nested else {"type": "response.create", **body}
)
with patch.object(
import_module("litellm.responses.main").base_llm_http_handler, "async_responses_websocket",
new_callable=AsyncMock,
) as mock_ws:
await _aresponses_websocket(
model="openai/gpt-5.6",
websocket=MagicMock(),
api_key="sk-test",
litellm_logging_obj=MagicMock(),
input=list(_STRIPPED_WS_INPUT),
first_message=first_message,
)
forwarded = json.loads(mock_ws.call_args.kwargs["first_message"])
container = forwarded["response"] if nested else forwarded
assert container["input"] == _STRIPPED_WS_INPUT
assert container["store"] is False
assert container["model"] == "gpt-5.6"
assert forwarded["type"] == "response.create"
@pytest.mark.asyncio
async def test_aresponses_websocket_forwards_the_first_frame_verbatim_when_routing_left_the_input_alone(): # test-quality-ok: the relay kwargs are the boundary; byte-identical passthrough is only observable there
from unittest.mock import MagicMock
from litellm.responses.main import _aresponses_websocket
first_message = '{"type": "response.create", "model": "gpt-5.6", "input": [{"role": "user", "content": "hi"}]}'
with patch.object(
import_module("litellm.responses.main").base_llm_http_handler, "async_responses_websocket",
new_callable=AsyncMock,
) as mock_ws:
await _aresponses_websocket(
model="openai/gpt-5.6",
websocket=MagicMock(),
api_key="sk-test",
litellm_logging_obj=MagicMock(),
input=list(_STRIPPED_WS_INPUT),
first_message=first_message,
)
assert mock_ws.call_args.kwargs["first_message"] == first_message
_INJECTION_POINT_INPUT = [{"role": "system", "content": "You are terse."}, {"role": "user", "content": "hi"}]
_SYSTEM_POINT = {"location": "message", "role": "system"}
_USER_POINT = {"location": "message", "role": "user"}

407
tests/test_litellm/responses/test_responses_websocket_all_providers.py
View file

@ -1502,6 +1502,34 @@ class TestNativeWebSocketDeploymentDefaults:
assert dict(request_defaults.fill_missing) == {"reasoning": {"effort": "high"}, "service_tier": "priority"}
assert dict(request_defaults.overrides) == {"provider_default": "configured"}
@pytest.mark.asyncio
async def test_aresponses_websocket_keeps_first_frame_routing_hints_out_of_the_defaults(
self, monkeypatch: pytest.MonkeyPatch
):
import importlib
from unittest.mock import AsyncMock
responses_main = importlib.import_module("litellm.responses.main")
stub = MagicMock()
stub.async_responses_websocket = AsyncMock()
monkeypatch.setattr(responses_main, "base_llm_http_handler", stub)
await responses_main._aresponses_websocket.__wrapped__(
model="openai/gpt-5-pro",
websocket=MagicMock(),
api_key="sk-test",
litellm_logging_obj=MagicMock(),
reasoning_effort="high",
input=[{"id": "encitem_abc", "type": "reasoning", "encrypted_content": "litellm_enc:abc"}],
previous_response_id="resp_first_turn",
)
call_kwargs = stub.async_responses_websocket.call_args.kwargs
assert dict(call_kwargs["request_defaults"].fill_missing) == {"reasoning": {"effort": "high"}}
assert "input" not in call_kwargs
assert "previous_response_id" not in call_kwargs
class TestNativeWebSocketGuardrails:
@pytest.mark.asyncio
@ -2927,3 +2955,382 @@ class TestNativeWebSocketUrlConstruction:
mock_config.get_websocket_url.assert_called_once()
_, call_kwargs = mock_config.get_websocket_url.call_args
assert call_kwargs["litellm_params"]["api_version"] == "2025-04-01-preview"
_AFFINITY_METADATA = {
"model_info": {"id": "dep-1"},
"encrypted_content_affinity_enabled": True,
}
def _wrapped_reasoning_item():
from litellm.responses.utils import ResponsesAPIRequestUtils
return {
"type": "reasoning",
"id": ResponsesAPIRequestUtils._build_encrypted_item_id("dep-1", "rs_orig"),
"encrypted_content": ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAA-blob", "dep-1"),
"summary": [],
}
class TestNativeWebSocketEncryptedContentAffinity:
@pytest.mark.asyncio
@pytest.mark.parametrize("nested", [False, True])
async def test_client_to_backend_restores_wrapped_ids(self, nested: bool):
from unittest.mock import AsyncMock
from litellm.responses.utils import ResponsesAPIRequestUtils
wrapped_previous = ResponsesAPIRequestUtils._build_responses_api_response_id(
custom_llm_provider="openai", model_id="dep-1", response_id="resp_orig"
)
payload = {
"input": [_wrapped_reasoning_item(), {"type": "message", "role": "user", "content": "hi"}],
"previous_response_id": wrapped_previous,
}
frame = {"type": "response.create", "response": payload} if nested else {"type": "response.create", **payload}
backend_ws = MagicMock()
backend_ws.send = AsyncMock()
websocket = MagicMock()
websocket.receive_text = AsyncMock(side_effect=[json.dumps(frame), Exception("stop")])
handler = _make_streaming(websocket=websocket, backend_ws=backend_ws, request_data={})
await handler.client_to_backend()
sent = json.loads(backend_ws.send.await_args_list[0][0][0])
body = sent["response"] if nested else sent
assert body["input"][0]["id"] == "rs_orig"
assert body["input"][0]["encrypted_content"] == "gAAAA-blob"
assert body["input"][1] == {"type": "message", "role": "user", "content": "hi"}
assert body["previous_response_id"] == "resp_orig"
@pytest.mark.asyncio
async def test_client_to_backend_leaves_unwrapped_frames_untouched(self):
from unittest.mock import AsyncMock
frame = json.dumps({"type": "response.create", "input": "hello", "previous_response_id": "resp_raw"})
backend_ws = MagicMock()
backend_ws.send = AsyncMock()
websocket = MagicMock()
websocket.receive_text = AsyncMock(side_effect=[frame, Exception("stop")])
handler = _make_streaming(websocket=websocket, backend_ws=backend_ws, request_data={})
await handler.client_to_backend()
assert backend_ws.send.await_args_list[0][0][0] == frame
@pytest.mark.asyncio
async def test_backend_to_client_wraps_ids_when_affinity_is_enabled(self):
import asyncio
from unittest.mock import AsyncMock
import websockets.exceptions # noqa: F401 (lazy submodule must be importable)
from litellm.responses.utils import ResponsesAPIRequestUtils
reasoning_item = {"type": "reasoning", "id": "rs_1", "encrypted_content": "gAAAA-blob", "summary": []}
websocket = MagicMock()
websocket.send_text = AsyncMock()
backend_ws = MagicMock()
backend_ws.recv = AsyncMock(
side_effect=[
json.dumps({"type": "response.output_item.done", "output_index": 0, "item": dict(reasoning_item)}),
json.dumps(
{
"type": "response.completed",
"response": {"id": "resp_1", "output": [dict(reasoning_item)], "usage": {"total_tokens": 3}},
}
),
Exception("stop"),
]
)
logging_obj = MagicMock()
logging_obj.dispatch_success_handlers = AsyncMock()
handler = _make_streaming(
websocket=websocket,
backend_ws=backend_ws,
logging_obj=logging_obj,
request_data={"litellm_metadata": dict(_AFFINITY_METADATA)},
custom_llm_provider="openai",
)
await handler.backend_to_client()
wrapped_content = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAA-blob", "dep-1")
item_done = json.loads(websocket.send_text.await_args_list[0][0][0])
assert item_done["item"]["encrypted_content"] == wrapped_content
completed = json.loads(websocket.send_text.await_args_list[1][0][0])
assert completed["response"]["id"] == ResponsesAPIRequestUtils._build_responses_api_response_id(
custom_llm_provider="openai", model_id="dep-1", response_id="resp_1"
)
assert completed["response"]["output"][0]["id"] == ResponsesAPIRequestUtils._build_encrypted_item_id(
"dep-1", "rs_1"
)
assert completed["response"]["output"][0]["encrypted_content"] == wrapped_content
await asyncio.sleep(0)
logged = logging_obj.dispatch_success_handlers.await_args[0][0]
assert logged[0]["response"]["id"] == completed["response"]["id"]
@pytest.mark.asyncio
async def test_backend_to_client_wraps_only_response_id_without_affinity(self):
from unittest.mock import AsyncMock
import websockets.exceptions # noqa: F401 (lazy submodule must be importable)
from litellm.responses.utils import ResponsesAPIRequestUtils
reasoning_item = {"type": "reasoning", "id": "rs_1", "encrypted_content": "gAAAA-blob", "summary": []}
websocket = MagicMock()
websocket.send_text = AsyncMock()
backend_ws = MagicMock()
backend_ws.recv = AsyncMock(
side_effect=[
json.dumps({"type": "response.output_item.done", "output_index": 0, "item": dict(reasoning_item)}),
json.dumps({"type": "response.completed", "response": {"id": "resp_1", "output": [dict(reasoning_item)]}}),
Exception("stop"),
]
)
logging_obj = MagicMock()
logging_obj.dispatch_success_handlers = AsyncMock()
handler = _make_streaming(
websocket=websocket,
backend_ws=backend_ws,
logging_obj=logging_obj,
request_data={"litellm_metadata": {"model_info": {"id": "dep-1"}}},
custom_llm_provider="openai",
)
await handler.backend_to_client()
item_done = json.loads(websocket.send_text.await_args_list[0][0][0])
assert item_done["item"] == reasoning_item
completed = json.loads(websocket.send_text.await_args_list[1][0][0])
assert completed["response"]["id"] == ResponsesAPIRequestUtils._build_responses_api_response_id(
custom_llm_provider="openai", model_id="dep-1", response_id="resp_1"
)
assert completed["response"]["output"][0] == reasoning_item
@pytest.mark.asyncio
@pytest.mark.parametrize(
"failure_frame, expected_status",
[
(
{
"type": "error",
"error": {
"type": "invalid_request_error",
"code": "invalid_encrypted_content",
"message": "The encrypted content for item rs_1 could not be verified.",
},
},
400,
),
(
{
"type": "response.failed",
"response": {
"id": "resp_1",
"status": "failed",
"error": {"code": "server_error", "message": "upstream blew up"},
},
},
500,
),
],
)
async def test_backend_to_client_books_failure_frames_as_failures(
self, failure_frame: dict[str, object], expected_status: int
):
import asyncio
from unittest.mock import AsyncMock
import websockets.exceptions # noqa: F401 (lazy submodule must be importable)
websocket = MagicMock()
websocket.send_text = AsyncMock()
backend_ws = MagicMock()
backend_ws.recv = AsyncMock(
side_effect=[
json.dumps({"type": "response.created", "response": {"id": "resp_1", "status": "in_progress"}}),
json.dumps(failure_frame),
Exception("stop"),
]
)
logging_obj = MagicMock()
logging_obj.dispatch_success_handlers = AsyncMock()
logging_obj.dispatch_failure_handlers = AsyncMock()
logging_obj._response_cost_calculator = MagicMock(return_value=0.0)
handler = _make_streaming(
websocket=websocket,
backend_ws=backend_ws,
logging_obj=logging_obj,
request_data={},
authorized_model="gpt-5.6",
custom_llm_provider="openai",
)
await handler.backend_to_client()
await asyncio.sleep(0)
logging_obj.dispatch_success_handlers.assert_not_awaited()
logging_obj.dispatch_failure_handlers.assert_awaited_once()
exception = logging_obj.dispatch_failure_handlers.await_args[0][0]
assert exception.status_code == expected_status
assert failure_frame.get("error", failure_frame.get("response", {}).get("error"))["message"] in str(exception)
@pytest.mark.asyncio
async def test_backend_to_client_bills_completed_turns_before_a_failure(self):
import asyncio
from unittest.mock import AsyncMock
import websockets.exceptions # noqa: F401 (lazy submodule must be importable)
websocket = MagicMock()
websocket.send_text = AsyncMock()
backend_ws = MagicMock()
backend_ws.recv = AsyncMock(
side_effect=[
json.dumps(
{
"type": "response.completed",
"response": {
"id": "resp_1",
"status": "completed",
"output": [],
"usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15},
},
}
),
json.dumps({"type": "error", "error": {"type": "invalid_request_error", "message": "bad turn"}}),
Exception("stop"),
]
)
logging_obj = MagicMock()
logging_obj.dispatch_success_handlers = AsyncMock()
logging_obj.dispatch_failure_handlers = AsyncMock()
logging_obj._response_cost_calculator = MagicMock(return_value=0.01)
handler = _make_streaming(websocket=websocket, backend_ws=backend_ws, logging_obj=logging_obj, request_data={})
await handler.backend_to_client()
await asyncio.sleep(0)
logging_obj.record_partial_usage_for_failure.assert_called_once()
usage, response_cost = logging_obj.record_partial_usage_for_failure.call_args[0]
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (10, 5, 15)
assert response_cost == 0.01
logging_obj.dispatch_success_handlers.assert_not_awaited()
logging_obj.dispatch_failure_handlers.assert_awaited_once()
@pytest.mark.asyncio
async def test_bidirectional_forward_returns_the_provider_failure(self):
import asyncio
from unittest.mock import AsyncMock
import websockets.exceptions # noqa: F401 (lazy submodule must be importable)
backend_drained = asyncio.Event()
backend_events = [
json.dumps({"type": "response.created", "response": {"id": "resp_1", "status": "in_progress"}}),
json.dumps(
{
"type": "error",
"status": 400,
"error": {
"type": "invalid_request_error",
"code": "invalid_encrypted_content",
"message": "could not be verified",
},
}
),
]
async def recv(decode=False):
if backend_events:
return backend_events.pop(0)
backend_drained.set()
raise Exception("stop")
async def receive_text():
await backend_drained.wait()
raise Exception("client gone")
websocket = MagicMock()
websocket.send_text = AsyncMock()
websocket.receive_text = receive_text
backend_ws = MagicMock()
backend_ws.recv = recv
backend_ws.send = AsyncMock()
backend_ws.close = AsyncMock()
logging_obj = MagicMock()
logging_obj.dispatch_success_handlers = AsyncMock()
logging_obj.dispatch_failure_handlers = AsyncMock()
logging_obj._response_cost_calculator = MagicMock(return_value=0.0)
handler = _make_streaming(
websocket=websocket,
backend_ws=backend_ws,
logging_obj=logging_obj,
request_data={},
authorized_model="gpt-5.6",
custom_llm_provider="openai",
)
failure = await handler.bidirectional_forward()
assert isinstance(failure, Exception)
assert failure.status_code == 400
assert "could not be verified" in str(failure)
@pytest.mark.asyncio
async def test_bidirectional_forward_returns_none_after_a_completed_turn(self):
import asyncio
from unittest.mock import AsyncMock
import websockets.exceptions # noqa: F401 (lazy submodule must be importable)
backend_drained = asyncio.Event()
backend_events = [
json.dumps(
{
"type": "response.completed",
"response": {
"id": "resp_1",
"status": "completed",
"output": [],
"usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15},
},
}
),
]
async def recv(decode=False):
if backend_events:
return backend_events.pop(0)
backend_drained.set()
raise Exception("stop")
async def receive_text():
await backend_drained.wait()
raise Exception("client gone")
websocket = MagicMock()
websocket.send_text = AsyncMock()
websocket.receive_text = receive_text
backend_ws = MagicMock()
backend_ws.recv = recv
backend_ws.send = AsyncMock()
backend_ws.close = AsyncMock()
logging_obj = MagicMock()
logging_obj.dispatch_success_handlers = AsyncMock()
logging_obj.dispatch_failure_handlers = AsyncMock()
handler = _make_streaming(
websocket=websocket,
backend_ws=backend_ws,
logging_obj=logging_obj,
request_data={},
authorized_model="gpt-5.6",
custom_llm_provider="openai",
)
assert await handler.bidirectional_forward() is None

12
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -24767,10 +24767,16 @@ export interface components {
redirect_uri?: string;
/** Refresh Token */
refresh_token?: string | null;
/** Requested Token Type */
requested_token_type?: string | null;
/** Resource */
resource?: string | null;
/** Scope */
scope?: string | null;
/** Subject Token */
subject_token?: string | null;
/** Subject Token Type */
subject_token_type?: string | null;
};
/** Body_token_endpoint_token_post */
Body_token_endpoint_token_post: {
@ -24788,10 +24794,16 @@ export interface components {
redirect_uri?: string;
/** Refresh Token */
refresh_token?: string | null;
/** Requested Token Type */
requested_token_type?: string | null;
/** Resource */
resource?: string | null;
/** Scope */
scope?: string | null;
/** Subject Token */
subject_token?: string | null;
/** Subject Token Type */
subject_token_type?: string | null;
};
/** Body_upload_logo_upload_logo_post */
Body_upload_logo_upload_logo_post: {