scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_/api-reference-dark-mode-4ec67f

Browse source
This commit is contained in:
Yuneng Jiang 2026-08-29 10:25:55 -07:00
commit 8f07a12726
No known key found for this signature in database
16 changed files with 1110 additions and 130 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

71
enterprise/litellm_enterprise/proxy/hooks/managed_files.py
View file

@ -473,19 +473,56 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
)
page_size: Final = min(limit or 20, 100)
cursor_args: _CursorPageArgs = {"cursor": {"unified_object_id": after}, "skip": 1} if after else {}
batches = await _managed_object_table(self.prisma_client).find_many(
where=where_clause,
take=page_size + 1,
order=[{"created_at": "desc"}, {"unified_object_id": "desc"}],
**cursor_args,
matches: Final = await self._collect_listed_batches(
where_clause=where_clause,
after=after,
wanted=page_size + 1,
user_api_key_dict=user_api_key_dict,
)
return build_list_page(list(matches[:page_size]), has_more=len(matches) > page_size)
has_more = len(batches) > page_size
async def _collect_listed_batches(
self,
where_clause: Mapping[str, object],
after: Optional[str],
wanted: int,
user_api_key_dict: UserAPIKeyAuth,
) -> tuple[LiteLLMBatch, ...]:
"""Read chunks newest-first until ``wanted`` batches survive parsing and
file-id resolution or the caller's rows run out, so a run of rows that will
not parse refills the page instead of emptying it. The first chunk is
``wanted`` rows, so a healthy page still costs one query; a scan that has to
continue widens to ``FILE_LIST_CONTINUATION_CHUNK_SIZE`` like ``afile_list``,
and every chunk advances the keyset cursor, so the walk ends once the
caller's rows are exhausted."""
matches: tuple[LiteLLMBatch, ...] = () # rebind-ok: accumulates survivors across chunks
cursor_id: Optional[str] = after # rebind-ok: keyset cursor advances to each chunk's last row
chunk_size: int = wanted # rebind-ok: widens once a scan has to continue past the first chunk
while len(matches) < wanted:
cursor_args: _CursorPageArgs = {"cursor": {"unified_object_id": cursor_id}, "skip": 1} if cursor_id else {}
chunk = await _managed_object_table(self.prisma_client).find_many(
where=where_clause,
take=chunk_size,
order=[{"created_at": "desc"}, {"unified_object_id": "desc"}],
**cursor_args,
)
matches = matches + await self._resolve_listed_rows(
rows=chunk, wanted=wanted - len(matches), user_api_key_dict=user_api_key_dict
)
if len(chunk) < chunk_size:
break
cursor_id = chunk[-1].unified_object_id
chunk_size = max(chunk_size, FILE_LIST_CONTINUATION_CHUNK_SIZE)
return matches
async def _resolve_listed_rows(
self,
rows: "Sequence[PrismaManagedObjectRow]",
wanted: int,
user_api_key_dict: UserAPIKeyAuth,
) -> tuple[LiteLLMBatch, ...]:
parsed_rows: Final = tuple(
(row, batch_obj) for row in batches[:page_size] if (batch_obj := _parse_managed_batch_row(row)) is not None
(row, batch_obj) for row in rows if (batch_obj := _parse_managed_batch_row(row)) is not None
)
unified_id_by_raw_id: Final = await map_raw_file_ids_to_unified(
raw_file_ids=frozenset(
@ -496,19 +533,19 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
),
prisma_client=self.prisma_client,
)
resolved_batches: Final = [
await self._resolve_listed_batch(
resolved: Final[list[LiteLLMBatch]] = [] # mutable-ok: resolution stops as soon as the page is full
for row, batch_obj in parsed_rows:
if len(resolved) == wanted:
break
resolved_batch = await self._resolve_listed_batch(
row=row,
batch_obj=batch_obj,
unified_id_by_raw_id=unified_id_by_raw_id,
user_api_key_dict=user_api_key_dict,
)
for row, batch_obj in parsed_rows
]
return build_list_page(
[batch_obj for batch_obj in resolved_batches if batch_obj is not None],
has_more=has_more,
)
if resolved_batch is not None:
resolved.append(resolved_batch)
return tuple(resolved)
async def _resolve_listed_batch(
self,

8
litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
View file

@ -903,8 +903,9 @@ class MCPRequestHandler:
NotSessionBearer,
SessionBearerAdmitted,
SessionBearerInvalid,
SessionSigningConfigError,
active_session_signing_keys,
resolve_session_bearer,
session_keys_from_master_key,
)
from litellm.proxy.proxy_server import master_key
@ -913,7 +914,10 @@ class MCPRequestHandler:
await MCPRequestHandler._run_pre_db_read_auth_checks(request=request, route=route)
keys: Final = session_keys_from_master_key(master_key)
keys: Final = active_session_signing_keys(master_key)
if isinstance(keys, SessionSigningConfigError):
verbose_logger.error("mcp gateway session admission rejected: %s", keys.detail)
raise HTTPException(status_code=500, detail="Server misconfigured: mcp_session_token_signing is invalid")
result: Final = resolve_session_bearer(authorization_value, keys, datetime.now(timezone.utc))
match result:
case SessionBearerAdmitted():

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

@ -65,15 +65,16 @@ from litellm.proxy._experimental.mcp_server.oauth_utils import (
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import (
SessionRefreshOpened,
SessionSigningConfigError,
active_session_signing_keys,
open_session_refresh_bearer,
session_keys_from_master_key,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import (
SESSION_REFRESH_TTL_SECONDS,
MintedSessionToken,
SessionAudience,
SessionKeys,
SessionPrincipal,
SessionSigningKeys,
mint_session_refresh_token,
mint_session_token,
)
@ -885,7 +886,7 @@ class _SingleUseGuard:
return "first" if count == 1 else "replayed"
def _session_token_pair(principal: SessionPrincipal, keys: SessionKeys, now: datetime) -> Response:
def _session_token_pair(principal: SessionPrincipal, keys: SessionSigningKeys, now: datetime) -> Response:
access: Final = mint_session_token(principal, keys, now)
refresh: Final = mint_session_refresh_token(principal, keys, now)
if not isinstance(access, MintedSessionToken) or not isinstance(refresh, MintedSessionToken):
@ -912,7 +913,7 @@ class _ProxyCredentialTokenResponse(TypedDict):
def _proxy_credential_response(
minted: MintedProxyCredential, principal: SessionPrincipal, keys: SessionKeys, now: datetime
minted: MintedProxyCredential, principal: SessionPrincipal, keys: SessionSigningKeys, now: datetime
) -> 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
@ -998,7 +999,10 @@ async def aggregate_token(
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")
keys: Final = session_keys_from_master_key(master_key)
keys: Final = active_session_signing_keys(master_key)
if isinstance(keys, SessionSigningConfigError):
verbose_logger.error("mcp_gateway_dcr token grant rejected: %s", keys.detail)
return _oauth_error(500, "server_error", "the gateway session signing configuration is invalid")
now: Final = datetime.now(timezone.utc)
issue: Final = _GrantIssuer(
request=request,
@ -1043,7 +1047,7 @@ class _GrantIssuer:
self,
request: Request,
resource: str | None,
keys: SessionKeys,
keys: SessionSigningKeys,
now: datetime,
reload_user: ReloadUser,
mint_proxy_credential: MintProxyCredential,
@ -1146,7 +1150,7 @@ async def _refresh_token_grant(
refresh_token: str | None,
client_id: str,
resource: str | None,
keys: SessionKeys,
keys: SessionSigningKeys,
now: datetime,
issue: _GrantIssuer,
) -> Response:
@ -1182,7 +1186,10 @@ async def revoke_refresh_token(token: str, client_id: str, master_key: str | Non
if master_key is None:
verbose_logger.error("mcp_gateway_dcr revoke rejected: no master_key configured")
return _oauth_error(500, "server_error", "the gateway has no master key configured")
keys: Final = session_keys_from_master_key(master_key)
keys: Final = active_session_signing_keys(master_key)
if isinstance(keys, SessionSigningConfigError):
verbose_logger.error("mcp_gateway_dcr revoke rejected: %s", keys.detail)
return _oauth_error(500, "server_error", "the gateway session signing configuration is invalid")
now: Final = datetime.now(timezone.utc)
opened: Final = open_session_refresh_bearer(token, keys, now, expected_client_id=client_id)
if isinstance(opened, SessionRefreshOpened):

102
litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py
View file

@ -20,13 +20,16 @@ from datetime import datetime
from functools import lru_cache
from typing import Final, Literal, TypeAlias
from pydantic import BaseModel, ConfigDict, SecretStr
from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError
from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import (
AsymmetricSessionKeys,
OpenedSessionToken,
SessionExpired,
SessionKeys,
SessionPrincipal,
SessionRotatedPublicKey,
SessionSigningKeys,
is_session_refresh_token,
is_session_token,
open_session_refresh_token,
@ -68,6 +71,99 @@ def session_keys_from_master_key(master_key: str) -> SessionKeys:
return SessionKeys(signing_key=SecretStr(signing))
class SessionSigningPreviousKey(BaseModel):
"""One retired key in ``mcp_session_token_signing.previous_public_keys``: its ``kid``
and the PEM public half (inline or an ``os.environ/`` reference)."""
model_config = ConfigDict(frozen=True, extra="forbid")
kid: str = Field(min_length=1)
public_key: str = Field(min_length=1)
class MCPSessionTokenSigningSettings(BaseModel):
"""The ``general_settings.mcp_session_token_signing`` block: opt-in asymmetric signing
for the gateway session tokens. Absent, the gateway keeps the backward-compatible
HS256 key derived from ``master_key``. ``private_key`` and each ``public_key`` accept
a PEM string inline or an ``os.environ/<NAME>`` (or secret manager) reference."""
model_config = ConfigDict(frozen=True, extra="forbid")
algorithm: Literal["RS256"]
kid: str = Field(min_length=1)
private_key: str = Field(min_length=1)
previous_public_keys: tuple[SessionSigningPreviousKey, ...] = ()
class SessionSigningConfigError(BaseModel):
"""``mcp_session_token_signing`` is present but unusable (bad shape, unresolvable
secret reference, or a key that is not a loadable RSA PEM); the caller fails closed
with a server error instead of silently falling back to HS256."""
model_config = ConfigDict(frozen=True)
tag: Literal["session_signing_config_error"] = "session_signing_config_error"
detail: str
def _resolve_key_material(value: str) -> str | None:
if not value.startswith("os.environ/"):
return value
from litellm.secret_managers.main import get_secret_str # noqa: PLC0415 # heavy import kept off the pure path
return get_secret_str(value)
def resolve_session_signing_keys(
master_key: str,
raw_settings: object | None,
) -> SessionSigningKeys | SessionSigningConfigError:
"""Turn the operator's ``mcp_session_token_signing`` setting into signing key material.
``None`` (the setting absent) keeps the backward-compatible HS256 key derived from
``master_key``. A present setting must fully validate into RS256 material; any defect
is a ``SessionSigningConfigError`` value so token issuance and admission fail closed
rather than minting under a key the operator did not intend.
"""
if raw_settings is None:
return session_keys_from_master_key(master_key)
try:
settings: Final = MCPSessionTokenSigningSettings.model_validate(raw_settings)
except ValidationError as exc:
return SessionSigningConfigError(detail=f"mcp_session_token_signing is malformed: {exc}")
private_pem: Final = _resolve_key_material(settings.private_key)
if private_pem is None:
return SessionSigningConfigError(detail="mcp_session_token_signing.private_key reference did not resolve")
resolved_previous: Final = tuple(
(previous.kid, _resolve_key_material(previous.public_key)) for previous in settings.previous_public_keys
)
unresolved: Final = tuple(kid for kid, pem in resolved_previous if pem is None)
if unresolved:
return SessionSigningConfigError(
detail=f"mcp_session_token_signing.previous_public_keys reference did not resolve for kid(s): {', '.join(unresolved)}"
)
try:
return AsymmetricSessionKeys(
private_key_pem=SecretStr(private_pem),
kid=settings.kid,
previous_public_keys=tuple(
SessionRotatedPublicKey(kid=kid, public_key_pem=pem)
for kid, pem in resolved_previous
if pem is not None
),
)
except ValidationError as exc:
return SessionSigningConfigError(
detail=f"mcp_session_token_signing keys are not usable RSA PEM material: {exc}"
)
def active_session_signing_keys(master_key: str) -> SessionSigningKeys | SessionSigningConfigError:
"""Wiring helper for the token endpoint and the admission edge: resolve the signing
keys from the live ``general_settings.mcp_session_token_signing`` block, or derive the
default HS256 key from ``master_key`` when the block is absent."""
from litellm.proxy.proxy_server import general_settings # noqa: PLC0415 # circular import at module load
return resolve_session_signing_keys(master_key, general_settings.get("mcp_session_token_signing"))
class NotSessionBearer(BaseModel):
"""The bearer is not session-shaped; admission continues on its normal path."""
@ -116,7 +212,7 @@ def is_session_bearer_shaped(authorization_value: str) -> bool:
def resolve_session_bearer(
authorization_value: str,
keys: SessionKeys,
keys: SessionSigningKeys,
now: datetime,
) -> SessionBearerResult:
"""Classify an ``Authorization`` value presented at the aggregate MCP edge.
@ -166,7 +262,7 @@ SessionRefreshResult: TypeAlias = SessionRefreshOpened | SessionRefreshInvalid
def open_session_refresh_bearer(
refresh_value: str,
keys: SessionKeys,
keys: SessionSigningKeys,
now: datetime,
expected_client_id: str,
) -> SessionRefreshResult:

181
litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py
View file

@ -8,8 +8,11 @@ is therefore a stable REFERENCE, not an authorization: admission reloads the liv
record and policy on every request, so deactivating the user (or their team) kills
outstanding sessions immediately without a revocation store.
Wire shape: ``llm_session_`` (access) / ``llm_srefresh_`` (refresh) + an HS256 JWT,
the same signing approach as :mod:`.envelope`. Claims are ``iss``/``iat``/``exp``
Wire shape: ``llm_session_`` (access) / ``llm_srefresh_`` (refresh) + a JWT signed with
the injected key material: HS256 under the default master-key-derived secret (the same
signing approach as :mod:`.envelope`), or RS256 under an operator-provided RSA private
key (:class:`AsymmetricSessionKeys`) so downstream validators hold only the public half.
Claims are ``iss``/``iat``/``exp``
plus ``jti`` (per-mint uniqueness, so two tokens minted in the same second never
collide and a future revocation list has a stable handle), ``kind``, ``user_id``, and
``client_id``; ``client_id`` binds the refresh token
@ -31,11 +34,16 @@ injected ``now``); the strict pydantic claims model is the sole, total type gate
from __future__ import annotations
import secrets
from collections import Counter
from datetime import datetime, timedelta
from functools import lru_cache
from typing import Final, Literal, TypeAlias
import jwt
from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError
from cryptography.exceptions import UnsupportedAlgorithm
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError, field_validator, model_validator
SESSION_TOKEN_PREFIX: Final = "llm_session_"
"""Marker prefix on every serialized session ACCESS token so the admission edge can cheaply
@ -71,6 +79,11 @@ limits while bounding hostile input before JWT parsing."""
_SESSION_JWT_ALGORITHM: Final = "HS256"
_SESSION_RSA_ALGORITHM: Final = "RS256"
_MIN_RSA_KEY_BITS: Final = 2048
"""RFC 7518 section 3.3: RS256 requires a key of at least 2048 bits."""
SessionTokenKind = Literal["session", "session_refresh"]
"""Which credential a session token is. Stamped into the signed claims and required to match
on open, so a signature-valid token of one kind cannot be replayed as the other even if its
@ -120,6 +133,85 @@ class SessionKeys(BaseModel):
signing_key: SecretStr = Field(min_length=32)
class SessionRotatedPublicKey(BaseModel):
"""The public half of a retired signing key, kept verifiable under its ``kid`` during a
rotation window so tokens minted before the rotation stay valid until they expire."""
model_config = ConfigDict(frozen=True)
kid: str = Field(min_length=1)
public_key_pem: str = Field(min_length=1)
@field_validator("public_key_pem")
@classmethod
def _pem_is_an_rsa_public_key(cls, value: str) -> str:
try:
loaded: Final = serialization.load_pem_public_key(value.encode())
except (ValueError, TypeError, UnsupportedAlgorithm) as exc:
raise ValueError(f"public_key_pem is not a loadable PEM public key: {exc}") from exc
if not isinstance(loaded, rsa.RSAPublicKey):
raise ValueError("public_key_pem must be an RSA public key in PEM format") # noqa: TRY004 # pydantic validators must raise ValueError
if loaded.key_size < _MIN_RSA_KEY_BITS:
raise ValueError(f"public_key_pem must be an RSA key of at least {_MIN_RSA_KEY_BITS} bits")
return value
class AsymmetricSessionKeys(BaseModel):
"""Injected RS256 key material: the issuer-held RSA private key and the stable ``kid``
stamped into every minted token's JOSE header, plus the public halves of previously
rotated keys that verification still accepts while their tokens age out. Downstream
validators never need the private key: :func:`session_public_key_pem` yields the
public half to distribute."""
model_config = ConfigDict(frozen=True)
private_key_pem: SecretStr
kid: str = Field(min_length=1)
previous_public_keys: tuple[SessionRotatedPublicKey, ...] = ()
@field_validator("private_key_pem")
@classmethod
def _pem_is_a_strong_rsa_private_key(cls, value: SecretStr) -> SecretStr:
try:
loaded: Final = serialization.load_pem_private_key(value.get_secret_value().encode(), password=None)
except (ValueError, TypeError, UnsupportedAlgorithm) as exc:
raise ValueError(f"private_key_pem is not a loadable unencrypted PEM private key: {exc}") from exc
if not isinstance(loaded, rsa.RSAPrivateKey):
raise ValueError("private_key_pem must be an unencrypted RSA private key in PEM format") # noqa: TRY004 # pydantic validators must raise ValueError
if loaded.key_size < _MIN_RSA_KEY_BITS:
raise ValueError(f"private_key_pem must be an RSA key of at least {_MIN_RSA_KEY_BITS} bits")
return value
@model_validator(mode="after")
def _kids_are_unique(self) -> AsymmetricSessionKeys:
kids: Final = (self.kid, *(previous.kid for previous in self.previous_public_keys))
duplicates: Final = tuple(kid for kid, count in Counter(kids).items() if count > 1)
if duplicates:
raise ValueError(
f"every kid must be unique across the current and previous keys; duplicated: {', '.join(duplicates)}"
)
return self
SessionSigningKeys: TypeAlias = SessionKeys | AsymmetricSessionKeys
"""Every key material shape the mints and openers accept: the default master-key-derived
HS256 secret, or operator-configured RS256 RSA keys."""
@lru_cache(maxsize=8)
def _public_key_pem_from_private(private_key_pem: str) -> str:
loaded: Final = serialization.load_pem_private_key(private_key_pem.encode(), password=None)
return (
loaded.public_key()
.public_bytes(serialization.Encoding.PEM, serialization.PublicFormat.SubjectPublicKeyInfo)
.decode()
)
def session_public_key_pem(keys: AsymmetricSessionKeys) -> str:
"""The PEM public half of the current RS256 signing key: the only material a downstream
validator (an external gateway verifying ``kid``-matched tokens) ever needs."""
return _public_key_pem_from_private(keys.private_key_pem.get_secret_value())
class MintedSessionToken(BaseModel):
"""A minted session token: the client-held bearer value and when it expires."""
@ -221,7 +313,7 @@ def is_session_refresh_token(candidate: str) -> bool:
def mint_session_token(
principal: SessionPrincipal,
keys: SessionKeys,
keys: SessionSigningKeys,
now: datetime,
) -> MintedSessionToken | SessionTokenMintError:
"""Mint the short-lived session ACCESS token for ``principal``.
@ -241,7 +333,7 @@ def mint_session_token(
def mint_session_refresh_token(
principal: SessionPrincipal,
keys: SessionKeys,
keys: SessionSigningKeys,
now: datetime,
) -> MintedSessionToken | SessionTokenMintError:
"""Mint the long-lived session REFRESH token for ``principal``.
@ -262,7 +354,7 @@ def mint_session_refresh_token(
def open_session_token(
candidate: str,
keys: SessionKeys,
keys: SessionSigningKeys,
now: datetime,
) -> OpenedSessionToken | SessionTokenOpenError:
"""Validate a session ACCESS ``candidate`` and recover the principal.
@ -275,7 +367,7 @@ def open_session_token(
def open_session_refresh_token(
candidate: str,
keys: SessionKeys,
keys: SessionSigningKeys,
now: datetime,
) -> OpenedSessionToken | SessionTokenOpenError:
"""Validate a session REFRESH ``candidate`` and recover the principal.
@ -292,7 +384,7 @@ def _mint(
prefix: str,
principal: SessionPrincipal,
expires_at: datetime,
keys: SessionKeys,
keys: SessionSigningKeys,
now: datetime,
) -> MintedSessionToken | SessionTokenTooLarge:
"""Sign the claims for either token kind and enforce the size cap. Shared by both mints
@ -309,20 +401,33 @@ def _mint(
audience=principal.audience,
team_id=principal.team_id,
)
token: Final = prefix + jwt.encode(
claims.model_dump(exclude_none=True), keys.signing_key.get_secret_value(), algorithm=_SESSION_JWT_ALGORITHM
)
token: Final = prefix + _sign_claims(claims, keys)
size_bytes: Final = len(token.encode("utf-8"))
if size_bytes > MAX_SESSION_TOKEN_BYTES:
return SessionTokenTooLarge(size_bytes=size_bytes, max_bytes=MAX_SESSION_TOKEN_BYTES)
return MintedSessionToken(token=SecretStr(token), expires_at=expires_at)
def _sign_claims(claims: _SessionClaims, keys: SessionSigningKeys) -> str:
"""Sign the claim set under whichever key material was injected: RS256 with the ``kid``
in the JOSE header (so a validator can pick the right public key), or the default
HS256 secret with no header extras (byte-compatible with every pre-RS256 token)."""
payload: Final = claims.model_dump(exclude_none=True)
if isinstance(keys, AsymmetricSessionKeys):
return jwt.encode(
payload,
keys.private_key_pem.get_secret_value(),
algorithm=_SESSION_RSA_ALGORITHM,
headers={"kid": keys.kid},
)
return jwt.encode(payload, keys.signing_key.get_secret_value(), algorithm=_SESSION_JWT_ALGORITHM)
def _open(
candidate: str,
prefix: str,
expected_kind: SessionTokenKind,
keys: SessionKeys,
keys: SessionSigningKeys,
now: datetime,
) -> OpenedSessionToken | SessionTokenOpenError:
"""Prefix-route, size-bound, signature-verify, kind-check, and expiry-check an
@ -337,7 +442,7 @@ def _open(
return SessionMalformed()
if len(candidate.encode("utf-8", "surrogatepass")) > MAX_SESSION_TOKEN_BYTES:
return SessionMalformed()
claims: Final = _decode_claims(candidate.removeprefix(prefix), keys.signing_key)
claims: Final = _decode_claims(candidate.removeprefix(prefix), keys)
if not isinstance(claims, _SessionClaims):
return claims
if claims.kind != expected_kind:
@ -356,14 +461,51 @@ def _open(
)
class _VerificationMaterial(BaseModel):
model_config = ConfigDict(frozen=True)
key: SecretStr
algorithm: Literal["HS256", "RS256"]
def _verification_material(
compact: str,
keys: SessionSigningKeys,
) -> _VerificationMaterial | SessionBadSignature | SessionMalformed:
"""Pick the single key and algorithm the candidate is allowed to verify under.
HS256 mode has exactly one secret. RS256 mode routes by the JOSE header ``kid``: the
current key's derived public half, or a retired key's stored public half during a
rotation window. An unknown or missing ``kid`` is ``SessionBadSignature`` (a foreign
key), and an undecodable header is ``SessionMalformed``. The algorithm is pinned per
key shape, never read from the header, so an HS256 token can never be verified
against a public key or vice versa.
"""
if isinstance(keys, SessionKeys):
return _VerificationMaterial(key=keys.signing_key, algorithm=_SESSION_JWT_ALGORITHM)
try:
header: Final = jwt.get_unverified_header(compact)
except jwt.InvalidTokenError:
return SessionMalformed()
kid: Final = header.get("kid")
if kid == keys.kid:
return _VerificationMaterial(key=SecretStr(session_public_key_pem(keys)), algorithm=_SESSION_RSA_ALGORITHM)
for previous in keys.previous_public_keys:
if previous.kid == kid:
return _VerificationMaterial(key=SecretStr(previous.public_key_pem), algorithm=_SESSION_RSA_ALGORITHM)
return SessionBadSignature()
def _decode_claims(
compact: str,
signing_key: SecretStr,
keys: SessionSigningKeys,
) -> _SessionClaims | SessionBadSignature | SessionMalformed:
"""Verify the HS256 signature and shape of an attacker-controlled compact JWT.
"""Verify the signature and shape of an attacker-controlled compact JWT.
``compact`` is fully hostile and bounded to ``MAX_SESSION_TOKEN_BYTES`` by the caller.
PyJWT's ``iat``/``nbf``/``exp`` validators are disabled: they raise on hostile claim
The accepted algorithm is pinned by :func:`_verification_material` from the injected
key shape, so ``alg`` confusion (``none``, or HS256 signed with a public key as the
secret) fails before or at signature verification. PyJWT's ``iat``/``nbf``/``exp``
validators are disabled: they raise on hostile claim
types and, for ``iat``/``nbf``, compare against the wall clock rather than the injected
``now`` (``exp`` is checked by the caller against ``now``). Apart from a signature
mismatch, every decode failure is ``SessionMalformed``: a non-UTF-8 candidate surfaces
@ -371,11 +513,14 @@ def _decode_claims(
``TypeError`` from PyJWT's claim validators, and a wrong issuer or structurally invalid
token as an ``InvalidTokenError``. ``_SessionClaims`` is the total type gate.
"""
material: Final = _verification_material(compact, keys)
if not isinstance(material, _VerificationMaterial):
return material
try:
payload: Final = jwt.decode(
compact,
signing_key.get_secret_value(),
algorithms=[_SESSION_JWT_ALGORITHM],
material.key.get_secret_value(),
algorithms=[material.algorithm],
issuer=SESSION_ISSUER,
options={
"verify_exp": False,

46
litellm/proxy/health_endpoints/_health_endpoints.py
View file

@ -1372,8 +1372,25 @@ class DBHealthCache(TypedDict):
db_health_cache: DBHealthCache = {"status": "unknown", "last_updated": datetime.now()}
# Bounds each DB round-trip on the probe path so a hung connection during a
# failover cannot make the probe fail by timeout (k8s default timeoutSeconds: 5).
DB_READINESS_CHECK_TIMEOUT_SECONDS: Final = 2.0
# One deadline for the whole probe-path DB check (initial check + reconnect +
# re-check, including reconnect lock waits), kept under timeoutSeconds: 5.
DB_READINESS_PROBE_DEADLINE_SECONDS: Final = 4.0
async def _db_health_readiness_check():
async def _db_health_readiness_check() -> DBHealthCache:
try:
return await asyncio.wait_for(
_db_health_readiness_check_unbounded(),
timeout=DB_READINESS_PROBE_DEADLINE_SECONDS,
)
except asyncio.TimeoutError:
return {"status": "disconnected", "last_updated": db_health_cache["last_updated"]}
async def _db_health_readiness_check_unbounded() -> DBHealthCache:
from litellm.proxy.proxy_server import prisma_client
global db_health_cache
@ -1387,7 +1404,7 @@ async def _db_health_readiness_check():
db_health_cache = {"status": "disconnected", "last_updated": datetime.now()}
return db_health_cache
await prisma_client.health_check()
await asyncio.wait_for(prisma_client.health_check(), timeout=DB_READINESS_CHECK_TIMEOUT_SECONDS)
db_health_cache = {"status": "connected", "last_updated": datetime.now()}
return db_health_cache
except Exception as e:
@ -1395,8 +1412,15 @@ async def _db_health_readiness_check():
if PrismaDBExceptionHandler.is_database_transport_error(e):
try:
verbose_proxy_logger.warning("_db_health_readiness_check: health_check failed, attempting reconnect")
await prisma_client.attempt_db_reconnect(reason="health_readiness_check")
await prisma_client.health_check()
await prisma_client.attempt_db_reconnect(
reason="health_readiness_check",
timeout_seconds=DB_READINESS_CHECK_TIMEOUT_SECONDS,
lock_timeout_seconds=DB_READINESS_CHECK_TIMEOUT_SECONDS,
)
await asyncio.wait_for(
prisma_client.health_check(),
timeout=DB_READINESS_CHECK_TIMEOUT_SECONDS,
)
verbose_proxy_logger.info("_db_health_readiness_check: reconnect succeeded")
db_health_cache = {
"status": "connected",
@ -1580,7 +1604,14 @@ async def _get_health_readiness_details(
# serve requests that depend on persisted state (keys, budgets,
# spend logs). Return 503 so orchestrators take this pod out of
# rotation; "Not connected" (no DB configured at all) stays 200.
if response is not None and db_health_status["status"] != "connected":
# With allow_requests_on_db_unavailable the proxy keeps serving
# during a DB outage, so the pod must stay in rotation (200) and
# report the DB state through the body instead.
if (
response is not None
and db_health_status["status"] != "connected"
and not PrismaDBExceptionHandler.should_allow_request_on_db_unavailable()
):
response.status_code = status.HTTP_503_SERVICE_UNAVAILABLE
return {
"status": "healthy",
@ -1671,7 +1702,10 @@ async def _resolve_public_readiness_db(response: Response) -> str:
return "Not connected"
db_health_status: Final = await _db_health_readiness_check()
if db_health_status["status"] != "connected":
if (
db_health_status["status"] != "connected"
and not PrismaDBExceptionHandler.should_allow_request_on_db_unavailable()
):
response.status_code = status.HTTP_503_SERVICE_UNAVAILABLE
return db_health_status["status"]

7
litellm/proxy/litellm_pre_call_utils.py
View file

@ -1948,6 +1948,13 @@ async def add_litellm_data_to_request(
for key, value in data["litellm_metadata"].items():
if key not in data[_metadata_variable_name]:
data[_metadata_variable_name][key] = value
if _metadata_variable_name == "metadata":
data["metadata"]["tags"] = LiteLLMProxyRequestSetup._merge_tags( # pyright: ignore[reportPrivateUsage] # same-module helper, budget blocks the unsuppressed idiom sibling call sites use
request_tags=data["metadata"].get("tags"),
tags_to_add=data["litellm_metadata"].get("tags"),
)
if _metadata_variable_name == "metadata":
data.pop("litellm_metadata", None)
data = LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata(
data=data,

20
litellm/proxy/utils.py
View file

@ -5580,12 +5580,8 @@ class PrismaClient:
return True
acquire_task: Final = asyncio.create_task(_acquire_reconnect_lock())
done, _pending = await asyncio.wait(
{acquire_task},
timeout=lock_timeout_seconds,
return_when=asyncio.FIRST_COMPLETED,
)
if acquire_task not in done:
async def _abandon_acquire_task() -> None:
acquire_task.cancel()
try:
await acquire_task
@ -5600,6 +5596,18 @@ class PrismaClient:
self._db_reconnect_lock.release()
except RuntimeError:
pass
try:
done, _pending = await asyncio.wait(
{acquire_task},
timeout=lock_timeout_seconds,
return_when=asyncio.FIRST_COMPLETED,
)
except asyncio.CancelledError:
await asyncio.shield(_abandon_acquire_task())
raise
if acquire_task not in done:
await _abandon_acquire_task()
verbose_proxy_logger.debug(
"Skipping DB reconnect attempt due to lock acquisition timeout. reason=%s timeout=%ss",
reason,

87
tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py
View file

@ -2635,6 +2635,93 @@ async def test_list_batches_unparseable_row_does_not_truncate_pagination():
assert len(seen) == len(set(seen))
@pytest.mark.asyncio
async def test_list_batches_fills_a_page_past_a_full_page_of_unparseable_rows():
"""A page whose rows all fail to parse must still let the caller advance.
``has_more`` came from the raw fetch while ``last_id`` came from the parsed
survivors, so a full page of corrupt rows answered ``data: []``,
``last_id: None``, ``has_more: True``, and a client following ``last_id``
could not move past them.
"""
from litellm.proxy._types import UserAPIKeyAuth
rows = [_managed_batch_row(i) for i in range(5)]
for corrupt_row in rows[2:4]:
corrupt_row.file_object = "{ not valid json"
prisma_client = _fake_managed_object_table(rows)
proxy_managed_files = _PROXY_LiteLLMManagedFiles(
DualCache(), prisma_client=prisma_client
)
pages = await _walk_batch_pages(
proxy_managed_files, UserAPIKeyAuth(user_id="test-user"), limit=1
)
assert [[batch.id for batch in page["data"]] for page in pages] == [
[rows[4].unified_object_id],
[rows[1].unified_object_id],
[rows[0].unified_object_id],
]
assert [page["has_more"] for page in pages] == [True, True, False]
_DEEP_BATCH_SCAN_ROW_COUNT = 2000
_DEEP_BATCH_SCAN_QUERY_BUDGET = 10
@pytest.mark.asyncio
async def test_list_batches_bounds_the_queries_a_deep_unparseable_run_costs():
"""A tiny limit behind thousands of corrupt rows must not turn one request into thousands of queries."""
from litellm.proxy._types import UserAPIKeyAuth
rows = [_managed_batch_row(0)] + [
_managed_batch_row(index, file_object="{ not valid json")
for index in range(1, _DEEP_BATCH_SCAN_ROW_COUNT + 1)
]
prisma_client = _fake_managed_object_table(rows)
proxy_managed_files = _PROXY_LiteLLMManagedFiles(
DualCache(), prisma_client=prisma_client
)
page = await proxy_managed_files.list_user_batches(
user_api_key_dict=UserAPIKeyAuth(user_id="test-user"), limit=1
)
assert [batch.id for batch in page["data"]] == [rows[0].unified_object_id]
assert page["has_more"] is False
assert (
prisma_client.db.litellm_managedobjecttable.find_many.call_count
<= _DEEP_BATCH_SCAN_QUERY_BUDGET
)
@pytest.mark.asyncio
async def test_list_batches_reads_one_chunk_when_the_first_one_fills_the_page():
"""The widened chunk must stay off the common path, where the newest rows already fill the page."""
from litellm.proxy._types import UserAPIKeyAuth
rows = [_managed_batch_row(index) for index in range(_DEEP_BATCH_SCAN_ROW_COUNT)]
prisma_client = _fake_managed_object_table(rows)
proxy_managed_files = _PROXY_LiteLLMManagedFiles(
DualCache(), prisma_client=prisma_client
)
page = await proxy_managed_files.list_user_batches(
user_api_key_dict=UserAPIKeyAuth(user_id="test-user"), limit=2
)
assert [batch.id for batch in page["data"]] == [
rows[-1].unified_object_id,
rows[-2].unified_object_id,
]
assert page["has_more"] is True
assert prisma_client.db.litellm_managedobjecttable.find_many.call_count == 1
@pytest.mark.asyncio
async def test_return_unified_file_id_includes_expires_at():
from litellm.types.llms.openai import OpenAIFileObject

91
tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py
View file

@ -3,6 +3,9 @@
from datetime import datetime, timedelta, timezone
import pytest
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from pydantic import SecretStr
from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import (
envelope_keys_from_master_key,
@ -13,17 +16,22 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credent
SessionBearerInvalid,
SessionRefreshInvalid,
SessionRefreshOpened,
SessionSigningConfigError,
is_session_bearer_shaped,
open_session_refresh_bearer,
resolve_session_bearer,
resolve_session_signing_keys,
session_keys_from_master_key,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import (
SESSION_TTL_SECONDS,
AsymmetricSessionKeys,
MintedSessionToken,
SessionKeys,
SessionPrincipal,
mint_session_refresh_token,
mint_session_token,
session_public_key_pem,
)
NOW = datetime(2026, 1, 1, 12, 0, 0, tzinfo=timezone.utc)
@ -133,3 +141,86 @@ def test_refresh_grant_rejects_a_different_client():
def test_refresh_grant_rejects_access_token_presented_as_refresh():
result = open_session_refresh_bearer(_access_token(), KEYS, NOW, expected_client_id="llm_client_abc")
assert isinstance(result, SessionRefreshInvalid)
def _rsa_private_pem() -> str:
key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
return key.private_bytes(
serialization.Encoding.PEM,
serialization.PrivateFormat.PKCS8,
serialization.NoEncryption(),
).decode()
def test_absent_signing_setting_keeps_the_master_key_hs256_default():
resolved = resolve_session_signing_keys(MASTER_KEY, None)
assert isinstance(resolved, SessionKeys)
assert resolved.signing_key.get_secret_value() == KEYS.signing_key.get_secret_value()
def test_rs256_signing_setting_resolves_inline_pem_material():
pem = _rsa_private_pem()
resolved = resolve_session_signing_keys(
MASTER_KEY,
{"algorithm": "RS256", "kid": "2026-01", "private_key": pem},
)
assert isinstance(resolved, AsymmetricSessionKeys)
assert resolved.kid == "2026-01"
minted = mint_session_token(PRINCIPAL, resolved, NOW)
assert isinstance(minted, MintedSessionToken)
admitted = resolve_session_bearer(f"Bearer {minted.token.get_secret_value()}", resolved, NOW)
assert isinstance(admitted, SessionBearerAdmitted)
def test_rs256_signing_setting_resolves_env_reference(monkeypatch):
monkeypatch.setenv("MCP_SESSION_PRIVATE_KEY", _rsa_private_pem())
resolved = resolve_session_signing_keys(
MASTER_KEY,
{"algorithm": "RS256", "kid": "2026-01", "private_key": "os.environ/MCP_SESSION_PRIVATE_KEY"},
)
assert isinstance(resolved, AsymmetricSessionKeys)
def test_rs256_signing_setting_resolves_previous_public_keys():
old_pem = _rsa_private_pem()
old_keys = AsymmetricSessionKeys(private_key_pem=SecretStr(old_pem), kid="2025-06")
resolved = resolve_session_signing_keys(
MASTER_KEY,
{
"algorithm": "RS256",
"kid": "2026-01",
"private_key": _rsa_private_pem(),
"previous_public_keys": [{"kid": "2025-06", "public_key": session_public_key_pem(old_keys)}],
},
)
assert isinstance(resolved, AsymmetricSessionKeys)
minted = mint_session_token(PRINCIPAL, old_keys, NOW)
assert isinstance(minted, MintedSessionToken)
admitted = resolve_session_bearer(f"Bearer {minted.token.get_secret_value()}", resolved, NOW)
assert isinstance(admitted, SessionBearerAdmitted)
@pytest.mark.parametrize(
"raw",
[
{"algorithm": "HS512", "kid": "k", "private_key": "irrelevant"},
{"algorithm": "RS256", "kid": "k"},
{"algorithm": "RS256", "kid": "k", "private_key": "not a pem"},
{"algorithm": "RS256", "kid": "k", "private_key": "os.environ/UNSET_MCP_SESSION_KEY_VAR"},
{"algorithm": "RS256", "kid": "k", "private_key": "x", "unexpected": True},
"not-a-mapping",
],
)
def test_defective_signing_setting_fails_closed_never_falls_back_to_hs256(raw):
resolved = resolve_session_signing_keys(MASTER_KEY, raw)
assert isinstance(resolved, SessionSigningConfigError)
def test_signing_config_error_detail_never_leaks_key_material():
pem = _rsa_private_pem()
resolved = resolve_session_signing_keys(
MASTER_KEY,
{"algorithm": "RS256", "kid": "k", "private_key": pem, "unexpected": True},
)
assert isinstance(resolved, SessionSigningConfigError)
assert pem.splitlines()[1] not in resolved.detail

187
tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_token.py
View file

@ -4,6 +4,8 @@ from datetime import datetime, timedelta, timezone
import jwt
import pytest
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from pydantic import SecretStr, ValidationError
from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import (
@ -13,6 +15,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token i
SESSION_REFRESH_TTL_SECONDS,
SESSION_TOKEN_PREFIX,
SESSION_TTL_SECONDS,
AsymmetricSessionKeys,
MintedSessionToken,
NotASessionToken,
OpenedSessionToken,
@ -21,6 +24,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token i
SessionKeys,
SessionMalformed,
SessionPrincipal,
SessionRotatedPublicKey,
SessionTokenTooLarge,
is_session_refresh_token,
is_session_token,
@ -28,8 +32,22 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token i
mint_session_token,
open_session_refresh_token,
open_session_token,
session_public_key_pem,
)
def _rsa_private_pem(bits: int = 2048) -> str:
key = rsa.generate_private_key(public_exponent=65537, key_size=bits)
return key.private_bytes(
serialization.Encoding.PEM,
serialization.PrivateFormat.PKCS8,
serialization.NoEncryption(),
).decode()
_RSA_PEM_A = _rsa_private_pem()
_RSA_PEM_B = _rsa_private_pem()
NOW = datetime(2026, 1, 1, 12, 0, 0, tzinfo=timezone.utc)
KEYS = SessionKeys(signing_key=SecretStr("k" * 32))
OTHER_KEYS = SessionKeys(signing_key=SecretStr("x" * 32))
@ -264,3 +282,172 @@ def test_signed_claims_with_a_non_string_team_are_rejected():
def test_principal_rejects_an_unknown_audience_at_construction():
with pytest.raises(ValidationError):
SessionPrincipal(user_id="user-123", client_id="llm_client_abc", audience="mcp")
RSA_KEYS = AsymmetricSessionKeys(private_key_pem=SecretStr(_RSA_PEM_A), kid="2026-01")
OTHER_RSA_KEYS = AsymmetricSessionKeys(private_key_pem=SecretStr(_RSA_PEM_B), kid="2025-06")
def test_rs256_access_round_trip_with_kid_and_alg_pinned_in_header():
minted = mint_session_token(PRINCIPAL, RSA_KEYS, NOW)
assert isinstance(minted, MintedSessionToken)
token = minted.token.get_secret_value()
header = jwt.get_unverified_header(token.removeprefix(SESSION_TOKEN_PREFIX))
assert header["alg"] == "RS256"
assert header["kid"] == "2026-01"
opened = open_session_token(token, RSA_KEYS, NOW)
assert isinstance(opened, OpenedSessionToken)
assert opened.principal == PRINCIPAL
def test_rs256_refresh_round_trip():
minted = mint_session_refresh_token(PRINCIPAL, RSA_KEYS, NOW)
assert isinstance(minted, MintedSessionToken)
token = minted.token.get_secret_value()
opened = open_session_refresh_token(token, RSA_KEYS, NOW)
assert isinstance(opened, OpenedSessionToken)
assert opened.principal == PRINCIPAL
def test_rs256_token_verifies_with_public_key_only():
minted = mint_session_token(PRINCIPAL, RSA_KEYS, NOW)
assert isinstance(minted, MintedSessionToken)
public_pem = session_public_key_pem(RSA_KEYS)
assert "PUBLIC KEY" in public_pem
assert "PRIVATE" not in public_pem
claims = jwt.decode(
minted.token.get_secret_value().removeprefix(SESSION_TOKEN_PREFIX),
public_pem,
algorithms=["RS256"],
issuer=SESSION_ISSUER,
options={"verify_exp": False},
)
assert claims["user_id"] == "user-123"
def test_rs256_tampered_signature_is_bad_signature():
minted = mint_session_token(PRINCIPAL, RSA_KEYS, NOW)
assert isinstance(minted, MintedSessionToken)
token = minted.token.get_secret_value()
tampered = token[:-2] + ("aa" if not token.endswith("aa") else "bb")
assert isinstance(open_session_token(tampered, RSA_KEYS, NOW), SessionBadSignature)
def test_rs256_expired_token_is_expired():
minted = mint_session_token(PRINCIPAL, RSA_KEYS, NOW)
assert isinstance(minted, MintedSessionToken)
after = NOW + timedelta(seconds=SESSION_TTL_SECONDS + 1)
assert isinstance(open_session_token(minted.token.get_secret_value(), RSA_KEYS, after), SessionExpired)
def test_hs256_token_is_rejected_in_rs256_mode():
assert isinstance(open_session_token(_mint_access(), RSA_KEYS, NOW), SessionBadSignature)
def test_hs256_token_claiming_the_current_kid_is_rejected_by_alg_pinning():
token = SESSION_TOKEN_PREFIX + jwt.encode(
_valid_claims(),
KEYS.signing_key.get_secret_value(),
algorithm="HS256",
headers={"kid": RSA_KEYS.kid},
)
assert isinstance(open_session_token(token, RSA_KEYS, NOW), SessionMalformed)
def test_rs256_token_is_rejected_in_hs256_mode():
minted = mint_session_token(PRINCIPAL, RSA_KEYS, NOW)
assert isinstance(minted, MintedSessionToken)
assert isinstance(open_session_token(minted.token.get_secret_value(), KEYS, NOW), SessionMalformed)
def test_rs256_token_from_an_unknown_kid_is_bad_signature():
minted = mint_session_token(PRINCIPAL, OTHER_RSA_KEYS, NOW)
assert isinstance(minted, MintedSessionToken)
assert isinstance(open_session_token(minted.token.get_secret_value(), RSA_KEYS, NOW), SessionBadSignature)
def test_rs256_token_signed_by_a_foreign_key_claiming_the_current_kid_is_bad_signature():
token = SESSION_TOKEN_PREFIX + jwt.encode(
_valid_claims(),
_RSA_PEM_B,
algorithm="RS256",
headers={"kid": RSA_KEYS.kid},
)
assert isinstance(open_session_token(token, RSA_KEYS, NOW), SessionBadSignature)
def test_alg_none_token_with_the_current_kid_is_rejected_in_rs256_mode():
unsigned = jwt.api_jws.encode(
b'{"iss":"litellm-mcp-gateway"}', key=None, algorithm="none", headers={"kid": RSA_KEYS.kid}
)
assert isinstance(open_session_token(SESSION_TOKEN_PREFIX + unsigned, RSA_KEYS, NOW), SessionMalformed)
def test_rotation_previous_public_key_still_verifies_until_removed():
minted = mint_session_token(PRINCIPAL, OTHER_RSA_KEYS, NOW)
assert isinstance(minted, MintedSessionToken)
token = minted.token.get_secret_value()
rotated = AsymmetricSessionKeys(
private_key_pem=SecretStr(_RSA_PEM_A),
kid="2026-01",
previous_public_keys=(
SessionRotatedPublicKey(kid="2025-06", public_key_pem=session_public_key_pem(OTHER_RSA_KEYS)),
),
)
opened = open_session_token(token, rotated, NOW)
assert isinstance(opened, OpenedSessionToken)
assert opened.principal == PRINCIPAL
assert isinstance(open_session_token(token, RSA_KEYS, NOW), SessionBadSignature)
def test_rotation_window_still_enforces_expiry_and_tamper_on_the_previous_key():
minted = mint_session_token(PRINCIPAL, OTHER_RSA_KEYS, NOW)
assert isinstance(minted, MintedSessionToken)
token = minted.token.get_secret_value()
rotated = AsymmetricSessionKeys(
private_key_pem=SecretStr(_RSA_PEM_A),
kid="2026-01",
previous_public_keys=(
SessionRotatedPublicKey(kid="2025-06", public_key_pem=session_public_key_pem(OTHER_RSA_KEYS)),
),
)
after = NOW + timedelta(seconds=SESSION_TTL_SECONDS + 1)
assert isinstance(open_session_token(token, rotated, after), SessionExpired)
tampered = token[:-2] + ("aa" if not token.endswith("aa") else "bb")
assert isinstance(open_session_token(tampered, rotated, NOW), SessionBadSignature)
def test_weak_or_garbage_private_key_pem_rejected_at_construction():
with pytest.raises(ValidationError):
AsymmetricSessionKeys(private_key_pem=SecretStr(_rsa_private_pem(bits=1024)), kid="weak")
with pytest.raises(ValidationError):
AsymmetricSessionKeys(private_key_pem=SecretStr("not a pem"), kid="junk")
with pytest.raises(ValidationError):
SessionRotatedPublicKey(kid="junk", public_key_pem="not a pem")
with pytest.raises(ValidationError):
SessionRotatedPublicKey(kid="private-half", public_key_pem=_RSA_PEM_A)
def test_weak_rotated_public_key_rejected_at_construction():
weak_public = (
serialization.load_pem_private_key(_rsa_private_pem(bits=1024).encode(), password=None)
.public_key()
.public_bytes(serialization.Encoding.PEM, serialization.PublicFormat.SubjectPublicKeyInfo)
.decode()
)
with pytest.raises(ValidationError):
SessionRotatedPublicKey(kid="2024-01", public_key_pem=weak_public)
def test_duplicate_kids_rejected_at_construction():
previous = SessionRotatedPublicKey(kid="2025-06", public_key_pem=session_public_key_pem(OTHER_RSA_KEYS))
with pytest.raises(ValidationError):
AsymmetricSessionKeys(private_key_pem=SecretStr(_RSA_PEM_A), kid="2025-06", previous_public_keys=(previous,))
with pytest.raises(ValidationError):
AsymmetricSessionKeys(
private_key_pem=SecretStr(_RSA_PEM_A), kid="2026-01", previous_public_keys=(previous, previous)
)
def test_asymmetric_keys_repr_never_leaks_the_private_key():
assert _RSA_PEM_A not in repr(RSA_KEYS)

176
tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py
View file

@ -1,10 +1,10 @@
import asyncio
import json
import time
from datetime import datetime, timedelta
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
import respx
@ -15,7 +15,6 @@ from prisma.errors import ClientNotConnectedError, HTTPClientClosedError, Prisma
import litellm
import litellm.proxy.health_endpoints._health_endpoints as _health_endpoints_module
from litellm.litellm_core_utils.health_check_helpers import TEST_IMAGE_BASE64
from litellm.proxy._types import LitellmUserRoles, ProxyException, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.health_endpoints._health_endpoints import (
@ -145,7 +144,11 @@ async def test_db_health_transport_error_never_raises(transport_error):
result = await _db_health_readiness_check()
assert result["status"] == "disconnected"
mock_prisma.attempt_db_reconnect.assert_called_once_with(reason="health_readiness_check")
mock_prisma.attempt_db_reconnect.assert_called_once_with(
reason="health_readiness_check",
timeout_seconds=_health_endpoints_module.DB_READINESS_CHECK_TIMEOUT_SECONDS,
lock_timeout_seconds=_health_endpoints_module.DB_READINESS_CHECK_TIMEOUT_SECONDS,
)
@pytest.mark.asyncio
@ -175,7 +178,11 @@ async def test_db_health_transport_error_reconnect_succeeds(transport_error):
result = await _db_health_readiness_check()
assert result["status"] == "connected"
mock_prisma.attempt_db_reconnect.assert_called_once_with(reason="health_readiness_check")
mock_prisma.attempt_db_reconnect.assert_called_once_with(
reason="health_readiness_check",
timeout_seconds=_health_endpoints_module.DB_READINESS_CHECK_TIMEOUT_SECONDS,
lock_timeout_seconds=_health_endpoints_module.DB_READINESS_CHECK_TIMEOUT_SECONDS,
)
assert mock_prisma.health_check.call_count == 2
@ -2276,6 +2283,159 @@ async def test_health_readiness_returns_503_when_db_disconnected():
assert result == {"status": "healthy", "db": "disconnected"}
@pytest.mark.asyncio
async def test_health_readiness_returns_200_when_db_down_and_allow_requests_on_db_unavailable():
"""
Regression test for https://github.com/BerriAI/litellm/issues/34934.
allow_requests_on_db_unavailable keeps the proxy serving through a DB
outage, so the readiness probe must keep the pod in rotation (200) and
report the DB state through the body, not the status code. Otherwise
K8s pulls every replica before the request-layer fail-open can run.
"""
from fastapi import Response
from litellm.proxy.health_endpoints._health_endpoints import health_readiness
mock_prisma = MagicMock()
mock_prisma.health_check = AsyncMock(side_effect=PrismaError("nope"))
mock_prisma.attempt_db_reconnect = AsyncMock(side_effect=Exception("still nope"))
_health_endpoints_module.db_health_cache = {
"status": "unknown",
"last_updated": datetime.now() - timedelta(seconds=60),
}
response = Response()
with (
patch( # test-quality-ok: the readiness path reads the proxy-global DB client; it has no injection seam
"litellm.proxy.proxy_server.prisma_client", mock_prisma
),
patch.dict( # test-quality-ok: the fail-open flag lives in the proxy-global general_settings; no injection seam
"litellm.proxy.proxy_server.general_settings",
{"allow_requests_on_db_unavailable": True},
),
):
result = await health_readiness(response=response)
assert response.status_code == 200
assert result == {"status": "healthy", "db": "disconnected"}
@pytest.mark.asyncio
async def test_health_readiness_details_returns_200_when_db_down_and_allow_requests_on_db_unavailable():
"""
The detailed readiness payload (public via
allow_public_health_readiness_details, or /health/readiness/details)
must honor the same flag so probes pointed at it also stay 200.
"""
from fastapi import Response
from litellm.proxy.health_endpoints._health_endpoints import (
_get_health_readiness_details,
)
mock_prisma = MagicMock()
mock_prisma.health_check = AsyncMock(side_effect=PrismaError("nope"))
mock_prisma.attempt_db_reconnect = AsyncMock(side_effect=Exception("still nope"))
_health_endpoints_module.db_health_cache = {
"status": "unknown",
"last_updated": datetime.now() - timedelta(seconds=60),
}
response = Response()
with (
patch( # test-quality-ok: the readiness path reads the proxy-global DB client; it has no injection seam
"litellm.proxy.proxy_server.prisma_client", mock_prisma
),
patch.dict( # test-quality-ok: the fail-open flag lives in the proxy-global general_settings; no injection seam
"litellm.proxy.proxy_server.general_settings",
{"allow_requests_on_db_unavailable": True},
),
):
result = await _get_health_readiness_details(response=response)
assert response.status_code == 200
assert result["db"] == "disconnected"
@pytest.mark.asyncio
async def test_db_health_readiness_check_bounds_hung_health_check():
"""
A connection that hangs mid-failover must not stall the probe past the
kubelet's timeoutSeconds; the DB round-trip is bounded and reported as
disconnected instead.
"""
from litellm.proxy.health_endpoints._health_endpoints import (
_db_health_readiness_check,
)
async def hang():
await asyncio.sleep(60)
mock_prisma = MagicMock()
mock_prisma.health_check = hang
mock_prisma.attempt_db_reconnect = AsyncMock(side_effect=Exception("still down"))
_health_endpoints_module.db_health_cache = {
"status": "unknown",
"last_updated": datetime.now() - timedelta(seconds=60),
}
with patch( # test-quality-ok: lowers the module-level probe timeout so the hung-call test finishes fast
"litellm.proxy.health_endpoints._health_endpoints.DB_READINESS_CHECK_TIMEOUT_SECONDS",
0.05,
):
start = time.monotonic()
with patch( # test-quality-ok: the readiness path reads the proxy-global DB client; it has no injection seam
"litellm.proxy.proxy_server.prisma_client", mock_prisma
):
result = await _db_health_readiness_check()
elapsed = time.monotonic() - start
assert result["status"] == "disconnected"
assert elapsed < 5
@pytest.mark.asyncio
async def test_db_health_readiness_check_overall_deadline_bounds_hung_reconnect():
"""
The whole probe-path DB check (initial check + reconnect + re-check,
including reconnect lock waits) runs under one deadline, so a reconnect
that hangs on the lock still returns disconnected within the deadline.
"""
from litellm.proxy.health_endpoints._health_endpoints import (
_db_health_readiness_check,
)
async def hang(**kwargs):
await asyncio.sleep(60)
mock_prisma = MagicMock()
mock_prisma.health_check = AsyncMock(side_effect=httpx.ConnectError("down"))
mock_prisma.attempt_db_reconnect = hang
_health_endpoints_module.db_health_cache = {
"status": "unknown",
"last_updated": datetime.now() - timedelta(seconds=60),
}
with patch( # test-quality-ok: lowers the module-level probe timeout so the hung-call test finishes fast
"litellm.proxy.health_endpoints._health_endpoints.DB_READINESS_PROBE_DEADLINE_SECONDS",
0.05,
):
start = time.monotonic()
with patch( # test-quality-ok: the readiness path reads the proxy-global DB client; it has no injection seam
"litellm.proxy.proxy_server.prisma_client", mock_prisma
):
result = await _db_health_readiness_check()
elapsed = time.monotonic() - start
assert result["status"] == "disconnected"
assert elapsed < 5
@pytest.mark.asyncio
async def test_health_readiness_returns_200_when_db_connected():
"""Happy path: connected DB keeps the legacy 200."""
@ -2746,13 +2906,13 @@ def test_test_model_connection_accepts_image_edit_mode(monkeypatch):
app = FastAPI()
app.include_router(_health_endpoints_module.router)
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN
)
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
client = TestClient(app)
with (
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), # test-quality-ok: the endpoint reads the proxy-global DB client and 500s when it is None; it has no injection seam
patch( # test-quality-ok: the endpoint reads the proxy-global DB client and 500s when it is None; it has no injection seam
"litellm.proxy.proxy_server.prisma_client", MagicMock()
),
respx.mock(assert_all_called=True) as respx_mock,
):
respx_mock.post(host="api.openai.com", path="/v1/images/edits").respond(

94
tests/test_litellm/proxy/test_litellm_pre_call_utils.py
View file

@ -1009,11 +1009,10 @@ async def test_add_litellm_data_to_request_strips_user_control_fields():
"_guardrail_pipelines",
"_pipeline_managed_guardrails",
}
for metadata_key in ("metadata", "litellm_metadata"):
cleaned_metadata = updated.get(metadata_key) or {}
for stripped_key in stripped_keys:
assert stripped_key not in cleaned_metadata
assert cleaned_metadata.get("safe_user_metadata") == "kept"
assert "litellm_metadata" not in updated
for stripped_key in stripped_keys:
assert stripped_key not in updated["metadata"]
assert updated["metadata"]["safe_user_metadata"] == "kept"
requester_metadata = updated["metadata"]["requester_metadata"]
for stripped_key in stripped_keys:
@ -1576,10 +1575,7 @@ async def test_add_litellm_data_to_request_allows_redaction_opt_out_with_admin_o
header.lower()
for header in updated["proxy_server_request"]["body"]["metadata"]["headers"]
}
assert "litellm-disable-message-redaction" in {
header.lower()
for header in (updated.get("litellm_metadata") or {}).get("headers", {})
}
assert "litellm_metadata" not in updated
@pytest.mark.asyncio
@ -6658,9 +6654,9 @@ async def test_add_litellm_data_to_request_strips_caller_supplied_callback_crede
assert "gcs_bucket_name" not in updated
assert updated["dd_api_key"] == "team-dd-key"
assert updated[TRUSTED_CALLBACK_VARS_FIELD] == {"dd_api_key": "team-dd-key"}
for metadata_key in ("metadata", "litellm_metadata"):
assert "dd_site" not in updated[metadata_key]
assert "dd_agent_host" not in updated[metadata_key]
assert "litellm_metadata" not in updated
assert "dd_site" not in updated["metadata"]
assert "dd_agent_host" not in updated["metadata"]
assert "dd_site" not in updated["litellm_params"]["metadata"]
assert updated["metadata"]["safe_user_metadata"] == "kept"
@ -7510,10 +7506,10 @@ async def test_add_litellm_data_to_request_strips_router_reserved_stamps_from_bo
version="test-version",
)
for bucket in ("metadata", "litellm_metadata"):
assert "attempted_fallbacks" not in updated[bucket]
assert "original_model_group" not in updated[bucket]
assert updated[bucket]["client_key"] == "client_value"
assert "litellm_metadata" not in updated
assert "attempted_fallbacks" not in updated["metadata"]
assert "original_model_group" not in updated["metadata"]
assert updated["metadata"]["client_key"] == "client_value"
@pytest.mark.asyncio
@ -7535,10 +7531,10 @@ async def test_add_litellm_data_to_request_strips_router_reserved_stamps_from_js
version="test-version",
)
assert isinstance(updated["litellm_metadata"], dict)
assert "attempted_fallbacks" not in updated["litellm_metadata"]
assert "original_model_group" not in updated["litellm_metadata"]
assert updated["litellm_metadata"]["client_key"] == "client_value"
assert "litellm_metadata" not in updated
assert "attempted_fallbacks" not in updated["metadata"]
assert "original_model_group" not in updated["metadata"]
assert updated["metadata"]["client_key"] == "client_value"
@pytest.mark.asyncio
@ -7562,9 +7558,10 @@ async def test_add_litellm_data_to_request_strips_router_reserved_stamps_despite
version="test-version",
)
assert updated["litellm_metadata"]["model_info"] == {"input_cost_per_token": 0.0}
assert "attempted_fallbacks" not in updated["litellm_metadata"]
assert "original_model_group" not in updated["litellm_metadata"]
assert "litellm_metadata" not in updated
assert updated["metadata"]["model_info"] == {"input_cost_per_token": 0.0}
assert "attempted_fallbacks" not in updated["metadata"]
assert "original_model_group" not in updated["metadata"]
@pytest.mark.asyncio
@ -7601,7 +7598,8 @@ async def test_router_keeps_proxy_metadata_bucket_identity_after_reserved_stamp_
litellm_metadata made the router hand downstream a scrubbed copy, so the proxy's
post_call write-backs (guardrail telemetry, applied guardrails) landed in a dict the
spend row never read. After the boundary strip plus the in-place scrub, the object the
router forwards is the proxy's own request_data bucket."""
router forwards is the proxy's own request_data bucket; on chat routes that bucket is
``metadata``, since the boundary folds client ``litellm_metadata`` into it."""
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
data = {
@ -7617,7 +7615,9 @@ async def test_router_keeps_proxy_metadata_bucket_identity_after_reserved_stamp_
general_settings={},
version="test-version",
)
proxy_bucket = request_data["litellm_metadata"]
proxy_bucket = request_data["metadata"]
assert "attempted_fallbacks" not in proxy_bucket
assert "original_model_group" not in proxy_bucket
router = litellm.Router(
model_list=[
{
@ -7630,7 +7630,7 @@ async def test_router_keeps_proxy_metadata_bucket_identity_after_reserved_stamp_
original_acompletion = router._acompletion
async def _spy(*args, **spy_kwargs):
forwarded_buckets.append(spy_kwargs["litellm_metadata"])
forwarded_buckets.append(spy_kwargs["metadata"])
return await original_acompletion(*args, **spy_kwargs)
router._acompletion = _spy
@ -7639,7 +7639,45 @@ async def test_router_keeps_proxy_metadata_bucket_identity_after_reserved_stamp_
assert forwarded_buckets == [proxy_bucket]
assert forwarded_buckets[0] is proxy_bucket
assert "attempted_fallbacks" not in proxy_bucket
assert "original_model_group" not in proxy_bucket
assert proxy_bucket["attempted_fallbacks"] == 0
assert proxy_bucket.get("original_model_group") != "spoofed-group"
proxy_bucket["standard_logging_guardrail_information"] = [{"guardrail_name": "postcall-guard"}]
assert forwarded_buckets[0]["standard_logging_guardrail_information"] == [{"guardrail_name": "postcall-guard"}]
@pytest.mark.asyncio
async def test_add_litellm_data_to_request_folds_litellm_metadata_into_metadata_on_chat_routes():
data = {
"model": "gpt-3.5-turbo",
"metadata": {"tags": ["from-metadata"]},
"litellm_metadata": {"trace_id": "abc", "tags": ["from-litellm-metadata"]},
}
updated = await add_litellm_data_to_request(
data=data,
request=_make_chat_request_mock(),
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
proxy_config=MagicMock(),
general_settings={},
version="test-version",
)
assert "litellm_metadata" not in updated
assert updated["metadata"]["trace_id"] == "abc"
assert updated["metadata"]["tags"] == ["from-metadata", "from-litellm-metadata"]
@pytest.mark.asyncio
async def test_add_litellm_data_to_request_keeps_litellm_metadata_on_litellm_metadata_routes():
data = {"model": "claude-sonnet-5", "litellm_metadata": {"trace_id": "abc"}}
updated = await add_litellm_data_to_request(
data=data,
request=_make_request_mock("/v1/messages", {"Content-Type": "application/json"}),
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
proxy_config=MagicMock(),
general_settings={},
version="test-version",
)
assert updated["litellm_metadata"]["trace_id"] == "abc"

9
tests/test_litellm/proxy/test_pricing_field_strip.py
View file

@ -285,8 +285,8 @@ async def test_add_litellm_data_to_request_skips_strip_with_key_opt_in():
async def test_add_litellm_data_to_request_strips_json_string_litellm_metadata():
"""``litellm_metadata`` may arrive as a JSON-encoded string (multipart/
form-data or ``extra_body``). The strip has to run after the proxy parses
it into a dict; otherwise the ``isinstance(dict)`` guard skips the field
and ``model_info`` survives the strip via the string path.
it into a dict but before the chat-route fold into ``metadata``; otherwise
``model_info`` survives via the string path and lands in the folded bucket.
"""
import json
@ -305,9 +305,8 @@ async def test_add_litellm_data_to_request_strips_json_string_litellm_metadata()
version="test-version",
)
parsed_metadata = updated.get("litellm_metadata")
assert isinstance(parsed_metadata, dict)
assert "model_info" not in parsed_metadata
assert "litellm_metadata" not in updated
assert "model_info" not in updated["metadata"]
@pytest.mark.asyncio

63
tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py
View file

@ -89,9 +89,7 @@ async def test_run_reconnect_cycle_direct_path_recreates_when_probe_fails(
prisma_client._cleanup_engine_watcher = MagicMock()
writer = MagicMock()
writer.query_raw = AsyncMock(
side_effect=[ConnectionError("probe failed"), [{"?column?": 1}]]
)
writer.query_raw = AsyncMock(side_effect=[ConnectionError("probe failed"), [{"?column?": 1}]])
monkeypatch.setattr(
PrismaClient,
"writer_db",
@ -171,9 +169,7 @@ async def test_run_reconnect_cycle_passes_writer_generation_to_recreate(
writer = MagicMock()
writer._engine_generation = 7
writer.query_raw = AsyncMock(
side_effect=[ConnectionError("probe failed"), [{"?column?": 1}]]
)
writer.query_raw = AsyncMock(side_effect=[ConnectionError("probe failed"), [{"?column?": 1}]])
monkeypatch.setattr(
PrismaClient,
"writer_db",
@ -229,9 +225,7 @@ async def test_attempt_reconnect_inside_lock_runs_cycle_and_resets_counter(
prisma_client._consecutive_reconnect_failures = 2
prisma_client._run_reconnect_cycle = AsyncMock()
ok = await prisma_client._attempt_reconnect_inside_lock(
force=True, reason="test", timeout_seconds=1
)
ok = await prisma_client._attempt_reconnect_inside_lock(force=True, reason="test", timeout_seconds=1)
pinned = {
"returned": ok,
"cycle_called": prisma_client._run_reconnect_cycle.await_count,
@ -254,9 +248,7 @@ async def test_attempt_reconnect_inside_lock_skips_when_in_cooldown(
prisma_client._db_last_reconnect_attempt_ts = time.time()
prisma_client._run_reconnect_cycle = AsyncMock()
ok = await prisma_client._attempt_reconnect_inside_lock(
force=False, reason="test", timeout_seconds=1
)
ok = await prisma_client._attempt_reconnect_inside_lock(force=False, reason="test", timeout_seconds=1)
assert ok is False
assert prisma_client._run_reconnect_cycle.await_count == 0
@ -269,9 +261,7 @@ async def test_attempt_reconnect_inside_lock_increments_failure_counter_on_error
prisma_client._consecutive_reconnect_failures = 0
prisma_client._run_reconnect_cycle = AsyncMock(side_effect=RuntimeError("boom"))
ok = await prisma_client._attempt_reconnect_inside_lock(
force=True, reason="failing_test", timeout_seconds=1
)
ok = await prisma_client._attempt_reconnect_inside_lock(force=True, reason="failing_test", timeout_seconds=1)
assert ok is False
assert prisma_client._consecutive_reconnect_failures == 1
@ -316,9 +306,7 @@ async def test_attempt_db_reconnect_lock_timeout_returns_false(
by replacing ``asyncio.wait`` with a callable that returns the loser
task as still-pending after it's already been completed elsewhere.
"""
completed_task: asyncio.Task[bool] = asyncio.get_running_loop().create_task(
_no_op_returning_true()
)
completed_task: asyncio.Task[bool] = asyncio.get_running_loop().create_task(_no_op_returning_true())
# Ensure the inner task has finished before attempt_db_reconnect sees it.
await completed_task
@ -329,7 +317,7 @@ async def test_attempt_db_reconnect_lock_timeout_returns_false(
monkeypatch.setattr(
asyncio,
"create_task",
lambda coro, *a, **kw: (coro.close() or completed_task),
lambda coro, *a, **kw: coro.close() or completed_task,
)
prisma_client._db_last_reconnect_attempt_ts = 0.0
@ -465,9 +453,7 @@ async def test_db_health_watchdog_loop_triggers_reconnect_on_timeout(
await prisma_client._db_health_watchdog_loop()
pinned = {
"reconnect_called": prisma_client.attempt_db_reconnect.await_count,
"reconnect_reason": prisma_client.attempt_db_reconnect.await_args.kwargs[
"reason"
],
"reconnect_reason": prisma_client.attempt_db_reconnect.await_args.kwargs["reason"],
"wait_for_calls": call_count["n"],
"loop_exited_clean": True,
}
@ -522,10 +508,7 @@ async def test_iam_refresh_racing_reconnect_recreates_engine_only_once(
from litellm.proxy.db.prisma_client import PrismaWrapper
def token_db_url(created: datetime) -> str:
token = (
f"host/?X-Amz-Date={created.strftime('%Y%m%dT%H%M%SZ')}"
f"&X-Amz-Expires=900&X-Amz-Signature=abc"
)
token = f"host/?X-Amz-Date={created.strftime('%Y%m%dT%H%M%SZ')}&X-Amz-Expires=900&X-Amz-Signature=abc"
return f"postgresql://user:{urllib.parse.quote(token, safe='')}@host:5432/db"
# Old engine (PID 111) carries an expired token; in-flight queries on it
@ -577,9 +560,7 @@ async def test_iam_refresh_racing_reconnect_recreates_engine_only_once(
# In-flight transport-error path fires while the refresh holds the
# wrapper's reconnection lock mid-recreate.
reconnect_task = asyncio.create_task(
prisma_client.attempt_db_reconnect(
reason="in_flight_transport_error", force=True
)
prisma_client.attempt_db_reconnect(reason="in_flight_transport_error", force=True)
)
await asyncio.sleep(0.05)
release_connect.set()
@ -1096,3 +1077,27 @@ async def test_unrelated_reconnect_failure_does_not_erase_the_burst_record(
"cycles_after": prisma_client._run_reconnect_cycle.await_count,
}
assert pinned == {"cycles_before": 2, "cycles_after": 2}
@pytest.mark.asyncio
async def test_attempt_db_reconnect_cancelled_while_waiting_does_not_strand_lock(
prisma_client: PrismaClient,
) -> None:
"""A reconnect cancelled while waiting on the lock (e.g. the readiness
probe deadline firing) must abandon its lock-acquisition task instead of
leaving it to grab the lock later with no owner to release it."""
prisma_client._db_last_reconnect_attempt_ts = 0.0
prisma_client._attempt_reconnect_inside_lock = AsyncMock(return_value=True)
await prisma_client._db_reconnect_lock.acquire()
waiting_reconnect: Final = asyncio.create_task(
prisma_client.attempt_db_reconnect(reason="probe_deadline", lock_timeout_seconds=30.0)
)
await asyncio.sleep(0.05)
waiting_reconnect.cancel()
with pytest.raises(asyncio.CancelledError):
await waiting_reconnect
prisma_client._db_reconnect_lock.release()
await asyncio.sleep(0.05)
assert prisma_client._db_reconnect_lock.locked() is False

75
tests/test_litellm/router_strategy/test_router_tag_routing.py
View file

@ -3072,3 +3072,78 @@ async def test_non_router_tags_still_pick_the_matching_tier_deployment():
)
assert response._hidden_params["model_id"] == "tier-gemini-flash-us"
def _chat_completions_request_mock():
from unittest.mock import MagicMock
from fastapi import Request
request_mock = MagicMock(spec=Request)
request_mock.url = MagicMock()
request_mock.url.path = "/v1/chat/completions"
request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions"
request_mock.method = "POST"
request_mock.query_params = {}
request_mock.headers = {"Content-Type": "application/json"}
request_mock.client = MagicMock()
request_mock.client.host = "127.0.0.1"
return request_mock
def _team_a_and_default_router():
return litellm.Router(
model_list=[
{
"model_name": "gpt-5.4-mini",
"litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "mock", "tags": ["team-a"]},
"model_info": {"id": "team-a-deployment"},
},
{
"model_name": "gpt-5.4-mini",
"litellm_params": {"model": "openai/gpt-5.4-nano", "api_key": "mock", "tags": ["default"]},
"model_info": {"id": "default-deployment"},
},
],
enable_tag_filtering=True,
)
@pytest.mark.asyncio()
@pytest.mark.parametrize(
"team_metadata,body_extra",
[
({"tags": ["team-a"]}, {}),
({}, {"tags": ["team-a"]}),
],
ids=["team-tags", "body-tags"],
)
async def test_chat_request_carrying_litellm_metadata_still_routes_on_proxy_merged_tags(team_metadata, body_extra):
from unittest.mock import MagicMock
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
router = _team_a_and_default_router()
data = {
"model": "gpt-5.4-mini",
"messages": [{"role": "user", "content": "hi"}],
"litellm_metadata": {"trace_id": "abc"},
**body_extra,
}
request_kwargs = await add_litellm_data_to_request(
data=data,
request=_chat_completions_request_mock(),
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key", metadata={}, team_metadata=team_metadata),
proxy_config=MagicMock(),
general_settings={},
version="test-version",
)
deployment = await router.async_get_available_deployment(
model="gpt-5.4-mini",
request_kwargs=request_kwargs,
messages=request_kwargs["messages"],
)
assert deployment["model_info"]["id"] == "team-a-deployment"