From 88b3cfd789b5726af3acb7387dc7acb54dbacd85 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 29 Aug 2026 00:55:43 +0000 Subject: [PATCH 01/33] fix(mcp): bind per-user OAuth credentials to the authenticated LiteLLM caller Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/discoverable_endpoints.py | 20 +- .../mcp_server/mcp_server_manager.py | 2 + .../mcp_server/oauth_identity_binding.py | 249 ++++++++++++++++++ .../mcp_management_endpoints.py | 22 ++ .../types/mcp_server/mcp_server_manager.py | 21 ++ .../mcp_server/test_oauth_identity_binding.py | 240 +++++++++++++++++ .../test_mcp_management_endpoints.py | 61 +++++ 7 files changed, 614 insertions(+), 1 deletion(-) create mode 100644 litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_identity_binding.py diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 93b85edd88d..370dee77083 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -54,6 +54,9 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( relative_request_url, revoke_refresh_token, ) +from litellm.proxy._experimental.mcp_server.oauth_identity_binding import ( + enforce_oauth_identity_binding, +) from litellm.proxy._experimental.mcp_server.oauth_utils import ( TOKEN_NO_CACHE_HEADERS, build_upstream_oauth2_token_request, @@ -1133,11 +1136,26 @@ async def exchange_token_with_server( server_id=resolved_server.server_id, ) + # Bind the exchanged token to the LiteLLM caller BEFORE it is returned, stored, or cached, so a + # token minted for a different upstream principal never becomes usable under the caller's user_id. + resolved_user_id: Final = ( + await _extract_user_id_from_request(request) + if resolved_server.needs_user_oauth_token or resolved_server.oauth_identity_binding is not None + else None + ) + if isinstance(token_response, dict): + await enforce_oauth_identity_binding( + server=resolved_server, + token_response=token_response, + litellm_user_id=resolved_user_id, + grant_type=grant_type, + ) + # Store server-side when the server is configured for per-user OAuth and # the calling client has provided a valid LiteLLM identity. # Errors are non-fatal: the token is still returned to the client. if resolved_server.needs_user_oauth_token: - user_id: Final = await _extract_user_id_from_request(request) + user_id: Final = resolved_user_id if user_id: try: await _store_per_user_token_server_side( diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 2330120adad..1d61aba8bc5 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2174,6 +2174,8 @@ class MCPServerManager: allow_elicitation=bool(server_config.get("allow_elicitation", False)), timeout=server_config.get("timeout", None), max_concurrent_requests=server_config.get("max_concurrent_requests", None), + token_validation=server_config.get("token_validation", None), + oauth_identity_binding=server_config.get("oauth_identity_binding", None), ) self._assign_unique_short_prefix(new_server) _warn_internal_delegate_pkce_if_applicable(new_server, source="config") diff --git a/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py b/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py new file mode 100644 index 00000000000..f132f59ae99 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py @@ -0,0 +1,249 @@ +"""Per-user OAuth identity binding: verify the upstream OIDC principal matches the LiteLLM caller. + +Closes the confused-deputy gap where a browser authenticated upstream as one principal produces a +token that the relay stores under a different, LiteLLM-authenticated principal: before the token +endpoint returns, stores, or caches an exchanged token for an identity-bound server, the id_token +is validated (signature via the pinned issuer's JWKS, issuer, audience, expiry) and its principal +claim is compared to the caller's trusted LiteLLM identity. Mismatches fail closed in enforce mode +and are logged in audit mode. +""" + +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass +from typing import Final, Literal + +import jwt +from fastapi import HTTPException + +from litellm._logging import verbose_logger +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.types.llms.custom_http import httpxSpecialProvider +from litellm.types.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer + +_ALLOWED_ID_TOKEN_ALGORITHMS: Final = ( + "RS256", + "RS384", + "RS512", + "ES256", + "ES384", + "ES512", + "PS256", + "PS384", + "PS512", +) +_JWKS_CACHE_TTL_SECONDS: Final = 3600 +_jwks_cache: Final = InMemoryCache(default_ttl=_JWKS_CACHE_TTL_SECONDS) + +JwksFetcher = Callable[[MCPOAuthIdentityBinding], Awaitable[list[Mapping[str, object]]]] +CallerPrincipalLoader = Callable[[str, MCPOAuthIdentityBinding], Awaitable[str | None]] + +_RejectionCode = Literal["oauth_principal_mismatch", "oauth_identity_binding_failed"] + + +@dataclass(frozen=True, slots=True) +class _BindingRejection: + code: _RejectionCode + description: str + + +async def _fetch_issuer_jwks(binding: MCPOAuthIdentityBinding) -> list[Mapping[str, object]]: + jwks_url: Final[str] = binding.jwks_url or await _discover_jwks_url(binding.issuer) + cached: Final = await _jwks_cache.async_get_cache(jwks_url) + if isinstance(cached, list): + return cached + client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check) + response: Final = await client.get(jwks_url) + response.raise_for_status() + document: Final = response.json() + keys: Final = document.get("keys") if isinstance(document, dict) else None + if not isinstance(keys, list): + raise TypeError(f"JWKS document at {jwks_url} has no 'keys' array") + await _jwks_cache.async_set_cache(jwks_url, keys, ttl=_JWKS_CACHE_TTL_SECONDS) + return keys + + +async def _discover_jwks_url(issuer: str) -> str: + discovery_url: Final = f"{issuer.rstrip('/')}/.well-known/openid-configuration" + client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check) + response: Final = await client.get(discovery_url) + response.raise_for_status() + metadata: Final = response.json() + jwks_uri: Final = metadata.get("jwks_uri") if isinstance(metadata, dict) else None + if not isinstance(jwks_uri, str) or not jwks_uri: + raise ValueError(f"OIDC discovery at {discovery_url} returned no jwks_uri") + return jwks_uri + + +def _select_signing_key(id_token: str, keys: list[Mapping[str, object]]) -> "jwt.PyJWK | _BindingRejection": + header: Final = jwt.get_unverified_header(id_token) + kid: Final = header.get("kid") + for key in keys: + if kid is None or key.get("kid") == kid: + return jwt.PyJWK(dict(key)) + return _BindingRejection( + code="oauth_identity_binding_failed", + description=f"id_token signing key (kid={kid!r}) not found in the issuer's JWKS", + ) + + +def _decode_id_token( + id_token: str, + binding: MCPOAuthIdentityBinding, + signing_key: "jwt.PyJWK", +) -> "Mapping[str, object] | _BindingRejection": + try: + return jwt.decode( + id_token, + signing_key.key, + algorithms=list(_ALLOWED_ID_TOKEN_ALGORITHMS), + issuer=binding.issuer, + audience=binding.audiences if binding.audiences else None, + options={ + "require": ["iss", "exp"], + "verify_aud": bool(binding.audiences), + }, + ) + except jwt.InvalidTokenError as exc: + return _BindingRejection( + code="oauth_identity_binding_failed", + description=f"id_token validation failed: {exc}", + ) + + +def _upstream_principal( + claims: Mapping[str, object], + binding: MCPOAuthIdentityBinding, +) -> "str | _BindingRejection": + principal: Final = claims.get(binding.principal_claim) + if not isinstance(principal, str) or not principal: + return _BindingRejection( + code="oauth_identity_binding_failed", + description=f"id_token has no usable '{binding.principal_claim}' claim", + ) + if binding.principal_claim == "email" and binding.require_email_verified and claims.get("email_verified") is not True: + return _BindingRejection( + code="oauth_identity_binding_failed", + description="id_token email is not verified (email_verified is not true)", + ) + return principal + + +async def _load_caller_principal(litellm_user_id: str, binding: MCPOAuthIdentityBinding) -> str | None: + if binding.caller_field == "user_id": + return litellm_user_id + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( # noqa: PLC0415 # inline import avoids a module-load circular import + load_active_user_by_id, + ) + + loaded: Final = await load_active_user_by_id(litellm_user_id) + if isinstance(loaded, str): + return None + return loaded.user_email + + +def _principals_match(upstream: str, caller: str, binding: MCPOAuthIdentityBinding) -> bool: + if binding.principal_claim == "email" or binding.caller_field == "user_email": + return upstream.strip().casefold() == caller.strip().casefold() + return upstream == caller + + +async def _evaluate_binding( + binding: MCPOAuthIdentityBinding, + token_response: Mapping[str, object], + litellm_user_id: str | None, + grant_type: str, + jwks_fetcher: JwksFetcher, + caller_principal_loader: CallerPrincipalLoader, +) -> _BindingRejection | None: + id_token: Final = token_response.get("id_token") + if not isinstance(id_token, str) or not id_token: + if grant_type == "refresh_token": + return None + return _BindingRejection( + code="oauth_identity_binding_failed", + description="the upstream token response carries no id_token to bind the credential to a principal", + ) + if not litellm_user_id: + return _BindingRejection( + code="oauth_identity_binding_failed", + description="the request carries no resolvable LiteLLM user identity to bind the credential to", + ) + try: + keys: Final = await jwks_fetcher(binding) + except Exception as exc: # noqa: BLE001 # a JWKS fetch failure must fail closed, not surface as a 500 + return _BindingRejection( + code="oauth_identity_binding_failed", + description=f"could not fetch the issuer's JWKS: {exc}", + ) + signing_key: Final = _select_signing_key(id_token, keys) + if isinstance(signing_key, _BindingRejection): + return signing_key + claims: Final = _decode_id_token(id_token, binding, signing_key) + if isinstance(claims, _BindingRejection): + return claims + upstream: Final = _upstream_principal(claims, binding) + if isinstance(upstream, _BindingRejection): + return upstream + caller: Final = await caller_principal_loader(litellm_user_id, binding) + if not caller: + return _BindingRejection( + code="oauth_identity_binding_failed", + description=f"the LiteLLM user has no '{binding.caller_field}' to compare the upstream principal against", + ) + if not _principals_match(upstream, caller, binding): + return _BindingRejection( + code="oauth_principal_mismatch", + description="The browser account does not match the selected credential owner.", + ) + return None + + +async def enforce_oauth_identity_binding( + server: MCPServer, + token_response: Mapping[str, object], + litellm_user_id: str | None, + grant_type: str, + jwks_fetcher: JwksFetcher = _fetch_issuer_jwks, + caller_principal_loader: CallerPrincipalLoader = _load_caller_principal, +) -> None: + """Validate the exchanged token's upstream principal against the LiteLLM caller. + + No-op when the server has no binding or it is disabled. In enforce mode a failure raises 403 + before the caller returns, stores, or caches the token; in audit mode failures are logged only. + A refresh_token grant without an id_token is allowed in both modes: the stored credential keeps + the binding established at the original authorization_code exchange. + """ + binding: Final = server.oauth_identity_binding + if binding is None or binding.mode == "disabled": + return + rejection: Final = await _evaluate_binding( + binding=binding, + token_response=token_response, + litellm_user_id=litellm_user_id, + grant_type=grant_type, + jwks_fetcher=jwks_fetcher, + caller_principal_loader=caller_principal_loader, + ) + if rejection is None: + return + if binding.mode == "audit": + verbose_logger.warning( + "oauth_identity_binding audit: server=%s user=%s grant=%s rejected=%s (%s)", + server.server_id, + litellm_user_id, + grant_type, + rejection.code, + rejection.description, + ) + return + raise HTTPException( + status_code=403, + detail={ + "error": rejection.code, + "error_description": rejection.description, + "server_id": server.server_id, + "credential_owner": "caller", + "credential_stored": False, + }, + ) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 556a30d0b29..5e68e2440da 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -2135,6 +2135,28 @@ if MCP_AVAILABLE: """Persist the OAuth2 access token obtained by the calling user.""" prisma_client: Final = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") await _authorize_and_fetch_mcp_server(prisma_client, user_api_key_dict, server_id) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 + global_mcp_server_manager as _manager, + ) + + # This endpoint accepts an opaque token with no upstream identity validation, so it must be + # closed for identity-bound servers or it becomes a bypass of the token-relay binding check. + registry_server: Final = _manager.get_mcp_server_by_id(server_id) + binding: Final = registry_server.oauth_identity_binding if registry_server else None + if binding is not None and binding.mode == "enforce": + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ + "error": "oauth_identity_binding_enforced", + "error_description": ( + "Direct credential storage is disabled for this server: its OAuth identity " + "binding is enforced and this endpoint cannot validate the token's principal. " + "Complete the OAuth flow through the gateway instead." + ), + "server_id": server_id, + "credential_stored": False, + }, + ) user_id: Final = user_api_key_dict.user_id or "" if not user_id: raise HTTPException( diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 9bf3acc601c..cc3f6c8e8d8 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -38,6 +38,26 @@ class MCPOAuthMetadata(BaseModel): usable in memory but must never be persisted as configuration.""" +class MCPOAuthIdentityBinding(BaseModel): + """Per-server policy binding stored per-user OAuth credentials to the authenticated LiteLLM caller. + + When enabled for an interactive oauth2 server, the token relay validates the upstream OIDC + ``id_token`` (signature via the pinned issuer's JWKS, issuer, audience, expiry) and compares its + principal claim to the LiteLLM caller's trusted identity before the token is returned, stored, + or cached. ``audit`` logs mismatches without changing behavior; ``enforce`` fails closed with + 403 ``oauth_principal_mismatch`` and disables the direct ``oauth-user-credential`` POST, which + would otherwise bypass validation with an arbitrary opaque token. + """ + + mode: Literal["disabled", "audit", "enforce"] = "disabled" + issuer: str + jwks_url: str | None = None + audiences: list[str] = [] + principal_claim: str = "email" + caller_field: Literal["user_email", "user_id"] = "user_email" + require_email_verified: bool = True + + class MCPServer(BaseModel): server_id: str name: str @@ -172,6 +192,7 @@ class MCPServer(BaseModel): # response (supports dot-notation for nested fields, e.g. "team.enterprise_id"). # Tokens that fail validation are rejected before storage. token_validation: dict[str, Any] | None = None + oauth_identity_binding: MCPOAuthIdentityBinding | None = None # Optional TTL override (seconds) for the Redis per-user token cache, capped # at the token's expires_in minus the expiry buffer so a cached entry never # outlives the token. Defaults to the token's expires_in minus the expiry diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_identity_binding.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_identity_binding.py new file mode 100644 index 00000000000..69eaabeda4a --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_identity_binding.py @@ -0,0 +1,240 @@ +import time +from collections.abc import Mapping +from typing import Final + +import jwt +import pytest +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import rsa +from fastapi import HTTPException + +from litellm.proxy._experimental.mcp_server.oauth_identity_binding import ( + enforce_oauth_identity_binding, +) +from litellm.types.mcp import MCPTransport +from litellm.types.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer + +ISSUER: Final = "https://idp.example.com" +AUDIENCE: Final = "litellm-client" +KID: Final = "test-key" + +_PRIVATE_KEY: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) +_PRIVATE_PEM: Final = _PRIVATE_KEY.private_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=serialization.NoEncryption(), +) +_PUBLIC_JWK: Final = { + **jwt.algorithms.RSAAlgorithm.to_jwk(_PRIVATE_KEY.public_key(), as_dict=True), + "kid": KID, + "alg": "RS256", + "use": "sig", +} + + +def _sign_id_token(claims: Mapping[str, object]) -> str: + payload: Final = { + "iss": ISSUER, + "aud": AUDIENCE, + "exp": int(time.time()) + 300, + "iat": int(time.time()), + **claims, + } + return jwt.encode(payload, _PRIVATE_PEM, algorithm="RS256", headers={"kid": KID}) + + +async def _jwks_fetcher(_binding: MCPOAuthIdentityBinding) -> list[Mapping[str, object]]: + return [_PUBLIC_JWK] + + +def _caller_loader(email: str | None): + async def load(_user_id: str, _binding: MCPOAuthIdentityBinding) -> str | None: + return email + + return load + + +def _server(mode: str = "enforce", **binding_overrides: object) -> MCPServer: + return MCPServer( + server_id="srv-1", + name="srv-1", + url="https://mcp.example.com", + transport=MCPTransport.http, + oauth_identity_binding=MCPOAuthIdentityBinding( + mode=mode, + issuer=ISSUER, + audiences=[AUDIENCE], + **binding_overrides, + ), + ) + + +@pytest.mark.asyncio +async def test_matching_principal_passes(): + token: Final = _sign_id_token({"email": "Alice@Example.com", "email_verified": True}) + result: Final = await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at", "id_token": token}, + litellm_user_id="user-a", + grant_type="authorization_code", + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + ) + assert result is None + + +@pytest.mark.asyncio +async def test_mismatched_principal_rejected(): + token: Final = _sign_id_token({"email": "mallory@example.com", "email_verified": True}) + with pytest.raises(HTTPException) as exc_info: + await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at", "id_token": token}, + litellm_user_id="user-a", + grant_type="authorization_code", + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + ) + assert exc_info.value.status_code == 403 + assert exc_info.value.detail["error"] == "oauth_principal_mismatch" + assert exc_info.value.detail["credential_stored"] is False + + +@pytest.mark.asyncio +async def test_missing_id_token_rejected_on_authorization_code(): + with pytest.raises(HTTPException) as exc_info: + await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at"}, + litellm_user_id="user-a", + grant_type="authorization_code", + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + ) + assert exc_info.value.status_code == 403 + assert exc_info.value.detail["error"] == "oauth_identity_binding_failed" + + +@pytest.mark.asyncio +async def test_refresh_without_id_token_allowed(): + result: Final = await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at"}, + litellm_user_id="user-a", + grant_type="refresh_token", + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + ) + assert result is None + + +@pytest.mark.asyncio +async def test_refresh_with_mismatched_id_token_rejected(): + token: Final = _sign_id_token({"email": "mallory@example.com", "email_verified": True}) + with pytest.raises(HTTPException): + await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at", "id_token": token}, + litellm_user_id="user-a", + grant_type="refresh_token", + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + ) + + +@pytest.mark.asyncio +async def test_audit_mode_logs_but_does_not_reject(): + token: Final = _sign_id_token({"email": "mallory@example.com", "email_verified": True}) + result: Final = await enforce_oauth_identity_binding( + server=_server(mode="audit"), + token_response={"access_token": "at", "id_token": token}, + litellm_user_id="user-a", + grant_type="authorization_code", + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + ) + assert result is None + + +@pytest.mark.asyncio +async def test_unverified_email_rejected(): + token: Final = _sign_id_token({"email": "alice@example.com", "email_verified": False}) + with pytest.raises(HTTPException) as exc_info: + await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at", "id_token": token}, + litellm_user_id="user-a", + grant_type="authorization_code", + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + ) + assert exc_info.value.detail["error"] == "oauth_identity_binding_failed" + + +@pytest.mark.asyncio +async def test_wrong_issuer_rejected(): + payload: Final = { + "iss": "https://evil.example.com", + "aud": AUDIENCE, + "exp": int(time.time()) + 300, + "email": "alice@example.com", + "email_verified": True, + } + token: Final = jwt.encode(payload, _PRIVATE_PEM, algorithm="RS256", headers={"kid": KID}) + with pytest.raises(HTTPException) as exc_info: + await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at", "id_token": token}, + litellm_user_id="user-a", + grant_type="authorization_code", + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + ) + assert exc_info.value.detail["error"] == "oauth_identity_binding_failed" + + +@pytest.mark.asyncio +async def test_no_litellm_identity_rejected(): + token: Final = _sign_id_token({"email": "alice@example.com", "email_verified": True}) + with pytest.raises(HTTPException) as exc_info: + await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at", "id_token": token}, + litellm_user_id=None, + grant_type="authorization_code", + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + ) + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_disabled_binding_is_noop(): + result: Final = await enforce_oauth_identity_binding( + server=_server(mode="disabled"), + token_response={"access_token": "at"}, + litellm_user_id=None, + grant_type="authorization_code", + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader(None), + ) + assert result is None + + +@pytest.mark.asyncio +async def test_no_binding_is_noop(): + server: Final = MCPServer( + server_id="srv-2", + name="srv-2", + url="https://mcp.example.com", + transport=MCPTransport.http, + ) + result: Final = await enforce_oauth_identity_binding( + server=server, + token_response={"access_token": "at"}, + litellm_user_id=None, + grant_type="authorization_code", + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader(None), + ) + assert result is None diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index ceb44de5576..1cb36918f11 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -4638,6 +4638,67 @@ async def test_store_mcp_oauth_user_credential_returns_status(): assert result.expires_at == "2099-01-01T00:00:00+00:00" +@pytest.mark.asyncio +async def test_store_mcp_oauth_user_credential_blocked_when_identity_binding_enforced(): + """The direct opaque-token POST must be closed for enforce-mode identity-bound servers, + otherwise it bypasses the token-relay principal check.""" + from litellm.proxy._types import MCPOAuthUserCredentialRequest + from litellm.types.mcp import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer + + if not mgmt_endpoints.MCP_AVAILABLE: + pytest.skip("MCP module not installed") + + from litellm.proxy._experimental.mcp_server import mcp_server_manager as manager_module + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + store_mcp_oauth_user_credential, + ) + + server_id = "srv-binding-1" + bound_server = MCPServer( + server_id=server_id, + name=server_id, + url="https://mcp.example.com", + transport=MCPTransport.http, + oauth_identity_binding=MCPOAuthIdentityBinding(mode="enforce", issuer="https://idp.example.com"), + ) + store_mock = AsyncMock(return_value=None) + + with ( + patch( # test-quality-ok: mirrors the existing store-credential tests in this file + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=_make_prisma_client(), + ), + patch( # test-quality-ok: mirrors the existing store-credential tests in this file + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", + new=AsyncMock(return_value=generate_mock_mcp_server_db_record(server_id=server_id)), + ), + patch( # test-quality-ok: mirrors the existing store-credential tests in this file + "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", + return_value=True, + ), + patch.object( # test-quality-ok: registry is a module-level singleton; injecting it would change the endpoint signature + manager_module.global_mcp_server_manager, + "get_mcp_server_by_id", + return_value=bound_server, + ), + patch( # test-quality-ok: asserting the DB write is never reached is the point of the test + "litellm.proxy.management_endpoints.mcp_management_endpoints.store_user_oauth_credential", + new=store_mock, + ), + ): + with pytest.raises(HTTPException) as exc_info: + await store_mcp_oauth_user_credential( + server_id=server_id, + payload=MCPOAuthUserCredentialRequest(access_token="opaque-tok", expires_in=3600), + user_api_key_dict=_make_user_auth("user-123"), + ) + + assert exc_info.value.status_code == 403 + assert exc_info.value.detail["error"] == "oauth_identity_binding_enforced" + store_mock.assert_not_called() + + @pytest.mark.asyncio async def test_delete_mcp_oauth_user_credential_only_deletes_oauth(): """delete_mcp_oauth_user_credential only deletes OAuth2 credentials, not BYOK.""" From c60c60e6fe4a860377d9fcd9e611d71234b91e36 Mon Sep 17 00:00:00 2001 From: michelligabriele Date: Wed, 2 Sep 2026 10:35:49 +0200 Subject: [PATCH 02/33] fix(proxy): resolve /v1/models limits from the deployment, not the alias --- litellm/proxy/utils.py | 108 +++++++++++---- litellm/router.py | 42 ++++-- litellm/types/router.py | 17 +++ litellm/utils.py | 27 ++-- .../test_model_management_endpoints.py | 4 +- .../proxy/proxy_server/test_routes_models.py | 8 +- .../test_team_model_name_translation.py | 12 +- tests/test_litellm/proxy/test_proxy_utils.py | 127 +++++++++++++++++- tests/test_litellm/test_router.py | 99 ++++++++++++++ 9 files changed, 388 insertions(+), 56 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 051d36c4d0f..b86e6d857cc 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -7481,6 +7481,64 @@ async def get_available_models_for_user( return all_models +def _safe_get_model_info(model: str, get_model_info: Callable[[str], ModelInfo]) -> ModelInfo | None: + try: + return get_model_info(model) + except Exception as e: + verbose_proxy_logger.debug( + "create_model_info_response: cost map lookup failed for %s: %s", + model, + e, + ) + return None + + +def _resolve_listing_model_info( + deployment_model: str | None, + listed_model: str, + get_model_info: Callable[[str], ModelInfo], +) -> tuple[ModelInfo, ...]: + """ + Cost-map entries describing a listed model, best source first. + + The name a model is listed under is an arbitrary public alias, so it often misses the + cost map and lands on a fallback-generalization rule that answers with a conservative + family baseline instead of the real model's limits; the deployment's underlying model + is what the request actually reaches. Both names are kept because either can + generalize, and because a deployment's own model is registered into the cost map as a + stub that carries no limits of its own. Exact entries are consulted before generalized + ones, and each field is then taken from the first entry that has it. + """ + listed_info: Final = _safe_get_model_info(listed_model, get_model_info) + + # Fast path, and the only one a wildcard-expanded name takes: with a single name + # there is nothing to order, so skip the generalization test entirely. This keeps + # the per-model cost of the listing on the hot path #33721 exists to protect. + if deployment_model is None or deployment_model == listed_model: + return () if listed_info is None else (listed_info,) + + deployment_info: Final = _safe_get_model_info(deployment_model, get_model_info) + if deployment_info is None: + return () if listed_info is None else (listed_info,) + if listed_info is None: + return (deployment_info,) + + from litellm.utils import is_generalized_model_info + + # Both names resolved: the deployment's model leads unless it only generalized + # while the listed name is an exact cost-map entry. + if is_generalized_model_info(deployment_info) and not is_generalized_model_info(listed_info): + return (listed_info, deployment_info) + return (deployment_info, listed_info) + + +def _first_token_limit(candidates: tuple[ModelInfo, ...], field: str) -> int | None: + return next( + (limit for limit in (coerce_token_limit(info.get(field)) for info in candidates) if limit is not None), + None, + ) + + def create_model_info_response( model_id: str, provider: str, @@ -7505,31 +7563,35 @@ def create_model_info_response( "owned_by": provider, } - try: - model_cost_info: ModelInfo | None = get_model_info(model_id) - except Exception as e: - verbose_proxy_logger.debug( - "create_model_info_response: cost map lookup failed for %s: %s", - model_id, - e, - ) - model_cost_info = None + listing_info: Final = llm_router.get_model_listing_info(model_id) if llm_router is not None else None - max_input_tokens: int | None = None - max_output_tokens: int | None = None - if model_cost_info is not None: - max_input_tokens = coerce_token_limit(model_cost_info.get("max_input_tokens")) - max_output_tokens = coerce_token_limit(model_cost_info.get("max_output_tokens")) - mode: Final = model_cost_info.get("mode") - if isinstance(mode, str): - base["mode"] = mode + candidates: Final = _resolve_listing_model_info( + deployment_model=listing_info.cost_map_key if listing_info is not None else None, + listed_model=model_id, + get_model_info=get_model_info, + ) - if llm_router is not None: - configured_input, configured_output = llm_router.get_configured_token_limits(model_id) - if configured_input is not None: - max_input_tokens = configured_input - if configured_output is not None: - max_output_tokens = configured_output + max_input_tokens: int | None = _first_token_limit(candidates, "max_input_tokens") + max_output_tokens: int | None = _first_token_limit(candidates, "max_output_tokens") + mode: Final = next( + ( + m + for m in ( + cast("Mapping[str, object]", info).get("mode") # cast-ok: an entry need not carry "mode" + for info in candidates + ) + if isinstance(m, str) + ), + None, + ) + if mode is not None: + base["mode"] = mode + + if listing_info is not None: + if listing_info.max_input_tokens is not None: + max_input_tokens = listing_info.max_input_tokens + if listing_info.max_output_tokens is not None: + max_output_tokens = listing_info.max_output_tokens if max_input_tokens is not None: base["max_input_tokens"] = max_input_tokens diff --git a/litellm/router.py b/litellm/router.py index 23d8907fb49..7bc6a6c714a 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -200,6 +200,7 @@ from litellm.types.router import ( CredentialLiteLLMParams, CustomRoutingStrategyBase, Deployment, + DeploymentModelListingInfo, DeploymentTypedDict, FallbackAccessCheck, GuardrailTypedDict, @@ -9726,25 +9727,44 @@ class Router: return None return Deployment(**first_usable) if isinstance(first_usable, dict) else first_usable + def get_model_listing_info(self, model_name: str) -> DeploymentModelListingInfo | None: + """ + Return what the concrete deployment behind model_name contributes to its + /v1/models entry: the cost-map key for its underlying model, plus any token + limits explicitly configured in its model_info. Resolved via O(1) index lookup. + + Returns None for wildcard-expanded or unknown names, where the listed name is + the real model name and no deployment-specific information exists, and treats a + malformed configured limit as absent rather than failing the listing. Unlike + get_model_group_info, this never triggers pattern matching or deep copies, so it + is safe to call per listed model on the /v1/models hot path. + """ + deployment: Final = self.get_deployment_by_model_group_name(model_group_name=model_name) + if deployment is None: + return None + + model_info: Final = deployment.model_info + # base_model is a declared field, so read it as one: an unset or blank value + # means the deployment's own model name is the cost-map key. + base_model: Final = model_info.base_model + return DeploymentModelListingInfo( + cost_map_key=base_model or deployment.litellm_params.model, + max_input_tokens=coerce_token_limit(model_info.get("max_input_tokens")), + max_output_tokens=coerce_token_limit(model_info.get("max_output_tokens")), + ) + def get_configured_token_limits(self, model_name: str) -> "tuple[int | None, int | None]": """ Return (max_input_tokens, max_output_tokens) explicitly configured in a concrete deployment's model_info for model_name, via O(1) index lookup. Returns (None, None) for wildcard-expanded or unknown names, and treats a - malformed configured value as absent rather than failing the listing. Unlike - get_model_group_info, this never triggers pattern matching or deep copies, so it - is safe to call per listed model on the /v1/models hot path. + malformed configured value as absent rather than failing the caller. """ - deployment: Final = self.get_deployment_by_model_group_name(model_group_name=model_name) - if deployment is None: + listing_info: Final = self.get_model_listing_info(model_name=model_name) + if listing_info is None: return (None, None) - - model_info: Final = deployment.model_info - return ( - coerce_token_limit(model_info.get("max_input_tokens")), - coerce_token_limit(model_info.get("max_output_tokens")), - ) + return (listing_info.max_input_tokens, listing_info.max_output_tokens) def get_deployment_credentials_with_provider( self, model_id: str, team_id: str | None = None diff --git a/litellm/types/router.py b/litellm/types/router.py index e0957383aac..6f9fb4bf6ee 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -575,6 +575,23 @@ class Deployment(BaseModel): setattr(self, key, value) +@dataclass(frozen=True, slots=True) +class DeploymentModelListingInfo: + """What a concrete deployment contributes to its OpenAI-compatible listing entry. + + ``cost_map_key`` is the name the deployment's underlying model is known by in + ``litellm.model_cost`` (``model_info.base_model`` when set, else + ``litellm_params.model``), which is what the request actually reaches; the public + model name it is listed under is an arbitrary alias and often absent from the cost + map. The token limits are the ones explicitly set in ``model_info``, which outrank + anything the cost map says. + """ + + cost_map_key: str + max_input_tokens: int | None + max_output_tokens: int | None + + class RouterErrors(enum.Enum): """ Enum for router specific errors with common codes diff --git a/litellm/utils.py b/litellm/utils.py index 252b6756937..80d2399c518 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2948,24 +2948,33 @@ def _resolve_builtin_model_cost_entry(key: str, provider: str) -> dict[str, obje return None +def is_generalized_model_info(model_info: ModelInfo) -> bool: + """Whether ``model_info`` came from a fallback-generalization capability rule. + + Detected as the resolved key missing ``litellm.model_cost`` while matching a + capability rule. A rule-derived entry carries no pricing and only a conservative + family-baseline context window, so callers holding a second candidate name should + prefer an exact cost-map entry from that name over this one. + """ + key: Final = cast("Mapping[str, object]", model_info).get("key") # cast-ok: partial dicts may omit "key" + if not isinstance(key, str): + return False + return key not in litellm.model_cost and match_capability_generalizations(key) is not None + + def _get_builtin_model_info_for_registration(model: str) -> ModelInfo | None: """Resolve ``model`` to its built-in cost-map entry for registration merging. Returns ``None`` when the lookup raises or when it resolved via a - fallback-generalization capability rule, detected as the resolved key missing - ``litellm.model_cost`` while matching a capability rule. A rule-derived entry - carries no pricing, so treating it as a hit would skip the built-in - cache-pricing inheritance for prefix-mangled keys. + fallback-generalization capability rule. A rule-derived entry carries no + pricing, so treating it as a hit would skip the built-in cache-pricing + inheritance for prefix-mangled keys. """ try: info: Final = get_model_info(model=model) except Exception: return None - if info["key"] in litellm.model_cost: - return info - if match_capability_generalizations(info["key"]) is None: - return info - return None + return None if is_generalized_model_info(info) else info _runtime_registered_model_cost: Final[dict[str, dict[str, object]]] = {} # mutable-ok: replayed on reload diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 4661cc17dbc..55e195ff14b 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -1944,7 +1944,7 @@ class TestModelInfoEndpoint: ): mock_router.get_fully_blocked_model_names.return_value = set() mock_router.get_model_list.return_value = [] - mock_router.get_configured_token_limits.return_value = (None, None) + mock_router.get_model_listing_info.return_value = None mock_router.get_deployment_by_model_group_name.return_value = Deployment( model_name="gpt-4", litellm_params=LiteLLM_Params(model="openai/gpt-4"), @@ -2021,7 +2021,7 @@ class TestModelInfoEndpoint: ): mock_router.get_fully_blocked_model_names.return_value = set() mock_router.get_model_list.return_value = [] - mock_router.get_configured_token_limits.return_value = (None, None) + mock_router.get_model_listing_info.return_value = None mock_router.get_deployment_by_model_group_name.return_value = Deployment( model_name="team-model-1", litellm_params=LiteLLM_Params(model="custom/team-model-1"), diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_models.py b/tests/test_litellm/proxy/proxy_server/test_routes_models.py index 2b126b1ea95..9063010b366 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_models.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_models.py @@ -17,6 +17,7 @@ import litellm from litellm.proxy import proxy_server from litellm.proxy import utils as proxy_utils from litellm.proxy.utils import create_model_info_response +from litellm.types.router import DeploymentModelListingInfo from .conftest import normalize # type: ignore[import-not-found] @@ -165,13 +166,16 @@ def test_anthropic_format_carries_router_configured_token_limits(client, auth_as built from another entry's lookup shows up as the wrong numbers.""" def _configured(model_name): - return (300000, 32000) if model_name == "gpt-4" else (500000, 4096) + max_input, max_output = (300000, 32000) if model_name == "gpt-4" else (500000, 4096) + return DeploymentModelListingInfo( + cost_map_key=model_name, max_input_tokens=max_input, max_output_tokens=max_output + ) def _cost_map_lookup(model_id): max_input, max_output = (200000, 64000) if model_id == "gpt-4" else (100000, 8000) return {"max_input_tokens": max_input, "max_output_tokens": max_output, "mode": "chat"} - patched_models.get_configured_token_limits = MagicMock(side_effect=_configured) + patched_models.get_model_listing_info = MagicMock(side_effect=_configured) def _resolved(**kwargs): return create_model_info_response(**kwargs, get_model_info=_cost_map_lookup) diff --git a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py index aa35fd64f18..6c1f2eb9bcf 100644 --- a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py +++ b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py @@ -877,7 +877,7 @@ async def test_v1_models_translates_team_model_for_access_group_key(monkeypatch) router.get_model_names.return_value = ["model_name_teamX_uuid9"] router.get_model_access_groups.return_value = {"grp-a": ["model_name_teamX_uuid9"]} router.get_fully_blocked_model_names.return_value = set() - router.get_configured_token_limits.return_value = (None, None) + router.get_model_listing_info.return_value = None router.model_list = [team_dep] router.get_model_list.return_value = [team_dep] @@ -919,7 +919,7 @@ async def test_v1_models_keeps_internal_names_when_public_name_flag_disabled( router.get_model_names.return_value = ["model_name_teamX_uuid9"] router.get_model_access_groups.return_value = {"grp-a": ["model_name_teamX_uuid9"]} router.get_fully_blocked_model_names.return_value = set() - router.get_configured_token_limits.return_value = (None, None) + router.get_model_listing_info.return_value = None router.model_list = [team_dep] router.get_model_list.return_value = [team_dep] @@ -954,7 +954,7 @@ async def test_v1_models_translates_team_model_with_metadata(monkeypatch): router.get_model_names.return_value = ["model_name_teamX_uuid9"] router.get_model_access_groups.return_value = {"grp-a": ["model_name_teamX_uuid9"]} router.get_fully_blocked_model_names.return_value = set() - router.get_configured_token_limits.return_value = (None, None) + router.get_model_listing_info.return_value = None router.model_list = [team_dep] router.get_model_list.return_value = [team_dep] router.get_model_group_info.return_value = None @@ -1000,7 +1000,7 @@ async def test_v1_models_metadata_fallbacks_use_internal_routing_key(monkeypatch router.get_model_names.return_value = ["model_name_teamX_uuid9"] router.get_model_access_groups.return_value = {"grp-a": ["model_name_teamX_uuid9"]} router.get_fully_blocked_model_names.return_value = set() - router.get_configured_token_limits.return_value = (None, None) + router.get_model_listing_info.return_value = None router.model_list = [team_dep] router.get_model_list.return_value = [team_dep] # Fallbacks are keyed on the internal routing name, as the router stores them. @@ -1057,7 +1057,7 @@ async def test_v1_models_metadata_does_not_leak_other_team_fallbacks(monkeypatch router.get_model_names.return_value = ["model_name_teamX_uuid9"] router.get_model_access_groups.return_value = {"grp-a": ["model_name_teamX_uuid9"]} router.get_fully_blocked_model_names.return_value = set() - router.get_configured_token_limits.return_value = (None, None) + router.get_model_listing_info.return_value = None router.model_list = [team_x, team_y] router.get_model_list.return_value = [team_x, team_y] router.fallbacks = [ @@ -1312,7 +1312,7 @@ def test_translate_team_model_names_for_listing_respects_legacy_flag(): def _public_named_router(*team_rows: dict) -> MagicMock: router = MagicMock() router.get_model_list.return_value = list(team_rows) - router.get_configured_token_limits.return_value = (None, None) + router.get_model_listing_info.return_value = None return router diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index dcaad968663..58f62bc4928 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -886,6 +886,7 @@ from typing import cast import litellm from litellm.proxy.utils import create_model_info_response +from litellm.types.router import DeploymentModelListingInfo from litellm.types.utils import ModelInfo @@ -913,7 +914,7 @@ def test_create_model_info_response_includes_max_tokens_from_lookup(): def test_create_model_info_response_does_not_call_router_group_info(): router = MagicMock() - router.get_configured_token_limits.return_value = (None, None) + router.get_model_listing_info.return_value = None response = create_model_info_response( model_id="some-model", @@ -928,7 +929,9 @@ def test_create_model_info_response_does_not_call_router_group_info(): def test_create_model_info_response_uses_deployment_limits_when_not_in_cost_map(): router = MagicMock() - router.get_configured_token_limits.return_value = (32000, 8000) + router.get_model_listing_info.return_value = DeploymentModelListingInfo( + cost_map_key="my-custom-deployment", max_input_tokens=32000, max_output_tokens=8000 + ) response = create_model_info_response( model_id="my-custom-deployment", @@ -944,7 +947,9 @@ def test_create_model_info_response_uses_deployment_limits_when_not_in_cost_map( def test_create_model_info_response_deployment_limits_override_cost_map(): router = MagicMock() - router.get_configured_token_limits.return_value = (200000, None) + router.get_model_listing_info.return_value = DeploymentModelListingInfo( + cost_map_key="gpt-4o", max_input_tokens=200000, max_output_tokens=None + ) response = create_model_info_response( model_id="gpt-4o", @@ -1878,3 +1883,119 @@ async def test_proxy_only_error_5xx_keeps_traceback_and_runs_sync_callbacks(monk Logging.failure_handler = orig_sync_failure assert "test_proxy_utils" in captured["async_traceback"] + + +def test_create_model_info_response_resolves_alias_to_deployment_model(): + """A public model name that is not itself a cost-map key must not be resolved through + the fallback-generalization rules: `bedrock-claude-opus-5` matches the generic + claude-family baseline (200k/64k) by substring, while the deployment it fronts really + accepts 1M/128k. Regression for the /v1/models alias resolution introduced in v1.94.0.""" + from litellm import Router + + saved_model_cost = dict(litellm.model_cost) + try: + router = Router( + model_list=[ + { + "model_name": "bedrock-claude-opus-5", + "litellm_params": { + "custom_llm_provider": "bedrock", + "model": "bedrock/eu.anthropic.claude-opus-5", + }, + "model_info": {"base_model": "eu.anthropic.claude-opus-5"}, + } + ] + ) + + response = create_model_info_response( + model_id="bedrock-claude-opus-5", provider="openai", llm_router=router + ) + finally: + litellm.model_cost.clear() + litellm.model_cost.update(saved_model_cost) + + assert response["max_input_tokens"] == 1000000 + assert response["max_output_tokens"] == 128000 + + +def test_create_model_info_response_keeps_exact_alias_over_generalized_deployment_model(): + """Mirror of the alias bug: when the deployment points at a custom backend name that + only matches a generalization rule, the listed name's exact cost-map entry is the + better answer and must win.""" + from litellm import Router + + saved_model_cost = dict(litellm.model_cost) + try: + router = Router( + model_list=[ + { + "model_name": "claude-opus-5", + "litellm_params": { + "custom_llm_provider": "bedrock", + "model": "bedrock/my-claude-opus-5-provisioned", + }, + } + ] + ) + + response = create_model_info_response( + model_id="claude-opus-5", provider="openai", llm_router=router + ) + finally: + litellm.model_cost.clear() + litellm.model_cost.update(saved_model_cost) + + assert response["max_input_tokens"] == 1000000 + + +def test_create_model_info_response_falls_back_to_alias_for_opaque_deployment_name(): + """An Azure deployment named after the resource rather than the model has no cost-map + entry; the listed name still does, and must keep answering.""" + from litellm import Router + + saved_model_cost = dict(litellm.model_cost) + try: + router = Router( + model_list=[ + { + "model_name": "gpt-4o", + "litellm_params": {"model": "azure/my-gpt4o-deployment"}, + } + ] + ) + + response = create_model_info_response( + model_id="gpt-4o", provider="openai", llm_router=router + ) + finally: + litellm.model_cost.clear() + litellm.model_cost.update(saved_model_cost) + + assert response["max_input_tokens"] == 128000 + assert response["max_output_tokens"] == 16384 + + +def test_create_model_info_response_resolves_mode_through_deployment_model(): + """`mode` is derived from the same lookup, so an aliased embedding deployment + currently reports no mode at all; it must report `embedding`.""" + from litellm import Router + + saved_model_cost = dict(litellm.model_cost) + try: + router = Router( + model_list=[ + { + "model_name": "my-embeddings", + "litellm_params": {"model": "openai/text-embedding-3-small"}, + } + ] + ) + + response = create_model_info_response( + model_id="my-embeddings", provider="openai", llm_router=router + ) + finally: + litellm.model_cost.clear() + litellm.model_cost.update(saved_model_cost) + + assert response["mode"] == "embedding" diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 84f6344be35..37805436f8d 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -7271,6 +7271,105 @@ def test_get_configured_token_limits_coerces_numeric_strings(): assert router.get_configured_token_limits("quoted-limits-model") == (32000, 8000) +def test_get_model_listing_info_prefers_base_model_over_litellm_params_model(): + """The cost-map key comes from base_model when set, so a deployment pointing at an + opaque backend name still resolves the real catalog entry.""" + router = litellm.Router( + model_list=[ + { + "model_name": "bedrock-claude-opus-5", + "litellm_params": {"model": "bedrock/eu.anthropic.claude-opus-5"}, + "model_info": {"base_model": "eu.anthropic.claude-opus-5"}, + } + ] + ) + + info = router.get_model_listing_info("bedrock-claude-opus-5") + assert info is not None + assert info.cost_map_key == "eu.anthropic.claude-opus-5" + + +def test_get_model_listing_info_falls_back_to_litellm_params_model(): + router = litellm.Router( + model_list=[ + { + "model_name": "bedrock-claude-opus-5", + "litellm_params": {"model": "bedrock/eu.anthropic.claude-opus-5"}, + } + ] + ) + + info = router.get_model_listing_info("bedrock-claude-opus-5") + assert info is not None + assert info.cost_map_key == "bedrock/eu.anthropic.claude-opus-5" + + +def test_get_model_listing_info_ignores_blank_base_model(): + """A base_model set to an empty string is absent, not a cost-map key.""" + router = litellm.Router( + model_list=[ + { + "model_name": "bedrock-claude-opus-5", + "litellm_params": {"model": "bedrock/eu.anthropic.claude-opus-5"}, + "model_info": {"base_model": ""}, + } + ] + ) + + info = router.get_model_listing_info("bedrock-claude-opus-5") + assert info is not None + assert info.cost_map_key == "bedrock/eu.anthropic.claude-opus-5" + + +def test_get_model_listing_info_returns_none_for_unknown_name(): + router = litellm.Router( + model_list=[ + { + "model_name": "no-limits-model", + "litellm_params": {"model": "openai/some-unmapped-model"}, + } + ] + ) + + assert router.get_model_listing_info("not-a-real-model") is None + + +def test_get_model_listing_info_carries_configured_limits(): + router = litellm.Router( + model_list=[ + { + "model_name": "my-custom-model", + "litellm_params": {"model": "openai/some-unmapped-model"}, + "model_info": {"max_input_tokens": 32000, "max_output_tokens": 8000}, + } + ] + ) + + info = router.get_model_listing_info("my-custom-model") + assert info is not None + assert (info.max_input_tokens, info.max_output_tokens) == (32000, 8000) + + +def test_get_model_listing_info_skips_wildcard_pattern_matching(): + router = litellm.Router( + model_list=[ + { + "model_name": "bedrock/*", + "litellm_params": {"model": "bedrock/*"}, + "model_info": {"max_input_tokens": 12345}, + } + ] + ) + + with patch.object( + router.pattern_router, "route", side_effect=AssertionError("pattern route called") + ): + assert ( + router.get_model_listing_info("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") + is None + ) + + @pytest.mark.asyncio async def test_acreate_batch_disable_fallbacks_surfaces_owning_provider_error(): router = litellm.Router( From d5270890c409fe5942ff98c5b3a69370fc0bff81 Mon Sep 17 00:00:00 2001 From: michelligabriele Date: Wed, 2 Sep 2026 12:03:36 +0200 Subject: [PATCH 03/33] fix(proxy): report the widest window across a model group, not the first deployment's --- litellm/proxy/utils.py | 45 +++++++-- litellm/router.py | 79 +++++++++++----- litellm/types/router.py | 17 ++-- .../proxy/proxy_server/test_routes_models.py | 2 +- tests/test_litellm/proxy/test_proxy_utils.py | 52 ++++++++++- tests/test_litellm/test_router.py | 91 ++++++++++++++++++- 6 files changed, 242 insertions(+), 44 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index b86e6d857cc..874dca4238f 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -7496,10 +7496,11 @@ def _safe_get_model_info(model: str, get_model_info: Callable[[str], ModelInfo]) def _resolve_listing_model_info( deployment_model: str | None, listed_model: str, + listed_info: ModelInfo | None, get_model_info: Callable[[str], ModelInfo], ) -> tuple[ModelInfo, ...]: """ - Cost-map entries describing a listed model, best source first. + Cost-map entries describing one deployment behind a listed model, best source first. The name a model is listed under is an arbitrary public alias, so it often misses the cost map and lands on a fallback-generalization rule that answers with a conservative @@ -7508,9 +7509,10 @@ def _resolve_listing_model_info( generalize, and because a deployment's own model is registered into the cost map as a stub that carries no limits of its own. Exact entries are consulted before generalized ones, and each field is then taken from the first entry that has it. - """ - listed_info: Final = _safe_get_model_info(listed_model, get_model_info) + ``listed_info`` is resolved once by the caller, since a group with several distinct + underlying models resolves the same alias for each of them. + """ # Fast path, and the only one a wildcard-expanded name takes: with a single name # there is nothing to order, so skip the generalization test entirely. This keeps # the per-model cost of the listing on the hot path #33721 exists to protect. @@ -7539,6 +7541,20 @@ def _first_token_limit(candidates: tuple[ModelInfo, ...], field: str) -> int | N ) +def _group_token_limit(candidate_sets: tuple[tuple[ModelInfo, ...], ...], field: str) -> int | None: + """The widest limit any deployment behind the listed name declares for ``field``. + + A model group is normally one model behind several interchangeable deployments, so + there is a single value to report. When a group genuinely mixes models, reporting the + widest window keeps the listing independent of config order and agreeing with + ``/model_group/info``, which aggregates the same way for the Admin UI. + """ + limits: Final = tuple( + limit for limit in (_first_token_limit(candidates, field) for candidates in candidate_sets) if limit is not None + ) + return max(limits) if limits else None + + def create_model_info_response( model_id: str, provider: str, @@ -7565,19 +7581,30 @@ def create_model_info_response( listing_info: Final = llm_router.get_model_listing_info(model_id) if llm_router is not None else None - candidates: Final = _resolve_listing_model_info( - deployment_model=listing_info.cost_map_key if listing_info is not None else None, - listed_model=model_id, - get_model_info=get_model_info, + # One entry per distinct model behind the listed name; (None,) when the router knows + # nothing about it, so the listed name is resolved on its own as before. + deployment_models: Final[tuple[str | None, ...]] = ( + listing_info.cost_map_keys if listing_info is not None and listing_info.cost_map_keys else (None,) + ) + listed_info: Final = _safe_get_model_info(model_id, get_model_info) + candidate_sets: Final = tuple( + _resolve_listing_model_info( + deployment_model=deployment_model, + listed_model=model_id, + listed_info=listed_info, + get_model_info=get_model_info, + ) + for deployment_model in deployment_models ) - max_input_tokens: int | None = _first_token_limit(candidates, "max_input_tokens") - max_output_tokens: int | None = _first_token_limit(candidates, "max_output_tokens") + max_input_tokens: int | None = _group_token_limit(candidate_sets, "max_input_tokens") + max_output_tokens: int | None = _group_token_limit(candidate_sets, "max_output_tokens") mode: Final = next( ( m for m in ( cast("Mapping[str, object]", info).get("mode") # cast-ok: an entry need not carry "mode" + for candidates in candidate_sets for info in candidates ) if isinstance(m, str) diff --git a/litellm/router.py b/litellm/router.py index 7bc6a6c714a..7011f7f504c 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -9729,29 +9729,57 @@ class Router: def get_model_listing_info(self, model_name: str) -> DeploymentModelListingInfo | None: """ - Return what the concrete deployment behind model_name contributes to its - /v1/models entry: the cost-map key for its underlying model, plus any token - limits explicitly configured in its model_info. Resolved via O(1) index lookup. + Return what the concrete deployments behind model_name contribute to its + /v1/models entry: the cost-map keys for their underlying models, plus the widest + token limits explicitly configured in their model_info. Resolved via O(1) index + lookup. - Returns None for wildcard-expanded or unknown names, where the listed name is - the real model name and no deployment-specific information exists, and treats a - malformed configured limit as absent rather than failing the listing. Unlike - get_model_group_info, this never triggers pattern matching or deep copies, so it - is safe to call per listed model on the /v1/models hot path. + Returns None for wildcard-expanded or unknown names, where the listed name is the + real model name and no deployment-specific information exists, and treats a + malformed configured limit as absent rather than failing the listing. + + The whole group is read rather than just its first deployment, so a group that + mixes models does not advertise a window that depends on config order; the widest + one is reported, which is what get_model_group_info already shows the Admin UI. + Keys are deduplicated, so the ordinary group of interchangeable deployments of one + model still costs the caller a single cost-map lookup. Unlike get_model_group_info, + this never triggers pattern matching or deep copies, so it is safe to call per + listed model on the /v1/models hot path. """ - deployment: Final = self.get_deployment_by_model_group_name(model_group_name=model_name) - if deployment is None: + indices: Final = self.model_name_to_deployment_indices.get(model_name) + if not indices: return None - model_info: Final = deployment.model_info - # base_model is a declared field, so read it as one: an unset or blank value - # means the deployment's own model name is the cost-map key. - base_model: Final = model_info.base_model - return DeploymentModelListingInfo( - cost_map_key=base_model or deployment.litellm_params.model, - max_input_tokens=coerce_token_limit(model_info.get("max_input_tokens")), - max_output_tokens=coerce_token_limit(model_info.get("max_output_tokens")), + deployments: Final = tuple(self.model_list[index] for index in indices) + model_infos: Final = tuple(deployment.get("model_info") or MappingProxyType({}) for deployment in deployments) + params: Final = tuple(deployment.get("litellm_params") or MappingProxyType({}) for deployment in deployments) + # base_model resolution mirrors get_router_model_info: unset or blank means the + # deployment's own model name is the cost-map key. + cost_map_keys: Final = tuple( + dict.fromkeys( # deduplicates while preserving config order + key + for key in ( + model_info.get("base_model") or litellm_params.get("base_model") or litellm_params.get("model") + for model_info, litellm_params in zip(model_infos, params) + ) + if isinstance(key, str) and key + ) ) + return DeploymentModelListingInfo( + cost_map_keys=cost_map_keys, + max_input_tokens=self._widest_configured_limit(model_infos, "max_input_tokens"), + max_output_tokens=self._widest_configured_limit(model_infos, "max_output_tokens"), + ) + + @staticmethod + def _widest_configured_limit(model_infos: Sequence[Mapping[str, Any]], field: str) -> int | None: + """The largest usable value of ``field`` across a group's configured model_info blocks.""" + limits: Final = tuple( + limit + for limit in (coerce_token_limit(model_info.get(field)) for model_info in model_infos) + if limit is not None + ) + return max(limits) if limits else None def get_configured_token_limits(self, model_name: str) -> "tuple[int | None, int | None]": """ @@ -9760,11 +9788,20 @@ class Router: Returns (None, None) for wildcard-expanded or unknown names, and treats a malformed configured value as absent rather than failing the caller. + + Deliberately reads one deployment rather than aggregating the group the way + get_model_listing_info does: its caller truncates an embedding input to this + value, so the widest window in a mixed group would be the wrong answer there. """ - listing_info: Final = self.get_model_listing_info(model_name=model_name) - if listing_info is None: + deployment: Final = self.get_deployment_by_model_group_name(model_group_name=model_name) + if deployment is None: return (None, None) - return (listing_info.max_input_tokens, listing_info.max_output_tokens) + + model_info: Final = deployment.model_info + return ( + coerce_token_limit(model_info.get("max_input_tokens")), + coerce_token_limit(model_info.get("max_output_tokens")), + ) def get_deployment_credentials_with_provider( self, model_id: str, team_id: str | None = None diff --git a/litellm/types/router.py b/litellm/types/router.py index 6f9fb4bf6ee..a773959f5f1 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -577,17 +577,18 @@ class Deployment(BaseModel): @dataclass(frozen=True, slots=True) class DeploymentModelListingInfo: - """What a concrete deployment contributes to its OpenAI-compatible listing entry. + """What the deployments behind a model name contribute to its OpenAI-compatible listing entry. - ``cost_map_key`` is the name the deployment's underlying model is known by in - ``litellm.model_cost`` (``model_info.base_model`` when set, else - ``litellm_params.model``), which is what the request actually reaches; the public - model name it is listed under is an arbitrary alias and often absent from the cost - map. The token limits are the ones explicitly set in ``model_info``, which outrank - anything the cost map says. + ``cost_map_keys`` are the names those deployments' underlying models are known by in + ``litellm.model_cost`` (``base_model`` when set, else ``litellm_params.model``), which + is what a request actually reaches; the public model name they are listed under is an + arbitrary alias and often absent from the cost map. Keys are deduplicated in config + order, so the ordinary group -- several interchangeable deployments of one model -- + carries exactly one. The token limits are the widest explicitly set in any + deployment's ``model_info``, which outrank anything the cost map says. """ - cost_map_key: str + cost_map_keys: tuple[str, ...] max_input_tokens: int | None max_output_tokens: int | None diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_models.py b/tests/test_litellm/proxy/proxy_server/test_routes_models.py index 9063010b366..54768fef2ce 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_models.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_models.py @@ -168,7 +168,7 @@ def test_anthropic_format_carries_router_configured_token_limits(client, auth_as def _configured(model_name): max_input, max_output = (300000, 32000) if model_name == "gpt-4" else (500000, 4096) return DeploymentModelListingInfo( - cost_map_key=model_name, max_input_tokens=max_input, max_output_tokens=max_output + cost_map_keys=(model_name,), max_input_tokens=max_input, max_output_tokens=max_output ) def _cost_map_lookup(model_id): diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index 58f62bc4928..90e7431a22d 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -930,7 +930,7 @@ def test_create_model_info_response_does_not_call_router_group_info(): def test_create_model_info_response_uses_deployment_limits_when_not_in_cost_map(): router = MagicMock() router.get_model_listing_info.return_value = DeploymentModelListingInfo( - cost_map_key="my-custom-deployment", max_input_tokens=32000, max_output_tokens=8000 + cost_map_keys=("my-custom-deployment",), max_input_tokens=32000, max_output_tokens=8000 ) response = create_model_info_response( @@ -948,7 +948,7 @@ def test_create_model_info_response_uses_deployment_limits_when_not_in_cost_map( def test_create_model_info_response_deployment_limits_override_cost_map(): router = MagicMock() router.get_model_listing_info.return_value = DeploymentModelListingInfo( - cost_map_key="gpt-4o", max_input_tokens=200000, max_output_tokens=None + cost_map_keys=("gpt-4o",), max_input_tokens=200000, max_output_tokens=None ) response = create_model_info_response( @@ -962,6 +962,54 @@ def test_create_model_info_response_deployment_limits_override_cost_map(): assert response["max_output_tokens"] == 16384 +def test_create_model_info_response_reports_widest_window_in_a_mixed_group(): + """A group mixing models advertises the widest window, not whichever is listed first.""" + limits = { + "small-model": _fake_model_info(max_input_tokens=200000, max_output_tokens=4096, mode="chat"), + "large-model": _fake_model_info(max_input_tokens=1000000, max_output_tokens=128000, mode="chat"), + } + + for keys in (("small-model", "large-model"), ("large-model", "small-model")): + router = MagicMock() + router.get_model_listing_info.return_value = DeploymentModelListingInfo( + cost_map_keys=keys, max_input_tokens=None, max_output_tokens=None + ) + + response = create_model_info_response( + model_id="house-claude", + provider="openai", + llm_router=router, + get_model_info=lambda model: limits[model], + ) + + assert response["max_input_tokens"] == 1000000, keys + assert response["max_output_tokens"] == 128000, keys + + +def test_create_model_info_response_resolves_alias_once_per_listing(): + """The alias is the same for every deployment in the group, so it is looked up once.""" + seen: list[str] = [] + + def _tracking_get_model_info(model: str) -> ModelInfo: + seen.append(model) + return _fake_model_info(max_input_tokens=128000) + + router = MagicMock() + router.get_model_listing_info.return_value = DeploymentModelListingInfo( + cost_map_keys=("model-a", "model-b"), max_input_tokens=None, max_output_tokens=None + ) + + create_model_info_response( + model_id="house-model", + provider="openai", + llm_router=router, + get_model_info=_tracking_get_model_info, + ) + + assert seen.count("house-model") == 1 + assert sorted(seen) == ["house-model", "model-a", "model-b"] + + def test_create_model_info_response_survives_malformed_configured_limits(): from litellm import Router diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 37805436f8d..588fd98abd1 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -7286,7 +7286,7 @@ def test_get_model_listing_info_prefers_base_model_over_litellm_params_model(): info = router.get_model_listing_info("bedrock-claude-opus-5") assert info is not None - assert info.cost_map_key == "eu.anthropic.claude-opus-5" + assert info.cost_map_keys == ("eu.anthropic.claude-opus-5",) def test_get_model_listing_info_falls_back_to_litellm_params_model(): @@ -7301,7 +7301,7 @@ def test_get_model_listing_info_falls_back_to_litellm_params_model(): info = router.get_model_listing_info("bedrock-claude-opus-5") assert info is not None - assert info.cost_map_key == "bedrock/eu.anthropic.claude-opus-5" + assert info.cost_map_keys == ("bedrock/eu.anthropic.claude-opus-5",) def test_get_model_listing_info_ignores_blank_base_model(): @@ -7318,7 +7318,7 @@ def test_get_model_listing_info_ignores_blank_base_model(): info = router.get_model_listing_info("bedrock-claude-opus-5") assert info is not None - assert info.cost_map_key == "bedrock/eu.anthropic.claude-opus-5" + assert info.cost_map_keys == ("bedrock/eu.anthropic.claude-opus-5",) def test_get_model_listing_info_returns_none_for_unknown_name(): @@ -7350,6 +7350,91 @@ def test_get_model_listing_info_carries_configured_limits(): assert (info.max_input_tokens, info.max_output_tokens) == (32000, 8000) +def test_get_model_listing_info_dedupes_interchangeable_deployments(): + """The ordinary group is N deployments of one model, so it yields exactly one key.""" + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4o", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-a"}, + }, + { + "model_name": "gpt-4o", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-b"}, + }, + ] + ) + + info = router.get_model_listing_info("gpt-4o") + assert info is not None + assert info.cost_map_keys == ("openai/gpt-4o",) + + +def test_get_model_listing_info_collects_every_model_in_a_mixed_group(): + router = litellm.Router( + model_list=[ + { + "model_name": "house-claude", + "litellm_params": {"model": "anthropic/claude-3-haiku-20240307"}, + }, + { + "model_name": "house-claude", + "litellm_params": {"model": "bedrock/eu.anthropic.claude-opus-5"}, + }, + ] + ) + + info = router.get_model_listing_info("house-claude") + assert info is not None + assert info.cost_map_keys == ( + "anthropic/claude-3-haiku-20240307", + "bedrock/eu.anthropic.claude-opus-5", + ) + + +def test_get_model_listing_info_reports_widest_configured_limits_in_a_mixed_group(): + """Matches how get_model_group_info aggregates for the Admin UI, so the two agree.""" + router = litellm.Router( + model_list=[ + { + "model_name": "house-model", + "litellm_params": {"model": "openai/some-unmapped-model"}, + "model_info": {"max_input_tokens": 32000, "max_output_tokens": 4096}, + }, + { + "model_name": "house-model", + "litellm_params": {"model": "openai/another-unmapped-model"}, + "model_info": {"max_input_tokens": 128000, "max_output_tokens": 16384}, + }, + ] + ) + + info = router.get_model_listing_info("house-model") + assert info is not None + assert (info.max_input_tokens, info.max_output_tokens) == (128000, 16384) + + +def test_get_model_listing_info_reads_base_model_from_litellm_params(): + """base_model resolution mirrors get_router_model_info, which also accepts it there.""" + router = litellm.Router( + model_list=[ + { + "model_name": "azure-deployment", + "litellm_params": { + "model": "azure/my-azure-deployment-name", + "base_model": "azure/gpt-4o", + "api_key": "sk-a", + "api_base": "https://example.openai.azure.com", + }, + } + ] + ) + + info = router.get_model_listing_info("azure-deployment") + assert info is not None + assert info.cost_map_keys == ("azure/gpt-4o",) + + def test_get_model_listing_info_skips_wildcard_pattern_matching(): router = litellm.Router( model_list=[ From f0f6ff8de4c5bd09b092206003f1f1aa00b7bb41 Mon Sep 17 00:00:00 2001 From: michelligabriele Date: Wed, 2 Sep 2026 14:12:23 +0200 Subject: [PATCH 04/33] test(router): cover _widest_configured_limit directly for the coverage gate --- litellm/proxy/utils.py | 14 +++++++++++--- tests/test_litellm/test_router.py | 14 ++++++++++++++ 2 files changed, 25 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 874dca4238f..f891c85320a 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -7545,9 +7545,17 @@ def _group_token_limit(candidate_sets: tuple[tuple[ModelInfo, ...], ...], field: """The widest limit any deployment behind the listed name declares for ``field``. A model group is normally one model behind several interchangeable deployments, so - there is a single value to report. When a group genuinely mixes models, reporting the - widest window keeps the listing independent of config order and agreeing with - ``/model_group/info``, which aggregates the same way for the Admin UI. + there is a single value to report and the choice of aggregate does not arise. + + When a group genuinely mixes models no single number is right, and the widest is the + deliberate pick over the narrowest for two reasons. It is what ``/model_group/info`` + has long reported to the Admin UI, so the two surfaces agree; disagreeing is the very + complaint this resolution path exists to fix. And of the two ways to be wrong, + under-advertising is worse: a client that trusts a narrowed window silently refuses + prompts the group would have served, while an over-long prompt that reaches a smaller + deployment comes back as a legible context-length error -- and does not reach one at + all when ``enable_pre_call_checks`` is set, which filters deployments the prompt does + not fit. """ limits: Final = tuple( limit for limit in (_first_token_limit(candidates, field) for candidates in candidate_sets) if limit is not None diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 588fd98abd1..eba9a85feb6 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -7350,6 +7350,20 @@ def test_get_model_listing_info_carries_configured_limits(): assert (info.max_input_tokens, info.max_output_tokens) == (32000, 8000) +def test_widest_configured_limit_ignores_absent_and_malformed_values(): + model_infos = ( + {"max_input_tokens": 32000}, + {}, + {"max_input_tokens": "not-a-number"}, + {"max_input_tokens": "128000"}, + {"max_output_tokens": 4096}, + ) + + assert litellm.Router._widest_configured_limit(model_infos, "max_input_tokens") == 128000 + assert litellm.Router._widest_configured_limit(model_infos, "max_output_tokens") == 4096 + assert litellm.Router._widest_configured_limit((), "max_input_tokens") is None + + def test_get_model_listing_info_dedupes_interchangeable_deployments(): """The ordinary group is N deployments of one model, so it yields exactly one key.""" router = litellm.Router( From 50874f9fc9cdc9e925d0fc7abda06a178a303980 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 2 Sep 2026 15:52:47 +0000 Subject: [PATCH 05/33] fix(mcp): require audiences and prove refresh-token ownership for identity binding Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/discoverable_endpoints.py | 9 ++ .../mcp_server/oauth_identity_binding.py | 93 +++++++++++--- .../types/mcp_server/mcp_server_manager.py | 4 +- .../mcp_server/test_oauth_identity_binding.py | 117 +++++++++++++++++- .../test_mcp_management_endpoints.py | 6 +- 5 files changed, 211 insertions(+), 18 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 92128a82856..88f2c47f5fc 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -56,6 +56,8 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( revoke_refresh_token, ) from litellm.proxy._experimental.mcp_server.oauth_identity_binding import ( + RefreshOwnershipProven, + RefreshTokenPresented, enforce_oauth_identity_binding, ) from litellm.proxy._experimental.mcp_server.oauth_utils import ( @@ -1046,7 +1048,13 @@ async def exchange_token_with_server( refresh_request_scope = scope or bridge_upstream_scope if refresh_request_scope: token_data["scope"] = refresh_request_scope + refresh_ownership = ( + RefreshOwnershipProven() + if bridge_upstream_refresh is not None + else RefreshTokenPresented(upstream_refresh_token) + ) else: + refresh_ownership = None if not code: raise HTTPException( status_code=400, @@ -1145,6 +1153,7 @@ async def exchange_token_with_server( token_response=token_response, litellm_user_id=resolved_user_id, grant_type=grant_type, + refresh_ownership=refresh_ownership, ) # Store server-side when the server is configured for per-user OAuth and diff --git a/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py b/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py index f132f59ae99..d4339c15e1a 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py @@ -14,6 +14,7 @@ from typing import Final, Literal import jwt from fastapi import HTTPException +from typing_extensions import assert_never from litellm._logging import verbose_logger from litellm.caching.in_memory_cache import InMemoryCache @@ -37,6 +38,7 @@ _jwks_cache: Final = InMemoryCache(default_ttl=_JWKS_CACHE_TTL_SECONDS) JwksFetcher = Callable[[MCPOAuthIdentityBinding], Awaitable[list[Mapping[str, object]]]] CallerPrincipalLoader = Callable[[str, MCPOAuthIdentityBinding], Awaitable[str | None]] +StoredRefreshTokenLoader = Callable[[str, str], Awaitable[str | None]] _RejectionCode = Literal["oauth_principal_mismatch", "oauth_identity_binding_failed"] @@ -47,6 +49,19 @@ class _BindingRejection: description: str +@dataclass(frozen=True, slots=True) +class RefreshOwnershipProven: + """The gateway itself unwrapped the upstream refresh token from a sealed per-user envelope.""" + + +@dataclass(frozen=True, slots=True) +class RefreshTokenPresented: + refresh_token: str + + +RefreshOwnership = RefreshOwnershipProven | RefreshTokenPresented | None + + async def _fetch_issuer_jwks(binding: MCPOAuthIdentityBinding) -> list[Mapping[str, object]]: jwks_url: Final[str] = binding.jwks_url or await _discover_jwks_url(binding.issuer) cached: Final = await _jwks_cache.async_get_cache(jwks_url) @@ -98,11 +113,8 @@ def _decode_id_token( signing_key.key, algorithms=list(_ALLOWED_ID_TOKEN_ALGORITHMS), issuer=binding.issuer, - audience=binding.audiences if binding.audiences else None, - options={ - "require": ["iss", "exp"], - "verify_aud": bool(binding.audiences), - }, + audience=binding.audiences, + options={"require": ["iss", "exp"]}, ) except jwt.InvalidTokenError as exc: return _BindingRejection( @@ -121,7 +133,11 @@ def _upstream_principal( code="oauth_identity_binding_failed", description=f"id_token has no usable '{binding.principal_claim}' claim", ) - if binding.principal_claim == "email" and binding.require_email_verified and claims.get("email_verified") is not True: + if ( + binding.principal_claim == "email" + and binding.require_email_verified + and claims.get("email_verified") is not True + ): return _BindingRejection( code="oauth_identity_binding_failed", description="id_token email is not verified (email_verified is not true)", @@ -142,6 +158,26 @@ async def _load_caller_principal(litellm_user_id: str, binding: MCPOAuthIdentity return loaded.user_email +async def _load_stored_refresh_token(litellm_user_id: str, server_id: str) -> str | None: + try: + from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 + get_user_oauth_credential, + ) + from litellm.proxy.utils import get_prisma_client_or_throw # noqa: PLC0415 + + prisma_client: Final = get_prisma_client_or_throw( + "Database not connected. Cannot verify OAuth refresh token ownership." + ) + cred: Final = await get_user_oauth_credential( + prisma_client=prisma_client, + user_id=litellm_user_id, + server_id=server_id, + ) + return cred.get("refresh_token") if cred else None + except Exception: # noqa: BLE001 # a credential lookup failure must fail closed + return None + + def _principals_match(upstream: str, caller: str, binding: MCPOAuthIdentityBinding) -> bool: if binding.principal_claim == "email" or binding.caller_field == "user_email": return upstream.strip().casefold() == caller.strip().casefold() @@ -153,17 +189,41 @@ async def _evaluate_binding( token_response: Mapping[str, object], litellm_user_id: str | None, grant_type: str, + server_id: str, + refresh_ownership: RefreshOwnership, jwks_fetcher: JwksFetcher, caller_principal_loader: CallerPrincipalLoader, + stored_refresh_token_loader: StoredRefreshTokenLoader, ) -> _BindingRejection | None: id_token: Final = token_response.get("id_token") if not isinstance(id_token, str) or not id_token: - if grant_type == "refresh_token": - return None - return _BindingRejection( - code="oauth_identity_binding_failed", - description="the upstream token response carries no id_token to bind the credential to a principal", - ) + if grant_type != "refresh_token": + return _BindingRejection( + code="oauth_identity_binding_failed", + description="the upstream token response carries no id_token to bind the credential to a principal", + ) + match refresh_ownership: + case RefreshOwnershipProven(): + return None + case None: + return _BindingRejection( + code="oauth_identity_binding_failed", + description="refresh_token grant without an id_token carries no refresh token to prove ownership of", + ) + case RefreshTokenPresented(refresh_token): + if not litellm_user_id: + return _BindingRejection( + code="oauth_identity_binding_failed", + description="the request carries no resolvable LiteLLM user identity to bind the credential to", + ) + stored: Final = await stored_refresh_token_loader(litellm_user_id, server_id) + if stored is None or stored != refresh_token: + return _BindingRejection( + code="oauth_identity_binding_failed", + description="the presented refresh_token is not the caller's stored credential for this server", + ) + return None + assert_never(refresh_ownership) if not litellm_user_id: return _BindingRejection( code="oauth_identity_binding_failed", @@ -204,15 +264,17 @@ async def enforce_oauth_identity_binding( token_response: Mapping[str, object], litellm_user_id: str | None, grant_type: str, + refresh_ownership: RefreshOwnership, jwks_fetcher: JwksFetcher = _fetch_issuer_jwks, caller_principal_loader: CallerPrincipalLoader = _load_caller_principal, + stored_refresh_token_loader: StoredRefreshTokenLoader = _load_stored_refresh_token, ) -> None: """Validate the exchanged token's upstream principal against the LiteLLM caller. No-op when the server has no binding or it is disabled. In enforce mode a failure raises 403 before the caller returns, stores, or caches the token; in audit mode failures are logged only. - A refresh_token grant without an id_token is allowed in both modes: the stored credential keeps - the binding established at the original authorization_code exchange. + A refresh_token grant without an id_token is allowed only when the presented refresh token matches + the caller's stored credential or a sealed bridge envelope already proved ownership. """ binding: Final = server.oauth_identity_binding if binding is None or binding.mode == "disabled": @@ -222,8 +284,11 @@ async def enforce_oauth_identity_binding( token_response=token_response, litellm_user_id=litellm_user_id, grant_type=grant_type, + server_id=server.server_id, + refresh_ownership=refresh_ownership, jwks_fetcher=jwks_fetcher, caller_principal_loader=caller_principal_loader, + stored_refresh_token_loader=stored_refresh_token_loader, ) if rejection is None: return diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index cc3f6c8e8d8..cebc7decac4 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -1,7 +1,7 @@ from datetime import datetime from typing import Any, Final, Literal -from pydantic import BaseModel, ConfigDict, field_validator +from pydantic import BaseModel, ConfigDict, Field, field_validator from litellm.types.mcp import ( DEFAULT_SUBJECT_TOKEN_TYPE, @@ -52,7 +52,7 @@ class MCPOAuthIdentityBinding(BaseModel): mode: Literal["disabled", "audit", "enforce"] = "disabled" issuer: str jwks_url: str | None = None - audiences: list[str] = [] + audiences: list[str] = Field(min_length=1) principal_claim: str = "email" caller_field: Literal["user_email", "user_id"] = "user_email" require_email_verified: bool = True diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_identity_binding.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_identity_binding.py index 69eaabeda4a..de59020c44c 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_identity_binding.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_identity_binding.py @@ -7,8 +7,11 @@ import pytest from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import rsa from fastapi import HTTPException +from pydantic import ValidationError from litellm.proxy._experimental.mcp_server.oauth_identity_binding import ( + RefreshOwnershipProven, + RefreshTokenPresented, enforce_oauth_identity_binding, ) from litellm.types.mcp import MCPTransport @@ -54,6 +57,13 @@ def _caller_loader(email: str | None): return load +def _stored_refresh_token_loader(refresh_token: str | None): + async def load(_user_id: str, _server_id: str) -> str | None: + return refresh_token + + return load + + def _server(mode: str = "enforce", **binding_overrides: object) -> MCPServer: return MCPServer( server_id="srv-1", @@ -77,6 +87,7 @@ async def test_matching_principal_passes(): token_response={"access_token": "at", "id_token": token}, litellm_user_id="user-a", grant_type="authorization_code", + refresh_ownership=None, jwks_fetcher=_jwks_fetcher, caller_principal_loader=_caller_loader("alice@example.com"), ) @@ -92,6 +103,7 @@ async def test_mismatched_principal_rejected(): token_response={"access_token": "at", "id_token": token}, litellm_user_id="user-a", grant_type="authorization_code", + refresh_ownership=None, jwks_fetcher=_jwks_fetcher, caller_principal_loader=_caller_loader("alice@example.com"), ) @@ -108,6 +120,7 @@ async def test_missing_id_token_rejected_on_authorization_code(): token_response={"access_token": "at"}, litellm_user_id="user-a", grant_type="authorization_code", + refresh_ownership=None, jwks_fetcher=_jwks_fetcher, caller_principal_loader=_caller_loader("alice@example.com"), ) @@ -116,18 +129,89 @@ async def test_missing_id_token_rejected_on_authorization_code(): @pytest.mark.asyncio -async def test_refresh_without_id_token_allowed(): +async def test_refresh_without_id_token_allowed_when_presented_token_matches_stored_credential(): result: Final = await enforce_oauth_identity_binding( server=_server(), token_response={"access_token": "at"}, litellm_user_id="user-a", grant_type="refresh_token", + refresh_ownership=RefreshTokenPresented("rt-1"), jwks_fetcher=_jwks_fetcher, caller_principal_loader=_caller_loader("alice@example.com"), + stored_refresh_token_loader=_stored_refresh_token_loader("rt-1"), ) assert result is None +@pytest.mark.asyncio +async def test_refresh_without_id_token_rejects_different_presented_token(): + with pytest.raises(HTTPException) as exc_info: + await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at"}, + litellm_user_id="user-a", + grant_type="refresh_token", + refresh_ownership=RefreshTokenPresented("rt-stolen"), + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + stored_refresh_token_loader=_stored_refresh_token_loader("rt-1"), + ) + assert exc_info.value.status_code == 403 + assert exc_info.value.detail["error"] == "oauth_identity_binding_failed" + + +@pytest.mark.asyncio +async def test_refresh_without_id_token_rejects_missing_stored_token(): + with pytest.raises(HTTPException) as exc_info: + await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at"}, + litellm_user_id="user-a", + grant_type="refresh_token", + refresh_ownership=RefreshTokenPresented("rt-1"), + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + stored_refresh_token_loader=_stored_refresh_token_loader(None), + ) + assert exc_info.value.status_code == 403 + assert exc_info.value.detail["error"] == "oauth_identity_binding_failed" + + +@pytest.mark.asyncio +async def test_refresh_without_id_token_passes_when_bridge_proves_ownership(): + async def fail_if_called(_user_id: str, _server_id: str) -> str | None: + raise AssertionError("stored refresh token loader should not be called") + + result: Final = await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at"}, + litellm_user_id="user-a", + grant_type="refresh_token", + refresh_ownership=RefreshOwnershipProven(), + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + stored_refresh_token_loader=fail_if_called, + ) + assert result is None + + +@pytest.mark.asyncio +async def test_refresh_without_id_token_rejects_without_ownership_proof(): + with pytest.raises(HTTPException) as exc_info: + await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at"}, + litellm_user_id="user-a", + grant_type="refresh_token", + refresh_ownership=None, + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + stored_refresh_token_loader=_stored_refresh_token_loader("rt-1"), + ) + assert exc_info.value.status_code == 403 + assert exc_info.value.detail["error"] == "oauth_identity_binding_failed" + + @pytest.mark.asyncio async def test_refresh_with_mismatched_id_token_rejected(): token: Final = _sign_id_token({"email": "mallory@example.com", "email_verified": True}) @@ -137,6 +221,7 @@ async def test_refresh_with_mismatched_id_token_rejected(): token_response={"access_token": "at", "id_token": token}, litellm_user_id="user-a", grant_type="refresh_token", + refresh_ownership=None, jwks_fetcher=_jwks_fetcher, caller_principal_loader=_caller_loader("alice@example.com"), ) @@ -150,6 +235,7 @@ async def test_audit_mode_logs_but_does_not_reject(): token_response={"access_token": "at", "id_token": token}, litellm_user_id="user-a", grant_type="authorization_code", + refresh_ownership=None, jwks_fetcher=_jwks_fetcher, caller_principal_loader=_caller_loader("alice@example.com"), ) @@ -165,6 +251,7 @@ async def test_unverified_email_rejected(): token_response={"access_token": "at", "id_token": token}, litellm_user_id="user-a", grant_type="authorization_code", + refresh_ownership=None, jwks_fetcher=_jwks_fetcher, caller_principal_loader=_caller_loader("alice@example.com"), ) @@ -187,6 +274,7 @@ async def test_wrong_issuer_rejected(): token_response={"access_token": "at", "id_token": token}, litellm_user_id="user-a", grant_type="authorization_code", + refresh_ownership=None, jwks_fetcher=_jwks_fetcher, caller_principal_loader=_caller_loader("alice@example.com"), ) @@ -202,6 +290,7 @@ async def test_no_litellm_identity_rejected(): token_response={"access_token": "at", "id_token": token}, litellm_user_id=None, grant_type="authorization_code", + refresh_ownership=None, jwks_fetcher=_jwks_fetcher, caller_principal_loader=_caller_loader("alice@example.com"), ) @@ -215,6 +304,7 @@ async def test_disabled_binding_is_noop(): token_response={"access_token": "at"}, litellm_user_id=None, grant_type="authorization_code", + refresh_ownership=None, jwks_fetcher=_jwks_fetcher, caller_principal_loader=_caller_loader(None), ) @@ -234,7 +324,32 @@ async def test_no_binding_is_noop(): token_response={"access_token": "at"}, litellm_user_id=None, grant_type="authorization_code", + refresh_ownership=None, jwks_fetcher=_jwks_fetcher, caller_principal_loader=_caller_loader(None), ) assert result is None + + +def test_identity_binding_requires_non_empty_audiences(): + with pytest.raises(ValidationError): + MCPOAuthIdentityBinding(mode="enforce", issuer=ISSUER, audiences=[]) + with pytest.raises(ValidationError): + MCPOAuthIdentityBinding(mode="enforce", issuer=ISSUER) + + +@pytest.mark.asyncio +async def test_wrong_audience_rejected(): + token: Final = _sign_id_token({"aud": "other-client", "email": "alice@example.com", "email_verified": True}) + with pytest.raises(HTTPException) as exc_info: + await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at", "id_token": token}, + litellm_user_id="user-a", + grant_type="authorization_code", + refresh_ownership=None, + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + ) + assert exc_info.value.status_code == 403 + assert exc_info.value.detail["error"] == "oauth_identity_binding_failed" diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index bfbbf816e69..f3510e7e0a8 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -4674,7 +4674,11 @@ async def test_store_mcp_oauth_user_credential_blocked_when_identity_binding_enf name=server_id, url="https://mcp.example.com", transport=MCPTransport.http, - oauth_identity_binding=MCPOAuthIdentityBinding(mode="enforce", issuer="https://idp.example.com"), + oauth_identity_binding=MCPOAuthIdentityBinding( + mode="enforce", + issuer="https://idp.example.com", + audiences=["litellm-client"], + ), ) store_mock = AsyncMock(return_value=None) From 35db36ee936643d36391eab1cb3c9f5031284b0d Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 2 Sep 2026 16:04:53 +0000 Subject: [PATCH 06/33] fix(mcp): satisfy type discipline lint gate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/discoverable_endpoints.py | 4 +- .../mcp_server/oauth_identity_binding.py | 39 ++++++++++++------- .../mcp_management_endpoints.py | 6 +-- .../types/mcp_server/mcp_server_manager.py | 2 +- 4 files changed, 31 insertions(+), 20 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 88f2c47f5fc..5f54bbb4958 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -1048,13 +1048,13 @@ async def exchange_token_with_server( refresh_request_scope = scope or bridge_upstream_scope if refresh_request_scope: token_data["scope"] = refresh_request_scope - refresh_ownership = ( + refresh_ownership = ( # rebind-ok: grant-specific branches assign one ownership value RefreshOwnershipProven() if bridge_upstream_refresh is not None else RefreshTokenPresented(upstream_refresh_token) ) else: - refresh_ownership = None + refresh_ownership = None # rebind-ok: grant-specific branches assign one ownership value if not code: raise HTTPException( status_code=400, diff --git a/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py b/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py index d4339c15e1a..b73227f3413 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py @@ -8,12 +8,13 @@ claim is compared to the caller's trusted LiteLLM identity. Mismatches fail clos and are logged in audit mode. """ -from collections.abc import Awaitable, Callable, Mapping +from collections.abc import Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass -from typing import Final, Literal +from typing import Final, Literal, TypeAlias import jwt from fastapi import HTTPException +from jwt.types import Options from typing_extensions import assert_never from litellm._logging import verbose_logger @@ -36,11 +37,20 @@ _ALLOWED_ID_TOKEN_ALGORITHMS: Final = ( _JWKS_CACHE_TTL_SECONDS: Final = 3600 _jwks_cache: Final = InMemoryCache(default_ttl=_JWKS_CACHE_TTL_SECONDS) -JwksFetcher = Callable[[MCPOAuthIdentityBinding], Awaitable[list[Mapping[str, object]]]] -CallerPrincipalLoader = Callable[[str, MCPOAuthIdentityBinding], Awaitable[str | None]] -StoredRefreshTokenLoader = Callable[[str, str], Awaitable[str | None]] +JwksFetcher: TypeAlias = Callable[ + [MCPOAuthIdentityBinding], # mutable-ok: Callable parameter syntax requires a list + Awaitable[Sequence[Mapping[str, object]]], +] +CallerPrincipalLoader: TypeAlias = Callable[ + [str, MCPOAuthIdentityBinding], # mutable-ok: Callable parameter syntax requires a list + Awaitable[str | None], +] +StoredRefreshTokenLoader: TypeAlias = Callable[ + [str, str], # mutable-ok: Callable parameter syntax requires a list + Awaitable[str | None], +] -_RejectionCode = Literal["oauth_principal_mismatch", "oauth_identity_binding_failed"] +_RejectionCode: TypeAlias = Literal["oauth_principal_mismatch", "oauth_identity_binding_failed"] @dataclass(frozen=True, slots=True) @@ -59,10 +69,10 @@ class RefreshTokenPresented: refresh_token: str -RefreshOwnership = RefreshOwnershipProven | RefreshTokenPresented | None +RefreshOwnership: TypeAlias = RefreshOwnershipProven | RefreshTokenPresented | None -async def _fetch_issuer_jwks(binding: MCPOAuthIdentityBinding) -> list[Mapping[str, object]]: +async def _fetch_issuer_jwks(binding: MCPOAuthIdentityBinding) -> Sequence[Mapping[str, object]]: jwks_url: Final[str] = binding.jwks_url or await _discover_jwks_url(binding.issuer) cached: Final = await _jwks_cache.async_get_cache(jwks_url) if isinstance(cached, list): @@ -90,12 +100,12 @@ async def _discover_jwks_url(issuer: str) -> str: return jwks_uri -def _select_signing_key(id_token: str, keys: list[Mapping[str, object]]) -> "jwt.PyJWK | _BindingRejection": +def _select_signing_key(id_token: str, keys: Sequence[Mapping[str, object]]) -> "jwt.PyJWK | _BindingRejection": header: Final = jwt.get_unverified_header(id_token) kid: Final = header.get("kid") for key in keys: if kid is None or key.get("kid") == kid: - return jwt.PyJWK(dict(key)) + return jwt.PyJWK(dict(key)) # mutable-ok: PyJWT requires a concrete JWK dictionary return _BindingRejection( code="oauth_identity_binding_failed", description=f"id_token signing key (kid={kid!r}) not found in the issuer's JWKS", @@ -108,13 +118,14 @@ def _decode_id_token( signing_key: "jwt.PyJWK", ) -> "Mapping[str, object] | _BindingRejection": try: + decode_options: Final[Options] = {"require": ("iss", "exp")} return jwt.decode( id_token, signing_key.key, - algorithms=list(_ALLOWED_ID_TOKEN_ALGORITHMS), + algorithms=_ALLOWED_ID_TOKEN_ALGORITHMS, issuer=binding.issuer, audience=binding.audiences, - options={"require": ["iss", "exp"]}, + options=decode_options, ) except jwt.InvalidTokenError as exc: return _BindingRejection( @@ -160,10 +171,10 @@ async def _load_caller_principal(litellm_user_id: str, binding: MCPOAuthIdentity async def _load_stored_refresh_token(litellm_user_id: str, server_id: str) -> str | None: try: - from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 + from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 # keep database imports lazy get_user_oauth_credential, ) - from litellm.proxy.utils import get_prisma_client_or_throw # noqa: PLC0415 + from litellm.proxy.utils import get_prisma_client_or_throw # noqa: PLC0415 # keep database imports lazy prisma_client: Final = get_prisma_client_or_throw( "Database not connected. Cannot verify OAuth refresh token ownership." diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 92a6997abff..109030a2df5 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -259,7 +259,7 @@ if MCP_AVAILABLE: if normalize_upstream_header_name(raw) is None: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, - detail={ + detail={ # mutable-ok: FastAPI exception detail requires a JSON-serializable dictionary "error": ( f"Invalid upstream_token_header {raw!r}: must be a valid HTTP header name " "(RFC 7230 token, e.g. 'esb-oauth')" @@ -2256,7 +2256,7 @@ if MCP_AVAILABLE: """Persist the OAuth2 access token obtained by the calling user.""" prisma_client: Final = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") await _authorize_and_fetch_mcp_server(prisma_client, user_api_key_dict, server_id) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # keep manager import lazy global_mcp_server_manager as _manager, ) @@ -2267,7 +2267,7 @@ if MCP_AVAILABLE: if binding is not None and binding.mode == "enforce": raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, - detail={ + detail={ # mutable-ok: FastAPI exception detail requires a JSON-serializable dictionary "error": "oauth_identity_binding_enforced", "error_description": ( "Direct credential storage is disabled for this server: its OAuth identity " diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index cebc7decac4..c38f1c366b1 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -52,7 +52,7 @@ class MCPOAuthIdentityBinding(BaseModel): mode: Literal["disabled", "audit", "enforce"] = "disabled" issuer: str jwks_url: str | None = None - audiences: list[str] = Field(min_length=1) + audiences: list[str] = Field(min_length=1) # mutable-ok: public Pydantic schema requires list values principal_claim: str = "email" caller_field: Literal["user_email", "user_id"] = "user_email" require_email_verified: bool = True From 22e3fcba54c2409a46e04ba81bd6885cbfc0442a Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 2 Sep 2026 16:05:37 +0000 Subject: [PATCH 07/33] fix(mcp): restore unrelated exception formatting Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/management_endpoints/mcp_management_endpoints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 109030a2df5..fd6294fd77e 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -259,7 +259,7 @@ if MCP_AVAILABLE: if normalize_upstream_header_name(raw) is None: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, - detail={ # mutable-ok: FastAPI exception detail requires a JSON-serializable dictionary + detail={ "error": ( f"Invalid upstream_token_header {raw!r}: must be a valid HTTP header name " "(RFC 7230 token, e.g. 'esb-oauth')" From 7ca9f28575021d026d28bb0bfb11a667ba3add77 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 2 Sep 2026 16:33:59 +0000 Subject: [PATCH 08/33] test(mcp): cover OAuth identity binding paths Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/oauth_identity_binding.py | 2 +- .../mcp_server/test_discoverable_endpoints.py | 126 ++++++++++ .../mcp_server/test_oauth_identity_binding.py | 229 ++++++++++++++++++ 3 files changed, 356 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py b/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py index b73227f3413..58ec3cfdb4f 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py @@ -234,7 +234,7 @@ async def _evaluate_binding( description="the presented refresh_token is not the caller's stored credential for this server", ) return None - assert_never(refresh_ownership) + assert_never(refresh_ownership) # pragma: no cover if not litellm_user_id: return _BindingRejection( code="oauth_identity_binding_failed", diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 598e9276423..4ae1b105dbf 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -7816,6 +7816,132 @@ async def test_token_exchange_pairs_client_secret_with_server_client_id(): assert "client_secret" not in sent +@pytest.mark.asyncio +async def test_token_exchange_refresh_passes_presented_refresh_ownership(): + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + exchange_token_with_server, + ) + from litellm.proxy._experimental.mcp_server.oauth_identity_binding import RefreshTokenPresented + from litellm.proxy._types import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer + + server = MCPServer( + server_id="srv-1", + name="srv-1", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="cid", + token_url="https://provider.example/token", + oauth_identity_binding=MCPOAuthIdentityBinding( + mode="enforce", + issuer="https://provider.example", + audiences=["cid"], + ), + ) + request = MagicMock(spec=Request) + request.base_url = "https://litellm.example.com/" + request.headers = {} + response = MagicMock() + response.json.return_value = {"access_token": "at"} + response.raise_for_status = MagicMock() + client = MagicMock() + client.post = AsyncMock(return_value=response) + enforce = AsyncMock() + + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=client, + ), + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._extract_user_id_from_request", + new=AsyncMock(return_value="user-a"), + ), + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.enforce_oauth_identity_binding", + new=enforce, + ), + ): + await exchange_token_with_server( + request=request, + mcp_server=server, + grant_type="refresh_token", + code=None, + redirect_uri=None, + client_id="cid", + client_secret=None, + code_verifier=None, + refresh_token="rt-1", + ) + + ownership = enforce.await_args.kwargs["refresh_ownership"] + assert isinstance(ownership, RefreshTokenPresented) + assert ownership.refresh_token == "rt-1" + + +@pytest.mark.asyncio +async def test_token_exchange_authorization_code_passes_no_refresh_ownership(): + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + exchange_token_with_server, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer + + server = MCPServer( + server_id="srv-1", + name="srv-1", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="cid", + token_url="https://provider.example/token", + oauth_identity_binding=MCPOAuthIdentityBinding( + mode="enforce", + issuer="https://provider.example", + audiences=["cid"], + ), + ) + request = MagicMock(spec=Request) + request.base_url = "https://litellm.example.com/" + request.headers = {} + response = MagicMock() + response.json.return_value = {"access_token": "at"} + response.raise_for_status = MagicMock() + client = MagicMock() + client.post = AsyncMock(return_value=response) + enforce = AsyncMock() + + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=client, + ), + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._extract_user_id_from_request", + new=AsyncMock(return_value="user-a"), + ), + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.enforce_oauth_identity_binding", + new=enforce, + ), + ): + await exchange_token_with_server( + request=request, + mcp_server=server, + grant_type="authorization_code", + code="auth-code", + redirect_uri="https://litellm.example.com/callback", + client_id="cid", + client_secret=None, + code_verifier=None, + ) + + assert enforce.await_args.kwargs["refresh_ownership"] is None + + def _upstream_token_response(status_code: int, *, json_body: object = None, text_body: str = "") -> "httpx.Response": import httpx diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_identity_binding.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_identity_binding.py index de59020c44c..c2662f9f298 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_identity_binding.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_identity_binding.py @@ -1,6 +1,8 @@ import time from collections.abc import Mapping +from types import SimpleNamespace from typing import Final +from unittest.mock import AsyncMock, MagicMock, patch import jwt import pytest @@ -12,6 +14,11 @@ from pydantic import ValidationError from litellm.proxy._experimental.mcp_server.oauth_identity_binding import ( RefreshOwnershipProven, RefreshTokenPresented, + _discover_jwks_url, + _fetch_issuer_jwks, + _load_caller_principal, + _load_stored_refresh_token, + _select_signing_key, enforce_oauth_identity_binding, ) from litellm.types.mcp import MCPTransport @@ -79,6 +86,121 @@ def _server(mode: str = "enforce", **binding_overrides: object) -> MCPServer: ) +@pytest.mark.asyncio +async def test_fetch_issuer_jwks_fetches_and_caches_keys(): + binding: Final = _server(jwks_url="https://idp.example.com/jwks-cache").oauth_identity_binding + assert binding is not None + response: Final = MagicMock() + response.json.return_value = {"keys": [_PUBLIC_JWK]} + client: Final = MagicMock() + client.get = AsyncMock(return_value=response) + + with patch( + "litellm.proxy._experimental.mcp_server.oauth_identity_binding.get_async_httpx_client", + return_value=client, + ): + first: Final = await _fetch_issuer_jwks(binding) + second: Final = await _fetch_issuer_jwks(binding) + + assert first == second == [_PUBLIC_JWK] + client.get.assert_awaited_once_with("https://idp.example.com/jwks-cache") + + +@pytest.mark.asyncio +async def test_fetch_issuer_jwks_rejects_malformed_document(): + binding: Final = _server(jwks_url="https://idp.example.com/jwks-invalid").oauth_identity_binding + assert binding is not None + response: Final = MagicMock() + response.json.return_value = {} + client: Final = MagicMock() + client.get = AsyncMock(return_value=response) + + with patch( + "litellm.proxy._experimental.mcp_server.oauth_identity_binding.get_async_httpx_client", + return_value=client, + ): + with pytest.raises(TypeError, match="has no 'keys' array"): + await _fetch_issuer_jwks(binding) + + +@pytest.mark.asyncio +async def test_discover_jwks_url_returns_provider_uri(): + response: Final = MagicMock() + response.json.return_value = {"jwks_uri": "https://idp.example.com/jwks"} + client: Final = MagicMock() + client.get = AsyncMock(return_value=response) + + with patch( + "litellm.proxy._experimental.mcp_server.oauth_identity_binding.get_async_httpx_client", + return_value=client, + ): + result: Final = await _discover_jwks_url("https://idp.example.com/") + + assert result == "https://idp.example.com/jwks" + client.get.assert_awaited_once_with("https://idp.example.com/.well-known/openid-configuration") + + +@pytest.mark.asyncio +async def test_discover_jwks_url_rejects_missing_provider_uri(): + response: Final = MagicMock() + response.json.return_value = {} + client: Final = MagicMock() + client.get = AsyncMock(return_value=response) + + with patch( + "litellm.proxy._experimental.mcp_server.oauth_identity_binding.get_async_httpx_client", + return_value=client, + ): + with pytest.raises(ValueError, match="returned no jwks_uri"): + await _discover_jwks_url("https://idp.example.com") + + +def test_select_signing_key_returns_matching_key_or_rejection(): + token: Final = _sign_id_token({"email": "alice@example.com", "email_verified": True}) + + selected: Final = _select_signing_key(token, [_PUBLIC_JWK]) + rejected: Final = _select_signing_key(token, []) + + assert selected.__class__.__name__ == "PyJWK" + assert rejected.code == "oauth_identity_binding_failed" + + +@pytest.mark.asyncio +async def test_load_caller_principal_supports_user_id_and_database_email(): + user_id_binding: Final = _server(caller_field="user_id").oauth_identity_binding + assert user_id_binding is not None + assert await _load_caller_principal("user-a", user_id_binding) == "user-a" + + with patch( + "litellm.proxy._experimental.mcp_server.bridge_token_flow.load_active_user_by_id", + new=AsyncMock(side_effect=["no_active_key", SimpleNamespace(user_email="alice@example.com")]), + ): + assert await _load_caller_principal("user-a", _server().oauth_identity_binding) is None + assert await _load_caller_principal("user-a", _server().oauth_identity_binding) == "alice@example.com" + + +@pytest.mark.asyncio +async def test_load_stored_refresh_token_returns_credential_and_fails_closed(): + get_credential: Final = AsyncMock(return_value={"refresh_token": "rt-1"}) + with ( + patch( + "litellm.proxy._experimental.mcp_server.db.get_user_oauth_credential", + new=get_credential, + ), + patch( + "litellm.proxy.utils.get_prisma_client_or_throw", + return_value="prisma", + ), + ): + assert await _load_stored_refresh_token("user-a", "srv-1") == "rt-1" + + with patch( + "litellm.proxy.utils.get_prisma_client_or_throw", + side_effect=RuntimeError("database unavailable"), + ): + assert await _load_stored_refresh_token("user-a", "srv-1") is None + + @pytest.mark.asyncio async def test_matching_principal_passes(): token: Final = _sign_id_token({"email": "Alice@Example.com", "email_verified": True}) @@ -112,6 +234,113 @@ async def test_mismatched_principal_rejected(): assert exc_info.value.detail["credential_stored"] is False +@pytest.mark.asyncio +async def test_missing_upstream_principal_rejected(): + token: Final = _sign_id_token({"email_verified": True}) + with pytest.raises(HTTPException) as exc_info: + await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at", "id_token": token}, + litellm_user_id="user-a", + grant_type="authorization_code", + refresh_ownership=None, + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + ) + assert exc_info.value.detail["error"] == "oauth_identity_binding_failed" + assert "no usable" in exc_info.value.detail["error_description"] + + +@pytest.mark.asyncio +async def test_jwks_fetch_failure_is_rejected(): + async def fail(_binding: MCPOAuthIdentityBinding) -> list[Mapping[str, object]]: + raise RuntimeError("jwks unavailable") + + token: Final = _sign_id_token({"email": "alice@example.com", "email_verified": True}) + with pytest.raises(HTTPException) as exc_info: + await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at", "id_token": token}, + litellm_user_id="user-a", + grant_type="authorization_code", + refresh_ownership=None, + jwks_fetcher=fail, + caller_principal_loader=_caller_loader("alice@example.com"), + ) + assert exc_info.value.detail["error"] == "oauth_identity_binding_failed" + assert "jwks unavailable" in exc_info.value.detail["error_description"] + + +@pytest.mark.asyncio +async def test_missing_signing_key_is_rejected(): + token: Final = _sign_id_token({"email": "alice@example.com", "email_verified": True}) + + async def no_keys(_binding: MCPOAuthIdentityBinding) -> tuple[Mapping[str, object], ...]: + return () + + with pytest.raises(HTTPException) as exc_info: + await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at", "id_token": token}, + litellm_user_id="user-a", + grant_type="authorization_code", + refresh_ownership=None, + jwks_fetcher=no_keys, + caller_principal_loader=_caller_loader("alice@example.com"), + ) + assert exc_info.value.detail["error"] == "oauth_identity_binding_failed" + assert "signing key" in exc_info.value.detail["error_description"] + + +@pytest.mark.asyncio +async def test_missing_caller_principal_is_rejected(): + token: Final = _sign_id_token({"email": "alice@example.com", "email_verified": True}) + with pytest.raises(HTTPException) as exc_info: + await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at", "id_token": token}, + litellm_user_id="user-a", + grant_type="authorization_code", + refresh_ownership=None, + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader(None), + ) + assert exc_info.value.detail["error"] == "oauth_identity_binding_failed" + assert "has no" in exc_info.value.detail["error_description"] + + +@pytest.mark.asyncio +async def test_refresh_without_id_token_requires_litellm_identity_for_presented_token(): + with pytest.raises(HTTPException) as exc_info: + await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at"}, + litellm_user_id=None, + grant_type="refresh_token", + refresh_ownership=RefreshTokenPresented("rt-1"), + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + stored_refresh_token_loader=_stored_refresh_token_loader("rt-1"), + ) + assert exc_info.value.detail["error"] == "oauth_identity_binding_failed" + assert "no resolvable LiteLLM user identity" in exc_info.value.detail["error_description"] + + +@pytest.mark.asyncio +async def test_user_id_principal_matching_uses_exact_comparison(): + token: Final = _sign_id_token({"sub": "user-a"}) + result: Final = await enforce_oauth_identity_binding( + server=_server(principal_claim="sub", caller_field="user_id"), + token_response={"access_token": "at", "id_token": token}, + litellm_user_id="user-a", + grant_type="authorization_code", + refresh_ownership=None, + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("user-a"), + ) + assert result is None + + @pytest.mark.asyncio async def test_missing_id_token_rejected_on_authorization_code(): with pytest.raises(HTTPException) as exc_info: From 540e528d2190e6ed7292174e1985673a602a3018 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 2 Sep 2026 16:43:18 +0000 Subject: [PATCH 09/33] test(mcp): satisfy test quality gate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/test_discoverable_endpoints.py | 16 +++++++++------- .../mcp_server/test_oauth_identity_binding.py | 16 ++++++++-------- 2 files changed, 17 insertions(+), 15 deletions(-) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 4ae1b105dbf..33b71f8d50f 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -7851,15 +7851,15 @@ async def test_token_exchange_refresh_passes_presented_refresh_ownership(): enforce = AsyncMock() with ( - patch( + patch( # test-quality-ok: no injection seam exists for the exchange's HTTP and identity collaborators "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", return_value=client, ), - patch( + patch( # test-quality-ok: no injection seam exists for request identity extraction "litellm.proxy._experimental.mcp_server.discoverable_endpoints._extract_user_id_from_request", new=AsyncMock(return_value="user-a"), ), - patch( + patch( # test-quality-ok: captures the ownership value at the exchange boundary "litellm.proxy._experimental.mcp_server.discoverable_endpoints.enforce_oauth_identity_binding", new=enforce, ), @@ -7915,20 +7915,20 @@ async def test_token_exchange_authorization_code_passes_no_refresh_ownership(): enforce = AsyncMock() with ( - patch( + patch( # test-quality-ok: no injection seam exists for the exchange's HTTP and identity collaborators "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", return_value=client, ), - patch( + patch( # test-quality-ok: no injection seam exists for request identity extraction "litellm.proxy._experimental.mcp_server.discoverable_endpoints._extract_user_id_from_request", new=AsyncMock(return_value="user-a"), ), - patch( + patch( # test-quality-ok: captures the ownership value at the exchange boundary "litellm.proxy._experimental.mcp_server.discoverable_endpoints.enforce_oauth_identity_binding", new=enforce, ), ): - await exchange_token_with_server( + result = await exchange_token_with_server( request=request, mcp_server=server, grant_type="authorization_code", @@ -7939,6 +7939,8 @@ async def test_token_exchange_authorization_code_passes_no_refresh_ownership(): code_verifier=None, ) + assert result.status_code == 200 + assert json.loads(result.body)["access_token"] == "at" assert enforce.await_args.kwargs["refresh_ownership"] is None diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_identity_binding.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_identity_binding.py index c2662f9f298..36364c32ae5 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_identity_binding.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_identity_binding.py @@ -95,7 +95,7 @@ async def test_fetch_issuer_jwks_fetches_and_caches_keys(): client: Final = MagicMock() client.get = AsyncMock(return_value=response) - with patch( + with patch( # test-quality-ok: no HTTP dependency injection seam exists for JWKS fetching "litellm.proxy._experimental.mcp_server.oauth_identity_binding.get_async_httpx_client", return_value=client, ): @@ -115,7 +115,7 @@ async def test_fetch_issuer_jwks_rejects_malformed_document(): client: Final = MagicMock() client.get = AsyncMock(return_value=response) - with patch( + with patch( # test-quality-ok: no HTTP dependency injection seam exists for JWKS fetching "litellm.proxy._experimental.mcp_server.oauth_identity_binding.get_async_httpx_client", return_value=client, ): @@ -130,7 +130,7 @@ async def test_discover_jwks_url_returns_provider_uri(): client: Final = MagicMock() client.get = AsyncMock(return_value=response) - with patch( + with patch( # test-quality-ok: no HTTP dependency injection seam exists for OIDC discovery "litellm.proxy._experimental.mcp_server.oauth_identity_binding.get_async_httpx_client", return_value=client, ): @@ -147,7 +147,7 @@ async def test_discover_jwks_url_rejects_missing_provider_uri(): client: Final = MagicMock() client.get = AsyncMock(return_value=response) - with patch( + with patch( # test-quality-ok: no HTTP dependency injection seam exists for OIDC discovery "litellm.proxy._experimental.mcp_server.oauth_identity_binding.get_async_httpx_client", return_value=client, ): @@ -171,7 +171,7 @@ async def test_load_caller_principal_supports_user_id_and_database_email(): assert user_id_binding is not None assert await _load_caller_principal("user-a", user_id_binding) == "user-a" - with patch( + with patch( # test-quality-ok: caller loading is a lazy database boundary without injection "litellm.proxy._experimental.mcp_server.bridge_token_flow.load_active_user_by_id", new=AsyncMock(side_effect=["no_active_key", SimpleNamespace(user_email="alice@example.com")]), ): @@ -183,18 +183,18 @@ async def test_load_caller_principal_supports_user_id_and_database_email(): async def test_load_stored_refresh_token_returns_credential_and_fails_closed(): get_credential: Final = AsyncMock(return_value={"refresh_token": "rt-1"}) with ( - patch( + patch( # test-quality-ok: stored-token loading is a lazy database boundary without injection "litellm.proxy._experimental.mcp_server.db.get_user_oauth_credential", new=get_credential, ), - patch( + patch( # test-quality-ok: stored-token loading is a lazy database boundary without injection "litellm.proxy.utils.get_prisma_client_or_throw", return_value="prisma", ), ): assert await _load_stored_refresh_token("user-a", "srv-1") == "rt-1" - with patch( + with patch( # test-quality-ok: stored-token loading is a lazy database boundary without injection "litellm.proxy.utils.get_prisma_client_or_throw", side_effect=RuntimeError("database unavailable"), ): From eddb29d90ed45849967c2318f15c2edd3e991d8f Mon Sep 17 00:00:00 2001 From: michelligabriele Date: Thu, 3 Sep 2026 11:33:44 +0200 Subject: [PATCH 10/33] fix(proxy): dedup latest health checks in SQL and gate the DB save per window --- litellm/constants.py | 1 + litellm/proxy/db/health_check_latest.py | 96 +++++++++ litellm/proxy/proxy_server.py | 58 +++++- litellm/proxy/utils.py | 52 +---- .../proxy/db/test_health_check_latest.py | 111 ++++++++++ .../proxy_server/test_background_health.py | 81 ++++++++ .../proxy/test_health_check_functions.py | 190 ++++++++++-------- .../test_prisma_client_health.py | 74 ++++--- 8 files changed, 502 insertions(+), 161 deletions(-) create mode 100644 litellm/proxy/db/health_check_latest.py create mode 100644 tests/test_litellm/proxy/db/test_health_check_latest.py diff --git a/litellm/constants.py b/litellm/constants.py index ef9329b9dfc..1f5bd673997 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1619,6 +1619,7 @@ CLOUDZERO_EXPORT_USAGE_DATA_JOB_NAME: Final = "cloudzero_export_usage_data" MAVVRIK_FOCUS_EXPORT_JOB_NAME: Final = "mavvrik_focus_export_usage_data" CLOUDZERO_MAX_FETCHED_DATA_RECORDS: Final = int(os.getenv("CLOUDZERO_MAX_FETCHED_DATA_RECORDS", 50000)) SPEND_LOG_CLEANUP_JOB_NAME: Final = "spend_log_cleanup" +BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME: Final = "background_health_check_db_save" KEY_ROTATION_JOB_NAME: Final = "litellm_key_rotation_job" EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME: Final = "litellm_expired_ui_session_key_cleanup_job" WEEKLY_SPEND_REPORT_JOB_ID: Final = "weekly_spend_report_job" diff --git a/litellm/proxy/db/health_check_latest.py b/litellm/proxy/db/health_check_latest.py new file mode 100644 index 00000000000..35bc838379c --- /dev/null +++ b/litellm/proxy/db/health_check_latest.py @@ -0,0 +1,96 @@ +""" +Latest health-check row per model, deduplicated by Postgres. + +prisma-client-py's ``find_many(distinct=...)`` dedups client-side: the emitted +SQL carries no DISTINCT, so the whole append-only history table streams to the +worker on every call. ``SELECT DISTINCT ON`` keeps the transfer at one row per +(model_id, model_name) and is served by the matching descending index. +""" + +from __future__ import annotations + +import json +from collections.abc import Sequence +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Final + +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, field_validator + +from litellm._logging import verbose_proxy_logger + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + +LATEST_HEALTH_CHECKS_SQL: Final = """ +SELECT DISTINCT ON ("model_id", "model_name") + "health_check_id", "model_name", "model_id", "status", + "healthy_count", "unhealthy_count", "error_message", + "response_time_ms", "details", "checked_by", + "checked_at", "created_at", "updated_at" +FROM "LiteLLM_HealthCheckTable" +ORDER BY "model_id" ASC, "model_name" ASC, "checked_at" DESC +""" + +LATEST_HEALTH_CHECKS_FOR_MODELS_SQL: Final = """ +SELECT DISTINCT ON ("model_id", "model_name") + "health_check_id", "model_name", "model_id", "status", + "healthy_count", "unhealthy_count", "error_message", + "response_time_ms", "details", "checked_by", + "checked_at", "created_at", "updated_at" +FROM "LiteLLM_HealthCheckTable" +WHERE "model_name" = ANY($1) +ORDER BY "model_id" ASC, "model_name" ASC, "checked_at" DESC +""" + + +class LatestHealthCheckRow(BaseModel): + model_config = ConfigDict(frozen=True, protected_namespaces=()) + + health_check_id: str + model_name: str + model_id: str | None = None + status: str + healthy_count: int = 0 + unhealthy_count: int = 0 + error_message: str | None = None + response_time_ms: float | None = None + details: JsonValue | None = None + checked_by: str | None = None + checked_at: datetime + created_at: datetime + updated_at: datetime + + @field_validator("details", mode="before") + @classmethod + def _decode_json_text(cls, value: object) -> object: + return json.loads(value) if isinstance(value, str) else value + + @field_validator("checked_at", "created_at", "updated_at") + @classmethod + def _assume_utc(cls, value: datetime) -> datetime: + return value.replace(tzinfo=timezone.utc) if value.tzinfo is None else value + + +_ROWS_ADAPTER: Final = TypeAdapter(tuple[LatestHealthCheckRow, ...]) + + +async def fetch_latest_health_checks(prisma_client: PrismaClient) -> tuple[LatestHealthCheckRow, ...]: + try: + rows: Final = await prisma_client.db.query_raw(LATEST_HEALTH_CHECKS_SQL) + return _ROWS_ADAPTER.validate_python(rows) + except Exception as query_err: # noqa: BLE001 # health decorates other reads; a driver error must not fail them + verbose_proxy_logger.error("Error getting all latest health checks: %s", query_err) + return () + + +async def fetch_latest_health_checks_for_models( + prisma_client: PrismaClient, model_names: Sequence[str] +) -> tuple[LatestHealthCheckRow, ...]: + if not model_names: + return () + try: + rows: Final = await prisma_client.db.query_raw(LATEST_HEALTH_CHECKS_FOR_MODELS_SQL, list(model_names)) + return _ROWS_ADAPTER.validate_python(rows) + except Exception as query_err: # noqa: BLE001 # a paged model list must not fail on its health decoration + verbose_proxy_logger.error("Error getting latest health checks for models: %s", query_err) + return () diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 27132c90e05..0e330272bdf 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -16,7 +16,16 @@ import threading import time import traceback import warnings -from collections.abc import AsyncGenerator, AsyncIterator, Callable, Collection, Mapping, MutableMapping, Sequence +from collections.abc import ( + AsyncGenerator, + AsyncIterator, + Awaitable, + Callable, + Collection, + Mapping, + MutableMapping, + Sequence, +) from datetime import datetime, timedelta, timezone from types import MappingProxyType, UnionType from typing import ( @@ -50,6 +59,7 @@ from litellm.constants import ( AIOHTTP_NEEDS_CLEANUP_CLOSED, AIOHTTP_TTL_DNS_CACHE, AUDIO_SPEECH_CHUNK_SIZE, + BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME, BASE_MCP_ROUTE, DAILY_TAG_SPEND_BATCH_MULTIPLIER, DEFAULT_MAX_RECURSE_DEPTH, @@ -224,7 +234,7 @@ def generate_feedback_box(): import contextlib from collections import defaultdict from contextlib import asynccontextmanager -from functools import lru_cache +from functools import lru_cache, partial import litellm import litellm._redis @@ -402,6 +412,7 @@ from litellm.proxy.config_resolvers.alerting import ( ) from litellm.proxy.container_endpoints.endpoints import router as container_router from litellm.proxy.credential_endpoints.endpoints import router as credential_router +from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import ( SPEND_LOG_CLEANUP_BOUND_SETTINGS, SpendLogCleanup, @@ -3580,12 +3591,35 @@ async def _run_direct_health_check_with_instrumentation( raise AssertionError("perform_health_check rejected every optional argument") +async def _window_gated_health_check_db_save( + save: Callable[[], Awaitable[None]], + pod_lock_manager: PodLockManager | None, + lock_ttl: int | None, +) -> None: + """ + Persist at most once per window fleet-wide: the lock is the "this window's save is + done" marker, so it is deliberately never released and expires with the interval. + """ + if pod_lock_manager is not None and pod_lock_manager.redis_cache is not None: + acquired: Final = await pod_lock_manager.acquire_lock( + cronjob_id=BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME, + ttl=lock_ttl, + allow_reentrant=False, + ) + if not acquired: + verbose_proxy_logger.debug("background_health_check_db_save_skipped another pod persisted this window") + return + await save() + + def _schedule_background_health_check_db_save( prisma_client, shared_health_manager, model_list: list, healthy_endpoints: list, unhealthy_endpoints: list, + pod_lock_manager: PodLockManager | None = None, + lock_ttl: int | None = None, ): """Fire-and-forget: persist health check results to DB if prisma is available.""" if prisma_client is None: @@ -3598,16 +3632,16 @@ def _schedule_background_health_check_db_save( checked_by: Final = shared_health_manager.pod_id if shared_health_manager is not None else "background_health_check" start_time: Final = time_module.time() - asyncio.create_task( - _save_background_health_checks_to_db( - prisma_client, - model_list, - healthy_endpoints, - unhealthy_endpoints, - start_time, - checked_by=checked_by, - ) + save: Final = partial( + _save_background_health_checks_to_db, + prisma_client, + model_list, + healthy_endpoints, + unhealthy_endpoints, + start_time, + checked_by=checked_by, ) + asyncio.create_task(_window_gated_health_check_db_save(save, pod_lock_manager, lock_ttl)) def _get_endpoint_exception_status(endpoint: dict, exceptions: dict) -> int: @@ -3914,6 +3948,8 @@ async def _run_background_health_check(): _llm_model_list, healthy_endpoints, unhealthy_endpoints, + pod_lock_manager=proxy_logging_obj.db_spend_update_writer.pod_lock_manager, + lock_ttl=health_check_interval, ) # Write health state to router cache for health-check-driven routing diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index cab2bd6d9db..01530762800 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -121,6 +121,11 @@ from litellm.proxy.db.exception_handler import ( PrismaDBExceptionHandler, call_with_db_reconnect_retry, ) +from litellm.proxy.db.health_check_latest import ( + LatestHealthCheckRow, + fetch_latest_health_checks, + fetch_latest_health_checks_for_models, +) from litellm.proxy.db.log_db_metrics import log_db_metrics from litellm.proxy.db.prisma_client import ( PrismaWrapper, @@ -5962,48 +5967,13 @@ class PrismaClient: verbose_proxy_logger.error("Error getting health check history: %s", e) return [] - async def get_all_latest_health_checks(self) -> "Sequence[prisma_models.LiteLLM_HealthCheckTable]": - """ - Get the latest health check for each model. + async def get_all_latest_health_checks(self) -> tuple[LatestHealthCheckRow, ...]: + """Latest health check per (model_id, model_name), deduplicated in Postgres.""" + return await fetch_latest_health_checks(self) - Uses DB-level DISTINCT ON (model_id, model_name) with ORDER BY checked_at DESC - (via Prisma ``distinct`` + ``order``) so we never load the full history into memory. - """ - try: - return await HealthCheckRepository(self).table.find_many( - distinct=["model_id", "model_name"], - order=[ - {"model_id": "asc"}, - {"model_name": "asc"}, - {"checked_at": "desc"}, - ], - ) - except Exception as e: - verbose_proxy_logger.error("Error getting all latest health checks: %s", e) - return [] - - async def get_latest_health_checks_for_models( - self, model_names: "Sequence[str]" - ) -> "Sequence[prisma_models.LiteLLM_HealthCheckTable]": - """ - Get the latest health check for each of the named models. - - Same DISTINCT ON as ``get_all_latest_health_checks``, bounded to the models asked - about, so a paged caller reads health for its page instead of for the whole table. - """ - if not model_names: - return () - latest_first: Final = (("model_id", "asc"), ("model_name", "asc"), ("checked_at", "desc")) - order: Final = [{field: direction} for field, direction in latest_first] # mutable-ok: prisma order is a list - try: - return await HealthCheckRepository(self).table.find_many( - where={"model_name": {"in": list(model_names)}}, # mutable-ok: prisma filters are dicts and lists - distinct=["model_id", "model_name"], # mutable-ok: prisma distinct takes a list - order=order, - ) - except Exception as e: # noqa: BLE001 # health decorates a list; a driver error must not fail the page - verbose_proxy_logger.error("Error getting latest health checks for models: %s", e) - return () + async def get_latest_health_checks_for_models(self, model_names: Sequence[str]) -> tuple[LatestHealthCheckRow, ...]: + """Same as ``get_all_latest_health_checks``, bounded to the named models.""" + return await fetch_latest_health_checks_for_models(self, model_names) ### HELPER FUNCTIONS ### diff --git a/tests/test_litellm/proxy/db/test_health_check_latest.py b/tests/test_litellm/proxy/db/test_health_check_latest.py new file mode 100644 index 00000000000..529e4d8f2e9 --- /dev/null +++ b/tests/test_litellm/proxy/db/test_health_check_latest.py @@ -0,0 +1,111 @@ +from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.proxy.db.health_check_latest import ( + LATEST_HEALTH_CHECKS_FOR_MODELS_SQL, + LATEST_HEALTH_CHECKS_SQL, + fetch_latest_health_checks, + fetch_latest_health_checks_for_models, +) + + +def _prisma(rows): + prisma = MagicMock() + prisma.db.query_raw = AsyncMock(return_value=rows) + return prisma + + +def _raw_row(**overrides): + row = { + "health_check_id": "hc-1", + "model_name": "gpt-4", + "model_id": "deployment-abc", + "status": "healthy", + "healthy_count": 1, + "unhealthy_count": 0, + "error_message": None, + "response_time_ms": 12.5, + "details": None, + "checked_by": "pod-1", + "checked_at": "2026-08-25T00:00:00+00:00", + "created_at": "2026-08-25T00:00:00+00:00", + "updated_at": "2026-08-25T00:00:00+00:00", + } + return {**row, **overrides} + + +@pytest.mark.asyncio +async def test_fetch_all_runs_one_distinct_on_query_with_no_parameters(): + """The dedup must be in the SQL: prisma find_many(distinct=...) streams the whole history table.""" + prisma = _prisma([]) + assert await fetch_latest_health_checks(prisma) == () + assert prisma.db.query_raw.await_args.args == (LATEST_HEALTH_CHECKS_SQL,) + assert 'DISTINCT ON ("model_id", "model_name")' in LATEST_HEALTH_CHECKS_SQL + assert '"checked_at" DESC' in LATEST_HEALTH_CHECKS_SQL + + +@pytest.mark.asyncio +async def test_raw_datetimes_come_back_tz_aware_with_or_without_an_offset(): + """The save path subtracts checked_at from datetime.now(timezone.utc); a naive value would TypeError.""" + naive = _raw_row(health_check_id="hc-naive", model_id=None, checked_at="2026-08-25T00:00:00") + aware = _raw_row(health_check_id="hc-aware", checked_at="2026-08-25T01:00:00+02:00") + rows = await fetch_latest_health_checks(_prisma([naive, aware])) + assert {row.health_check_id: (row.model_id, row.checked_at) for row in rows} == { + "hc-naive": (None, datetime(2026, 8, 25, 0, 0, tzinfo=timezone.utc)), + "hc-aware": ("deployment-abc", datetime(2026, 8, 24, 23, 0, tzinfo=timezone.utc)), + } + + +@pytest.mark.asyncio +async def test_json_details_decode_from_text_and_pass_through_as_dict(): + rows = await fetch_latest_health_checks( + _prisma( + [ + _raw_row(health_check_id="text", details='{"region": "eu"}'), + _raw_row(health_check_id="dict", details={"region": "us"}), + _raw_row(health_check_id="none", details=None), + ] + ) + ) + assert {row.health_check_id: row.details for row in rows} == { + "text": {"region": "eu"}, + "dict": {"region": "us"}, + "none": None, + } + + +@pytest.mark.asyncio +async def test_fetch_all_degrades_to_no_rows_when_the_query_fails(): + prisma = _prisma([]) + prisma.db.query_raw.side_effect = RuntimeError("db down") + assert await fetch_latest_health_checks(prisma) == () + + +@pytest.mark.asyncio +async def test_fetch_all_degrades_to_no_rows_for_a_malformed_row(): + assert await fetch_latest_health_checks(_prisma([{"unexpected": "shape"}])) == () + + +@pytest.mark.asyncio +async def test_fetch_for_models_binds_the_page_as_the_only_parameter(): + prisma = _prisma([_raw_row()]) + rows = await fetch_latest_health_checks_for_models(prisma, ("gpt-4", "claude-opus")) + assert prisma.db.query_raw.await_args.args == (LATEST_HEALTH_CHECKS_FOR_MODELS_SQL, ["gpt-4", "claude-opus"]) + assert [row.model_name for row in rows] == ["gpt-4"] + assert 'WHERE "model_name" = ANY($1)' in LATEST_HEALTH_CHECKS_FOR_MODELS_SQL + + +@pytest.mark.asyncio +async def test_fetch_for_models_skips_the_database_for_an_empty_page(): + prisma = _prisma([]) + assert await fetch_latest_health_checks_for_models(prisma, ()) == () + prisma.db.query_raw.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_fetch_for_models_degrades_to_no_rows_when_the_query_fails(): + prisma = _prisma([]) + prisma.db.query_raw.side_effect = RuntimeError("db down") + assert await fetch_latest_health_checks_for_models(prisma, ("gpt-4",)) == () diff --git a/tests/test_litellm/proxy/proxy_server/test_background_health.py b/tests/test_litellm/proxy/proxy_server/test_background_health.py index 990844369f7..4874f75752c 100644 --- a/tests/test_litellm/proxy/proxy_server/test_background_health.py +++ b/tests/test_litellm/proxy/proxy_server/test_background_health.py @@ -20,6 +20,7 @@ from unittest.mock import AsyncMock, MagicMock import pytest import litellm.proxy.proxy_server as proxy_server +from litellm.constants import BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME from litellm.proxy.proxy_server import ( _adaptive_router_flusher_loop, _get_endpoint_exception_status, @@ -245,6 +246,86 @@ async def test_schedule_background_health_check_db_save_invalid_no_event_loop_ra ) +def _lock_manager(redis_cache, acquired): + manager = MagicMock() + manager.redis_cache = redis_cache + manager.acquire_lock = AsyncMock(return_value=acquired) + manager.release_lock = AsyncMock() + return manager + + +def _capture_saves(monkeypatch): + saves = [] + + async def _fake_save(*_args, **kwargs): + saves.append(kwargs) + + import litellm.proxy.health_endpoints._health_endpoints as he + + monkeypatch.setattr(he, "_save_background_health_checks_to_db", _fake_save) + return saves + + +def _schedule_with(lock_manager): + _schedule_background_health_check_db_save( + prisma_client=MagicMock(), + shared_health_manager=None, + model_list=[], + healthy_endpoints=[], + unhealthy_endpoints=[], + pod_lock_manager=lock_manager, + lock_ttl=300, + ) + + +@pytest.mark.asyncio +async def test_schedule_background_health_check_db_save_skips_a_window_another_pod_persisted(monkeypatch): + saves = _capture_saves(monkeypatch) + lock_manager = _lock_manager(redis_cache=MagicMock(), acquired=False) + + _schedule_with(lock_manager) + await asyncio.sleep(0) + + assert saves == [] + + +@pytest.mark.asyncio +async def test_schedule_background_health_check_db_save_holds_the_window_lock_for_the_whole_interval(monkeypatch): + """The lock is the "saved this window" marker: never reentrant, TTL = interval, and never released.""" + saves = _capture_saves(monkeypatch) + lock_manager = _lock_manager(redis_cache=MagicMock(), acquired=True) + + _schedule_with(lock_manager) + await asyncio.sleep(0) + + assert normalize( + { + "saves": len(saves), + "lock_request": lock_manager.acquire_lock.await_args.kwargs, + "released": lock_manager.release_lock.await_count, + } + ) == { + "saves": 1, + "lock_request": { + "cronjob_id": BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME, + "ttl": 300, + "allow_reentrant": False, + }, + "released": 0, + } + + +@pytest.mark.asyncio +async def test_schedule_background_health_check_db_save_runs_ungated_without_redis(monkeypatch): + saves = _capture_saves(monkeypatch) + lock_manager = _lock_manager(redis_cache=None, acquired=True) + + _schedule_with(lock_manager) + await asyncio.sleep(0) + + assert (len(saves), lock_manager.acquire_lock.await_count) == (1, 0) + + # --------------------------------------------------------------------------- # _get_endpoint_exception_status # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/test_health_check_functions.py b/tests/test_litellm/proxy/test_health_check_functions.py index fdae11d517a..9adaddb032d 100644 --- a/tests/test_litellm/proxy/test_health_check_functions.py +++ b/tests/test_litellm/proxy/test_health_check_functions.py @@ -6,6 +6,8 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.db.health_check_latest import LatestHealthCheckRow from litellm.proxy.health_endpoints._health_endpoints import ( _aggregate_health_check_results, _build_model_param_to_info_mapping, @@ -13,6 +15,7 @@ from litellm.proxy.health_endpoints._health_endpoints import ( _save_background_health_checks_to_db, _save_health_check_results_if_changed, _save_health_check_to_db, + latest_health_checks_endpoint, ) from litellm.proxy.utils import PrismaClient @@ -451,96 +454,125 @@ async def test_save_background_health_checks_to_db_exception_handling(): # Function should complete without raising -@pytest.mark.asyncio -async def test_get_all_latest_health_checks_with_model_id(mock_prisma): - """Test get_all_latest_health_checks properly groups by model_id""" - mock_check2 = MagicMock() - mock_check2.model_id = "model-456" - mock_check2.model_name = "gpt-3.5-turbo" - mock_check2.checked_at = datetime.now(timezone.utc) - timedelta(minutes=5) - - mock_check3 = MagicMock() - mock_check3.model_id = "model-123" - mock_check3.model_name = "gpt-3.5-turbo" - mock_check3.checked_at = datetime.now(timezone.utc) - timedelta( - minutes=1 - ) # Latest for model-123 - - # Order by checked_at desc - mock_prisma.db.litellm_healthchecktable.find_many = AsyncMock( - return_value=[mock_check3, mock_check2] - ) - - result = await mock_prisma.get_all_latest_health_checks() - - # Should return 2 unique models (by model_id) - assert len(result) == 2 - - # Should have latest check for each model_id - model_ids = {check.model_id for check in result} - assert "model-123" in model_ids - assert "model-456" in model_ids - - # model-123 should have the latest check (1 minute ago) - model123_check = next(c for c in result if c.model_id == "model-123") - assert model123_check.checked_at == mock_check3.checked_at +def _raw_latest_row(model_name: str, model_id, checked_at: datetime) -> dict: + return { + "health_check_id": f"hc-{model_id or 'no-id'}-{model_name}", + "model_name": model_name, + "model_id": model_id, + "status": "healthy", + "healthy_count": 1, + "unhealthy_count": 0, + "error_message": None, + "response_time_ms": 10.0, + "details": None, + "checked_by": "pod-1", + "checked_at": checked_at.isoformat(), + "created_at": checked_at.isoformat(), + "updated_at": checked_at.isoformat(), + } @pytest.mark.asyncio -async def test_get_all_latest_health_checks_without_model_id(mock_prisma): - """Test get_all_latest_health_checks groups by model_name when model_id is None""" - mock_check2 = MagicMock() - mock_check2.model_id = None - mock_check2.model_name = "gpt-3.5-turbo" - mock_check2.checked_at = datetime.now(timezone.utc) - timedelta(minutes=1) # Latest - - mock_prisma.db.litellm_healthchecktable.find_many = AsyncMock( - return_value=[mock_check2] - ) - - result = await mock_prisma.get_all_latest_health_checks() - - # Should return 1 unique model (by model_name) - assert len(result) == 1 - assert result[0].model_name == "gpt-3.5-turbo" - assert result[0].checked_at == mock_check2.checked_at # Latest - - -@pytest.mark.asyncio -async def test_get_all_latest_health_checks_same_name_with_and_without_model_id( - mock_prisma, -): +async def test_get_all_latest_health_checks_keeps_every_distinct_group_with_its_own_checked_at(mock_prisma): """ - Same model_name can appear twice after DISTINCT ON: once keyed by (model_id, name) - and once by (NULL, name) — different Postgres groups than a single row with id. + Postgres owns the dedup. (id, name), (other id, name) and (NULL, name) are distinct groups and each row + must arrive typed, with its own checked_at, for the 1h re-save compare and the id-or-name lookup key. """ now = datetime.now(timezone.utc) - with_id = MagicMock() - with_id.model_id = "deployment-abc" - with_id.model_name = "gpt-4" - with_id.checked_at = now - timedelta(minutes=2) - - without_id = MagicMock() - without_id.model_id = None - without_id.model_name = "gpt-4" - without_id.checked_at = now - timedelta(minutes=1) - - mock_prisma.db.litellm_healthchecktable.find_many = AsyncMock( - return_value=[without_id, with_id] + mock_prisma.db.query_raw = AsyncMock( + return_value=[ + _raw_latest_row("gpt-3.5-turbo", "model-123", now - timedelta(minutes=1)), + _raw_latest_row("gpt-3.5-turbo", "model-456", now - timedelta(minutes=5)), + _raw_latest_row("gpt-4", "deployment-abc", now - timedelta(minutes=2)), + _raw_latest_row("gpt-4", None, now - timedelta(minutes=3)), + ] ) result = await mock_prisma.get_all_latest_health_checks() - assert len(result) == 2 - names = {r.model_name for r in result} - assert names == {"gpt-4"} - ids = {r.model_id for r in result} - assert "deployment-abc" in ids - assert None in ids + assert {(check.model_id, check.model_name): check.checked_at for check in result} == { + ("model-123", "gpt-3.5-turbo"): now - timedelta(minutes=1), + ("model-456", "gpt-3.5-turbo"): now - timedelta(minutes=5), + ("deployment-abc", "gpt-4"): now - timedelta(minutes=2), + (None, "gpt-4"): now - timedelta(minutes=3), + } - by_key = {(r.model_id, r.model_name): r for r in result} - assert by_key[("deployment-abc", "gpt-4")].checked_at == with_id.checked_at - assert by_key[(None, "gpt-4")].checked_at == without_id.checked_at + +@pytest.mark.asyncio +async def test_save_background_health_checks_compares_raw_checked_at_against_utc_now(mock_prisma): + """ + Raw rows carry ISO strings and the engine may omit the offset. A naive checked_at would TypeError + inside the 1h compare, be swallowed, and silently stop every save; a stale row must still re-save. + """ + stale = (datetime.now(timezone.utc) - timedelta(hours=2)).replace(tzinfo=None) + fresh = datetime.now(timezone.utc) - timedelta(minutes=5) + mock_prisma.db.query_raw = AsyncMock( + return_value=[ + _raw_latest_row("stale-model", "stale-id", stale), + _raw_latest_row("fresh-model", "fresh-id", fresh), + ] + ) + mock_prisma.save_health_check_result = AsyncMock() + model_list = [ + {"model_name": "stale-model", "model_info": {"id": "stale-id"}, "litellm_params": {"model": "openai/stale"}}, + {"model_name": "fresh-model", "model_info": {"id": "fresh-id"}, "litellm_params": {"model": "openai/fresh"}}, + ] + + await _save_background_health_checks_to_db( + mock_prisma, + model_list, + [{"model": "openai/stale"}, {"model": "openai/fresh"}], + [], + time.time(), + "pod-1", + ) + await asyncio.sleep(0) + + assert [call.kwargs["model_id"] for call in mock_prisma.save_health_check_result.await_args_list] == ["stale-id"] + + +@pytest.mark.asyncio +async def test_latest_health_checks_endpoint_serialises_raw_rows(monkeypatch): + row = LatestHealthCheckRow( + health_check_id="hc-1", + model_name="gpt-4", + model_id="deployment-abc", + status="healthy", + healthy_count=1, + unhealthy_count=0, + error_message=None, + response_time_ms=12.5, + details='{"region": "eu"}', + checked_by="pod-1", + checked_at=datetime(2026, 8, 25), + created_at=datetime(2026, 8, 25, tzinfo=timezone.utc), + updated_at=datetime(2026, 8, 25, tzinfo=timezone.utc), + ) + prisma = MagicMock() + prisma.get_all_latest_health_checks = AsyncMock(return_value=(row,)) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma) + + response = await latest_health_checks_endpoint(user_api_key_dict=UserAPIKeyAuth()) + + assert response == { + "latest_health_checks": { + "deployment-abc": { + "health_check_id": "hc-1", + "model_name": "gpt-4", + "model_id": "deployment-abc", + "status": "healthy", + "healthy_count": 1, + "unhealthy_count": 0, + "error_message": None, + "response_time_ms": 12.5, + "details": {"region": "eu"}, + "checked_by": "pod-1", + "checked_at": "2026-08-25T00:00:00+00:00", + "created_at": "2026-08-25T00:00:00+00:00", + } + }, + "total_models": 1, + } @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_health.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_health.py index 9f48ba68b4f..fbe9934f06c 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_health.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_health.py @@ -22,6 +22,10 @@ from unittest.mock import AsyncMock, MagicMock import pytest +from litellm.proxy.db.health_check_latest import ( + LATEST_HEALTH_CHECKS_FOR_MODELS_SQL, + LATEST_HEALTH_CHECKS_SQL, +) from litellm.proxy.utils import PrismaClient @@ -261,36 +265,49 @@ async def test_get_health_check_history_db_error_returns_empty_list( assert await prisma_client.get_health_check_history() == [] +def _raw_health_check_row(model_name: str = "gpt-4", model_id: str | None = "deployment-abc") -> dict[str, Any]: + return { + "health_check_id": f"hc-{model_name}", + "model_name": model_name, + "model_id": model_id, + "status": "healthy", + "healthy_count": 1, + "unhealthy_count": 0, + "error_message": None, + "response_time_ms": 12.5, + "details": None, + "checked_by": "pod-1", + "checked_at": "2026-08-25T00:00:00+00:00", + "created_at": "2026-08-25T00:00:00+00:00", + "updated_at": "2026-08-25T00:00:00+00:00", + } + + @pytest.mark.asyncio -async def test_get_all_latest_health_checks_uses_distinct( +async def test_get_all_latest_health_checks_dedups_in_postgres( prisma_client: PrismaClient, ) -> None: - rows = [MagicMock(name=f"row-{i}") for i in range(3)] - prisma_client.db.litellm_healthchecktable.find_many = AsyncMock(return_value=rows) + """A revert to prisma find_many(distinct=...) streams the whole history table; the SQL must own the DISTINCT.""" + prisma_client.db.query_raw = AsyncMock(return_value=[_raw_health_check_row()]) result = await prisma_client.get_all_latest_health_checks() - kwargs = prisma_client.db.litellm_healthchecktable.find_many.await_args.kwargs actual = { - "len": len(result), - "distinct": kwargs["distinct"], - "order_len": len(kwargs["order"]), - "first_order": kwargs["order"][0], + "query": prisma_client.db.query_raw.await_args.args, + "rows": [(row.model_id, row.model_name, row.status) for row in result], + "row_type": type(result[0]).__name__, } assert actual == { - "len": 3, - "distinct": ["model_id", "model_name"], - "order_len": 3, - "first_order": {"model_id": "asc"}, + "query": (LATEST_HEALTH_CHECKS_SQL,), + "rows": [("deployment-abc", "gpt-4", "healthy")], + "row_type": "LatestHealthCheckRow", } @pytest.mark.asyncio -async def test_get_all_latest_health_checks_db_error_returns_empty_list( +async def test_get_all_latest_health_checks_db_error_returns_no_rows( prisma_client: PrismaClient, ) -> None: - prisma_client.db.litellm_healthchecktable.find_many = AsyncMock( - side_effect=RuntimeError("oops") - ) - assert await prisma_client.get_all_latest_health_checks() == [] + prisma_client.db.query_raw = AsyncMock(side_effect=RuntimeError("oops")) + assert await prisma_client.get_all_latest_health_checks() == () @pytest.mark.asyncio @@ -298,18 +315,15 @@ async def test_get_latest_health_checks_for_models_bounds_the_query_to_those_mod prisma_client: PrismaClient, ) -> None: """A paged caller reads health for its page; an unbounded read is the bug this exists to avoid.""" - prisma_client.db.litellm_healthchecktable.find_many = AsyncMock(return_value=[]) - await prisma_client.get_latest_health_checks_for_models(["gpt-5", "claude-opus"]) - kwargs = prisma_client.db.litellm_healthchecktable.find_many.await_args.kwargs + prisma_client.db.query_raw = AsyncMock(return_value=[_raw_health_check_row(model_name="gpt-5")]) + result = await prisma_client.get_latest_health_checks_for_models(["gpt-5", "claude-opus"]) actual = { - "where": kwargs["where"], - "distinct": kwargs["distinct"], - "order": kwargs["order"], + "query": prisma_client.db.query_raw.await_args.args, + "rows": [row.model_name for row in result], } assert actual == { - "where": {"model_name": {"in": ["gpt-5", "claude-opus"]}}, - "distinct": ["model_id", "model_name"], - "order": [{"model_id": "asc"}, {"model_name": "asc"}, {"checked_at": "desc"}], + "query": (LATEST_HEALTH_CHECKS_FOR_MODELS_SQL, ["gpt-5", "claude-opus"]), + "rows": ["gpt-5"], } @@ -317,14 +331,14 @@ async def test_get_latest_health_checks_for_models_bounds_the_query_to_those_mod async def test_get_latest_health_checks_for_models_does_not_query_for_an_empty_page( prisma_client: PrismaClient, ) -> None: - prisma_client.db.litellm_healthchecktable.find_many = AsyncMock(return_value=[]) + prisma_client.db.query_raw = AsyncMock(return_value=[]) assert await prisma_client.get_latest_health_checks_for_models([]) == () - assert prisma_client.db.litellm_healthchecktable.find_many.await_count == 0 + assert prisma_client.db.query_raw.await_count == 0 @pytest.mark.asyncio -async def test_get_latest_health_checks_for_models_db_error_returns_empty_list( +async def test_get_latest_health_checks_for_models_db_error_returns_no_rows( prisma_client: PrismaClient, ) -> None: - prisma_client.db.litellm_healthchecktable.find_many = AsyncMock(side_effect=RuntimeError("oops")) + prisma_client.db.query_raw = AsyncMock(side_effect=RuntimeError("oops")) assert await prisma_client.get_latest_health_checks_for_models(["gpt-5"]) == () From f980b94923e25e5c01607fcdf2088bd0c232fc9e Mon Sep 17 00:00:00 2001 From: michelligabriele Date: Thu, 3 Sep 2026 12:36:04 +0200 Subject: [PATCH 11/33] fix(proxy): release the health check save window lock on failure or cancel --- .../health_endpoints/_health_endpoints.py | 105 +++++++----- litellm/proxy/proxy_server.py | 45 +++-- .../proxy_server/test_background_health.py | 109 +++++++----- .../proxy/test_health_check_functions.py | 155 +++++++++++++----- 4 files changed, 278 insertions(+), 136 deletions(-) diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 8b57bdca2fe..91825a5606c 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -7,7 +7,7 @@ import secrets import time import traceback from collections.abc import Iterable, Mapping -from datetime import datetime, timedelta +from datetime import datetime, timedelta, timezone from typing import Any, Final, Literal, TypedDict, cast import fastapi @@ -41,6 +41,7 @@ from litellm.proxy.auth.auth_utils import ( ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler +from litellm.proxy.db.health_check_latest import LatestHealthCheckRow from litellm.proxy.db.proxy_worker_heartbeat import count_live_proxy_workers from litellm.proxy.health_check import ( ADMIN_ONLY_HEALTH_DISPLAY_PARAMS, @@ -697,13 +698,42 @@ def _aggregate_health_check_results( return model_results +class _AggregatedHealthResult(TypedDict): + """One entry of ``_aggregate_health_check_results``: a model's counts for this cycle.""" + + model_name: ReadOnly[str] + model_id: ReadOnly[str | None] + healthy_count: ReadOnly[int] + unhealthy_count: ReadOnly[int] + error_message: ReadOnly[str | None] + + +def _new_health_status(result: _AggregatedHealthResult) -> str: + return "healthy" if result["healthy_count"] > 0 else "unhealthy" + + +def _should_persist_health_check_result( + result: _AggregatedHealthResult, latest_checks_map: Mapping[str, LatestHealthCheckRow] +) -> bool: + """ + True when this result has to be written: no previous row, the status changed, or the + previous row is older than one hour (periodic refresh while the status is stable). + """ + lookup_key: Final = result["model_id"] if result["model_id"] else result["model_name"] + last_check: Final = latest_checks_map.get(lookup_key) + if last_check is None or last_check.status != _new_health_status(result): + return True + time_since_last_check: Final = (datetime.now(timezone.utc) - last_check.checked_at).total_seconds() + return time_since_last_check >= 3600 # 1 hour threshold + + async def _save_health_check_results_if_changed( prisma_client, model_results: dict, latest_checks_map: dict, start_time: float, checked_by: str | None = None, -): +) -> bool: """ Save health check results to database, but only if status changed or >1 hour since last save. @@ -714,47 +744,39 @@ async def _save_health_check_results_if_changed( - Status changes: Immediate write (no delay) - Result: ~92% reduction in DB writes for stable systems, while maintaining real-time updates on changes + The writes are awaited rather than detached so the caller learns whether this cycle's + persistence completed. + Args: prisma_client: Database client model_results: Dictionary of aggregated health check results per model latest_checks_map: Dictionary mapping model_id/model_name to latest health check start_time: Start time of health check for calculating response time checked_by: Identifier for who/what performed the check + + Returns: + True when every row that needed writing was written (including when nothing needed + writing); False when any write failed. """ - for result in model_results.values(): - new_status = "healthy" if result["healthy_count"] > 0 else "unhealthy" - - # Check if we should save this result - should_save = True - lookup_key = result["model_id"] if result["model_id"] else result["model_name"] - if lookup_key in latest_checks_map: - last_check = latest_checks_map[lookup_key] - # Only save if status changed or if it's been a while since last check - if last_check.status == new_status: - # Check if last check was recent (within 1 hour) - if last_check.checked_at: - from datetime import datetime, timezone - - time_since_last_check = (datetime.now(timezone.utc) - last_check.checked_at).total_seconds() - # Only skip if status unchanged AND checked recently (within 1 hour) - # This ensures we still get periodic updates even if status is stable - if time_since_last_check < 3600: # 1 hour threshold - should_save = False - - if should_save: - asyncio.create_task( - prisma_client.save_health_check_result( - model_name=result["model_name"], - model_id=result["model_id"], - status=new_status, - healthy_count=result["healthy_count"], - unhealthy_count=result["unhealthy_count"], - error_message=result["error_message"], - response_time_ms=(time.time() - start_time) * 1000, - details=None, - checked_by=checked_by, - ) - ) + to_write: Final = tuple( + result for result in model_results.values() if _should_persist_health_check_result(result, latest_checks_map) + ) + writes: Final = tuple( + prisma_client.save_health_check_result( + model_name=result["model_name"], + model_id=result["model_id"], + status=_new_health_status(result), + healthy_count=result["healthy_count"], + unhealthy_count=result["unhealthy_count"], + error_message=result["error_message"], + response_time_ms=(time.time() - start_time) * 1000, + details=None, + checked_by=checked_by, + ) + for result in to_write + ) + rows: Final = await asyncio.gather(*writes) + return all(row is not None for row in rows) async def _save_background_health_checks_to_db( @@ -764,7 +786,7 @@ async def _save_background_health_checks_to_db( unhealthy_endpoints: list, start_time: float, checked_by: str | None = None, -): +) -> bool: """ Save background health check results to database for each model. @@ -773,9 +795,13 @@ async def _save_background_health_checks_to_db( OPTIMIZATION: Only saves to database if the status has changed from the last saved check. This dramatically reduces database writes when health status remains stable. + + Returns: + True when this cycle's persistence completed; False when it was skipped or any step + failed. Never raises: a database failure must not break the health check loop. """ if prisma_client is None: - return + return False try: # Step 1: Build mapping from model parameter to model info @@ -798,7 +824,7 @@ async def _save_background_health_checks_to_db( latest_checks_map[key] = check # Step 4: Save aggregated results, but only if status changed - await _save_health_check_results_if_changed( + return await _save_health_check_results_if_changed( prisma_client, model_results, latest_checks_map, @@ -808,6 +834,7 @@ async def _save_background_health_checks_to_db( except Exception as db_error: verbose_proxy_logger.warning("Failed to save background health checks to database: %s", db_error) # Continue execution - don't let database save failure break health checks + return False _PROXY_ADMIN_ROLES: Final = frozenset( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 0e330272bdf..88ab36b3d1e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -152,6 +152,7 @@ if TYPE_CHECKING: from prisma import models as prisma_models from litellm.integrations.opentelemetry import OpenTelemetry + from litellm.proxy.health_check_utils.shared_health_check_manager import SharedHealthCheckManager Span = _Span | Any else: @@ -3592,35 +3593,49 @@ async def _run_direct_health_check_with_instrumentation( async def _window_gated_health_check_db_save( - save: Callable[[], Awaitable[None]], + save: Callable[[], Awaitable[bool]], pod_lock_manager: PodLockManager | None, lock_ttl: int | None, ) -> None: """ - Persist at most once per window fleet-wide: the lock is the "this window's save is - done" marker, so it is deliberately never released and expires with the interval. + Persist at most once per window fleet-wide. A completed save keeps the lock as the + "this window's save is done" marker, so it is deliberately never released and expires + with the interval. A save that reports failure or is cancelled releases the lock so + another pod's cycle in the same window can retry, instead of the fleet going a whole + window without a write. """ - if pod_lock_manager is not None and pod_lock_manager.redis_cache is not None: - acquired: Final = await pod_lock_manager.acquire_lock( - cronjob_id=BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME, - ttl=lock_ttl, - allow_reentrant=False, + if pod_lock_manager is None or pod_lock_manager.redis_cache is None: + await save() + return + acquired: Final = await pod_lock_manager.acquire_lock( + cronjob_id=BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME, + ttl=lock_ttl, + allow_reentrant=False, + ) + if not acquired: + verbose_proxy_logger.debug("background_health_check_db_save_skipped another pod persisted this window") + return + try: + persisted: Final = await save() + except BaseException: + await pod_lock_manager.release_lock(cronjob_id=BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME) + raise + if not persisted: + verbose_proxy_logger.warning( + "background_health_check_db_save_incomplete released the window lock so another pod can retry" ) - if not acquired: - verbose_proxy_logger.debug("background_health_check_db_save_skipped another pod persisted this window") - return - await save() + await pod_lock_manager.release_lock(cronjob_id=BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME) def _schedule_background_health_check_db_save( - prisma_client, - shared_health_manager, + prisma_client: PrismaClient | None, + shared_health_manager: "SharedHealthCheckManager | None", model_list: list, healthy_endpoints: list, unhealthy_endpoints: list, pod_lock_manager: PodLockManager | None = None, lock_ttl: int | None = None, -): +) -> None: """Fire-and-forget: persist health check results to DB if prisma is available.""" if prisma_client is None: return diff --git a/tests/test_litellm/proxy/proxy_server/test_background_health.py b/tests/test_litellm/proxy/proxy_server/test_background_health.py index 4874f75752c..d99aa637955 100644 --- a/tests/test_litellm/proxy/proxy_server/test_background_health.py +++ b/tests/test_litellm/proxy/proxy_server/test_background_health.py @@ -112,13 +112,11 @@ async def test_run_direct_health_check_with_instrumentation_returns_results( lambda _gs: {}, ) - healthy, unhealthy, exceptions = ( - await _run_direct_health_check_with_instrumentation( - model_list=[{"model_name": "gpt-4"}], - details=False, - max_concurrency=1, - instrumentation_context={"source": "test"}, - ) + healthy, unhealthy, exceptions = await _run_direct_health_check_with_instrumentation( + model_list=[{"model_name": "gpt-4"}], + details=False, + max_concurrency=1, + instrumentation_context={"source": "test"}, ) assert normalize( @@ -254,11 +252,12 @@ def _lock_manager(redis_cache, acquired): return manager -def _capture_saves(monkeypatch): +def _capture_saves(monkeypatch, persisted=True): saves = [] async def _fake_save(*_args, **kwargs): saves.append(kwargs) + return persisted import litellm.proxy.health_endpoints._health_endpoints as he @@ -266,6 +265,15 @@ def _capture_saves(monkeypatch): return saves +def _cancel_during_save(monkeypatch): + async def _fake_save(*_args, **_kwargs): + raise asyncio.CancelledError() + + import litellm.proxy.health_endpoints._health_endpoints as he + + monkeypatch.setattr(he, "_save_background_health_checks_to_db", _fake_save) + + def _schedule_with(lock_manager): _schedule_background_health_check_db_save( prisma_client=MagicMock(), @@ -315,15 +323,56 @@ async def test_schedule_background_health_check_db_save_holds_the_window_lock_fo } +@pytest.mark.asyncio +async def test_schedule_background_health_check_db_save_releases_the_window_lock_when_the_save_reports_failure( + monkeypatch, +): + """A failed save must not burn the window: release the lock so another pod's cycle can retry.""" + saves = _capture_saves(monkeypatch, persisted=False) + lock_manager = _lock_manager(redis_cache=MagicMock(), acquired=True) + + _schedule_with(lock_manager) + await asyncio.sleep(0) + + assert normalize( + { + "saves": len(saves), + "release_request": lock_manager.release_lock.await_args.kwargs, + "release_count": lock_manager.release_lock.await_count, + } + ) == { + "saves": 1, + "release_request": {"cronjob_id": BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME}, + "release_count": 1, + } + + +@pytest.mark.asyncio +async def test_schedule_background_health_check_db_save_releases_the_window_lock_when_the_save_is_cancelled( + monkeypatch, +): + """A pod shutting down mid-save releases the lock instead of holding it until the TTL.""" + _cancel_during_save(monkeypatch) + lock_manager = _lock_manager(redis_cache=MagicMock(), acquired=True) + + _schedule_with(lock_manager) + await asyncio.sleep(0) + + assert ( + lock_manager.release_lock.await_args.kwargs, + lock_manager.release_lock.await_count, + ) == ({"cronjob_id": BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME}, 1) + + @pytest.mark.asyncio async def test_schedule_background_health_check_db_save_runs_ungated_without_redis(monkeypatch): - saves = _capture_saves(monkeypatch) + saves = _capture_saves(monkeypatch, persisted=False) lock_manager = _lock_manager(redis_cache=None, acquired=True) _schedule_with(lock_manager) await asyncio.sleep(0) - assert (len(saves), lock_manager.acquire_lock.await_count) == (1, 0) + assert (len(saves), lock_manager.acquire_lock.await_count, lock_manager.release_lock.await_count) == (1, 0, 0) # --------------------------------------------------------------------------- @@ -400,13 +449,9 @@ def test_write_health_state_to_router_cache_sets_states(monkeypatch): _write_health_state_to_router_cache(healthy, unhealthy, exceptions) - fake_router.health_state_cache.set_deployment_health_states.assert_called_once_with( - fake_states - ) + fake_router.health_state_cache.set_deployment_health_states.assert_called_once_with(fake_states) - call_args = fake_router.health_state_cache.set_deployment_health_states.call_args[ - 0 - ][0] + call_args = fake_router.health_state_cache.set_deployment_health_states.call_args[0][0] assert normalize( { "states_keys": sorted(call_args.keys()), @@ -448,9 +493,7 @@ def test_write_health_state_to_router_cache_populates_for_listing_filter(monkeyp fake_router.cooldown_time = 30 monkeypatch.setattr(proxy_server, "llm_router", fake_router) - monkeypatch.setattr( - proxy_server, "general_settings", {"model_list_healthy_only": True} - ) + monkeypatch.setattr(proxy_server, "general_settings", {"model_list_healthy_only": True}) fake_states = {"m1": {"is_healthy": True}, "m2": {"is_healthy": False}} @@ -484,9 +527,7 @@ def test_write_health_state_to_router_cache_populates_for_listing_filter(monkeyp {"m2": SimpleNamespace(status_code=500)}, ) - fake_router.health_state_cache.set_deployment_health_states.assert_called_once_with( - fake_states - ) + fake_router.health_state_cache.set_deployment_health_states.assert_called_once_with(fake_states) assert cooldowns == [] assert failures == [] @@ -496,9 +537,7 @@ def test_write_health_state_to_router_cache_swallows_internal_failures(monkeypat fake_router = MagicMock() fake_router.enable_health_check_routing = True fake_router.health_check_ignore_transient_errors = False - fake_router.health_state_cache.set_deployment_health_states.side_effect = ( - RuntimeError("cache exploded") - ) + fake_router.health_state_cache.set_deployment_health_states.side_effect = RuntimeError("cache exploded") monkeypatch.setattr(proxy_server, "llm_router", fake_router) @@ -528,9 +567,7 @@ async def test_adaptive_router_flusher_loop_flushes_each_router(monkeypatch): from litellm.types.router import TaggedPreRoutingStrategy fake_router = MagicMock() - fake_router.adaptive_routers = { - "alpha": [TaggedPreRoutingStrategy(tags=(), strategy=fake_ar)] - } + fake_router.adaptive_routers = {"alpha": [TaggedPreRoutingStrategy(tags=(), strategy=fake_ar)]} monkeypatch.setattr(proxy_server, "llm_router", fake_router) monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) @@ -628,12 +665,8 @@ async def test_run_background_health_check_runs_one_cycle_then_cancels(monkeypat "_run_direct_health_check_with_instrumentation", _fake_direct, ) - monkeypatch.setattr( - proxy_server, "_schedule_background_health_check_db_save", lambda *a, **kw: None - ) - monkeypatch.setattr( - proxy_server, "_write_health_state_to_router_cache", lambda *a, **kw: None - ) + monkeypatch.setattr(proxy_server, "_schedule_background_health_check_db_save", lambda *a, **kw: None) + monkeypatch.setattr(proxy_server, "_write_health_state_to_router_cache", lambda *a, **kw: None) monkeypatch.setattr( proxy_server, "health_check_filter_kwargs_from_general_settings", @@ -711,12 +744,8 @@ async def test_run_background_health_check_probes_only_listed_model_groups(monke "_run_direct_health_check_with_instrumentation", _fake_direct, ) - monkeypatch.setattr( - proxy_server, "_schedule_background_health_check_db_save", lambda *a, **kw: None - ) - monkeypatch.setattr( - proxy_server, "_write_health_state_to_router_cache", lambda *a, **kw: None - ) + monkeypatch.setattr(proxy_server, "_schedule_background_health_check_db_save", lambda *a, **kw: None) + monkeypatch.setattr(proxy_server, "_write_health_state_to_router_cache", lambda *a, **kw: None) monkeypatch.setattr( proxy_server, "health_check_filter_kwargs_from_general_settings", diff --git a/tests/test_litellm/proxy/test_health_check_functions.py b/tests/test_litellm/proxy/test_health_check_functions.py index 9adaddb032d..5aa80213134 100644 --- a/tests/test_litellm/proxy/test_health_check_functions.py +++ b/tests/test_litellm/proxy/test_health_check_functions.py @@ -24,12 +24,8 @@ from litellm.proxy.utils import PrismaClient def mock_prisma(): """Simplified mock PrismaClient with bound methods""" client = MagicMock() - client.db.litellm_healthchecktable.create = AsyncMock( - return_value={"id": "test-id"} - ) - client.db.litellm_healthchecktable.find_many = AsyncMock( - return_value=[{"id": "1", "model_name": "test"}] - ) + client.db.litellm_healthchecktable.create = AsyncMock(return_value={"id": "test-id"}) + client.db.litellm_healthchecktable.find_many = AsyncMock(return_value=[{"id": "1", "model_name": "test"}]) # Bind actual methods import types @@ -55,14 +51,10 @@ def mock_prisma(): ("healthy", 1, 0, False), # Database error case ], ) -async def test_save_health_check_result( - mock_prisma, status, healthy, unhealthy, should_succeed -): +async def test_save_health_check_result(mock_prisma, status, healthy, unhealthy, should_succeed): """Test health check result saving with various scenarios""" if not should_succeed: - mock_prisma.db.litellm_healthchecktable.create.side_effect = Exception( - "DB Error" - ) + mock_prisma.db.litellm_healthchecktable.create.side_effect = Exception("DB Error") result = await mock_prisma.save_health_check_result( model_name="test-model", @@ -190,9 +182,7 @@ def test_aggregate_health_check_results(): {"model": "gpt-4", "error": "Rate limit exceeded"}, ] - result = _aggregate_health_check_results( - model_param_to_info, healthy_endpoints, unhealthy_endpoints - ) + result = _aggregate_health_check_results(model_param_to_info, healthy_endpoints, unhealthy_endpoints) # Check gpt-3.5-turbo is healthy gpt35_key = ("model-123", "gpt-3.5-turbo") @@ -223,9 +213,7 @@ def test_aggregate_health_check_results_multiple_endpoints(): ] unhealthy_endpoints = [] - result = _aggregate_health_check_results( - model_param_to_info, healthy_endpoints, unhealthy_endpoints - ) + result = _aggregate_health_check_results(model_param_to_info, healthy_endpoints, unhealthy_endpoints) key = ("model-123", "gpt-3.5-turbo") assert result[key]["healthy_count"] == 2 @@ -401,7 +389,7 @@ async def test_save_background_health_checks_to_db(): start_time = 1234567890.0 - await _save_background_health_checks_to_db( + persisted = await _save_background_health_checks_to_db( mock_prisma, model_list, healthy_endpoints, @@ -410,7 +398,8 @@ async def test_save_background_health_checks_to_db(): "background_health_check", ) - # Should call get_all_latest_health_checks and save_health_check_result + # Should call get_all_latest_health_checks and save_health_check_result, and report completion + assert persisted is True mock_prisma.get_all_latest_health_checks.assert_called_once() mock_prisma.save_health_check_result.assert_called_once() @@ -421,22 +410,112 @@ async def test_save_background_health_checks_to_db(): assert call_kwargs["checked_by"] == "background_health_check" +def _two_model_results(): + return { + ("model-1", "gpt-4"): { + "model_name": "gpt-4", + "model_id": "model-1", + "healthy_count": 1, + "unhealthy_count": 0, + "error_message": None, + }, + ("model-2", "gpt-4o"): { + "model_name": "gpt-4o", + "model_id": "model-2", + "healthy_count": 0, + "unhealthy_count": 1, + "error_message": "boom", + }, + } + + +@pytest.mark.asyncio +async def test_save_health_check_results_if_changed_awaits_every_write_and_reports_success(): + """Writes are awaited, not detached, so the caller can tell the cycle's persistence completed.""" + mock_prisma = MagicMock() + mock_prisma.save_health_check_result = AsyncMock(return_value={"id": "row"}) + + persisted = await _save_health_check_results_if_changed( + mock_prisma, _two_model_results(), {}, 1234567890.0, "background_health_check" + ) + + assert (persisted, mock_prisma.save_health_check_result.await_count) == (True, 2) + + +@pytest.mark.asyncio +async def test_save_health_check_results_if_changed_reports_failure_when_a_write_returns_none(): + """save_health_check_result swallows DB errors and returns None; that must surface as False.""" + mock_prisma = MagicMock() + mock_prisma.save_health_check_result = AsyncMock(side_effect=[{"id": "row"}, None]) + + persisted = await _save_health_check_results_if_changed( + mock_prisma, _two_model_results(), {}, 1234567890.0, "background_health_check" + ) + + assert (persisted, mock_prisma.save_health_check_result.await_count) == (False, 2) + + +@pytest.mark.asyncio +async def test_save_health_check_results_if_changed_reports_success_when_nothing_needed_writing(): + mock_prisma = MagicMock() + mock_prisma.save_health_check_result = AsyncMock() + model_results = { + ("model-1", "gpt-4"): { + "model_name": "gpt-4", + "model_id": "model-1", + "healthy_count": 1, + "unhealthy_count": 0, + "error_message": None, + }, + } + latest_checks_map = { + "model-1": MagicMock(status="healthy", checked_at=datetime.now(timezone.utc) - timedelta(minutes=5)), + } + + persisted = await _save_health_check_results_if_changed( + mock_prisma, model_results, latest_checks_map, 1234567890.0, "background_health_check" + ) + + assert (persisted, mock_prisma.save_health_check_result.await_count) == (True, 0) + + +def _one_model_setup(): + model_list = [ + { + "model_name": "gpt-3.5-turbo", + "model_info": {"id": "model-123"}, + "litellm_params": {"model": "gpt-3.5-turbo"}, + }, + ] + return model_list, [{"model": "gpt-3.5-turbo"}], [] + + +@pytest.mark.asyncio +async def test_save_background_health_checks_to_db_returns_false_when_a_write_fails(): + mock_prisma = MagicMock() + mock_prisma.get_all_latest_health_checks = AsyncMock(return_value=[]) + mock_prisma.save_health_check_result = AsyncMock(return_value=None) + model_list, healthy_endpoints, unhealthy_endpoints = _one_model_setup() + + persisted = await _save_background_health_checks_to_db( + mock_prisma, model_list, healthy_endpoints, unhealthy_endpoints, 1234567890.0, "background_health_check" + ) + + assert (persisted, mock_prisma.save_health_check_result.await_count) == (False, 1) + + @pytest.mark.asyncio async def test_save_background_health_checks_to_db_no_prisma(): """Test graceful handling when no prisma client""" - result = await _save_background_health_checks_to_db( - None, [], [], [], 0.0, "background_health_check" - ) - assert result is None + result = await _save_background_health_checks_to_db(None, [], [], [], 0.0, "background_health_check") + assert result is False @pytest.mark.asyncio async def test_save_background_health_checks_to_db_exception_handling(): """Test exception handling in background health check save""" mock_prisma = MagicMock() - mock_prisma.get_all_latest_health_checks = AsyncMock( - side_effect=Exception("DB Error") - ) + mock_prisma.get_all_latest_health_checks = AsyncMock(side_effect=Exception("DB Error")) model_list = [ { @@ -446,12 +525,13 @@ async def test_save_background_health_checks_to_db_exception_handling(): }, ] - # Should not raise exception, should handle gracefully - await _save_background_health_checks_to_db( + # Must not raise (the health check loop has to survive a DB outage) but must report + # the failure, so the window lock can be released for another pod to retry + persisted = await _save_background_health_checks_to_db( mock_prisma, model_list, [], [], 0.0, "background_health_check" ) - # Function should complete without raising + assert persisted is False def _raw_latest_row(model_name: str, model_id, checked_at: datetime) -> dict: @@ -660,12 +740,7 @@ def test_parse_background_health_check_model_groups_unset_returns_none(): assert parse_background_health_check_model_groups(None) is None assert parse_background_health_check_model_groups({}) is None - assert ( - parse_background_health_check_model_groups( - {"background_health_check_model_groups": None} - ) - is None - ) + assert parse_background_health_check_model_groups({"background_health_check_model_groups": None}) is None def test_parse_background_health_check_model_groups_list_returns_frozenset(): @@ -682,9 +757,7 @@ def test_parse_background_health_check_model_groups_malformed_raises(bad_value): from litellm.proxy.health_check import parse_background_health_check_model_groups with pytest.raises(ValueError, match="must be a list of model group names"): - parse_background_health_check_model_groups( - {"background_health_check_model_groups": bad_value} - ) + parse_background_health_check_model_groups({"background_health_check_model_groups": bad_value}) def test_filter_deployments_to_model_groups(): @@ -697,9 +770,7 @@ def test_filter_deployments_to_model_groups(): ] assert filter_deployments_to_model_groups(model_list, None) == tuple(model_list) - assert filter_deployments_to_model_groups( - model_list, frozenset({"prod-openai"}) - ) == (model_list[0], model_list[2]) + assert filter_deployments_to_model_groups(model_list, frozenset({"prod-openai"})) == (model_list[0], model_list[2]) assert filter_deployments_to_model_groups(model_list, frozenset()) == () From 7c1d060b7af3150320a2c88428871703d71f7a38 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 18:04:36 -0700 Subject: [PATCH 12/33] fix(mcp): enforce OAuth identity binding across credential lifetime --- litellm/proxy/_experimental/mcp_server/db.py | 34 +++- .../mcp_server/discoverable_endpoints.py | 94 +++++++--- .../mcp_server/oauth_identity_binding.py | 170 ++++++++++++++---- .../authz_code_refresher.py | 37 +++- .../per_user_oauth_store.py | 13 ++ .../types/mcp_server/mcp_server_manager.py | 13 +- .../test_authz_code_refresher.py | 34 ++++ .../test_per_user_oauth_store.py | 25 +++ .../mcp_server/test_db_credentials.py | 24 ++- .../mcp_server/test_discoverable_endpoints.py | 60 ++++++- .../mcp_server/test_oauth_identity_binding.py | 111 +++++++++--- .../test_mcp_management_endpoints.py | 3 +- 12 files changed, 526 insertions(+), 92 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 7379126983a..08f83aeb756 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -6,10 +6,17 @@ from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypedDict, cast +from typing_extensions import ReadOnly + from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.constants import MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.proxy._experimental.mcp_server.oauth_identity_binding import ( + RefreshTokenPresented, + credential_binding_matches, + enforce_oauth_identity_binding, +) from litellm.proxy._experimental.mcp_server.oauth_utils import build_upstream_oauth2_token_request from litellm.proxy._types import ( LiteLLM_MCPServerTable, @@ -117,6 +124,7 @@ class _OAuthCredentialAccessToken(TypedDict): class OAuthCredentialPayload(_OAuthCredentialAccessToken, total=False): + identity_binding_proof: ReadOnly[str] type: str refresh_token: str expires_at: str @@ -1393,6 +1401,7 @@ async def store_user_oauth_credential( expires_in: int | None = None, scopes: list[str] | None = None, skip_byok_guard: bool = False, + identity_binding_proof: str | None = None, ) -> None: """Persist an OAuth2 access token for a user+server pair. @@ -1409,6 +1418,7 @@ async def store_user_oauth_credential( "type": "oauth2", "access_token": access_token, "connected_at": datetime.now(timezone.utc).isoformat(), + **({"identity_binding_proof": identity_binding_proof} if identity_binding_proof else {}), } if refresh_token: payload["refresh_token"] = refresh_token @@ -1628,6 +1638,11 @@ async def refresh_user_oauth_token( warning and returns ``None`` — the caller is responsible for clearing the stale credential and triggering re-authentication. """ + binding: Final = server.oauth_identity_binding + if binding is not None and binding.mode == "enforce": + if not await credential_binding_matches(binding, user_id, server.server_id, cred): + return None + refresh_token: Final[str | None] = cred.get("refresh_token") token_url: Final[str | None] = getattr(server, "effective_token_url", None) or getattr(server, "token_url", None) server_id: Final[str] = getattr(server, "server_id", "") @@ -1677,6 +1692,14 @@ async def refresh_user_oauth_token( ) return None + binding_proof: Final = await enforce_oauth_identity_binding( + server=server, + token_response=body, + litellm_user_id=user_id, + grant_type="refresh_token", + refresh_ownership=RefreshTokenPresented(refresh_token), + ) + access_token: Final[str | None] = body.get("access_token") if not access_token: verbose_proxy_logger.warning( @@ -1709,6 +1732,7 @@ async def refresh_user_oauth_token( refresh_token=new_refresh_token, expires_in=expires_in, scopes=scopes, + identity_binding_proof=binding_proof, skip_byok_guard=True, # Row is already OAuth2; skip the extra find_unique check ) @@ -1742,6 +1766,10 @@ async def resolve_valid_user_oauth_token( grant: Final = oauth_grant_state(cred) if cred is None or grant == "absent": return None + binding: Final = server.oauth_identity_binding + if binding is not None and binding.mode == "enforce": + if not await credential_binding_matches(binding, user_id, server.server_id, cred): + return None if grant == "valid": return cred if prisma_client is None: @@ -1782,7 +1810,11 @@ async def resolve_user_oauth_access_token( mcp_per_user_token_cache, ) - if prefetched_creds is None: + binding: Final = server.oauth_identity_binding + enforce_binding: Final = binding is not None and binding.mode == "enforce" + if enforce_binding: + await mcp_per_user_token_cache.delete(user_id, server_id) + if prefetched_creds is None and not enforce_binding: cached_token: Final = await mcp_per_user_token_cache.get(user_id, server_id) if cached_token is not None: return cached_token diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 9f6bc1860e1..c8e02be308b 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -144,6 +144,7 @@ def encode_state_with_base_url( dcr_client_id: str | None = None, dcr_client_secret: str | None = None, dcr_token_endpoint_auth_method: MCPTokenEndpointAuthMethod | None = None, + oauth_nonce: str | None = None, ) -> str: """ Encode the base_url, original state, and PKCE parameters using encryption. @@ -154,9 +155,8 @@ def encode_state_with_base_url( code_challenge: PKCE code challenge from client code_challenge_method: PKCE code challenge method from client client_redirect_uri: Original redirect_uri from client - litellm_user_id: The SSO-authenticated litellm user captured at the bridge authorize - (interactive dcr_bridge oauth_delegate only); the callback seals it into the gateway - authorization code so the token mint can bind the envelope to this user + litellm_user_id: The authenticated user captured for bridge or identity-bound per-user OAuth; + the callback seals this credential owner into the authorization code mcp_server_id: The server the flow targets, sealed alongside litellm_user_id (bridge) or dcr_client_id (ephemeral mint) so the gateway code cannot be replayed against another server @@ -174,6 +174,7 @@ def encode_state_with_base_url( An encrypted string that encodes all values """ state_data: Final = { + "oauth_nonce": oauth_nonce, "base_url": base_url, "original_state": original_state, "code_challenge": code_challenge, @@ -215,10 +216,10 @@ _BRIDGE_AUTH_CODE_PREFIX: Final = "llm_bcode_" class _BridgeAuthorizationCode(BaseModel): - """The identity and upstream code the gateway seals into the authorization code it hands a DCR - client for an interactive dcr_bridge oauth_delegate sign-in, recovered at the token endpoint.""" + """Authenticated caller and upstream code sealed for bridge or identity-bound per-user OAuth.""" model_config = ConfigDict(frozen=True) + oauth_nonce: str | None = None upstream_code: str = Field(min_length=1) litellm_user_id: str = Field(min_length=1) mcp_server_id: str = Field(min_length=1) @@ -230,7 +231,12 @@ def is_bridge_authorization_code(code: str) -> bool: return code.startswith(_BRIDGE_AUTH_CODE_PREFIX) -def seal_bridge_authorization_code(upstream_code: str, litellm_user_id: str, mcp_server_id: str) -> str: +def seal_bridge_authorization_code( + upstream_code: str, + litellm_user_id: str, + mcp_server_id: str, + oauth_nonce: str | None = None, +) -> str: """Seal the upstream authorization code and the SSO-captured litellm user into a gateway authorization code. The DCR client only echoes this opaque value back at the token endpoint; the gateway decrypts it there to recover the user (to bind the envelope) and the upstream code (to @@ -239,7 +245,12 @@ def seal_bridge_authorization_code(upstream_code: str, litellm_user_id: str, mcp authenticated symmetric helper (the same family the OAuth state uses), so the client can neither read nor forge it.""" payload: Final = json.dumps( - {"upstream_code": upstream_code, "litellm_user_id": litellm_user_id, "mcp_server_id": mcp_server_id}, + { + "upstream_code": upstream_code, + "litellm_user_id": litellm_user_id, + "mcp_server_id": mcp_server_id, + "oauth_nonce": oauth_nonce, + }, sort_keys=True, ) return _BRIDGE_AUTH_CODE_PREFIX + encrypt_value_helper(payload) @@ -552,6 +563,7 @@ async def _store_per_user_token_server_side( server: MCPServer, user_id: str, token_response: dict[str, Any], + identity_binding_proof: str | None = None, ) -> None: """Persist the OAuth token server-side and warm the Redis cache. @@ -593,6 +605,7 @@ async def _store_per_user_token_server_side( refresh_token=refresh_token, expires_in=expires_in, scopes=scopes, + identity_binding_proof=identity_binding_proof, ) verbose_logger.info( "_store_per_user_token_server_side: stored token for user=%s server=%s", @@ -859,6 +872,11 @@ async def authorize_with_server( ), ) + binding: Final = resolved_server.oauth_identity_binding + enforce_binding: Final = binding is not None and binding.mode == "enforce" + if enforce_binding: + _require_s256_pkce(code_challenge, code_challenge_method) + if resolved_server.is_dcr_bridge: # Enforce S256 PKCE on both bridge arms. The relay arm forwards the validated, # now-non-optional pair to the upstream authorize; the short-circuit arm keeps @@ -889,19 +907,16 @@ async def authorize_with_server( base_url: Final = urlunparse(parsed._replace(query="")) request_base_url: Final = get_request_base_url(request) - # Interactive dcr_bridge oauth_delegate sign-in: this arm runs the gateway /callback and /token in - # the loop, so the gateway can capture the litellm user here (from the browser's UI session) and - # carry it to the back-channel token mint. Seal the SSO user and the target server into the state; - # the callback reads them back to mint the gateway authorization code. A DCR client cannot present a - # litellm key, so the browser session is the only identity source; without one there is nothing to - # bind, so send the user through login first. Every other oauth2 server keeps the identity-less state. + # Seal the authenticated caller into state so the token exchange cannot select another credential owner. litellm_user_id: str | None = None - if resolved_server.is_dcr_bridge and resolved_server.is_oauth_delegate: + if enforce_binding or (resolved_server.is_dcr_bridge and resolved_server.is_oauth_delegate): from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import ( # noqa: PLC0415 # inline import avoids a module-load circular import _user_id_from_session_cookie, ) - litellm_user_id = _user_id_from_session_cookie(request) + litellm_user_id = ( + await _extract_user_id_from_request(request) if enforce_binding else None + ) or _user_id_from_session_cookie(request) if litellm_user_id is None: return _redirect_to_litellm_login(request) denial: Final = await _bridge_authorize_access_denial( @@ -913,9 +928,11 @@ async def authorize_with_server( if denial is not None: return denial + oauth_nonce: Final = secrets.token_urlsafe(32) if enforce_binding else None encoded_state: Final = encode_state_with_base_url( base_url=base_url, original_state=state, + oauth_nonce=oauth_nonce, code_challenge=code_challenge, code_challenge_method=code_challenge_method, client_redirect_uri=redirect_uri, @@ -935,11 +952,16 @@ async def authorize_with_server( "state": relay_state, "response_type": response_type or "code", } + if oauth_nonce: + params["nonce"] = oauth_nonce if scope: params["scope"] = scope elif resolved_server.scopes: params["scope"] = " ".join(resolved_server.scopes) + if enforce_binding and "openid" not in params.get("scope", "").split(): + params["scope"] = f"openid {params.get('scope', '')}".strip() + if code_challenge: params["code_challenge"] = code_challenge if code_challenge_method: @@ -1020,6 +1042,12 @@ async def exchange_token_with_server( except TokenEndpointAuthConfigError as exc: raise HTTPException(status_code=400, detail=str(exc)) from exc + request_user_id: Final = ( + await _extract_user_id_from_request(request) + if resolved_server.needs_user_oauth_token or resolved_server.oauth_identity_binding is not None + else None + ) + bridge_identity: _BridgeAuthorizationCode | None = None bridge_mint_ready: _BridgeMintReady | None = None bridge_upstream_refresh: SecretStr | None = None @@ -1081,6 +1109,14 @@ async def exchange_token_with_server( detail="Authorization code was issued for a different MCP server", ) code = bridge_identity.upstream_code + binding: Final = resolved_server.oauth_identity_binding + if binding is not None and binding.mode == "enforce": + if bridge_identity is None or not bridge_identity.oauth_nonce: + raise HTTPException(status_code=403, detail={"error": "oauth_identity_binding_failed"}) + if request_user_id is not None and request_user_id != bridge_identity.litellm_user_id: + raise HTTPException(status_code=403, detail={"error": "oauth_principal_mismatch"}) + if not code_verifier: + raise HTTPException(status_code=403, detail={"error": "oauth_identity_binding_failed"}) bridge_token_relay: Final = _dcr_bridge_relays_client_registration(resolved_server) if bridge_token_relay and not redirect_uri: raise HTTPException( @@ -1108,6 +1144,16 @@ async def exchange_token_with_server( return _bridge_mint_error_response(prepared) bridge_mint_ready = prepared + refresh_binding: Final = resolved_server.oauth_identity_binding + if grant_type == "refresh_token" and refresh_binding is not None and refresh_binding.mode == "enforce": + await enforce_oauth_identity_binding( + server=resolved_server, + token_response={}, + litellm_user_id=request_user_id, + grant_type=grant_type, + refresh_ownership=refresh_ownership, + ) + async_client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check) try: response: Final = await async_client.post( @@ -1150,19 +1196,19 @@ async def exchange_token_with_server( # Bind the exchanged token to the LiteLLM caller BEFORE it is returned, stored, or cached, so a # token minted for a different upstream principal never becomes usable under the caller's user_id. - resolved_user_id: Final = ( - await _extract_user_id_from_request(request) - if resolved_server.needs_user_oauth_token or resolved_server.oauth_identity_binding is not None - else None - ) - if isinstance(token_response, dict): + resolved_user_id: Final = bridge_identity.litellm_user_id if bridge_identity else request_user_id + binding_proof: Final = ( await enforce_oauth_identity_binding( server=resolved_server, token_response=token_response, litellm_user_id=resolved_user_id, grant_type=grant_type, refresh_ownership=refresh_ownership, + expected_nonce=bridge_identity.oauth_nonce if bridge_identity else None, ) + if isinstance(token_response, dict) + else None + ) # Store server-side when the server is configured for per-user OAuth and # the calling client has provided a valid LiteLLM identity. @@ -1175,6 +1221,7 @@ async def exchange_token_with_server( server=resolved_server, user_id=user_id, token_response=token_response, + identity_binding_proof=binding_proof, ) except Exception as exc: verbose_logger.warning( @@ -2161,7 +2208,10 @@ async def callback( forwarded_code = code if isinstance(litellm_user_id, str) and litellm_user_id and isinstance(mcp_server_id, str) and mcp_server_id: forwarded_code = seal_bridge_authorization_code( - upstream_code=code, litellm_user_id=litellm_user_id, mcp_server_id=mcp_server_id + upstream_code=code, + litellm_user_id=litellm_user_id, + mcp_server_id=mcp_server_id, + oauth_nonce=state_data.get("oauth_nonce"), ) elif isinstance(dcr_client_id, str) and dcr_client_id and isinstance(mcp_server_id, str) and mcp_server_id: forwarded_code = seal_passthrough_authorization_code( diff --git a/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py b/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py index 58ec3cfdb4f..7ca9e6a23df 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py @@ -8,9 +8,12 @@ claim is compared to the caller's trusted LiteLLM identity. Mismatches fail clos and are logged in audit mode. """ +import hashlib +import hmac +import json from collections.abc import Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass -from typing import Final, Literal, TypeAlias +from typing import Final, Literal, Protocol, TypeAlias import jwt from fastapi import HTTPException @@ -45,9 +48,17 @@ CallerPrincipalLoader: TypeAlias = Callable[ [str, MCPOAuthIdentityBinding], # mutable-ok: Callable parameter syntax requires a list Awaitable[str | None], ] + + +@dataclass(frozen=True, slots=True) +class VerifiedRefreshToken: + refresh_token: str + binding_proof: str + + StoredRefreshTokenLoader: TypeAlias = Callable[ - [str, str], # mutable-ok: Callable parameter syntax requires a list - Awaitable[str | None], + [str, str, MCPOAuthIdentityBinding], # mutable-ok: Callable parameter syntax requires a list + Awaitable[VerifiedRefreshToken | None], ] _RejectionCode: TypeAlias = Literal["oauth_principal_mismatch", "oauth_identity_binding_failed"] @@ -72,6 +83,18 @@ class RefreshTokenPresented: RefreshOwnership: TypeAlias = RefreshOwnershipProven | RefreshTokenPresented | None +class BindingValidator(Protocol): + async def __call__( + self, + *, + server: MCPServer, + token_response: Mapping[str, object], + litellm_user_id: str | None, + grant_type: str, + refresh_ownership: RefreshOwnership, + ) -> str | None: ... + + async def _fetch_issuer_jwks(binding: MCPOAuthIdentityBinding) -> Sequence[Mapping[str, object]]: jwks_url: Final[str] = binding.jwks_url or await _discover_jwks_url(binding.issuer) cached: Final = await _jwks_cache.async_get_cache(jwks_url) @@ -118,7 +141,7 @@ def _decode_id_token( signing_key: "jwt.PyJWK", ) -> "Mapping[str, object] | _BindingRejection": try: - decode_options: Final[Options] = {"require": ("iss", "exp")} + decode_options: Final[Options] = {"require": ("iss", "exp", "aud", "sub", "iat")} return jwt.decode( id_token, signing_key.key, @@ -169,7 +192,9 @@ async def _load_caller_principal(litellm_user_id: str, binding: MCPOAuthIdentity return loaded.user_email -async def _load_stored_refresh_token(litellm_user_id: str, server_id: str) -> str | None: +async def _load_stored_refresh_token( + litellm_user_id: str, server_id: str, binding: MCPOAuthIdentityBinding +) -> VerifiedRefreshToken | None: try: from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 # keep database imports lazy get_user_oauth_credential, @@ -184,17 +209,89 @@ async def _load_stored_refresh_token(litellm_user_id: str, server_id: str) -> st user_id=litellm_user_id, server_id=server_id, ) - return cred.get("refresh_token") if cred else None + if not cred or not await credential_binding_matches(binding, litellm_user_id, server_id, cred): + return None + refresh_token: Final = cred.get("refresh_token") + proof: Final = cred.get("identity_binding_proof") + return VerifiedRefreshToken(refresh_token, proof) if refresh_token and proof else None except Exception: # noqa: BLE001 # a credential lookup failure must fail closed return None +async def current_binding_proof( + binding: MCPOAuthIdentityBinding, + user_id: str, + server_id: str, + caller_principal_loader: CallerPrincipalLoader = _load_caller_principal, +) -> str | None: + principal: Final = await caller_principal_loader(user_id, binding) + if not principal: + return None + return _binding_proof(binding, user_id, server_id, principal) + + +def _binding_proof(binding: MCPOAuthIdentityBinding, user_id: str, server_id: str, principal: str) -> str: + payload: Final = json.dumps( + ("oidc-nonce-v1", server_id, user_id, principal, binding.model_dump(mode="json")), + sort_keys=True, + ) + return hashlib.sha256(payload.encode()).hexdigest() + + +async def credential_binding_matches( + binding: MCPOAuthIdentityBinding, + user_id: str, + server_id: str, + credential: Mapping[str, object], + caller_principal_loader: CallerPrincipalLoader = _load_caller_principal, +) -> bool: + stored: Final = credential.get("identity_binding_proof") + if not isinstance(stored, str) or not stored: + return False + expected: Final = await current_binding_proof(binding, user_id, server_id, caller_principal_loader) + return expected is not None and hmac.compare_digest(stored, expected) + + def _principals_match(upstream: str, caller: str, binding: MCPOAuthIdentityBinding) -> bool: if binding.principal_claim == "email" or binding.caller_field == "user_email": return upstream.strip().casefold() == caller.strip().casefold() return upstream == caller +async def _evaluate_refresh_ownership( + binding: MCPOAuthIdentityBinding, + litellm_user_id: str | None, + server_id: str, + refresh_ownership: RefreshOwnership, + stored_refresh_token_loader: StoredRefreshTokenLoader, +) -> _BindingRejection | str: + match refresh_ownership: + case RefreshOwnershipProven(): + return _BindingRejection( + code="oauth_identity_binding_failed", + description="an identity envelope alone does not prove upstream principal binding", + ) + case None: + return _BindingRejection( + code="oauth_identity_binding_failed", + description="refresh_token grant without an id_token carries no refresh token to prove ownership of", + ) + case RefreshTokenPresented(refresh_token): + if not litellm_user_id: + return _BindingRejection( + code="oauth_identity_binding_failed", + description="the request carries no resolvable LiteLLM user identity to bind the credential to", + ) + stored: Final = await stored_refresh_token_loader(litellm_user_id, server_id, binding) + if stored is None or not hmac.compare_digest(stored.refresh_token, refresh_token): + return _BindingRejection( + code="oauth_identity_binding_failed", + description="the presented refresh_token is not the caller's stored credential for this server", + ) + return stored.binding_proof + assert_never(refresh_ownership) # pragma: no cover + + async def _evaluate_binding( binding: MCPOAuthIdentityBinding, token_response: Mapping[str, object], @@ -205,7 +302,8 @@ async def _evaluate_binding( jwks_fetcher: JwksFetcher, caller_principal_loader: CallerPrincipalLoader, stored_refresh_token_loader: StoredRefreshTokenLoader, -) -> _BindingRejection | None: + expected_nonce: str | None, +) -> _BindingRejection | str: id_token: Final = token_response.get("id_token") if not isinstance(id_token, str) or not id_token: if grant_type != "refresh_token": @@ -213,28 +311,13 @@ async def _evaluate_binding( code="oauth_identity_binding_failed", description="the upstream token response carries no id_token to bind the credential to a principal", ) - match refresh_ownership: - case RefreshOwnershipProven(): - return None - case None: - return _BindingRejection( - code="oauth_identity_binding_failed", - description="refresh_token grant without an id_token carries no refresh token to prove ownership of", - ) - case RefreshTokenPresented(refresh_token): - if not litellm_user_id: - return _BindingRejection( - code="oauth_identity_binding_failed", - description="the request carries no resolvable LiteLLM user identity to bind the credential to", - ) - stored: Final = await stored_refresh_token_loader(litellm_user_id, server_id) - if stored is None or stored != refresh_token: - return _BindingRejection( - code="oauth_identity_binding_failed", - description="the presented refresh_token is not the caller's stored credential for this server", - ) - return None - assert_never(refresh_ownership) # pragma: no cover + return await _evaluate_refresh_ownership( + binding, + litellm_user_id, + server_id, + refresh_ownership, + stored_refresh_token_loader, + ) if not litellm_user_id: return _BindingRejection( code="oauth_identity_binding_failed", @@ -247,12 +330,25 @@ async def _evaluate_binding( code="oauth_identity_binding_failed", description=f"could not fetch the issuer's JWKS: {exc}", ) - signing_key: Final = _select_signing_key(id_token, keys) + try: + signing_key: Final = _select_signing_key(id_token, keys) + except (jwt.PyJWTError, ValueError, TypeError): + return _BindingRejection( + code="oauth_identity_binding_failed", + description="invalid id_token header or issuer signing key", + ) if isinstance(signing_key, _BindingRejection): return signing_key claims: Final = _decode_id_token(id_token, binding, signing_key) if isinstance(claims, _BindingRejection): return claims + if grant_type == "authorization_code": + nonce: Final = claims.get("nonce") + if not expected_nonce or not isinstance(nonce, str) or not hmac.compare_digest(nonce, expected_nonce): + return _BindingRejection( + code="oauth_identity_binding_failed", + description="id_token nonce does not match the authenticated authorization transaction", + ) upstream: Final = _upstream_principal(claims, binding) if isinstance(upstream, _BindingRejection): return upstream @@ -267,7 +363,7 @@ async def _evaluate_binding( code="oauth_principal_mismatch", description="The browser account does not match the selected credential owner.", ) - return None + return _binding_proof(binding, litellm_user_id, server_id, caller) async def enforce_oauth_identity_binding( @@ -279,16 +375,17 @@ async def enforce_oauth_identity_binding( jwks_fetcher: JwksFetcher = _fetch_issuer_jwks, caller_principal_loader: CallerPrincipalLoader = _load_caller_principal, stored_refresh_token_loader: StoredRefreshTokenLoader = _load_stored_refresh_token, -) -> None: + expected_nonce: str | None = None, +) -> str | None: """Validate the exchanged token's upstream principal against the LiteLLM caller. No-op when the server has no binding or it is disabled. In enforce mode a failure raises 403 before the caller returns, stores, or caches the token; in audit mode failures are logged only. A refresh_token grant without an id_token is allowed only when the presented refresh token matches - the caller's stored credential or a sealed bridge envelope already proved ownership. + the caller's stored credential and that credential was previously identity-validated. """ binding: Final = server.oauth_identity_binding - if binding is None or binding.mode == "disabled": + if binding is None or binding.mode not in ("audit", "enforce"): return rejection: Final = await _evaluate_binding( binding=binding, @@ -300,9 +397,10 @@ async def enforce_oauth_identity_binding( jwks_fetcher=jwks_fetcher, caller_principal_loader=caller_principal_loader, stored_refresh_token_loader=stored_refresh_token_loader, + expected_nonce=expected_nonce, ) - if rejection is None: - return + if isinstance(rejection, str): + return rejection if binding.mode == "enforce" else None if binding.mode == "audit": verbose_logger.warning( "oauth_identity_binding audit: server=%s user=%s grant=%s rejected=%s (%s)", diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/authz_code_refresher.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/authz_code_refresher.py index 92bd30694af..824459fd9b9 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/authz_code_refresher.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/authz_code_refresher.py @@ -18,6 +18,11 @@ from litellm._logging import verbose_logger from litellm.proxy._experimental.mcp_server.auth.token_endpoint_auth import ( TokenEndpointAuthConfigError, ) +from litellm.proxy._experimental.mcp_server.oauth_identity_binding import ( + BindingValidator, + RefreshTokenPresented, + enforce_oauth_identity_binding, +) from litellm.proxy._experimental.mcp_server.oauth_utils import build_upstream_oauth2_token_request from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( OAuthToken, @@ -39,6 +44,7 @@ class CredentialPersist(Protocol): refresh_token: str | None, expires_in: int | None, scopes: tuple[str, ...] | None, + identity_binding_proof: str | None = None, ) -> None: ... @@ -78,11 +84,13 @@ class AuthorizationCodeRefresher: persist: CredentialPersist, *, clock: Callable[[], float] = time.time, + identity_validator: BindingValidator = enforce_oauth_identity_binding, ) -> None: self._server_lookup = server_lookup self._token_endpoint = token_endpoint self._persist = persist self._clock = clock + self._identity_validator = identity_validator async def refresh(self, user_id: str, server_id: str, token: OAuthToken) -> OAuthToken | None: if token.refresh_token is None: @@ -104,6 +112,15 @@ class AuthorizationCodeRefresher: except TokenEndpointAuthConfigError as exc: verbose_logger.warning("MCP OAuth refresh misconfigured for server %s: %s", server_id, exc) return None + binding: Final = server.oauth_identity_binding + if binding is not None and binding.mode == "enforce": + await self._identity_validator( + server=server, + token_response={}, + litellm_user_id=user_id, + grant_type="refresh_token", + refresh_ownership=RefreshTokenPresented(token.refresh_token), + ) form: Final = { "grant_type": "refresh_token", "refresh_token": token.refresh_token, @@ -116,12 +133,30 @@ class AuthorizationCodeRefresher: if not isinstance(access_token, str) or not access_token: return None + binding_proof: Final = await self._identity_validator( + server=server, + token_response=body, + litellm_user_id=user_id, + grant_type="refresh_token", + refresh_ownership=RefreshTokenPresented(token.refresh_token), + ) rotated: Final = body.get("refresh_token") new_refresh: Final = rotated if isinstance(rotated, str) and rotated else token.refresh_token expires_in: Final = _parse_expires_in(body.get("expires_in")) scopes: Final = _parse_scopes(body.get("scope")) or token.scopes - await self._persist(user_id, server_id, access_token, new_refresh, expires_in, scopes or None) + if binding_proof is not None: + await self._persist( + user_id, + server_id, + access_token, + new_refresh, + expires_in, + scopes or None, + identity_binding_proof=binding_proof, + ) + else: + await self._persist(user_id, server_id, access_token, new_refresh, expires_in, scopes or None) return OAuthToken( access_token=access_token, expires_at=self._clock() + expires_in if expires_in is not None else None, diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py index 5e28396dcfb..efcc613c539 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py @@ -15,6 +15,7 @@ from collections.abc import Callable, Mapping from typing import TYPE_CHECKING, Final from litellm._logging import verbose_logger +from litellm.proxy._experimental.mcp_server.oauth_identity_binding import credential_binding_matches from litellm.proxy._experimental.mcp_server.outbound_credentials.authz_code_refresher import ( AuthorizationCodeRefresher, ) @@ -38,6 +39,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.token_cache_cod OAuthTokenCacheCodec, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.v2_token_store import ( + CredentialReader, V2PerUserTokenStore, ) @@ -69,6 +71,7 @@ async def _persist_credential( refresh_token: str | None, expires_in: int | None, scopes: tuple[str, ...] | None, + identity_binding_proof: str | None = None, ) -> None: from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 store_user_oauth_credential, @@ -86,6 +89,7 @@ async def _persist_credential( expires_in=expires_in, scopes=list(scopes) if scopes else None, skip_byok_guard=True, + identity_binding_proof=identity_binding_proof, ) @@ -172,8 +176,10 @@ class LazyPerUserOAuthTokenStore: *, store_builder: StoreBuilder = _build_per_user_oauth_token_store, redis_available: Callable[[], bool] = _redis_cache_is_available, + credential_reader: CredentialReader = _read_credential, ) -> None: self._server_lookup = server_lookup + self._credential_reader = credential_reader self._store_builder = store_builder self._redis_available = redis_available self._store: InvalidatableOAuthTokenStore | None = None @@ -182,6 +188,13 @@ class LazyPerUserOAuthTokenStore: self._local_fetches = 0 async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None: + server: Final = self._server_lookup(server_id) + binding: Final = server.oauth_identity_binding if server else None + if binding is not None and binding.mode == "enforce": + credential: Final = await self._credential_reader(user_id, server_id) + await self.invalidate(user_id, server_id) + if credential is None or not await credential_binding_matches(binding, user_id, server_id, credential): + return None if self._uses_redis: store = self._store if store is not None: diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 117d95a09de..f6b5441a626 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -1,7 +1,8 @@ from datetime import datetime from typing import Any, Final, Literal -from pydantic import BaseModel, ConfigDict, Field, field_validator +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator +from typing_extensions import Self from litellm.types.mcp import ( DEFAULT_SUBJECT_TOKEN_TYPE, @@ -42,7 +43,7 @@ class MCPOAuthIdentityBinding(BaseModel): """Per-server policy binding stored per-user OAuth credentials to the authenticated LiteLLM caller. When enabled for an interactive oauth2 server, the token relay validates the upstream OIDC - ``id_token`` (signature via the pinned issuer's JWKS, issuer, audience, expiry) and compares its + ``id_token`` (signature via the pinned issuer's JWKS, issuer, audience, expiry, nonce) and compares its principal claim to the LiteLLM caller's trusted identity before the token is returned, stored, or cached. ``audit`` logs mismatches without changing behavior; ``enforce`` fails closed with 403 ``oauth_principal_mismatch`` and disables the direct ``oauth-user-credential`` POST, which @@ -246,6 +247,14 @@ class MCPServer(BaseModel): """ return self.oauth2_flow == "client_credentials" + @model_validator(mode="after") + def validate_identity_binding_mode(self) -> Self: + binding: Final = self.oauth_identity_binding + if binding is not None and binding.mode != "disabled": + if not self.needs_user_oauth_token or self.delegate_auth_to_upstream: + raise ValueError("oauth_identity_binding requires gateway-managed per-user OAuth2 credentials") + return self + @property def needs_user_oauth_token(self) -> bool: """True if this is an OAuth2 server that relies on per-user tokens (no client_credentials).""" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_authz_code_refresher.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_authz_code_refresher.py index bb2f2ff8b02..16707f6ace8 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_authz_code_refresher.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_authz_code_refresher.py @@ -22,6 +22,7 @@ class _Server: server_id="srv", configured_token_url=None, ): + self.oauth_identity_binding = None self.token_url = token_url self.configured_token_url = configured_token_url self.client_id = client_id @@ -287,3 +288,36 @@ async def test_refresh_uses_admin_entered_token_url_when_issuer_yield_empties_re assert token is not None assert token.access_token == "new-at" assert posted[0][0] == "https://idp.example.com/token" + + +@pytest.mark.asyncio +async def test_identity_rejection_never_persists_or_returns_refreshed_token(): + from unittest.mock import AsyncMock + from fastapi import HTTPException + + validator = AsyncMock(side_effect=HTTPException(status_code=403, detail="oauth_principal_mismatch")) + persist = AsyncMock() + refresher = AuthorizationCodeRefresher( + _lookup(_Server()), _endpoint({"access_token": "foreign-token", "id_token": "foreign-identity"}), + persist, identity_validator=validator, + ) + with pytest.raises(HTTPException) as error: + await refresher.refresh("alice", "srv", OAuthToken(access_token="old", refresh_token="old-rt")) + assert error.value.status_code == 403 + persist.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_verified_refresh_preserves_binding_proof_in_storage(): + from unittest.mock import AsyncMock + + validator = AsyncMock(return_value="verified-binding") + persist = AsyncMock() + refresher = AuthorizationCodeRefresher( + _lookup(_Server()), _endpoint({"access_token": "new", "refresh_token": "rotated"}), + persist, identity_validator=validator, + ) + token = await refresher.refresh("alice", "srv", OAuthToken(access_token="old", refresh_token="old-rt")) + assert token.access_token == "new" + assert token.refresh_token == "rotated" + assert persist.await_args.kwargs["identity_binding_proof"] == "verified-binding" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_per_user_oauth_store.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_per_user_oauth_store.py index ca32cf2bb8d..bd190b09c82 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_per_user_oauth_store.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_per_user_oauth_store.py @@ -246,3 +246,28 @@ async def test_lazy_store_invalidate_works_after_redis_chain_is_built() -> None: assert build_calls == 1 assert redis_store.invalidations == [("u", "s")] + + +@pytest.mark.asyncio +async def test_enforcement_invalidates_cached_legacy_credentials_before_use(): + from litellm.types.mcp import MCPAuth, MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer + + server = MCPServer( + server_id="srv", name="srv", transport=MCPTransport.http, auth_type=MCPAuth.oauth2, + oauth_identity_binding=MCPOAuthIdentityBinding( + mode="enforce", issuer="https://idp.example.com", audiences=["client"], + ), + ) + cached = _RecordingStore("belongs-to-bob") + + async def read_legacy(user_id, server_id): + return {"access_token": "belongs-to-bob", "refresh_token": "bobs-refresh"} + + store = LazyPerUserOAuthTokenStore( + lambda server_id: server, store_builder=lambda lookup: (cached, False), + redis_available=lambda: False, credential_reader=read_legacy, + ) + assert await store.fetch("alice", "srv") is None + assert cached.calls == [] + assert cached.invalidations == [("alice", "srv")] diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py index 355b3bfd30e..8e2adbeb06e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py @@ -628,6 +628,7 @@ async def test_oauth_round_trip_returns_payload(): access_token, refresh_token="rfr-xyz", scopes=["a", "b"], + identity_binding_proof="verified-proof", ) stored = _stored_value(prisma) @@ -642,6 +643,7 @@ async def test_oauth_round_trip_returns_payload(): assert result["access_token"] == access_token assert result["refresh_token"] == "rfr-xyz" assert result["scopes"] == ["a", "b"] + assert result["identity_binding_proof"] == "verified-proof" @pytest.mark.asyncio @@ -1184,6 +1186,7 @@ class _RefreshResponse: def _refresh_server(**overrides): base = dict( + oauth_identity_binding=None, token_url="https://idp.example.com/token", server_id="srv-1", client_id="cid", @@ -1309,7 +1312,7 @@ async def test_refresh_user_oauth_token_uses_client_secret_basic(monkeypatch): sends HTTP Basic and keeps the secret out of the body.""" import litellm.proxy._experimental.mcp_server.db as db_mod - server = MagicMock() + server = MagicMock(oauth_identity_binding=None) server.token_url = "https://idp.example.com/oauth2/token" server.server_id = "srv" server.client_id = "cid" @@ -1348,7 +1351,7 @@ async def test_refresh_user_oauth_token_defaults_to_client_secret_post(monkeypat the body (client_secret_post) and sends no Authorization header.""" import litellm.proxy._experimental.mcp_server.db as db_mod - server = MagicMock() + server = MagicMock(oauth_identity_binding=None) server.token_url = "https://idp.example.com/oauth2/token" server.server_id = "srv" server.client_id = "cid" @@ -1609,3 +1612,20 @@ def test_partial_update_defers_omitted_eligibility_fields_to_the_stored_row(): request = UpdateMCPServerRequest(server_id="relay-update", per_server_oauth_discovery=True) assert request.per_server_oauth_discovery is True + + +@pytest.mark.asyncio +async def test_enforcement_rejects_preexisting_unverified_credential(): + from litellm.types.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer + + server = MCPServer( + server_id="srv-1", name="srv-1", url="https://mcp.example.com", transport=MCPTransport.http, auth_type=MCPAuth.oauth2, + oauth_identity_binding=MCPOAuthIdentityBinding( + mode="enforce", issuer="https://idp.example.com", audiences=["client"], + ), + ) + result = await resolve_valid_user_oauth_token( + user_id="alice", server=server, + cred={"access_token": "belongs-to-bob", "refresh_token": "bobs-refresh-token"}, + ) + assert result is None diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 590d6606154..42187cb1771 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -8194,15 +8194,17 @@ async def test_token_exchange_refresh_passes_presented_refresh_ownership(): @pytest.mark.asyncio -async def test_token_exchange_authorization_code_passes_no_refresh_ownership(): +async def test_token_exchange_authorization_code_passes_no_refresh_ownership(monkeypatch): from fastapi import Request from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( exchange_token_with_server, + seal_bridge_authorization_code, ) from litellm.proxy._types import MCPTransport from litellm.types.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer + monkeypatch.setenv("LITELLM_SALT_KEY", "identity-binding-test-salt") server = MCPServer( server_id="srv-1", name="srv-1", @@ -8244,16 +8246,17 @@ async def test_token_exchange_authorization_code_passes_no_refresh_ownership(): request=request, mcp_server=server, grant_type="authorization_code", - code="auth-code", + code=seal_bridge_authorization_code("auth-code", "user-a", "srv-1", "login-nonce"), redirect_uri="https://litellm.example.com/callback", client_id="cid", client_secret=None, - code_verifier=None, + code_verifier="test-verifier", ) assert result.status_code == 200 assert json.loads(result.body)["access_token"] == "at" assert enforce.await_args.kwargs["refresh_ownership"] is None + assert enforce.await_args.kwargs["expected_nonce"] == "login-nonce" def _upstream_token_response(status_code: int, *, json_body: object = None, text_body: str = "") -> "httpx.Response": @@ -11049,3 +11052,54 @@ def test_introspect_route_answers_for_authenticated_caller(monkeypatch): assert active.status_code == 200 assert active.json()["active"] is True assert active.json()["sub"] == "u1" + + +@pytest.mark.asyncio +async def test_identity_bound_authorization_carries_nonce_and_caller_through_callback(monkeypatch): + from http.cookies import SimpleCookie + from urllib.parse import parse_qs, urlparse + from fastapi import Request + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _oauth_state_cookie_name, authorize_with_server, callback, open_bridge_authorization_code, + ) + from litellm.types.mcp import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer + + monkeypatch.setenv("LITELLM_SALT_KEY", "identity-binding-test-salt") + server = MCPServer( + server_id="srv", name="srv", transport=MCPTransport.http, auth_type=MCPAuth.oauth2, + client_id="client", authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + oauth_identity_binding=MCPOAuthIdentityBinding( + mode="enforce", issuer="https://idp.example.com", audiences=["client"], + ), + ) + request = Request({"type": "http", "scheme": "https", "server": ("proxy.example.com", 443), + "path": "/authorize", "query_string": b"", "headers": []}) + with ( + patch( # test-quality-ok: isolate authenticated request resolution from the real encrypted OAuth round trip + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._extract_user_id_from_request", + new=AsyncMock(return_value="alice")), + patch( # test-quality-ok: isolate user access lookup while testing nonce and caller preservation + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._bridge_authorize_access_denial", + new=AsyncMock(return_value=None)), + ): + authorized = await authorize_with_server( + request, server, "client", "http://127.0.0.1:6274/callback", state="client-state", + code_challenge="pkce-challenge", code_challenge_method="S256", + ) + query = parse_qs(urlparse(authorized.headers["location"]).query) + assert len(query["nonce"][0]) >= 32 + cookies = SimpleCookie() + cookies.load(authorized.headers["set-cookie"]) + name = _oauth_state_cookie_name(query["state"][0]) + callback_request = Request({**request.scope, "path": "/callback", + "headers": [(b"cookie", f"{name}={cookies[name].value}".encode())]}) + completed = await callback(callback_request, code="upstream-code", state=query["state"][0]) + returned = parse_qs(urlparse(completed.headers["location"]).query) + sealed = open_bridge_authorization_code(returned["code"][0]) + assert sealed.litellm_user_id == "alice" + assert sealed.mcp_server_id == "srv" + assert sealed.upstream_code == "upstream-code" + assert sealed.oauth_nonce == query["nonce"][0] + assert returned["state"] == ["client-state"] diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_identity_binding.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_identity_binding.py index 36364c32ae5..6f291865e20 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_identity_binding.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_identity_binding.py @@ -13,7 +13,9 @@ from pydantic import ValidationError from litellm.proxy._experimental.mcp_server.oauth_identity_binding import ( RefreshOwnershipProven, + VerifiedRefreshToken, RefreshTokenPresented, + current_binding_proof, _discover_jwks_url, _fetch_issuer_jwks, _load_caller_principal, @@ -21,7 +23,7 @@ from litellm.proxy._experimental.mcp_server.oauth_identity_binding import ( _select_signing_key, enforce_oauth_identity_binding, ) -from litellm.types.mcp import MCPTransport +from litellm.types.mcp import MCPAuth, MCPTransport from litellm.types.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer ISSUER: Final = "https://idp.example.com" @@ -48,6 +50,8 @@ def _sign_id_token(claims: Mapping[str, object]) -> str: "aud": AUDIENCE, "exp": int(time.time()) + 300, "iat": int(time.time()), + "nonce": "test-nonce", + "sub": "upstream-user", **claims, } return jwt.encode(payload, _PRIVATE_PEM, algorithm="RS256", headers={"kid": KID}) @@ -65,8 +69,8 @@ def _caller_loader(email: str | None): def _stored_refresh_token_loader(refresh_token: str | None): - async def load(_user_id: str, _server_id: str) -> str | None: - return refresh_token + async def load(_user_id: str, _server_id: str, _binding: MCPOAuthIdentityBinding) -> VerifiedRefreshToken | None: + return VerifiedRefreshToken(refresh_token, "verified-binding") if refresh_token else None return load @@ -77,6 +81,7 @@ def _server(mode: str = "enforce", **binding_overrides: object) -> MCPServer: name="srv-1", url="https://mcp.example.com", transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, oauth_identity_binding=MCPOAuthIdentityBinding( mode=mode, issuer=ISSUER, @@ -181,7 +186,9 @@ async def test_load_caller_principal_supports_user_id_and_database_email(): @pytest.mark.asyncio async def test_load_stored_refresh_token_returns_credential_and_fails_closed(): - get_credential: Final = AsyncMock(return_value={"refresh_token": "rt-1"}) + binding: Final = _server(caller_field="user_id", principal_claim="sub").oauth_identity_binding + proof: Final = await current_binding_proof(binding, "user-a", "srv-1") + get_credential: Final = AsyncMock(return_value={"refresh_token": "rt-1", "identity_binding_proof": proof}) with ( patch( # test-quality-ok: stored-token loading is a lazy database boundary without injection "litellm.proxy._experimental.mcp_server.db.get_user_oauth_credential", @@ -192,13 +199,15 @@ async def test_load_stored_refresh_token_returns_credential_and_fails_closed(): return_value="prisma", ), ): - assert await _load_stored_refresh_token("user-a", "srv-1") == "rt-1" + assert await _load_stored_refresh_token("user-a", "srv-1", binding) == VerifiedRefreshToken("rt-1", proof) + get_credential.return_value = {"refresh_token": "rt-1"} + assert await _load_stored_refresh_token("user-a", "srv-1", binding) is None with patch( # test-quality-ok: stored-token loading is a lazy database boundary without injection "litellm.proxy.utils.get_prisma_client_or_throw", side_effect=RuntimeError("database unavailable"), ): - assert await _load_stored_refresh_token("user-a", "srv-1") is None + assert await _load_stored_refresh_token("user-a", "srv-1", binding) is None @pytest.mark.asyncio @@ -209,11 +218,12 @@ async def test_matching_principal_passes(): token_response={"access_token": "at", "id_token": token}, litellm_user_id="user-a", grant_type="authorization_code", + expected_nonce="test-nonce", refresh_ownership=None, jwks_fetcher=_jwks_fetcher, caller_principal_loader=_caller_loader("alice@example.com"), ) - assert result is None + assert result is not None @pytest.mark.asyncio @@ -225,6 +235,7 @@ async def test_mismatched_principal_rejected(): token_response={"access_token": "at", "id_token": token}, litellm_user_id="user-a", grant_type="authorization_code", + expected_nonce="test-nonce", refresh_ownership=None, jwks_fetcher=_jwks_fetcher, caller_principal_loader=_caller_loader("alice@example.com"), @@ -243,6 +254,7 @@ async def test_missing_upstream_principal_rejected(): token_response={"access_token": "at", "id_token": token}, litellm_user_id="user-a", grant_type="authorization_code", + expected_nonce="test-nonce", refresh_ownership=None, jwks_fetcher=_jwks_fetcher, caller_principal_loader=_caller_loader("alice@example.com"), @@ -263,6 +275,7 @@ async def test_jwks_fetch_failure_is_rejected(): token_response={"access_token": "at", "id_token": token}, litellm_user_id="user-a", grant_type="authorization_code", + expected_nonce="test-nonce", refresh_ownership=None, jwks_fetcher=fail, caller_principal_loader=_caller_loader("alice@example.com"), @@ -284,6 +297,7 @@ async def test_missing_signing_key_is_rejected(): token_response={"access_token": "at", "id_token": token}, litellm_user_id="user-a", grant_type="authorization_code", + expected_nonce="test-nonce", refresh_ownership=None, jwks_fetcher=no_keys, caller_principal_loader=_caller_loader("alice@example.com"), @@ -301,6 +315,7 @@ async def test_missing_caller_principal_is_rejected(): token_response={"access_token": "at", "id_token": token}, litellm_user_id="user-a", grant_type="authorization_code", + expected_nonce="test-nonce", refresh_ownership=None, jwks_fetcher=_jwks_fetcher, caller_principal_loader=_caller_loader(None), @@ -334,11 +349,12 @@ async def test_user_id_principal_matching_uses_exact_comparison(): token_response={"access_token": "at", "id_token": token}, litellm_user_id="user-a", grant_type="authorization_code", + expected_nonce="test-nonce", refresh_ownership=None, jwks_fetcher=_jwks_fetcher, caller_principal_loader=_caller_loader("user-a"), ) - assert result is None + assert result is not None @pytest.mark.asyncio @@ -349,6 +365,7 @@ async def test_missing_id_token_rejected_on_authorization_code(): token_response={"access_token": "at"}, litellm_user_id="user-a", grant_type="authorization_code", + expected_nonce="test-nonce", refresh_ownership=None, jwks_fetcher=_jwks_fetcher, caller_principal_loader=_caller_loader("alice@example.com"), @@ -369,7 +386,7 @@ async def test_refresh_without_id_token_allowed_when_presented_token_matches_sto caller_principal_loader=_caller_loader("alice@example.com"), stored_refresh_token_loader=_stored_refresh_token_loader("rt-1"), ) - assert result is None + assert result is not None @pytest.mark.asyncio @@ -407,21 +424,18 @@ async def test_refresh_without_id_token_rejects_missing_stored_token(): @pytest.mark.asyncio -async def test_refresh_without_id_token_passes_when_bridge_proves_ownership(): - async def fail_if_called(_user_id: str, _server_id: str) -> str | None: - raise AssertionError("stored refresh token loader should not be called") - - result: Final = await enforce_oauth_identity_binding( - server=_server(), - token_response={"access_token": "at"}, - litellm_user_id="user-a", - grant_type="refresh_token", - refresh_ownership=RefreshOwnershipProven(), - jwks_fetcher=_jwks_fetcher, - caller_principal_loader=_caller_loader("alice@example.com"), - stored_refresh_token_loader=fail_if_called, - ) - assert result is None +async def test_identity_envelope_does_not_prove_upstream_binding(): + with pytest.raises(HTTPException) as error: + await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at"}, + litellm_user_id="user-a", + grant_type="refresh_token", + refresh_ownership=RefreshOwnershipProven(), + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + ) + assert error.value.status_code == 403 @pytest.mark.asyncio @@ -464,6 +478,7 @@ async def test_audit_mode_logs_but_does_not_reject(): token_response={"access_token": "at", "id_token": token}, litellm_user_id="user-a", grant_type="authorization_code", + expected_nonce="test-nonce", refresh_ownership=None, jwks_fetcher=_jwks_fetcher, caller_principal_loader=_caller_loader("alice@example.com"), @@ -480,6 +495,7 @@ async def test_unverified_email_rejected(): token_response={"access_token": "at", "id_token": token}, litellm_user_id="user-a", grant_type="authorization_code", + expected_nonce="test-nonce", refresh_ownership=None, jwks_fetcher=_jwks_fetcher, caller_principal_loader=_caller_loader("alice@example.com"), @@ -503,6 +519,7 @@ async def test_wrong_issuer_rejected(): token_response={"access_token": "at", "id_token": token}, litellm_user_id="user-a", grant_type="authorization_code", + expected_nonce="test-nonce", refresh_ownership=None, jwks_fetcher=_jwks_fetcher, caller_principal_loader=_caller_loader("alice@example.com"), @@ -519,6 +536,7 @@ async def test_no_litellm_identity_rejected(): token_response={"access_token": "at", "id_token": token}, litellm_user_id=None, grant_type="authorization_code", + expected_nonce="test-nonce", refresh_ownership=None, jwks_fetcher=_jwks_fetcher, caller_principal_loader=_caller_loader("alice@example.com"), @@ -533,6 +551,7 @@ async def test_disabled_binding_is_noop(): token_response={"access_token": "at"}, litellm_user_id=None, grant_type="authorization_code", + expected_nonce="test-nonce", refresh_ownership=None, jwks_fetcher=_jwks_fetcher, caller_principal_loader=_caller_loader(None), @@ -553,6 +572,7 @@ async def test_no_binding_is_noop(): token_response={"access_token": "at"}, litellm_user_id=None, grant_type="authorization_code", + expected_nonce="test-nonce", refresh_ownership=None, jwks_fetcher=_jwks_fetcher, caller_principal_loader=_caller_loader(None), @@ -576,9 +596,52 @@ async def test_wrong_audience_rejected(): token_response={"access_token": "at", "id_token": token}, litellm_user_id="user-a", grant_type="authorization_code", + expected_nonce="test-nonce", refresh_ownership=None, jwks_fetcher=_jwks_fetcher, caller_principal_loader=_caller_loader("alice@example.com"), ) assert exc_info.value.status_code == 403 assert exc_info.value.detail["error"] == "oauth_identity_binding_failed" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("nonce", [None, "another-login"]) +async def test_authorization_code_rejects_missing_or_foreign_nonce(nonce): + token = _sign_id_token({"email": "alice@example.com", "email_verified": True, "nonce": nonce}) + with pytest.raises(HTTPException) as error: + await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at", "id_token": token}, + litellm_user_id="user-a", + grant_type="authorization_code", + refresh_ownership=None, + expected_nonce="this-login", + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + ) + assert error.value.status_code == 403 + assert error.value.detail["error"] == "oauth_identity_binding_failed" + + +@pytest.mark.asyncio +async def test_binding_proof_rejects_changed_user_or_policy(): + from litellm.proxy._experimental.mcp_server.oauth_identity_binding import credential_binding_matches + + binding = _server(caller_field="user_id", principal_claim="sub").oauth_identity_binding + proof = await current_binding_proof(binding, "alice", "srv-1") + credential = {"identity_binding_proof": proof} + assert await credential_binding_matches(binding, "alice", "srv-1", credential) + assert not await credential_binding_matches(binding, "bob", "srv-1", credential) + assert not await credential_binding_matches(binding, "alice", "other-server", credential) + changed = binding.model_copy(update={"audiences": ["different-client"]}) + assert not await credential_binding_matches(changed, "alice", "srv-1", credential) + + +@pytest.mark.parametrize("auth_type", [MCPAuth.oauth_delegate, MCPAuth.true_passthrough]) +def test_identity_binding_rejects_modes_without_gateway_credential_custody(auth_type): + with pytest.raises(ValidationError, match="gateway-managed per-user"): + MCPServer( + server_id="srv", name="srv", transport=MCPTransport.http, auth_type=auth_type, + oauth_identity_binding=_server().oauth_identity_binding, + ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index 6a20688c09d..5e00e7d75be 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -4807,7 +4807,7 @@ async def test_store_mcp_oauth_user_credential_blocked_when_identity_binding_enf """The direct opaque-token POST must be closed for enforce-mode identity-bound servers, otherwise it bypasses the token-relay principal check.""" from litellm.proxy._types import MCPOAuthUserCredentialRequest - from litellm.types.mcp import MCPTransport + from litellm.types.mcp import MCPAuth, MCPTransport from litellm.types.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer if not mgmt_endpoints.MCP_AVAILABLE: @@ -4824,6 +4824,7 @@ async def test_store_mcp_oauth_user_credential_blocked_when_identity_binding_enf name=server_id, url="https://mcp.example.com", transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, oauth_identity_binding=MCPOAuthIdentityBinding( mode="enforce", issuer="https://idp.example.com", From 1273e46b8bf4821643eb4b948517c95de34bf632 Mon Sep 17 00:00:00 2001 From: eugene-yao-zocdoc Date: Wed, 26 Aug 2026 15:45:25 -0400 Subject: [PATCH 13/33] feat(redis): add ElastiCache IAM authentication --- litellm/_redis.py | 57 ++++- litellm/_redis_credential_provider.py | 85 +++++++- litellm/proxy/_types.py | 4 + .../cache_settings_endpoints.py | 36 ++++ .../coordination_redis_endpoints.py | 29 +++ tests/test_litellm/test_redis.py | 196 +++++++++++++++++- .../test_redis_credential_provider.py | 122 +++++++++++ 7 files changed, 520 insertions(+), 9 deletions(-) create mode 100644 tests/test_litellm/test_redis_credential_provider.py diff --git a/litellm/_redis.py b/litellm/_redis.py index 3e68d50cf16..e289cc3b6c3 100644 --- a/litellm/_redis.py +++ b/litellm/_redis.py @@ -24,6 +24,7 @@ from redis.credentials import CredentialProvider from litellm import get_secret, get_secret_str from litellm._redis_credential_provider import ( AzureADCredentialProvider, + ElastiCacheIAMCredentialProvider, GCPIAMCredentialProvider, _generate_gcp_iam_access_token, ) @@ -75,6 +76,10 @@ def _get_redis_kwargs(): "azure_client_id", "azure_tenant_id", "azure_client_secret", + "aws_iam_auth", + "aws_iam_user_name", + "aws_iam_cache_name", + "aws_iam_region", } available_args: Final = {x for x in _unwrapped_init_args(redis.Redis) if x not in exclude_args} | include_args @@ -270,6 +275,32 @@ def _redis_kwargs_from_environment(): return return_dict +def _is_true(value: object | None) -> bool: + return value is True or (isinstance(value, str) and value.lower() == "true") + + +def _build_elasticache_iam_provider(redis_kwargs: dict) -> ElastiCacheIAMCredentialProvider | None: + if not _is_true(redis_kwargs.get("aws_iam_auth")): + return None + + required_settings: Final = { + "aws_iam_user_name": redis_kwargs.get("aws_iam_user_name"), + "aws_iam_cache_name": redis_kwargs.get("aws_iam_cache_name"), + "aws_iam_region": redis_kwargs.get("aws_iam_region") + or get_secret_str("AWS_REGION") + or get_secret_str("AWS_DEFAULT_REGION"), + } + missing_settings: Final = tuple(name for name, value in required_settings.items() if not value) + if missing_settings: + raise ValueError("AWS ElastiCache IAM Redis authentication requires: " + ", ".join(missing_settings)) + + return ElastiCacheIAMCredentialProvider( + user_name=str(required_settings["aws_iam_user_name"]), + cache_name=str(required_settings["aws_iam_cache_name"]), + region=str(required_settings["aws_iam_region"]), + ) + + def create_gcp_iam_redis_connect_func( service_account: str, ssl_ca_certs: str | None = None, @@ -540,6 +571,7 @@ def _get_redis_client_logic(**env_overrides): _azure_redis_ad_token: Final = redis_kwargs.get("azure_redis_ad_token") or get_secret("REDIS_AZURE_AD_TOKEN") _azure_ad_enabled: Final = _azure_redis_ad_token is not None and str(_azure_redis_ad_token).lower() == "true" + _aws_iam_enabled: Final = _is_true(redis_kwargs.get("aws_iam_auth")) if _azure_ad_enabled and _gcp_service_account is not None: verbose_logger.warning( @@ -560,13 +592,24 @@ def _get_redis_client_logic(**env_overrides): azure_tenant_id=_azure_tenant_id, azure_client_secret=_azure_client_secret, ) - # Marker for async paths to detect Azure AD auth. The live credential - # object is attached separately as `_azure_credential` by - # `create_azure_ad_redis_connect_func`; the raw client_id/tenant_id/secret - # are intentionally NOT exposed on the function to avoid leaking - # credentials via inspection or logging. redis_kwargs["redis_connect_func"]._azure_redis_ad_token = True + if _aws_iam_enabled and _gcp_service_account is not None: + verbose_logger.warning( + "Both GCP IAM (gcp_service_account) and AWS ElastiCache IAM (aws_iam_auth) are configured " + "for Redis. Using GCP IAM. Remove one to avoid misconfiguration." + ) + elif _aws_iam_enabled and _azure_ad_enabled: + verbose_logger.warning( + "Both Azure AD (azure_redis_ad_token) and AWS ElastiCache IAM (aws_iam_auth) are configured " + "for Redis. Using Azure AD. Remove one to avoid misconfiguration." + ) + elif _aws_iam_enabled: + aws_provider: Final = _build_elasticache_iam_provider(redis_kwargs) + if aws_provider is not None: + verbose_logger.debug("Setting up AWS ElastiCache IAM authentication for Redis.") + redis_kwargs["credential_provider"] = aws_provider + redis_kwargs.pop("gcp_service_account", None) redis_kwargs.pop("gcp_ssl_ca_certs", None) @@ -575,6 +618,10 @@ def _get_redis_client_logic(**env_overrides): redis_kwargs.pop("azure_client_id", None) redis_kwargs.pop("azure_tenant_id", None) redis_kwargs.pop("azure_client_secret", None) + redis_kwargs.pop("aws_iam_auth", None) + redis_kwargs.pop("aws_iam_user_name", None) + redis_kwargs.pop("aws_iam_cache_name", None) + redis_kwargs.pop("aws_iam_region", None) if redis_kwargs.get("credential_provider") is not None: redis_kwargs.pop("redis_connect_func", None) diff --git a/litellm/_redis_credential_provider.py b/litellm/_redis_credential_provider.py index ba0398789a6..a20bbfc530c 100644 --- a/litellm/_redis_credential_provider.py +++ b/litellm/_redis_credential_provider.py @@ -1,7 +1,8 @@ import asyncio import threading import time -from typing import Final, Protocol +from typing import Any, Final, Protocol +from urllib.parse import quote from redis.credentials import CredentialProvider @@ -117,6 +118,88 @@ class GCPIAMCredentialProvider(CredentialProvider): return (token,) +_ELASTICACHE_SERVICE_NAME: Final = "elasticache" +_ELASTICACHE_TOKEN_TTL_SECONDS: Final = 900 + + +class _FrozenBotocoreCredentials(Protocol): + access_key: str + secret_key: str + token: str | None + + +class _BotocoreCredentials(Protocol): + def get_frozen_credentials(self) -> _FrozenBotocoreCredentials: ... + + +class _BotocoreCredentialsResolver(Protocol): + def __call__(self) -> _BotocoreCredentials | None: ... + + +class ElastiCacheIAMCredentialProvider(CredentialProvider): + def __init__( + self, + user_name: str, + cache_name: str, + region: str, + credentials_resolver: _BotocoreCredentialsResolver | None = None, + token_lifetime_seconds: int = _ELASTICACHE_TOKEN_TTL_SECONDS, + ) -> None: + self._user_name = user_name + self._cache_name = cache_name + self._region = region + self._credentials_resolver = credentials_resolver or self._resolve_credentials + self._credentials: _BotocoreCredentials | None = None + self._token_lifetime_seconds = token_lifetime_seconds + + @staticmethod + def _resolve_credentials() -> Any: + try: + import botocore.session + except ImportError as e: + raise ImportError( + "botocore is required for ElastiCache IAM Redis authentication. Install it with: pip install boto3" + ) from e + + return botocore.session.get_session().get_credentials() + + def _get_credentials(self) -> tuple[str, str]: + credentials: Final = self._credentials if self._credentials is not None else self._credentials_resolver() + if credentials is None: + raise RuntimeError("Unable to resolve AWS credentials for ElastiCache IAM Redis authentication") + self._credentials = credentials + + frozen_credentials: Final = credentials.get_frozen_credentials() + if frozen_credentials is None: + raise RuntimeError("Unable to resolve AWS credentials for ElastiCache IAM Redis authentication") + + try: + from botocore.auth import SigV4QueryAuth + from botocore.awsrequest import AWSRequest + except ImportError as e: + raise ImportError( + "botocore is required for ElastiCache IAM Redis authentication. Install it with: pip install boto3" + ) from e + + request: Final = AWSRequest( + method="GET", + url=(f"https://{self._cache_name}/?Action=connect&User={quote(self._user_name, safe='')}"), + ) + SigV4QueryAuth( + frozen_credentials, + _ELASTICACHE_SERVICE_NAME, + self._region, + expires=self._token_lifetime_seconds, + ).add_auth(request) + return self._user_name, request.url.removeprefix("https://") + + def get_credentials(self) -> tuple[str, str]: + return self._get_credentials() + + async def get_credentials_async(self) -> tuple[str, str]: + return await asyncio.to_thread(self._get_credentials) + + class AzureADCredentialProvider(CredentialProvider): """ redis.credentials.CredentialProvider implementation that supplies Azure AD diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c22bb76629d..99103ff3706 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2427,6 +2427,10 @@ class CoordinationRedisParams(LiteLLMPydanticObjectBase): ) sentinel_password: str | None = Field(None, description="password for the sentinel nodes") service_name: str | None = Field(None, description="sentinel service name") + aws_iam_auth: bool | str | None = Field(None, description="enable AWS ElastiCache IAM authentication") + aws_iam_user_name: str | None = Field(None, description="AWS ElastiCache IAM user name") + aws_iam_cache_name: str | None = Field(None, description="AWS ElastiCache cache name") + aws_iam_region: str | None = Field(None, description="AWS region for ElastiCache IAM authentication") def has_connection_target(self) -> bool: return any(value is not None for value in (self.host, self.url, self.startup_nodes, self.sentinel_nodes)) diff --git a/litellm/types/management_endpoints/cache_settings_endpoints.py b/litellm/types/management_endpoints/cache_settings_endpoints.py index 32b32991449..10aece6efa4 100644 --- a/litellm/types/management_endpoints/cache_settings_endpoints.py +++ b/litellm/types/management_endpoints/cache_settings_endpoints.py @@ -237,4 +237,40 @@ CACHE_SETTINGS_FIELDS: Final[list[CacheSettingsField]] = [ ui_field_name="SSL Check Hostname", redis_type=None, ), + CacheSettingsField( + field_name="aws_iam_auth", + field_type="Boolean", + field_value=None, + field_description="Enable AWS ElastiCache IAM authentication", + field_default=False, + ui_field_name="AWS IAM Authentication", + redis_type=None, + ), + CacheSettingsField( + field_name="aws_iam_user_name", + field_type="String", + field_value=None, + field_description="AWS ElastiCache IAM user name", + field_default=None, + ui_field_name="AWS IAM User Name", + redis_type=None, + ), + CacheSettingsField( + field_name="aws_iam_cache_name", + field_type="String", + field_value=None, + field_description="AWS ElastiCache cache name", + field_default=None, + ui_field_name="AWS IAM Cache Name", + redis_type=None, + ), + CacheSettingsField( + field_name="aws_iam_region", + field_type="String", + field_value=None, + field_description="AWS region for ElastiCache IAM authentication", + field_default=None, + ui_field_name="AWS IAM Region", + redis_type=None, + ), ] diff --git a/litellm/types/management_endpoints/coordination_redis_endpoints.py b/litellm/types/management_endpoints/coordination_redis_endpoints.py index 30033346ed7..86b32752cbb 100644 --- a/litellm/types/management_endpoints/coordination_redis_endpoints.py +++ b/litellm/types/management_endpoints/coordination_redis_endpoints.py @@ -102,4 +102,33 @@ COORDINATION_REDIS_SETTINGS_FIELDS: Final[list[CoordinationRedisSettingsField]] ui_field_name="Service Name", section="sentinel", ), + CoordinationRedisSettingsField( + field_name="aws_iam_auth", + field_type="Boolean", + field_description="Enable AWS ElastiCache IAM authentication", + field_default=False, + ui_field_name="AWS IAM Authentication", + section="connection", + ), + CoordinationRedisSettingsField( + field_name="aws_iam_user_name", + field_type="String", + field_description="AWS ElastiCache IAM user name", + ui_field_name="AWS IAM User Name", + section="connection", + ), + CoordinationRedisSettingsField( + field_name="aws_iam_cache_name", + field_type="String", + field_description="AWS ElastiCache cache name", + ui_field_name="AWS IAM Cache Name", + section="connection", + ), + CoordinationRedisSettingsField( + field_name="aws_iam_region", + field_type="String", + field_description="AWS region for ElastiCache IAM authentication", + ui_field_name="AWS IAM Region", + section="connection", + ), ] diff --git a/tests/test_litellm/test_redis.py b/tests/test_litellm/test_redis.py index a96e8541e06..c99a3a074fb 100644 --- a/tests/test_litellm/test_redis.py +++ b/tests/test_litellm/test_redis.py @@ -24,6 +24,7 @@ from litellm._redis import ( ) from litellm._redis_credential_provider import ( AzureADCredentialProvider, + ElastiCacheIAMCredentialProvider, GCPIAMCredentialProvider, _token_cache, ) @@ -78,6 +79,8 @@ def clean_redis_environment(monkeypatch): "REDIS_URL", "REDIS_CLUSTER_NODES", "REDIS_SENTINEL_NODES", + "AWS_REGION", + "AWS_DEFAULT_REGION", *_get_redis_env_kwarg_mapping(), ): monkeypatch.delenv(var, raising=False) @@ -110,6 +113,17 @@ def test_credential_provider_is_not_environment_derived(): assert "credential_provider" not in mapping.values() +def test_aws_iam_settings_are_environment_derived(): + allowed = _get_redis_kwargs() + mapping = _get_redis_env_kwarg_mapping() + + assert {"aws_iam_auth", "aws_iam_user_name", "aws_iam_cache_name", "aws_iam_region"} <= allowed + assert mapping["REDIS_AWS_IAM_AUTH"] == "aws_iam_auth" + assert mapping["REDIS_AWS_IAM_USER_NAME"] == "aws_iam_user_name" + assert mapping["REDIS_AWS_IAM_CACHE_NAME"] == "aws_iam_cache_name" + assert mapping["REDIS_AWS_IAM_REGION"] == "aws_iam_region" + + def test_sync_direct_preserves_credential_provider_identity(clean_redis_environment): provider = _StubCredentialProvider() @@ -300,6 +314,180 @@ def test_gcp_kwargs_never_survive_client_logic(clean_redis_environment, override assert "gcp_ssl_ca_certs" not in redis_kwargs +def test_aws_iam_environment_settings_install_provider(clean_redis_environment, monkeypatch): + monkeypatch.setenv("REDIS_AWS_IAM_AUTH", "true") + monkeypatch.setenv("REDIS_AWS_IAM_USER_NAME", "iam-user") + monkeypatch.setenv("REDIS_AWS_IAM_CACHE_NAME", "cache.example.com") + monkeypatch.setenv("REDIS_AWS_IAM_REGION", "us-east-1") + + redis_kwargs = _get_redis_client_logic(host="cache.example.com", port=6379) + + assert isinstance(redis_kwargs["credential_provider"], ElastiCacheIAMCredentialProvider) + assert not {"aws_iam_auth", "aws_iam_user_name", "aws_iam_cache_name", "aws_iam_region"} & redis_kwargs.keys() + + +def test_aws_iam_settings_are_removed_for_url_and_static_credentials(clean_redis_environment): + redis_kwargs = _get_redis_client_logic( + url="rediss://url-user:url-pass@cache.example.com:6380", + aws_iam_auth=True, + aws_iam_user_name="iam-user", + aws_iam_cache_name="cache.example.com", + aws_iam_region="us-east-1", + username="static-user", + password="static-password", + ) + + assert isinstance(redis_kwargs["credential_provider"], ElastiCacheIAMCredentialProvider) + assert redis_kwargs["url"] == "rediss://cache.example.com:6380" + assert "username" not in redis_kwargs + assert "password" not in redis_kwargs + assert not {"aws_iam_auth", "aws_iam_user_name", "aws_iam_cache_name", "aws_iam_region"} & redis_kwargs.keys() + + +@pytest.mark.parametrize("missing", ["aws_iam_user_name", "aws_iam_cache_name", "aws_iam_region"]) +def test_aws_iam_missing_setting_fails_closed(clean_redis_environment, missing): + settings = { + "aws_iam_auth": True, + "aws_iam_user_name": "iam-user", + "aws_iam_cache_name": "cache.example.com", + "aws_iam_region": "us-east-1", + } + settings[missing] = None + + with pytest.raises(ValueError, match=missing): + _get_redis_client_logic(host="cache.example.com", port=6379, **settings) + + +@pytest.mark.parametrize("region_var", ["AWS_REGION", "AWS_DEFAULT_REGION"]) +def test_aws_iam_region_falls_back_to_environment(clean_redis_environment, monkeypatch, region_var): + monkeypatch.setenv(region_var, "sa-east-1") + + redis_kwargs = _get_redis_client_logic( + host="cache.example.com", + port=6379, + aws_iam_auth=True, + aws_iam_user_name="iam-user", + aws_iam_cache_name="cache.example.com", + ) + + provider = redis_kwargs["credential_provider"] + assert isinstance(provider, ElastiCacheIAMCredentialProvider) + assert provider._region == "sa-east-1" + + +def test_aws_iam_region_prefers_aws_region_over_default_region(clean_redis_environment, monkeypatch): + monkeypatch.setenv("AWS_REGION", "sa-east-1") + monkeypatch.setenv("AWS_DEFAULT_REGION", "eu-west-1") + + redis_kwargs = _get_redis_client_logic( + host="cache.example.com", + port=6379, + aws_iam_auth=True, + aws_iam_user_name="iam-user", + aws_iam_cache_name="cache.example.com", + ) + + assert redis_kwargs["credential_provider"]._region == "sa-east-1" + + +def test_aws_iam_region_prefers_explicit_over_environment(clean_redis_environment, monkeypatch): + monkeypatch.setenv("AWS_REGION", "sa-east-1") + + redis_kwargs = _get_redis_client_logic( + host="cache.example.com", + port=6379, + aws_iam_auth=True, + aws_iam_user_name="iam-user", + aws_iam_cache_name="cache.example.com", + aws_iam_region="explicit-region", + ) + + assert redis_kwargs["credential_provider"]._region == "explicit-region" + + +def test_aws_iam_settings_map_to_distinct_provider_fields(clean_redis_environment): + redis_kwargs = _get_redis_client_logic( + host="cache.example.com", + port=6379, + aws_iam_auth=True, + aws_iam_user_name="iam-user-value", + aws_iam_cache_name="iam-cache-value", + aws_iam_region="iam-region-value", + ) + + provider = redis_kwargs["credential_provider"] + assert isinstance(provider, ElastiCacheIAMCredentialProvider) + assert provider._user_name == "iam-user-value" + assert provider._cache_name == "iam-cache-value" + assert provider._region == "iam-region-value" + + +@pytest.mark.parametrize("aws_iam_auth", [False, "false"]) +def test_aws_iam_auth_disabled_does_not_install_provider(clean_redis_environment, aws_iam_auth): + redis_kwargs = _get_redis_client_logic( + host="cache.example.com", + port=6379, + aws_iam_auth=aws_iam_auth, + aws_iam_user_name="iam-user", + aws_iam_cache_name="cache.example.com", + aws_iam_region="us-east-1", + ) + + assert "credential_provider" not in redis_kwargs + assert not {"aws_iam_auth", "aws_iam_user_name", "aws_iam_cache_name", "aws_iam_region"} & redis_kwargs.keys() + + +def test_explicit_provider_wins_over_aws_iam(clean_redis_environment): + provider = _StubCredentialProvider() + + redis_kwargs = _get_redis_client_logic( + host="cache.example.com", + port=6379, + credential_provider=provider, + aws_iam_auth=True, + aws_iam_user_name="iam-user", + aws_iam_cache_name="cache.example.com", + aws_iam_region="us-east-1", + ) + + assert redis_kwargs["credential_provider"] is provider + assert "aws_iam_auth" not in redis_kwargs + + +def test_gcp_wins_over_aws_iam(clean_redis_environment): + with patch("litellm._redis.create_gcp_iam_redis_connect_func") as mock_gcp: + mock_gcp.return_value = _gcp_marker_callback() + redis_kwargs = _get_redis_client_logic( + host="cache.example.com", + port=6379, + gcp_service_account="sa@example.com", + aws_iam_auth=True, + aws_iam_user_name="iam-user", + aws_iam_cache_name="cache.example.com", + aws_iam_region="us-east-1", + ) + + assert "credential_provider" not in redis_kwargs + assert redis_kwargs["redis_connect_func"] is mock_gcp.return_value + + +def test_azure_wins_over_aws_iam(clean_redis_environment): + with patch("litellm._redis.create_azure_ad_redis_connect_func") as mock_azure: + mock_azure.return_value = MagicMock() + redis_kwargs = _get_redis_client_logic( + host="cache.example.com", + port=6379, + azure_redis_ad_token="true", + aws_iam_auth=True, + aws_iam_user_name="iam-user", + aws_iam_cache_name="cache.example.com", + aws_iam_region="us-east-1", + ) + + assert "credential_provider" not in redis_kwargs + assert redis_kwargs["redis_connect_func"] is mock_azure.return_value + + def test_provider_keeps_the_rest_of_the_url_intact(clean_redis_environment): provider = _StubCredentialProvider() @@ -1530,9 +1718,11 @@ def test_async_sentinel_keeps_the_credential_provider_off_the_monitors(markers, "redis_connect_func": SimpleNamespace(**markers), } - with patch("litellm._redis.async_redis.Sentinel") as mock_sentinel_cls: - with patch("litellm._redis._get_redis_client_logic", return_value=redis_kwargs): - get_redis_async_client() + with ( + patch("litellm._redis.async_redis.Sentinel") as mock_sentinel_cls, + patch("litellm._redis._get_redis_client_logic", return_value=redis_kwargs), + ): + get_redis_async_client() sentinel_kwargs = mock_sentinel_cls.call_args[1]["sentinel_kwargs"] assert sentinel_kwargs["password"] == sentinel_password diff --git a/tests/test_litellm/test_redis_credential_provider.py b/tests/test_litellm/test_redis_credential_provider.py new file mode 100644 index 00000000000..6299d4d42e0 --- /dev/null +++ b/tests/test_litellm/test_redis_credential_provider.py @@ -0,0 +1,122 @@ +import asyncio +from types import SimpleNamespace +from urllib.parse import parse_qs, urlsplit + +import pytest + +from litellm._redis_credential_provider import ( + ElastiCacheIAMCredentialProvider, + _BotocoreCredentials, +) + + +class _FakeCredentials: + def __init__(self, access_key: str) -> None: + self.access_key = access_key + self.secret_key = "synthetic-secret" + self.token = "synthetic-session-token" + + def get_frozen_credentials(self): + return self + + +class _FakeResolver: + def __init__(self, credentials: _BotocoreCredentials | None) -> None: + self.credentials = credentials + self.calls = 0 + + def __call__(self): + self.calls += 1 + return self.credentials + + +class _RotatingFakeCredentials: + def __init__(self) -> None: + self.calls = 0 + + def __bool__(self) -> bool: + return False + + def get_frozen_credentials(self): + self.calls += 1 + return SimpleNamespace( + access_key=f"AKIA-SYNTHETIC-{self.calls}", + secret_key="synthetic-secret", + token="synthetic-session-token", + ) + + +def test_elasticache_provider_signs_expected_query(): + resolver = _FakeResolver(_FakeCredentials("AKIA-SYNTHETIC")) + provider = ElastiCacheIAMCredentialProvider( + user_name="iam-user", + cache_name="cache.example.com", + region="us-east-1", + credentials_resolver=resolver, + ) + + user_name, token = provider.get_credentials() + parsed = urlsplit("https://" + token) + query = parse_qs(parsed.query) + + assert user_name == "iam-user" + assert parsed.netloc == "cache.example.com" + assert query["Action"] == ["connect"] + assert query["User"] == ["iam-user"] + assert query["X-Amz-Expires"] == ["900"] + assert "elasticache" in query["X-Amz-Credential"][0] + assert query["X-Amz-Credential"][0].split("/")[2] == "us-east-1" + assert not token.startswith("https://") + + +def test_elasticache_provider_resolves_credentials_once_but_refreshes_signature(): + rotating_credentials = _RotatingFakeCredentials() + resolver = _FakeResolver(rotating_credentials) + provider = ElastiCacheIAMCredentialProvider( + user_name="iam-user", + cache_name="cache.example.com", + region="us-east-1", + credentials_resolver=resolver, + ) + + first = provider.get_credentials() + second = provider.get_credentials() + async_result = asyncio.run(provider.get_credentials_async()) + + assert first[0] == second[0] == async_result[0] == "iam-user" + assert first[1] != second[1] + assert async_result[1] != second[1] + assert resolver.calls == 1 + assert rotating_credentials.calls == 3 + + +def test_elasticache_provider_reports_missing_credentials(): + provider = ElastiCacheIAMCredentialProvider( + user_name="iam-user", + cache_name="cache.example.com", + region="us-east-1", + credentials_resolver=_FakeResolver(None), + ) + + with pytest.raises(RuntimeError, match="Unable to resolve AWS credentials"): + provider.get_credentials() + + +def test_elasticache_provider_recovers_after_a_failed_resolution(): + resolver = _FakeResolver(None) + provider = ElastiCacheIAMCredentialProvider( + user_name="iam-user", + cache_name="cache.example.com", + region="us-east-1", + credentials_resolver=resolver, + ) + + with pytest.raises(RuntimeError, match="Unable to resolve AWS credentials"): + provider.get_credentials() + + resolver.credentials = _FakeCredentials("AKIA-SYNTHETIC") + user_name, token = provider.get_credentials() + + assert user_name == "iam-user" + assert token + assert resolver.calls == 2 From 604b9edc533ec236eaa3563e61e5fade4bfd6279 Mon Sep 17 00:00:00 2001 From: eugene-yao-zocdoc Date: Wed, 26 Aug 2026 15:58:02 -0400 Subject: [PATCH 14/33] test(redis): cover AWS IAM provider install on async cluster nodes No existing test combined aws_iam_auth with startup_nodes, the shape ElastiCache/Valkey Serverless deployments actually use. --- tests/test_litellm/test_redis.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/tests/test_litellm/test_redis.py b/tests/test_litellm/test_redis.py index c99a3a074fb..21c23ba7866 100644 --- a/tests/test_litellm/test_redis.py +++ b/tests/test_litellm/test_redis.py @@ -488,6 +488,20 @@ def test_azure_wins_over_aws_iam(clean_redis_environment): assert redis_kwargs["redis_connect_func"] is mock_azure.return_value +def test_async_cluster_installs_aws_iam_provider(clean_redis_environment): + startup_nodes = [{"host": "cluster-node", "port": 6379}] + + client = get_redis_async_client( + startup_nodes=startup_nodes, + aws_iam_auth=True, + aws_iam_user_name="iam-user", + aws_iam_cache_name="cache.example.com", + aws_iam_region="us-east-1", + ) + + assert isinstance(client.connection_kwargs["credential_provider"], ElastiCacheIAMCredentialProvider) + + def test_provider_keeps_the_rest_of_the_url_intact(clean_redis_environment): provider = _StubCredentialProvider() From f66891d80f4c7ce4d47a5ddb4ce499e5df21065a Mon Sep 17 00:00:00 2001 From: eugene-yao-zocdoc Date: Wed, 26 Aug 2026 19:19:01 -0400 Subject: [PATCH 15/33] fix(redis): type ElastiCache IAM configuration Generated with AI Co-Authored-By: Claude Code --- litellm/_redis.py | 41 ++++++++++--------- litellm/_redis_credential_provider.py | 30 ++++++-------- .../test_redis_credential_provider.py | 25 +++++------ ui/litellm-dashboard/src/lib/http/schema.d.ts | 20 +++++++++ 4 files changed, 65 insertions(+), 51 deletions(-) diff --git a/litellm/_redis.py b/litellm/_redis.py index e289cc3b6c3..2fd8a7100e3 100644 --- a/litellm/_redis.py +++ b/litellm/_redis.py @@ -279,25 +279,24 @@ def _is_true(value: object | None) -> bool: return value is True or (isinstance(value, str) and value.lower() == "true") -def _build_elasticache_iam_provider(redis_kwargs: dict) -> ElastiCacheIAMCredentialProvider | None: - if not _is_true(redis_kwargs.get("aws_iam_auth")): - return None - - required_settings: Final = { - "aws_iam_user_name": redis_kwargs.get("aws_iam_user_name"), - "aws_iam_cache_name": redis_kwargs.get("aws_iam_cache_name"), - "aws_iam_region": redis_kwargs.get("aws_iam_region") - or get_secret_str("AWS_REGION") - or get_secret_str("AWS_DEFAULT_REGION"), - } - missing_settings: Final = tuple(name for name, value in required_settings.items() if not value) +def _build_elasticache_iam_provider( + user_name: object | None, + cache_name: object | None, + region: object | None, +) -> ElastiCacheIAMCredentialProvider: + required_settings: Final = ( + ("aws_iam_user_name", user_name), + ("aws_iam_cache_name", cache_name), + ("aws_iam_region", region), + ) + missing_settings: Final = tuple(name for name, value in required_settings if not value) if missing_settings: raise ValueError("AWS ElastiCache IAM Redis authentication requires: " + ", ".join(missing_settings)) return ElastiCacheIAMCredentialProvider( - user_name=str(required_settings["aws_iam_user_name"]), - cache_name=str(required_settings["aws_iam_cache_name"]), - region=str(required_settings["aws_iam_region"]), + user_name=str(user_name), + cache_name=str(cache_name), + region=str(region), ) @@ -605,10 +604,14 @@ def _get_redis_client_logic(**env_overrides): "for Redis. Using Azure AD. Remove one to avoid misconfiguration." ) elif _aws_iam_enabled: - aws_provider: Final = _build_elasticache_iam_provider(redis_kwargs) - if aws_provider is not None: - verbose_logger.debug("Setting up AWS ElastiCache IAM authentication for Redis.") - redis_kwargs["credential_provider"] = aws_provider + verbose_logger.debug("Setting up AWS ElastiCache IAM authentication for Redis.") + redis_kwargs["credential_provider"] = _build_elasticache_iam_provider( + user_name=redis_kwargs.get("aws_iam_user_name"), + cache_name=redis_kwargs.get("aws_iam_cache_name"), + region=redis_kwargs.get("aws_iam_region") + or get_secret_str("AWS_REGION") + or get_secret_str("AWS_DEFAULT_REGION"), + ) redis_kwargs.pop("gcp_service_account", None) redis_kwargs.pop("gcp_ssl_ca_certs", None) diff --git a/litellm/_redis_credential_provider.py b/litellm/_redis_credential_provider.py index a20bbfc530c..85759af2a88 100644 --- a/litellm/_redis_credential_provider.py +++ b/litellm/_redis_credential_provider.py @@ -1,11 +1,19 @@ +from __future__ import annotations + import asyncio import threading import time -from typing import Any, Final, Protocol +from collections.abc import Callable +from typing import TYPE_CHECKING, Any, Final, Protocol from urllib.parse import quote from redis.credentials import CredentialProvider +if TYPE_CHECKING: + from botocore.credentials import Credentials +else: + Credentials = Any # rebind-ok: runtime alias for the type-checking-only botocore import + # Azure AD scope for Redis Cache for Azure. AZURE_REDIS_SCOPE: Final = "https://redis.azure.com/.default" @@ -122,38 +130,24 @@ _ELASTICACHE_SERVICE_NAME: Final = "elasticache" _ELASTICACHE_TOKEN_TTL_SECONDS: Final = 900 -class _FrozenBotocoreCredentials(Protocol): - access_key: str - secret_key: str - token: str | None - - -class _BotocoreCredentials(Protocol): - def get_frozen_credentials(self) -> _FrozenBotocoreCredentials: ... - - -class _BotocoreCredentialsResolver(Protocol): - def __call__(self) -> _BotocoreCredentials | None: ... - - class ElastiCacheIAMCredentialProvider(CredentialProvider): def __init__( self, user_name: str, cache_name: str, region: str, - credentials_resolver: _BotocoreCredentialsResolver | None = None, + credentials_resolver: Callable[[], Credentials | None] | None = None, token_lifetime_seconds: int = _ELASTICACHE_TOKEN_TTL_SECONDS, ) -> None: self._user_name = user_name self._cache_name = cache_name self._region = region self._credentials_resolver = credentials_resolver or self._resolve_credentials - self._credentials: _BotocoreCredentials | None = None + self._credentials: Credentials | None = None self._token_lifetime_seconds = token_lifetime_seconds @staticmethod - def _resolve_credentials() -> Any: + def _resolve_credentials() -> Credentials | None: try: import botocore.session except ImportError as e: diff --git a/tests/test_litellm/test_redis_credential_provider.py b/tests/test_litellm/test_redis_credential_provider.py index 6299d4d42e0..de4afb7c313 100644 --- a/tests/test_litellm/test_redis_credential_provider.py +++ b/tests/test_litellm/test_redis_credential_provider.py @@ -4,10 +4,7 @@ from urllib.parse import parse_qs, urlsplit import pytest -from litellm._redis_credential_provider import ( - ElastiCacheIAMCredentialProvider, - _BotocoreCredentials, -) +from litellm._redis_credential_provider import ElastiCacheIAMCredentialProvider class _FakeCredentials: @@ -20,16 +17,6 @@ class _FakeCredentials: return self -class _FakeResolver: - def __init__(self, credentials: _BotocoreCredentials | None) -> None: - self.credentials = credentials - self.calls = 0 - - def __call__(self): - self.calls += 1 - return self.credentials - - class _RotatingFakeCredentials: def __init__(self) -> None: self.calls = 0 @@ -46,6 +33,16 @@ class _RotatingFakeCredentials: ) +class _FakeResolver: + def __init__(self, credentials: _FakeCredentials | _RotatingFakeCredentials | None) -> None: + self.credentials = credentials + self.calls = 0 + + def __call__(self): + self.calls += 1 + return self.credentials + + def test_elasticache_provider_signs_expected_query(): resolver = _FakeResolver(_FakeCredentials("AKIA-SYNTHETIC")) provider = ElastiCacheIAMCredentialProvider( diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 6d62ce2b675..9ca0df396b7 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -26357,6 +26357,26 @@ export interface components { * independently of the response-cache backend in `litellm_settings.cache_params`. */ CoordinationRedisParams: { + /** + * Aws Iam Auth + * @description enable AWS ElastiCache IAM authentication + */ + aws_iam_auth?: boolean | string | null; + /** + * Aws Iam Cache Name + * @description AWS ElastiCache cache name + */ + aws_iam_cache_name?: string | null; + /** + * Aws Iam Region + * @description AWS region for ElastiCache IAM authentication + */ + aws_iam_region?: string | null; + /** + * Aws Iam User Name + * @description AWS ElastiCache IAM user name + */ + aws_iam_user_name?: string | null; /** * Host * @description Redis hostname From bb043fe29a7f680f8a5aee1ba5998f3fdb26e561 Mon Sep 17 00:00:00 2001 From: eugene-yao-zocdoc Date: Wed, 26 Aug 2026 19:48:03 -0400 Subject: [PATCH 16/33] test(redis): cover ElastiCache IAM failures Generated with AI Co-Authored-By: Claude Code --- tests/test_litellm/test_redis.py | 52 ++++++-------- .../test_redis_credential_provider.py | 70 +++++++++++++++++++ 2 files changed, 93 insertions(+), 29 deletions(-) diff --git a/tests/test_litellm/test_redis.py b/tests/test_litellm/test_redis.py index 21c23ba7866..43fb7563233 100644 --- a/tests/test_litellm/test_redis.py +++ b/tests/test_litellm/test_redis.py @@ -455,37 +455,33 @@ def test_explicit_provider_wins_over_aws_iam(clean_redis_environment): def test_gcp_wins_over_aws_iam(clean_redis_environment): - with patch("litellm._redis.create_gcp_iam_redis_connect_func") as mock_gcp: - mock_gcp.return_value = _gcp_marker_callback() - redis_kwargs = _get_redis_client_logic( - host="cache.example.com", - port=6379, - gcp_service_account="sa@example.com", - aws_iam_auth=True, - aws_iam_user_name="iam-user", - aws_iam_cache_name="cache.example.com", - aws_iam_region="us-east-1", - ) + redis_kwargs = _get_redis_client_logic( + host="cache.example.com", + port=6379, + gcp_service_account="sa@example.com", + aws_iam_auth=True, + aws_iam_user_name="iam-user", + aws_iam_cache_name="cache.example.com", + aws_iam_region="us-east-1", + ) assert "credential_provider" not in redis_kwargs - assert redis_kwargs["redis_connect_func"] is mock_gcp.return_value + assert redis_kwargs["redis_connect_func"]._gcp_service_account == "sa@example.com" def test_azure_wins_over_aws_iam(clean_redis_environment): - with patch("litellm._redis.create_azure_ad_redis_connect_func") as mock_azure: - mock_azure.return_value = MagicMock() - redis_kwargs = _get_redis_client_logic( - host="cache.example.com", - port=6379, - azure_redis_ad_token="true", - aws_iam_auth=True, - aws_iam_user_name="iam-user", - aws_iam_cache_name="cache.example.com", - aws_iam_region="us-east-1", - ) + redis_kwargs = _get_redis_client_logic( + host="cache.example.com", + port=6379, + azure_redis_ad_token="true", + aws_iam_auth=True, + aws_iam_user_name="iam-user", + aws_iam_cache_name="cache.example.com", + aws_iam_region="us-east-1", + ) assert "credential_provider" not in redis_kwargs - assert redis_kwargs["redis_connect_func"] is mock_azure.return_value + assert redis_kwargs["redis_connect_func"]._azure_redis_ad_token is True def test_async_cluster_installs_aws_iam_provider(clean_redis_environment): @@ -1732,11 +1728,9 @@ def test_async_sentinel_keeps_the_credential_provider_off_the_monitors(markers, "redis_connect_func": SimpleNamespace(**markers), } - with ( - patch("litellm._redis.async_redis.Sentinel") as mock_sentinel_cls, - patch("litellm._redis._get_redis_client_logic", return_value=redis_kwargs), - ): - get_redis_async_client() + with patch("litellm._redis.async_redis.Sentinel") as mock_sentinel_cls: + with patch("litellm._redis._get_redis_client_logic", return_value=redis_kwargs): + get_redis_async_client() sentinel_kwargs = mock_sentinel_cls.call_args[1]["sentinel_kwargs"] assert sentinel_kwargs["password"] == sentinel_password diff --git a/tests/test_litellm/test_redis_credential_provider.py b/tests/test_litellm/test_redis_credential_provider.py index de4afb7c313..336e30e12c6 100644 --- a/tests/test_litellm/test_redis_credential_provider.py +++ b/tests/test_litellm/test_redis_credential_provider.py @@ -1,4 +1,6 @@ import asyncio +import builtins +import sys from types import SimpleNamespace from urllib.parse import parse_qs, urlsplit @@ -87,6 +89,41 @@ def test_elasticache_provider_resolves_credentials_once_but_refreshes_signature( assert rotating_credentials.calls == 3 +def test_elasticache_provider_uses_botocore_session_credentials(monkeypatch): + credentials = _FakeCredentials("AKIA-SYNTHETIC") + monkeypatch.setattr("botocore.session.get_session", lambda: SimpleNamespace(get_credentials=lambda: credentials)) + provider = ElastiCacheIAMCredentialProvider( + user_name="iam-user", + cache_name="cache.example.com", + region="us-east-1", + ) + + user_name, token = provider.get_credentials() + + assert user_name == "iam-user" + assert "AKIA-SYNTHETIC" in token + + +def test_elasticache_provider_reports_missing_botocore(monkeypatch): + original_import = builtins.__import__ + + def import_without_botocore(name, *args, **kwargs): + if name == "botocore.session": + raise ImportError("synthetic missing dependency") + return original_import(name, *args, **kwargs) + + monkeypatch.delitem(sys.modules, "botocore.session", raising=False) + monkeypatch.setattr(builtins, "__import__", import_without_botocore) + provider = ElastiCacheIAMCredentialProvider( + user_name="iam-user", + cache_name="cache.example.com", + region="us-east-1", + ) + + with pytest.raises(ImportError, match="pip install boto3"): + provider.get_credentials() + + def test_elasticache_provider_reports_missing_credentials(): provider = ElastiCacheIAMCredentialProvider( user_name="iam-user", @@ -99,6 +136,39 @@ def test_elasticache_provider_reports_missing_credentials(): provider.get_credentials() +def test_elasticache_provider_reports_missing_frozen_credentials(): + provider = ElastiCacheIAMCredentialProvider( + user_name="iam-user", + cache_name="cache.example.com", + region="us-east-1", + credentials_resolver=_FakeResolver(SimpleNamespace(get_frozen_credentials=lambda: None)), + ) + + with pytest.raises(RuntimeError, match="Unable to resolve AWS credentials"): + provider.get_credentials() + + +def test_elasticache_provider_reports_missing_signing_dependency(monkeypatch): + original_import = builtins.__import__ + + def import_without_botocore_auth(name, *args, **kwargs): + if name == "botocore.auth": + raise ImportError("synthetic missing dependency") + return original_import(name, *args, **kwargs) + + monkeypatch.delitem(sys.modules, "botocore.auth", raising=False) + monkeypatch.setattr(builtins, "__import__", import_without_botocore_auth) + provider = ElastiCacheIAMCredentialProvider( + user_name="iam-user", + cache_name="cache.example.com", + region="us-east-1", + credentials_resolver=_FakeResolver(_FakeCredentials("AKIA-SYNTHETIC")), + ) + + with pytest.raises(ImportError, match="pip install boto3"): + provider.get_credentials() + + def test_elasticache_provider_recovers_after_a_failed_resolution(): resolver = _FakeResolver(None) provider = ElastiCacheIAMCredentialProvider( From 921c263876600776be8dd69980f683b2c9ea3940 Mon Sep 17 00:00:00 2001 From: eugene-yao-zocdoc Date: Wed, 26 Aug 2026 20:21:03 -0400 Subject: [PATCH 17/33] fix(proxy): validate coordination Redis mappings Generated with AI Co-Authored-By: Claude Code --- .../management_endpoints/coordination_redis_endpoints.py | 2 +- litellm/proxy/proxy_server.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/management_endpoints/coordination_redis_endpoints.py b/litellm/proxy/management_endpoints/coordination_redis_endpoints.py index 86ce336c7a3..88dc09ab001 100644 --- a/litellm/proxy/management_endpoints/coordination_redis_endpoints.py +++ b/litellm/proxy/management_endpoints/coordination_redis_endpoints.py @@ -140,7 +140,7 @@ def _merge_over_saved( def _validated_params(settings: Mapping[str, object]) -> CoordinationRedisParams: """Validate settings the way startup does: resolve env refs, then require a connection target.""" try: - params: Final = CoordinationRedisParams(**_resolve_env_refs(settings)) + params: Final = CoordinationRedisParams.model_validate(_resolve_env_refs(settings)) except ValidationError as e: invalid_fields: Final = sorted({str(error["loc"][0]) for error in e.errors() if error["loc"]}) raise HTTPException( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9269fd48e6c..c21df8ce931 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -4906,7 +4906,7 @@ class ProxyConfig: if not isinstance(raw_params, dict): raise ValueError("general_settings.coordination_redis must be a mapping of Redis connection params") - coordination_params: Final = CoordinationRedisParams(**_resolve_coordination_redis_env_refs(raw_params)) + coordination_params: Final = CoordinationRedisParams.model_validate(_resolve_coordination_redis_env_refs(raw_params)) if not coordination_params.has_connection_target(): raise ValueError( "general_settings.coordination_redis needs a connection target: " @@ -9234,7 +9234,7 @@ class ProxyStartupEvent: if persisted is None: return None - coordination_params: Final = CoordinationRedisParams(**_resolve_coordination_redis_env_refs(persisted)) + coordination_params: Final = CoordinationRedisParams.model_validate(_resolve_coordination_redis_env_refs(persisted)) if not coordination_params.has_connection_target(): verbose_proxy_logger.warning( "coordination_redis saved in the database names no connection target; ignoring it." From 9763bd80ae442533d234e733cf334e3e1565681e Mon Sep 17 00:00:00 2001 From: eugene-yao-zocdoc Date: Wed, 26 Aug 2026 20:26:26 -0400 Subject: [PATCH 18/33] style(proxy): format coordination Redis validation Generated with AI Co-Authored-By: Claude Code --- litellm/proxy/proxy_server.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c21df8ce931..de10fde45c0 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -4906,7 +4906,9 @@ class ProxyConfig: if not isinstance(raw_params, dict): raise ValueError("general_settings.coordination_redis must be a mapping of Redis connection params") - coordination_params: Final = CoordinationRedisParams.model_validate(_resolve_coordination_redis_env_refs(raw_params)) + coordination_params: Final = CoordinationRedisParams.model_validate( + _resolve_coordination_redis_env_refs(raw_params) + ) if not coordination_params.has_connection_target(): raise ValueError( "general_settings.coordination_redis needs a connection target: " @@ -9234,7 +9236,9 @@ class ProxyStartupEvent: if persisted is None: return None - coordination_params: Final = CoordinationRedisParams.model_validate(_resolve_coordination_redis_env_refs(persisted)) + coordination_params: Final = CoordinationRedisParams.model_validate( + _resolve_coordination_redis_env_refs(persisted) + ) if not coordination_params.has_connection_target(): verbose_proxy_logger.warning( "coordination_redis saved in the database names no connection target; ignoring it." From a7ad6023d6a52ed035b638c4b3ab867c27a3b65b Mon Sep 17 00:00:00 2001 From: eugene-yao-zocdoc Date: Wed, 26 Aug 2026 21:14:07 -0400 Subject: [PATCH 19/33] ci: retry CodSpeed result upload Generated with AI Co-Authored-By: Claude Code From 0217a137fe3f8793163140fdb4d2f1fa4d132bbf Mon Sep 17 00:00:00 2001 From: eugene-yao-zocdoc Date: Thu, 27 Aug 2026 13:49:12 -0400 Subject: [PATCH 20/33] fix(redis): require TLS for ElastiCache IAM auth --- litellm/_redis.py | 9 ++++ tests/test_litellm/test_redis.py | 78 ++++++++++++++++++++++++++++++++ 2 files changed, 87 insertions(+) diff --git a/litellm/_redis.py b/litellm/_redis.py index 2fd8a7100e3..50900b2977b 100644 --- a/litellm/_redis.py +++ b/litellm/_redis.py @@ -279,6 +279,13 @@ def _is_true(value: object | None) -> bool: return value is True or (isinstance(value, str) and value.lower() == "true") +def _uses_tls(redis_kwargs: Mapping[str, object]) -> bool: + if redis_kwargs.get("startup_nodes") is not None: + return _is_true(redis_kwargs.get("ssl")) + url: Final = redis_kwargs.get("url") + return urlsplit(url).scheme.lower() == "rediss" if isinstance(url, str) else _is_true(redis_kwargs.get("ssl")) + + def _build_elasticache_iam_provider( user_name: object | None, cache_name: object | None, @@ -604,6 +611,8 @@ def _get_redis_client_logic(**env_overrides): "for Redis. Using Azure AD. Remove one to avoid misconfiguration." ) elif _aws_iam_enabled: + if not _uses_tls(redis_kwargs): + raise ValueError("AWS ElastiCache IAM Redis authentication requires TLS") verbose_logger.debug("Setting up AWS ElastiCache IAM authentication for Redis.") redis_kwargs["credential_provider"] = _build_elasticache_iam_provider( user_name=redis_kwargs.get("aws_iam_user_name"), diff --git a/tests/test_litellm/test_redis.py b/tests/test_litellm/test_redis.py index 43fb7563233..f2eda9a5753 100644 --- a/tests/test_litellm/test_redis.py +++ b/tests/test_litellm/test_redis.py @@ -319,6 +319,7 @@ def test_aws_iam_environment_settings_install_provider(clean_redis_environment, monkeypatch.setenv("REDIS_AWS_IAM_USER_NAME", "iam-user") monkeypatch.setenv("REDIS_AWS_IAM_CACHE_NAME", "cache.example.com") monkeypatch.setenv("REDIS_AWS_IAM_REGION", "us-east-1") + monkeypatch.setenv("REDIS_SSL", "true") redis_kwargs = _get_redis_client_logic(host="cache.example.com", port=6379) @@ -326,6 +327,77 @@ def test_aws_iam_environment_settings_install_provider(clean_redis_environment, assert not {"aws_iam_auth", "aws_iam_user_name", "aws_iam_cache_name", "aws_iam_region"} & redis_kwargs.keys() +@pytest.mark.parametrize( + "transport", + [ + pytest.param({"host": "cache.example.com", "port": 6379}, id="host_without_ssl"), + pytest.param({"host": "cache.example.com", "port": 6379, "ssl": False}, id="host_ssl_false"), + pytest.param({"url": "redis://cache.example.com:6379", "ssl": True}, id="plaintext_url"), + pytest.param( + {"startup_nodes": [{"host": "cache.example.com", "port": 6379}]}, + id="cluster_without_ssl", + ), + pytest.param( + { + "url": "rediss://cache.example.com:6379", + "startup_nodes": [{"host": "cache.example.com", "port": 6379}], + }, + id="cluster_without_ssl_ignores_url_scheme", + ), + pytest.param( + { + "sentinel_nodes": [("sentinel.example.com", 26379)], + "service_name": "cache", + }, + id="sentinel_without_ssl", + ), + ], +) +def test_aws_iam_auth_rejects_non_tls_connections(clean_redis_environment, transport): + with pytest.raises(ValueError, match="requires TLS"): + _get_redis_client_logic( + **transport, + aws_iam_auth=True, + aws_iam_user_name="iam-user", + aws_iam_cache_name="cache.example.com", + aws_iam_region="us-east-1", + ) + + +@pytest.mark.parametrize( + "transport", + [ + pytest.param({"host": "cache.example.com", "port": 6379, "ssl": True}, id="host"), + pytest.param({"url": "rediss://cache.example.com:6379"}, id="url"), + pytest.param( + { + "startup_nodes": [{"host": "cache.example.com", "port": 6379}], + "ssl": True, + }, + id="cluster", + ), + pytest.param( + { + "sentinel_nodes": [("sentinel.example.com", 26379)], + "service_name": "cache", + "ssl": True, + }, + id="sentinel", + ), + ], +) +def test_aws_iam_auth_accepts_tls_connections(clean_redis_environment, transport): + redis_kwargs = _get_redis_client_logic( + **transport, + aws_iam_auth=True, + aws_iam_user_name="iam-user", + aws_iam_cache_name="cache.example.com", + aws_iam_region="us-east-1", + ) + + assert isinstance(redis_kwargs["credential_provider"], ElastiCacheIAMCredentialProvider) + + def test_aws_iam_settings_are_removed_for_url_and_static_credentials(clean_redis_environment): redis_kwargs = _get_redis_client_logic( url="rediss://url-user:url-pass@cache.example.com:6380", @@ -351,6 +423,7 @@ def test_aws_iam_missing_setting_fails_closed(clean_redis_environment, missing): "aws_iam_user_name": "iam-user", "aws_iam_cache_name": "cache.example.com", "aws_iam_region": "us-east-1", + "ssl": True, } settings[missing] = None @@ -365,6 +438,7 @@ def test_aws_iam_region_falls_back_to_environment(clean_redis_environment, monke redis_kwargs = _get_redis_client_logic( host="cache.example.com", port=6379, + ssl=True, aws_iam_auth=True, aws_iam_user_name="iam-user", aws_iam_cache_name="cache.example.com", @@ -382,6 +456,7 @@ def test_aws_iam_region_prefers_aws_region_over_default_region(clean_redis_envir redis_kwargs = _get_redis_client_logic( host="cache.example.com", port=6379, + ssl=True, aws_iam_auth=True, aws_iam_user_name="iam-user", aws_iam_cache_name="cache.example.com", @@ -396,6 +471,7 @@ def test_aws_iam_region_prefers_explicit_over_environment(clean_redis_environmen redis_kwargs = _get_redis_client_logic( host="cache.example.com", port=6379, + ssl=True, aws_iam_auth=True, aws_iam_user_name="iam-user", aws_iam_cache_name="cache.example.com", @@ -409,6 +485,7 @@ def test_aws_iam_settings_map_to_distinct_provider_fields(clean_redis_environmen redis_kwargs = _get_redis_client_logic( host="cache.example.com", port=6379, + ssl=True, aws_iam_auth=True, aws_iam_user_name="iam-user-value", aws_iam_cache_name="iam-cache-value", @@ -489,6 +566,7 @@ def test_async_cluster_installs_aws_iam_provider(clean_redis_environment): client = get_redis_async_client( startup_nodes=startup_nodes, + ssl=True, aws_iam_auth=True, aws_iam_user_name="iam-user", aws_iam_cache_name="cache.example.com", From f1cd2b03caac2c544866be984089039aa1228956 Mon Sep 17 00:00:00 2001 From: eugene-yao-zocdoc Date: Mon, 31 Aug 2026 13:49:55 -0400 Subject: [PATCH 21/33] fix(redis): type signed ElastiCache IAM URL --- litellm/_redis_credential_provider.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/_redis_credential_provider.py b/litellm/_redis_credential_provider.py index 85759af2a88..4d728ef21b6 100644 --- a/litellm/_redis_credential_provider.py +++ b/litellm/_redis_credential_provider.py @@ -4,7 +4,7 @@ import asyncio import threading import time from collections.abc import Callable -from typing import TYPE_CHECKING, Any, Final, Protocol +from typing import TYPE_CHECKING, Any, Final, Protocol, cast from urllib.parse import quote from redis.credentials import CredentialProvider @@ -185,7 +185,7 @@ class ElastiCacheIAMCredentialProvider(CredentialProvider): self._region, expires=self._token_lifetime_seconds, ).add_auth(request) - return self._user_name, request.url.removeprefix("https://") + return self._user_name, cast(str, request.url).removeprefix("https://") def get_credentials(self) -> tuple[str, str]: return self._get_credentials() From aca81c263ccd6660bffcf05115e4971eb92412c8 Mon Sep 17 00:00:00 2001 From: eugene-yao-zocdoc Date: Mon, 31 Aug 2026 14:07:04 -0400 Subject: [PATCH 22/33] fix(redis): validate signed ElastiCache IAM URL --- litellm/_redis_credential_provider.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/litellm/_redis_credential_provider.py b/litellm/_redis_credential_provider.py index 4d728ef21b6..4441d318373 100644 --- a/litellm/_redis_credential_provider.py +++ b/litellm/_redis_credential_provider.py @@ -4,7 +4,7 @@ import asyncio import threading import time from collections.abc import Callable -from typing import TYPE_CHECKING, Any, Final, Protocol, cast +from typing import TYPE_CHECKING, Any, Final, Protocol from urllib.parse import quote from redis.credentials import CredentialProvider @@ -185,7 +185,10 @@ class ElastiCacheIAMCredentialProvider(CredentialProvider): self._region, expires=self._token_lifetime_seconds, ).add_auth(request) - return self._user_name, cast(str, request.url).removeprefix("https://") + signed_url: Final = request.url + if signed_url is None: + raise RuntimeError("Unable to generate AWS ElastiCache IAM credentials") + return self._user_name, signed_url.removeprefix("https://") def get_credentials(self) -> tuple[str, str]: return self._get_credentials() From cd4073895ff9884d6fa77337fe6ac4c62dc029ec Mon Sep 17 00:00:00 2001 From: eugene-yao-zocdoc Date: Wed, 9 Sep 2026 10:03:26 -0400 Subject: [PATCH 23/33] refactor(redis): type botocore credentials without Any Deferred annotation evaluation keeps the type-checking-only botocore import off the runtime path, so the alias only reintroduced typing.Any, which the strict ruff budget now bans --- litellm/_redis_credential_provider.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/litellm/_redis_credential_provider.py b/litellm/_redis_credential_provider.py index 4441d318373..15f625dc8b4 100644 --- a/litellm/_redis_credential_provider.py +++ b/litellm/_redis_credential_provider.py @@ -4,15 +4,13 @@ import asyncio import threading import time from collections.abc import Callable -from typing import TYPE_CHECKING, Any, Final, Protocol +from typing import TYPE_CHECKING, Final, Protocol from urllib.parse import quote from redis.credentials import CredentialProvider if TYPE_CHECKING: from botocore.credentials import Credentials -else: - Credentials = Any # rebind-ok: runtime alias for the type-checking-only botocore import # Azure AD scope for Redis Cache for Azure. AZURE_REDIS_SCOPE: Final = "https://redis.azure.com/.default" From 2b0711e4f518dc796395972896aa5c3e6608b965 Mon Sep 17 00:00:00 2001 From: eugene-yao-zocdoc Date: Wed, 9 Sep 2026 10:06:22 -0400 Subject: [PATCH 24/33] style(redis): restore the Azure AD marker comment The comment documents that the raw Azure client id, tenant id and secret are deliberately kept off the connect function, so this branch should never have dropped it --- litellm/_redis.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/litellm/_redis.py b/litellm/_redis.py index 50900b2977b..d77043a6dfd 100644 --- a/litellm/_redis.py +++ b/litellm/_redis.py @@ -598,6 +598,11 @@ def _get_redis_client_logic(**env_overrides): azure_tenant_id=_azure_tenant_id, azure_client_secret=_azure_client_secret, ) + # Marker for async paths to detect Azure AD auth. The live credential + # object is attached separately as `_azure_credential` by + # `create_azure_ad_redis_connect_func`; the raw client_id/tenant_id/secret + # are intentionally NOT exposed on the function to avoid leaking + # credentials via inspection or logging. redis_kwargs["redis_connect_func"]._azure_redis_ad_token = True if _aws_iam_enabled and _gcp_service_account is not None: From bf73c49d737b0ee1d4274a36a1846a8628de3478 Mon Sep 17 00:00:00 2001 From: eugene-yao-zocdoc Date: Wed, 9 Sep 2026 10:17:53 -0400 Subject: [PATCH 25/33] refactor(redis): drop unreachable frozen credentials guard --- litellm/_redis_credential_provider.py | 2 -- tests/test_litellm/test_redis_credential_provider.py | 12 ------------ 2 files changed, 14 deletions(-) diff --git a/litellm/_redis_credential_provider.py b/litellm/_redis_credential_provider.py index 15f625dc8b4..dd49a6793ec 100644 --- a/litellm/_redis_credential_provider.py +++ b/litellm/_redis_credential_provider.py @@ -162,8 +162,6 @@ class ElastiCacheIAMCredentialProvider(CredentialProvider): self._credentials = credentials frozen_credentials: Final = credentials.get_frozen_credentials() - if frozen_credentials is None: - raise RuntimeError("Unable to resolve AWS credentials for ElastiCache IAM Redis authentication") try: from botocore.auth import SigV4QueryAuth diff --git a/tests/test_litellm/test_redis_credential_provider.py b/tests/test_litellm/test_redis_credential_provider.py index 336e30e12c6..9a58722b04e 100644 --- a/tests/test_litellm/test_redis_credential_provider.py +++ b/tests/test_litellm/test_redis_credential_provider.py @@ -136,18 +136,6 @@ def test_elasticache_provider_reports_missing_credentials(): provider.get_credentials() -def test_elasticache_provider_reports_missing_frozen_credentials(): - provider = ElastiCacheIAMCredentialProvider( - user_name="iam-user", - cache_name="cache.example.com", - region="us-east-1", - credentials_resolver=_FakeResolver(SimpleNamespace(get_frozen_credentials=lambda: None)), - ) - - with pytest.raises(RuntimeError, match="Unable to resolve AWS credentials"): - provider.get_credentials() - - def test_elasticache_provider_reports_missing_signing_dependency(monkeypatch): original_import = builtins.__import__ From c4fc20bcf933606d0ffc91a47036dc16df7059ea Mon Sep 17 00:00:00 2001 From: eugene-yao-zocdoc Date: Wed, 9 Sep 2026 14:43:58 -0400 Subject: [PATCH 26/33] fix(redis): accept every truthy flag and sign serverless ElastiCache caches The ElastiCache IAM gate read `aws_iam_auth` and `ssl` with a helper that only accepted the literal string "true", while the kwarg coercion that runs later accepts "true", "1" and "yes". Type coercion happens after the gate, so `REDIS_AWS_IAM_AUTH=1` silently skipped IAM auth and `REDIS_SSL=1` made the "requires TLS" check fail closed on a connection that was in fact TLS. Both helpers now share `_str_to_bool`. AWS signs serverless cache tokens with an extra `ResourceType=ServerlessCache` query parameter, so tokens minted for a serverless cache were rejected. Adds an `aws_iam_serverless` setting (`REDIS_AWS_IAM_SERVERLESS`) that puts the parameter into the signed URL, and lowercases the cache name because ElastiCache lowercases it at creation time. --- litellm/_redis.py | 51 ++++++------ litellm/_redis_credential_provider.py | 17 ++-- litellm/proxy/_types.py | 3 + .../cache_settings_endpoints.py | 9 ++ .../coordination_redis_endpoints.py | 8 ++ tests/test_litellm/test_redis.py | 82 +++++++++++++++++-- .../test_redis_credential_provider.py | 54 ++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 ++ 8 files changed, 192 insertions(+), 37 deletions(-) diff --git a/litellm/_redis.py b/litellm/_redis.py index d77043a6dfd..c5acdcb038b 100644 --- a/litellm/_redis.py +++ b/litellm/_redis.py @@ -39,6 +39,14 @@ from ._logging import verbose_logger AZURE_REDIS_SCOPE: Final = "https://redis.azure.com/.default" +_AWS_IAM_KWARG_NAMES: Final = ( + "aws_iam_auth", + "aws_iam_user_name", + "aws_iam_cache_name", + "aws_iam_region", + "aws_iam_serverless", +) + def _unwrapped_init_args(cls: type) -> frozenset[str]: """Every parameter on a single class's own ``__init__``, decorator-unwrapped. @@ -76,10 +84,7 @@ def _get_redis_kwargs(): "azure_client_id", "azure_tenant_id", "azure_client_secret", - "aws_iam_auth", - "aws_iam_user_name", - "aws_iam_cache_name", - "aws_iam_region", + *_AWS_IAM_KWARG_NAMES, } available_args: Final = {x for x in _unwrapped_init_args(redis.Redis) if x not in exclude_args} | include_args @@ -275,22 +280,25 @@ def _redis_kwargs_from_environment(): return return_dict -def _is_true(value: object | None) -> bool: - return value is True or (isinstance(value, str) and value.lower() == "true") +def _coerces_to_true(value: object | None) -> bool: + return _str_to_bool(value) if isinstance(value, str) else bool(value) def _uses_tls(redis_kwargs: Mapping[str, object]) -> bool: if redis_kwargs.get("startup_nodes") is not None: - return _is_true(redis_kwargs.get("ssl")) + return _coerces_to_true(redis_kwargs.get("ssl")) url: Final = redis_kwargs.get("url") - return urlsplit(url).scheme.lower() == "rediss" if isinstance(url, str) else _is_true(redis_kwargs.get("ssl")) + if isinstance(url, str): + return urlsplit(url).scheme.lower() == "rediss" + return _coerces_to_true(redis_kwargs.get("ssl")) -def _build_elasticache_iam_provider( - user_name: object | None, - cache_name: object | None, - region: object | None, -) -> ElastiCacheIAMCredentialProvider: +def _build_elasticache_iam_provider(redis_kwargs: Mapping[str, object]) -> ElastiCacheIAMCredentialProvider: + user_name: Final = redis_kwargs.get("aws_iam_user_name") + cache_name: Final = redis_kwargs.get("aws_iam_cache_name") + region: Final = ( + redis_kwargs.get("aws_iam_region") or get_secret_str("AWS_REGION") or get_secret_str("AWS_DEFAULT_REGION") + ) required_settings: Final = ( ("aws_iam_user_name", user_name), ("aws_iam_cache_name", cache_name), @@ -304,6 +312,7 @@ def _build_elasticache_iam_provider( user_name=str(user_name), cache_name=str(cache_name), region=str(region), + is_serverless=_coerces_to_true(redis_kwargs.get("aws_iam_serverless")), ) @@ -577,7 +586,7 @@ def _get_redis_client_logic(**env_overrides): _azure_redis_ad_token: Final = redis_kwargs.get("azure_redis_ad_token") or get_secret("REDIS_AZURE_AD_TOKEN") _azure_ad_enabled: Final = _azure_redis_ad_token is not None and str(_azure_redis_ad_token).lower() == "true" - _aws_iam_enabled: Final = _is_true(redis_kwargs.get("aws_iam_auth")) + _aws_iam_enabled: Final = _coerces_to_true(redis_kwargs.get("aws_iam_auth")) if _azure_ad_enabled and _gcp_service_account is not None: verbose_logger.warning( @@ -619,13 +628,7 @@ def _get_redis_client_logic(**env_overrides): if not _uses_tls(redis_kwargs): raise ValueError("AWS ElastiCache IAM Redis authentication requires TLS") verbose_logger.debug("Setting up AWS ElastiCache IAM authentication for Redis.") - redis_kwargs["credential_provider"] = _build_elasticache_iam_provider( - user_name=redis_kwargs.get("aws_iam_user_name"), - cache_name=redis_kwargs.get("aws_iam_cache_name"), - region=redis_kwargs.get("aws_iam_region") - or get_secret_str("AWS_REGION") - or get_secret_str("AWS_DEFAULT_REGION"), - ) + redis_kwargs["credential_provider"] = _build_elasticache_iam_provider(redis_kwargs) redis_kwargs.pop("gcp_service_account", None) redis_kwargs.pop("gcp_ssl_ca_certs", None) @@ -635,10 +638,8 @@ def _get_redis_client_logic(**env_overrides): redis_kwargs.pop("azure_client_id", None) redis_kwargs.pop("azure_tenant_id", None) redis_kwargs.pop("azure_client_secret", None) - redis_kwargs.pop("aws_iam_auth", None) - redis_kwargs.pop("aws_iam_user_name", None) - redis_kwargs.pop("aws_iam_cache_name", None) - redis_kwargs.pop("aws_iam_region", None) + for aws_iam_key in _AWS_IAM_KWARG_NAMES: + redis_kwargs.pop(aws_iam_key, None) if redis_kwargs.get("credential_provider") is not None: redis_kwargs.pop("redis_connect_func", None) diff --git a/litellm/_redis_credential_provider.py b/litellm/_redis_credential_provider.py index dd49a6793ec..7d90f944657 100644 --- a/litellm/_redis_credential_provider.py +++ b/litellm/_redis_credential_provider.py @@ -5,7 +5,7 @@ import threading import time from collections.abc import Callable from typing import TYPE_CHECKING, Final, Protocol -from urllib.parse import quote +from urllib.parse import urlencode from redis.credentials import CredentialProvider @@ -126,6 +126,7 @@ class GCPIAMCredentialProvider(CredentialProvider): _ELASTICACHE_SERVICE_NAME: Final = "elasticache" _ELASTICACHE_TOKEN_TTL_SECONDS: Final = 900 +_ELASTICACHE_SERVERLESS_RESOURCE_TYPE: Final = "ServerlessCache" class ElastiCacheIAMCredentialProvider(CredentialProvider): @@ -134,12 +135,14 @@ class ElastiCacheIAMCredentialProvider(CredentialProvider): user_name: str, cache_name: str, region: str, + is_serverless: bool = False, credentials_resolver: Callable[[], Credentials | None] | None = None, token_lifetime_seconds: int = _ELASTICACHE_TOKEN_TTL_SECONDS, ) -> None: self._user_name = user_name - self._cache_name = cache_name + self._cache_name = cache_name.lower() self._region = region + self._is_serverless = is_serverless self._credentials_resolver = credentials_resolver or self._resolve_credentials self._credentials: Credentials | None = None self._token_lifetime_seconds = token_lifetime_seconds @@ -171,10 +174,14 @@ class ElastiCacheIAMCredentialProvider(CredentialProvider): "botocore is required for ElastiCache IAM Redis authentication. Install it with: pip install boto3" ) from e - request: Final = AWSRequest( - method="GET", - url=(f"https://{self._cache_name}/?Action=connect&User={quote(self._user_name, safe='')}"), + query: Final = urlencode( + ( + ("Action", "connect"), + ("User", self._user_name), + *((("ResourceType", _ELASTICACHE_SERVERLESS_RESOURCE_TYPE),) if self._is_serverless else ()), + ) ) + request: Final = AWSRequest(method="GET", url=f"https://{self._cache_name}/?{query}") SigV4QueryAuth( frozen_credentials, _ELASTICACHE_SERVICE_NAME, diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 99103ff3706..07e4515059a 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2431,6 +2431,9 @@ class CoordinationRedisParams(LiteLLMPydanticObjectBase): aws_iam_user_name: str | None = Field(None, description="AWS ElastiCache IAM user name") aws_iam_cache_name: str | None = Field(None, description="AWS ElastiCache cache name") aws_iam_region: str | None = Field(None, description="AWS region for ElastiCache IAM authentication") + aws_iam_serverless: bool | str | None = Field( + None, description="the ElastiCache cache is serverless rather than a self-designed cluster" + ) def has_connection_target(self) -> bool: return any(value is not None for value in (self.host, self.url, self.startup_nodes, self.sentinel_nodes)) diff --git a/litellm/types/management_endpoints/cache_settings_endpoints.py b/litellm/types/management_endpoints/cache_settings_endpoints.py index 10aece6efa4..cb05f3a50ac 100644 --- a/litellm/types/management_endpoints/cache_settings_endpoints.py +++ b/litellm/types/management_endpoints/cache_settings_endpoints.py @@ -273,4 +273,13 @@ CACHE_SETTINGS_FIELDS: Final[list[CacheSettingsField]] = [ ui_field_name="AWS IAM Region", redis_type=None, ), + CacheSettingsField( + field_name="aws_iam_serverless", + field_type="Boolean", + field_value=None, + field_description="The ElastiCache cache is serverless rather than a self-designed cluster", + field_default=False, + ui_field_name="AWS IAM Serverless Cache", + redis_type=None, + ), ] diff --git a/litellm/types/management_endpoints/coordination_redis_endpoints.py b/litellm/types/management_endpoints/coordination_redis_endpoints.py index 86b32752cbb..d70a921a22f 100644 --- a/litellm/types/management_endpoints/coordination_redis_endpoints.py +++ b/litellm/types/management_endpoints/coordination_redis_endpoints.py @@ -131,4 +131,12 @@ COORDINATION_REDIS_SETTINGS_FIELDS: Final[list[CoordinationRedisSettingsField]] ui_field_name="AWS IAM Region", section="connection", ), + CoordinationRedisSettingsField( + field_name="aws_iam_serverless", + field_type="Boolean", + field_description="The ElastiCache cache is serverless rather than a self-designed cluster", + field_default=False, + ui_field_name="AWS IAM Serverless Cache", + section="connection", + ), ] diff --git a/tests/test_litellm/test_redis.py b/tests/test_litellm/test_redis.py index f2eda9a5753..5337ed825a7 100644 --- a/tests/test_litellm/test_redis.py +++ b/tests/test_litellm/test_redis.py @@ -113,15 +113,25 @@ def test_credential_provider_is_not_environment_derived(): assert "credential_provider" not in mapping.values() +_AWS_IAM_SETTINGS = { + "aws_iam_auth", + "aws_iam_user_name", + "aws_iam_cache_name", + "aws_iam_region", + "aws_iam_serverless", +} + + def test_aws_iam_settings_are_environment_derived(): allowed = _get_redis_kwargs() mapping = _get_redis_env_kwarg_mapping() - assert {"aws_iam_auth", "aws_iam_user_name", "aws_iam_cache_name", "aws_iam_region"} <= allowed + assert _AWS_IAM_SETTINGS <= allowed assert mapping["REDIS_AWS_IAM_AUTH"] == "aws_iam_auth" assert mapping["REDIS_AWS_IAM_USER_NAME"] == "aws_iam_user_name" assert mapping["REDIS_AWS_IAM_CACHE_NAME"] == "aws_iam_cache_name" assert mapping["REDIS_AWS_IAM_REGION"] == "aws_iam_region" + assert mapping["REDIS_AWS_IAM_SERVERLESS"] == "aws_iam_serverless" def test_sync_direct_preserves_credential_provider_identity(clean_redis_environment): @@ -319,12 +329,15 @@ def test_aws_iam_environment_settings_install_provider(clean_redis_environment, monkeypatch.setenv("REDIS_AWS_IAM_USER_NAME", "iam-user") monkeypatch.setenv("REDIS_AWS_IAM_CACHE_NAME", "cache.example.com") monkeypatch.setenv("REDIS_AWS_IAM_REGION", "us-east-1") - monkeypatch.setenv("REDIS_SSL", "true") + monkeypatch.setenv("REDIS_AWS_IAM_SERVERLESS", "1") + monkeypatch.setenv("REDIS_SSL", "1") redis_kwargs = _get_redis_client_logic(host="cache.example.com", port=6379) - assert isinstance(redis_kwargs["credential_provider"], ElastiCacheIAMCredentialProvider) - assert not {"aws_iam_auth", "aws_iam_user_name", "aws_iam_cache_name", "aws_iam_region"} & redis_kwargs.keys() + provider = redis_kwargs["credential_provider"] + assert isinstance(provider, ElastiCacheIAMCredentialProvider) + assert provider._is_serverless is True + assert not _AWS_IAM_SETTINGS & redis_kwargs.keys() @pytest.mark.parametrize( @@ -332,6 +345,9 @@ def test_aws_iam_environment_settings_install_provider(clean_redis_environment, [ pytest.param({"host": "cache.example.com", "port": 6379}, id="host_without_ssl"), pytest.param({"host": "cache.example.com", "port": 6379, "ssl": False}, id="host_ssl_false"), + pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "false"}, id="host_ssl_false_string"), + pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "0"}, id="host_ssl_zero_string"), + pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "no"}, id="host_ssl_no_string"), pytest.param({"url": "redis://cache.example.com:6379", "ssl": True}, id="plaintext_url"), pytest.param( {"startup_nodes": [{"host": "cache.example.com", "port": 6379}]}, @@ -368,6 +384,13 @@ def test_aws_iam_auth_rejects_non_tls_connections(clean_redis_environment, trans "transport", [ pytest.param({"host": "cache.example.com", "port": 6379, "ssl": True}, id="host"), + pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "true"}, id="host_ssl_true_string"), + pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "1"}, id="host_ssl_one_string"), + pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "yes"}, id="host_ssl_yes_string"), + pytest.param( + {"startup_nodes": [{"host": "cache.example.com", "port": 6379}], "ssl": "1"}, + id="cluster_ssl_one_string", + ), pytest.param({"url": "rediss://cache.example.com:6379"}, id="url"), pytest.param( { @@ -413,7 +436,7 @@ def test_aws_iam_settings_are_removed_for_url_and_static_credentials(clean_redis assert redis_kwargs["url"] == "rediss://cache.example.com:6380" assert "username" not in redis_kwargs assert "password" not in redis_kwargs - assert not {"aws_iam_auth", "aws_iam_user_name", "aws_iam_cache_name", "aws_iam_region"} & redis_kwargs.keys() + assert not _AWS_IAM_SETTINGS & redis_kwargs.keys() @pytest.mark.parametrize("missing", ["aws_iam_user_name", "aws_iam_cache_name", "aws_iam_region"]) @@ -499,7 +522,7 @@ def test_aws_iam_settings_map_to_distinct_provider_fields(clean_redis_environmen assert provider._region == "iam-region-value" -@pytest.mark.parametrize("aws_iam_auth", [False, "false"]) +@pytest.mark.parametrize("aws_iam_auth", [None, False, "", "false", "0", "no"]) def test_aws_iam_auth_disabled_does_not_install_provider(clean_redis_environment, aws_iam_auth): redis_kwargs = _get_redis_client_logic( host="cache.example.com", @@ -511,7 +534,52 @@ def test_aws_iam_auth_disabled_does_not_install_provider(clean_redis_environment ) assert "credential_provider" not in redis_kwargs - assert not {"aws_iam_auth", "aws_iam_user_name", "aws_iam_cache_name", "aws_iam_region"} & redis_kwargs.keys() + assert not _AWS_IAM_SETTINGS & redis_kwargs.keys() + + +@pytest.mark.parametrize("aws_iam_auth", [True, "true", "True", "TRUE", "1", "yes"]) +def test_aws_iam_auth_enabled_by_any_truthy_flag(clean_redis_environment, aws_iam_auth): + redis_kwargs = _get_redis_client_logic( + host="cache.example.com", + port=6379, + ssl=True, + aws_iam_auth=aws_iam_auth, + aws_iam_user_name="iam-user", + aws_iam_cache_name="cache-name", + aws_iam_region="us-east-1", + ) + + assert isinstance(redis_kwargs["credential_provider"], ElastiCacheIAMCredentialProvider) + + +@pytest.mark.parametrize( + "aws_iam_serverless, expected", + [ + pytest.param(None, False, id="unset"), + pytest.param(False, False, id="bool_false"), + pytest.param("false", False, id="string_false"), + pytest.param("0", False, id="string_zero"), + pytest.param(True, True, id="bool_true"), + pytest.param("true", True, id="string_true"), + pytest.param("1", True, id="string_one"), + ], +) +def test_aws_iam_serverless_flag_reaches_the_provider(clean_redis_environment, aws_iam_serverless, expected): + redis_kwargs = _get_redis_client_logic( + host="cache.example.com", + port=6379, + ssl=True, + aws_iam_auth=True, + aws_iam_user_name="iam-user", + aws_iam_cache_name="cache-name", + aws_iam_region="us-east-1", + aws_iam_serverless=aws_iam_serverless, + ) + + provider = redis_kwargs["credential_provider"] + assert isinstance(provider, ElastiCacheIAMCredentialProvider) + assert provider._is_serverless is expected + assert "aws_iam_serverless" not in redis_kwargs def test_explicit_provider_wins_over_aws_iam(clean_redis_environment): diff --git a/tests/test_litellm/test_redis_credential_provider.py b/tests/test_litellm/test_redis_credential_provider.py index 9a58722b04e..e71fc5d530d 100644 --- a/tests/test_litellm/test_redis_credential_provider.py +++ b/tests/test_litellm/test_redis_credential_provider.py @@ -175,3 +175,57 @@ def test_elasticache_provider_recovers_after_a_failed_resolution(): assert user_name == "iam-user" assert token assert resolver.calls == 2 + + +@pytest.mark.parametrize( + "provider_kwargs, expected_resource_type", + [ + pytest.param({}, None, id="default_is_self_designed"), + pytest.param({"is_serverless": False}, None, id="self_designed"), + pytest.param({"is_serverless": True}, ["ServerlessCache"], id="serverless"), + ], +) +def test_elasticache_provider_signs_resource_type_only_for_serverless(provider_kwargs, expected_resource_type): + provider = ElastiCacheIAMCredentialProvider( + user_name="iam-user", + cache_name="cache-name", + region="us-east-1", + credentials_resolver=_FakeResolver(_FakeCredentials("AKIA-SYNTHETIC")), + **provider_kwargs, + ) + + _, token = provider.get_credentials() + query = parse_qs(urlsplit("https://" + token).query) + + assert query.get("ResourceType") == expected_resource_type + assert query["X-Amz-Signature"] + + +def test_elasticache_provider_lowercases_the_cache_name(): + provider = ElastiCacheIAMCredentialProvider( + user_name="iam-user", + cache_name="Mixed-Case-Cache", + region="us-east-1", + credentials_resolver=_FakeResolver(_FakeCredentials("AKIA-SYNTHETIC")), + ) + + _, token = provider.get_credentials() + + assert urlsplit("https://" + token).netloc == "mixed-case-cache" + + +def test_elasticache_provider_encodes_reserved_characters_in_the_user_name(): + user_name = "iam user/with+reserved&chars" + provider = ElastiCacheIAMCredentialProvider( + user_name=user_name, + cache_name="cache-name", + region="us-east-1", + credentials_resolver=_FakeResolver(_FakeCredentials("AKIA-SYNTHETIC")), + ) + + returned_user_name, token = provider.get_credentials() + query = parse_qs(urlsplit("https://" + token).query) + + assert returned_user_name == user_name + assert query["User"] == [user_name] + assert query["Action"] == ["connect"] diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 9ca0df396b7..0448f8373c1 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -26372,6 +26372,11 @@ export interface components { * @description AWS region for ElastiCache IAM authentication */ aws_iam_region?: string | null; + /** + * Aws Iam Serverless + * @description the ElastiCache cache is serverless rather than a self-designed cluster + */ + aws_iam_serverless?: boolean | string | null; /** * Aws Iam User Name * @description AWS ElastiCache IAM user name From f0e3e031c3fd79acf04db3bca41c4f2f5d1f12ec Mon Sep 17 00:00:00 2001 From: eugene-yao-zocdoc Date: Wed, 9 Sep 2026 15:01:06 -0400 Subject: [PATCH 27/33] test(redis): pin ElastiCache IAM signing and TLS coercion invariants Strengthens the serverless test to assert ResourceType is signed rather than merely present, ties _uses_tls to the redis-py kwarg coercion so the two cannot drift, locks the stripped-kwarg name tuple to the test's expectations, and adds "off" and "True" sentinel flag values. Renames the provider builder's parameter to redis_settings. --- tests/test_litellm/test_redis.py | 15 ++++++++++++++- .../test_redis_credential_provider.py | 19 ++++++++++++------- 2 files changed, 26 insertions(+), 8 deletions(-) diff --git a/tests/test_litellm/test_redis.py b/tests/test_litellm/test_redis.py index 5337ed825a7..0e8c86d26df 100644 --- a/tests/test_litellm/test_redis.py +++ b/tests/test_litellm/test_redis.py @@ -10,13 +10,16 @@ from redis.credentials import CredentialProvider import litellm from litellm._redis import ( + _AWS_IAM_KWARG_NAMES, _async_auth_kwargs, + _coerce_redis_kwargs_types, _get_redis_client_logic, _get_redis_cluster_kwargs, _get_redis_env_kwarg_mapping, _get_redis_kwargs, _get_redis_url_kwargs, _pretty_print_redis_config, + _uses_tls, get_redis_async_client, get_redis_client, get_redis_connection_pool, @@ -31,6 +34,7 @@ from litellm._redis_credential_provider import ( from litellm.caching.redis_cache import RedisCache from litellm.caching.redis_cluster_cache import RedisClusterCache from litellm.constants import REDIS_CLUSTER_HEALTH_CHECK_INTERVAL +from litellm.proxy._types import CoordinationRedisParams class _StubCredentialProvider(CredentialProvider): @@ -127,6 +131,8 @@ def test_aws_iam_settings_are_environment_derived(): mapping = _get_redis_env_kwarg_mapping() assert _AWS_IAM_SETTINGS <= allowed + assert set(_AWS_IAM_KWARG_NAMES) == _AWS_IAM_SETTINGS + assert {f for f in CoordinationRedisParams.model_fields if f.startswith("aws_iam_")} == _AWS_IAM_SETTINGS assert mapping["REDIS_AWS_IAM_AUTH"] == "aws_iam_auth" assert mapping["REDIS_AWS_IAM_USER_NAME"] == "aws_iam_user_name" assert mapping["REDIS_AWS_IAM_CACHE_NAME"] == "aws_iam_cache_name" @@ -348,6 +354,7 @@ def test_aws_iam_environment_settings_install_provider(clean_redis_environment, pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "false"}, id="host_ssl_false_string"), pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "0"}, id="host_ssl_zero_string"), pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "no"}, id="host_ssl_no_string"), + pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "off"}, id="host_ssl_off_string"), pytest.param({"url": "redis://cache.example.com:6379", "ssl": True}, id="plaintext_url"), pytest.param( {"startup_nodes": [{"host": "cache.example.com", "port": 6379}]}, @@ -385,6 +392,7 @@ def test_aws_iam_auth_rejects_non_tls_connections(clean_redis_environment, trans [ pytest.param({"host": "cache.example.com", "port": 6379, "ssl": True}, id="host"), pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "true"}, id="host_ssl_true_string"), + pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "True"}, id="host_ssl_true_capitalized"), pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "1"}, id="host_ssl_one_string"), pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "yes"}, id="host_ssl_yes_string"), pytest.param( @@ -421,6 +429,11 @@ def test_aws_iam_auth_accepts_tls_connections(clean_redis_environment, transport assert isinstance(redis_kwargs["credential_provider"], ElastiCacheIAMCredentialProvider) +@pytest.mark.parametrize("ssl", ["true", "True", "TRUE", "1", "yes", "YES", "false", "0", "no", "off", "", "maybe"]) +def test_tls_detection_agrees_with_the_ssl_kwarg_coercion(ssl): + assert _uses_tls({"ssl": ssl}) is _coerce_redis_kwargs_types({"ssl": ssl})["ssl"] + + def test_aws_iam_settings_are_removed_for_url_and_static_credentials(clean_redis_environment): redis_kwargs = _get_redis_client_logic( url="rediss://url-user:url-pass@cache.example.com:6380", @@ -522,7 +535,7 @@ def test_aws_iam_settings_map_to_distinct_provider_fields(clean_redis_environmen assert provider._region == "iam-region-value" -@pytest.mark.parametrize("aws_iam_auth", [None, False, "", "false", "0", "no"]) +@pytest.mark.parametrize("aws_iam_auth", [None, False, "", "false", "0", "no", "off"]) def test_aws_iam_auth_disabled_does_not_install_provider(clean_redis_environment, aws_iam_auth): redis_kwargs = _get_redis_client_logic( host="cache.example.com", diff --git a/tests/test_litellm/test_redis_credential_provider.py b/tests/test_litellm/test_redis_credential_provider.py index e71fc5d530d..96b1933b0e3 100644 --- a/tests/test_litellm/test_redis_credential_provider.py +++ b/tests/test_litellm/test_redis_credential_provider.py @@ -178,14 +178,14 @@ def test_elasticache_provider_recovers_after_a_failed_resolution(): @pytest.mark.parametrize( - "provider_kwargs, expected_resource_type", + "provider_kwargs, expected_operation_params", [ - pytest.param({}, None, id="default_is_self_designed"), - pytest.param({"is_serverless": False}, None, id="self_designed"), - pytest.param({"is_serverless": True}, ["ServerlessCache"], id="serverless"), + pytest.param({}, frozenset({"Action", "User"}), id="default_is_self_designed"), + pytest.param({"is_serverless": False}, frozenset({"Action", "User"}), id="self_designed"), + pytest.param({"is_serverless": True}, frozenset({"Action", "User", "ResourceType"}), id="serverless"), ], ) -def test_elasticache_provider_signs_resource_type_only_for_serverless(provider_kwargs, expected_resource_type): +def test_elasticache_provider_signs_resource_type_only_for_serverless(provider_kwargs, expected_operation_params): provider = ElastiCacheIAMCredentialProvider( user_name="iam-user", cache_name="cache-name", @@ -195,9 +195,14 @@ def test_elasticache_provider_signs_resource_type_only_for_serverless(provider_k ) _, token = provider.get_credentials() - query = parse_qs(urlsplit("https://" + token).query) + query_string = urlsplit("https://" + token).query + param_names = tuple(pair.split("=", 1)[0] for pair in query_string.split("&")) + first_auth_param = next(i for i, name in enumerate(param_names) if name.startswith("X-Amz-")) + query = parse_qs(query_string) - assert query.get("ResourceType") == expected_resource_type + assert frozenset(param_names[:first_auth_param]) == expected_operation_params + assert all(name.startswith("X-Amz-") for name in param_names[first_auth_param:]) + assert query.get("ResourceType") == (["ServerlessCache"] if "ResourceType" in expected_operation_params else None) assert query["X-Amz-Signature"] From 55ab5ee53d48ef165886de31144b1e172f8208e1 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Thu, 10 Sep 2026 13:08:46 -0700 Subject: [PATCH 28/33] fix(mcp): preserve identity checks across cached OAuth credentials --- AGENTS.md | 2 + litellm/proxy/_experimental/mcp_server/db.py | 34 +++++--- .../mcp_server/oauth2_token_cache.py | 25 ++++-- .../mcp_server/oauth_identity_binding.py | 2 +- .../authz_code_refresher.py | 11 +++ .../outbound_credentials/oauth_token_store.py | 1 + .../per_user_oauth_store.py | 33 ++++++-- .../outbound_credentials/token_cache_codec.py | 36 ++++++-- .../outbound_credentials/v2_token_store.py | 2 + .../test_authz_code_refresher.py | 18 ++-- .../test_per_user_oauth_store.py | 82 +++++++++++++++++-- .../test_token_cache_codec.py | 14 ++++ .../mcp_server/test_db_credentials.py | 75 +++++++++++++++++ .../mcp_server/test_oauth_identity_binding.py | 30 +++++-- 14 files changed, 311 insertions(+), 54 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 41921fdff4d..a1e8f6f618d 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1 +1,3 @@ Read @CLAUDE.md for coding guidelines + +Before requesting maintainer review, verify the current PR tip passes required CI and code coverage, meets Greptile confidence of at least 4/5, and has acceptable Veria and Bugbot reviews. Inspect warnings and findings, fix actionable issues, and rerun the affected checks and reviewers after changes. Record evidence for any false positive or unavailable review; never treat a pending or missing bot result as a pass. Do not lower coverage thresholds or lint budgets to satisfy a check diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 08f83aeb756..789b2ffaef4 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -6,6 +6,7 @@ from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypedDict, cast +from fastapi import HTTPException from typing_extensions import ReadOnly from litellm._logging import verbose_proxy_logger @@ -1692,13 +1693,18 @@ async def refresh_user_oauth_token( ) return None - binding_proof: Final = await enforce_oauth_identity_binding( - server=server, - token_response=body, - litellm_user_id=user_id, - grant_type="refresh_token", - refresh_ownership=RefreshTokenPresented(refresh_token), - ) + try: + binding_proof: Final = await enforce_oauth_identity_binding( + server=server, + token_response=body, + litellm_user_id=user_id, + grant_type="refresh_token", + refresh_ownership=RefreshTokenPresented(refresh_token), + ) + except HTTPException as exc: + if exc.status_code != 403: + raise + return None access_token: Final[str | None] = body.get("access_token") if not access_token: @@ -1812,8 +1818,14 @@ async def resolve_user_oauth_access_token( binding: Final = server.oauth_identity_binding enforce_binding: Final = binding is not None and binding.mode == "enforce" - if enforce_binding: - await mcp_per_user_token_cache.delete(user_id, server_id) + if prefetched_creds is None and enforce_binding and binding is not None: + bound_token: Final = await mcp_per_user_token_cache.get_token(user_id, server_id) + if bound_token is not None: + if await credential_binding_matches( + binding, user_id, server_id, {"identity_binding_proof": bound_token.identity_binding_proof} + ): + return bound_token.access_token + await mcp_per_user_token_cache.delete(user_id, server_id) if prefetched_creds is None and not enforce_binding: cached_token: Final = await mcp_per_user_token_cache.get(user_id, server_id) if cached_token is not None: @@ -1848,7 +1860,9 @@ async def resolve_user_oauth_access_token( access_token: Final[str] = cred["access_token"] if prefetched_creds is None: ttl: Final = _compute_per_user_token_ttl(server, _remaining_token_seconds(cred.get("expires_at"))) - await mcp_per_user_token_cache.set(user_id, server_id, access_token, ttl) + await mcp_per_user_token_cache.set( + user_id, server_id, access_token, ttl, identity_binding_proof=cred.get("identity_binding_proof") + ) return access_token except Exception as e: verbose_proxy_logger.warning( diff --git a/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py b/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py index a4ef970b87a..42edc2999ab 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py +++ b/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py @@ -28,6 +28,8 @@ from litellm.proxy._experimental.mcp_server.oauth_utils import ( build_upstream_oauth2_token_request, resolve_upstream_resource, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import OAuthToken +from litellm.proxy._experimental.mcp_server.outbound_credentials.token_cache_codec import OAuthTokenCacheCodec from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, encrypt_value_helper, @@ -233,8 +235,17 @@ class MCPPerUserTokenCache: def _cache_key(self, user_id: str, server_id: str) -> str: return f"{MCP_PER_USER_TOKEN_REDIS_KEY_PREFIX}:{user_id}:{server_id}" + def _codec(self) -> OAuthTokenCacheCodec: + return OAuthTokenCacheCodec( + encrypt_value_helper, + lambda blob: decrypt_value_helper(blob, key="mcp_per_user_token", exception_type="debug"), + ) + async def get(self, user_id: str, server_id: str) -> str | None: - """Return the plaintext access_token, or None on miss/error.""" + token: Final = await self.get_token(user_id, server_id) + return token.access_token if token is not None else None + + async def get_token(self, user_id: str, server_id: str) -> OAuthToken | None: try: from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415 @@ -242,12 +253,7 @@ class MCPPerUserTokenCache: encrypted: Final = await user_api_key_cache.async_get_cache(key) if encrypted is None: return None - plaintext: Final = decrypt_value_helper( - encrypted, - key="mcp_per_user_token", - exception_type="debug", - ) - return plaintext or None + return self._codec().decode(encrypted) except Exception as exc: verbose_logger.debug( "MCPPerUserTokenCache.get failed for user=%s server=%s: %s", @@ -263,13 +269,16 @@ class MCPPerUserTokenCache: server_id: str, access_token: str, ttl: int, + identity_binding_proof: str | None = None, ) -> None: """Store NaCl-encrypted access_token in Redis with the given TTL.""" try: from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415 key: Final = self._cache_key(user_id, server_id) - encrypted: Final = encrypt_value_helper(access_token) + encrypted: Final = self._codec().encode( + OAuthToken(access_token=access_token, identity_binding_proof=identity_binding_proof) + ) await user_api_key_cache.async_set_cache(key, encrypted, ttl=ttl) verbose_logger.debug( "MCPPerUserTokenCache.set: cached token for user=%s server=%s ttl=%ds", diff --git a/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py b/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py index 7ca9e6a23df..f02e6c85d9b 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py @@ -342,7 +342,7 @@ async def _evaluate_binding( claims: Final = _decode_id_token(id_token, binding, signing_key) if isinstance(claims, _BindingRejection): return claims - if grant_type == "authorization_code": + if grant_type == "authorization_code" and (binding.mode == "enforce" or expected_nonce is not None): nonce: Final = claims.get("nonce") if not expected_nonce or not isinstance(nonce, str) or not hmac.compare_digest(nonce, expected_nonce): return _BindingRejection( diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/authz_code_refresher.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/authz_code_refresher.py index 824459fd9b9..af1e82eab82 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/authz_code_refresher.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/authz_code_refresher.py @@ -14,6 +14,8 @@ import time from collections.abc import Awaitable, Callable from typing import TYPE_CHECKING, Final, Protocol +from fastapi import HTTPException + from litellm._logging import verbose_logger from litellm.proxy._experimental.mcp_server.auth.token_endpoint_auth import ( TokenEndpointAuthConfigError, @@ -93,6 +95,14 @@ class AuthorizationCodeRefresher: self._identity_validator = identity_validator async def refresh(self, user_id: str, server_id: str, token: OAuthToken) -> OAuthToken | None: + try: + return await self._refresh(user_id, server_id, token) + except HTTPException as exc: + if exc.status_code != 403: + raise + return None + + async def _refresh(self, user_id: str, server_id: str, token: OAuthToken) -> OAuthToken | None: if token.refresh_token is None: return None server: Final = self._server_lookup(server_id) @@ -162,4 +172,5 @@ class AuthorizationCodeRefresher: expires_at=self._clock() + expires_in if expires_in is not None else None, refresh_token=new_refresh, scopes=scopes, + identity_binding_proof=binding_proof, ) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/oauth_token_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/oauth_token_store.py index e0cd9e8e5ed..0ac4296f498 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/oauth_token_store.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/oauth_token_store.py @@ -41,6 +41,7 @@ class OAuthToken: expires_at: float | None = None refresh_token: str | None = None scopes: tuple[str, ...] = () + identity_binding_proof: str | None = None def __repr__(self) -> str: has_refresh: Final = self.refresh_token is not None diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py index efcc613c539..2897c3e8e4a 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py @@ -12,6 +12,7 @@ from __future__ import annotations import asyncio from collections.abc import Callable, Mapping +from functools import partial from typing import TYPE_CHECKING, Final from litellm._logging import verbose_logger @@ -39,7 +40,6 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.token_cache_cod OAuthTokenCacheCodec, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.v2_token_store import ( - CredentialReader, V2PerUserTokenStore, ) @@ -146,12 +146,26 @@ def _runtime_backend_and_coordinator() -> tuple[TokenCacheBackend | None, Refres return backend, coordinator, True +async def _read_bound_credential( + server_lookup: ServerLookup, user_id: str, server_id: str +) -> Mapping[str, object] | None: + credential: Final = await _read_credential(user_id, server_id) + server: Final = server_lookup(server_id) + binding: Final = server.oauth_identity_binding if server else None + if credential is not None and binding is not None and binding.mode == "enforce": + if not await credential_binding_matches(binding, user_id, server_id, credential): + return None + return credential + + def _build_per_user_oauth_token_store( server_lookup: ServerLookup, ) -> tuple[CachedOAuthTokenStore, bool]: backend, coordinator, uses_redis = _runtime_backend_and_coordinator() refresher: Final = AuthorizationCodeRefresher(server_lookup, _post_token_endpoint, _persist_credential) - refreshing: Final = RefreshingTokenStore(V2PerUserTokenStore(_read_credential), refresher, coordinator=coordinator) + refreshing: Final = RefreshingTokenStore( + V2PerUserTokenStore(partial(_read_bound_credential, server_lookup)), refresher, coordinator=coordinator + ) return CachedOAuthTokenStore(refreshing, default_ttl_seconds=_DEFAULT_TTL_SECONDS, backend=backend), uses_redis @@ -176,10 +190,8 @@ class LazyPerUserOAuthTokenStore: *, store_builder: StoreBuilder = _build_per_user_oauth_token_store, redis_available: Callable[[], bool] = _redis_cache_is_available, - credential_reader: CredentialReader = _read_credential, ) -> None: self._server_lookup = server_lookup - self._credential_reader = credential_reader self._store_builder = store_builder self._redis_available = redis_available self._store: InvalidatableOAuthTokenStore | None = None @@ -188,13 +200,18 @@ class LazyPerUserOAuthTokenStore: self._local_fetches = 0 async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None: + token: Final = await self._fetch_token(user_id, server_id) server: Final = self._server_lookup(server_id) binding: Final = server.oauth_identity_binding if server else None - if binding is not None and binding.mode == "enforce": - credential: Final = await self._credential_reader(user_id, server_id) - await self.invalidate(user_id, server_id) - if credential is None or not await credential_binding_matches(binding, user_id, server_id, credential): + if token is not None and binding is not None and binding.mode == "enforce": + if not await credential_binding_matches( + binding, user_id, server_id, {"identity_binding_proof": token.identity_binding_proof} + ): + await self.invalidate(user_id, server_id) return None + return token + + async def _fetch_token(self, user_id: str, server_id: str) -> OAuthToken | None: if self._uses_redis: store = self._store if store is not None: diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_cache_codec.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_cache_codec.py index 14cc309b685..56dc3f19cdd 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_cache_codec.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_cache_codec.py @@ -1,24 +1,26 @@ """Serialize + encrypt boundary for caching an OAuth token in a shared (Redis) cache. -A cross-replica cache must serialize the token, and a plaintext bearer in Redis is a leak, so this -encrypts the value (NaCl in production via the injected ``encrypt``, identity in tests). It caches -**only** the ``access_token``: the hot path needs just the bearer, expiry is carried by the cache -entry's TTL (set from the token's ``expires_at`` by the cache), and the long-lived refresh_token stays -in the DB - the refresh path is always a cache miss that re-reads it - so it never reaches Redis. A -decoded token therefore carries only the bearer (``expires_at`` and ``refresh_token`` both None); the -TTL, not the value, bounds its life. An empty/undecryptable blob (e.g. master-key rotation) is a miss. +Shared cache values contain an encrypted access token and optional identity-binding proof. +Refresh tokens remain in the database; cache TTL bounds the access token's lifetime. +Legacy bearer-only entries decode without proof and cannot satisfy identity enforcement. """ from __future__ import annotations +import json from collections.abc import Callable from dataclasses import dataclass from typing import Final +from pydantic import TypeAdapter, ValidationError + from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( OAuthToken, ) +_BOUND_PREFIX: Final = "litellm-bound-oauth-v1:" +_BOUND_PAYLOAD: Final = TypeAdapter(dict[str, str]) + @dataclass(frozen=True, slots=True) class OAuthTokenCacheCodec: @@ -26,10 +28,30 @@ class OAuthTokenCacheCodec: decrypt: Callable[[str], str | None] def encode(self, token: OAuthToken) -> str: + if token.identity_binding_proof is not None: + return self.encrypt( + _BOUND_PREFIX + + json.dumps( + { + "access_token": token.access_token, + "identity_binding_proof": token.identity_binding_proof, + } + ) + ) return self.encrypt(token.access_token) def decode(self, blob: str) -> OAuthToken | None: access_token: Final = self.decrypt(blob) if not access_token: return None + if access_token.startswith(_BOUND_PREFIX): + try: + payload: Final = _BOUND_PAYLOAD.validate_json(access_token[len(_BOUND_PREFIX) :]) + except ValidationError: + return None + bearer: Final = payload.get("access_token") + proof: Final = payload.get("identity_binding_proof") + if not bearer or not proof: + return None + return OAuthToken(access_token=bearer, identity_binding_proof=proof) return OAuthToken(access_token=access_token, refresh_token=None) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/v2_token_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/v2_token_store.py index 4d88ec8b025..0f18931d118 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/v2_token_store.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/v2_token_store.py @@ -46,11 +46,13 @@ def _to_oauth_token(payload: Mapping[str, object]) -> OAuthToken | None: return None refresh_token: Final = payload.get("refresh_token") expires_at: Final = payload.get("expires_at") + binding_proof: Final = payload.get("identity_binding_proof") return OAuthToken( access_token=access_token, expires_at=_iso_to_epoch(expires_at) if isinstance(expires_at, str) else None, refresh_token=refresh_token if isinstance(refresh_token, str) else None, scopes=_to_scopes(payload.get("scopes")), + identity_binding_proof=binding_proof if isinstance(binding_proof, str) else None, ) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_authz_code_refresher.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_authz_code_refresher.py index 16707f6ace8..df068b60338 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_authz_code_refresher.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_authz_code_refresher.py @@ -293,17 +293,18 @@ async def test_refresh_uses_admin_entered_token_url_when_issuer_yield_empties_re @pytest.mark.asyncio async def test_identity_rejection_never_persists_or_returns_refreshed_token(): from unittest.mock import AsyncMock + from fastapi import HTTPException validator = AsyncMock(side_effect=HTTPException(status_code=403, detail="oauth_principal_mismatch")) persist = AsyncMock() refresher = AuthorizationCodeRefresher( - _lookup(_Server()), _endpoint({"access_token": "foreign-token", "id_token": "foreign-identity"}), - persist, identity_validator=validator, + _lookup(_Server()), + _endpoint({"access_token": "foreign-token", "id_token": "foreign-identity"}), + persist, + identity_validator=validator, ) - with pytest.raises(HTTPException) as error: - await refresher.refresh("alice", "srv", OAuthToken(access_token="old", refresh_token="old-rt")) - assert error.value.status_code == 403 + assert await refresher.refresh("alice", "srv", OAuthToken(access_token="old", refresh_token="old-rt")) is None persist.assert_not_awaited() @@ -314,10 +315,13 @@ async def test_verified_refresh_preserves_binding_proof_in_storage(): validator = AsyncMock(return_value="verified-binding") persist = AsyncMock() refresher = AuthorizationCodeRefresher( - _lookup(_Server()), _endpoint({"access_token": "new", "refresh_token": "rotated"}), - persist, identity_validator=validator, + _lookup(_Server()), + _endpoint({"access_token": "new", "refresh_token": "rotated"}), + persist, + identity_validator=validator, ) token = await refresher.refresh("alice", "srv", OAuthToken(access_token="old", refresh_token="old-rt")) assert token.access_token == "new" assert token.refresh_token == "rotated" + assert token.identity_binding_proof == "verified-binding" assert persist.await_args.kwargs["identity_binding_proof"] == "verified-binding" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_per_user_oauth_store.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_per_user_oauth_store.py index bd190b09c82..e7308884060 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_per_user_oauth_store.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_per_user_oauth_store.py @@ -254,20 +254,86 @@ async def test_enforcement_invalidates_cached_legacy_credentials_before_use(): from litellm.types.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer server = MCPServer( - server_id="srv", name="srv", transport=MCPTransport.http, auth_type=MCPAuth.oauth2, + server_id="srv", + name="srv", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, oauth_identity_binding=MCPOAuthIdentityBinding( - mode="enforce", issuer="https://idp.example.com", audiences=["client"], + mode="enforce", + issuer="https://idp.example.com", + audiences=["client"], ), ) cached = _RecordingStore("belongs-to-bob") - async def read_legacy(user_id, server_id): - return {"access_token": "belongs-to-bob", "refresh_token": "bobs-refresh"} - store = LazyPerUserOAuthTokenStore( - lambda server_id: server, store_builder=lambda lookup: (cached, False), - redis_available=lambda: False, credential_reader=read_legacy, + lambda server_id: server, + store_builder=lambda lookup: (cached, False), + redis_available=lambda: False, ) assert await store.fetch("alice", "srv") is None - assert cached.calls == [] + assert cached.calls == [("alice", "srv")] assert cached.invalidations == [("alice", "srv")] + + +@pytest.mark.asyncio +async def test_enforced_cache_hit_avoids_credential_read_and_rejects_changed_policy(monkeypatch): + from unittest.mock import AsyncMock + + from litellm.proxy._experimental.mcp_server.oauth_identity_binding import current_binding_proof + from litellm.proxy._experimental.mcp_server.outbound_credentials import per_user_oauth_store as module + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="srv", + name="srv", + transport="http", + auth_type="oauth2", + oauth_identity_binding={ + "mode": "enforce", + "issuer": "https://idp.example", + "audiences": ["client"], + "caller_field": "user_id", + "principal_claim": "sub", + }, + ) + proof = await current_binding_proof(server.oauth_identity_binding, "alice", "srv") + read = AsyncMock(return_value={"access_token": "alice-token", "identity_binding_proof": proof}) + monkeypatch.setattr(module, "_read_credential", read) + monkeypatch.setattr(module, "_runtime_backend_and_coordinator", lambda: (None, None, False)) + store = LazyPerUserOAuthTokenStore(lambda _: server, redis_available=lambda: False) + assert (await store.fetch("alice", "srv")).access_token == "alice-token" + assert (await store.fetch("alice", "srv")).access_token == "alice-token" + read.assert_awaited_once_with("alice", "srv") + server.oauth_identity_binding = server.oauth_identity_binding.model_copy(update={"audiences": ["changed"]}) + assert await store.fetch("alice", "srv") is None + read.assert_awaited_once() + assert await store.fetch("alice", "srv") is None + assert read.await_count == 2 + + +@pytest.mark.asyncio +async def test_expired_unverified_credential_never_reaches_refresh(monkeypatch): + from unittest.mock import AsyncMock + + from litellm.proxy._experimental.mcp_server.outbound_credentials import per_user_oauth_store as module + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="srv", + name="srv", + transport="http", + auth_type="oauth2", + token_url="https://idp.example/token", + oauth_identity_binding={"mode": "enforce", "issuer": "https://idp.example", "audiences": ["client"]}, + ) + read = AsyncMock( + return_value={"access_token": "bob", "refresh_token": "bob-refresh", "expires_at": "2000-01-01T00:00:00Z"} + ) + post = AsyncMock() + monkeypatch.setattr(module, "_read_credential", read) + monkeypatch.setattr(module, "_post_token_endpoint", post) + monkeypatch.setattr(module, "_runtime_backend_and_coordinator", lambda: (None, None, False)) + store = LazyPerUserOAuthTokenStore(lambda _: server, redis_available=lambda: False) + assert await store.fetch("alice", "srv") is None + post.assert_not_awaited() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_cache_codec.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_cache_codec.py index 17a7e13af03..12a04c4b4dc 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_cache_codec.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_cache_codec.py @@ -52,3 +52,17 @@ def test_undecryptable_blob_is_a_miss(): def test_empty_plaintext_is_a_miss(): codec = OAuthTokenCacheCodec(encrypt=lambda s: s, decrypt=lambda b: b) assert codec.decode("") is None + + +def test_bound_token_round_trip_preserves_proof_without_refresh_secret(): + codec = _wrapping_codec() + blob = codec.encode(OAuthToken("alice-token", refresh_token="private-refresh", identity_binding_proof="proof")) + assert "private-refresh" not in blob + decoded = codec.decode(blob) + assert decoded == OAuthToken("alice-token", identity_binding_proof="proof") + + +def test_malformed_bound_entries_fail_closed(): + codec = _wrapping_codec() + for payload in ("not-json", "{}", '{"access_token":"at"}', '{"access_token":1,"identity_binding_proof":"p"}'): + assert codec.decode("enc:litellm-bound-oauth-v1:" + payload) is None diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py index 8e2adbeb06e..60a5e1a22bb 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py @@ -1629,3 +1629,78 @@ async def test_enforcement_rejects_preexisting_unverified_credential(): cred={"access_token": "belongs-to-bob", "refresh_token": "bobs-refresh-token"}, ) assert result is None + + +@pytest.mark.asyncio +async def test_refresh_identity_rejection_returns_reauthentication_without_persisting(monkeypatch): + from fastapi import HTTPException + from litellm.proxy._experimental.mcp_server import db as module + + validator = AsyncMock(side_effect=HTTPException(status_code=403, detail="oauth_principal_mismatch")) + monkeypatch.setattr(module, "enforce_oauth_identity_binding", validator) + result, captured = await _run_refresh( + monkeypatch, _refresh_server(), {"access_token": "bob", "refresh_token": "rotated"} + ) + assert result is None + assert captured["data"]["grant_type"] == "refresh_token" + module.store_user_oauth_credential.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_verified_legacy_cache_reads_avoid_database_and_reject_policy_changes(monkeypatch): + from litellm.caching.caching import DualCache + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server import db as module + from litellm.proxy._experimental.mcp_server.oauth2_token_cache import mcp_per_user_token_cache + from litellm.proxy._experimental.mcp_server.oauth_identity_binding import current_binding_proof + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="srv", + name="srv", + transport="http", + auth_type="oauth2", + oauth_identity_binding={ + "mode": "enforce", + "issuer": "https://idp.example", + "audiences": ["client"], + "caller_field": "user_id", + "principal_claim": "sub", + }, + ) + proof = await current_binding_proof(server.oauth_identity_binding, "alice", "srv") + monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache()) + read = AsyncMock(return_value=None) + monkeypatch.setattr(module, "get_user_oauth_credential", read) + monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) + await mcp_per_user_token_cache.set("alice", "srv", "alice-token", 60, identity_binding_proof=proof) + assert await module.resolve_user_oauth_access_token("alice", server) == "alice-token" + assert await module.resolve_user_oauth_access_token("alice", server) == "alice-token" + read.assert_not_awaited() + server.oauth_identity_binding = server.oauth_identity_binding.model_copy(update={"audiences": ["changed"]}) + assert await module.resolve_user_oauth_access_token("alice", server) is None + read.assert_awaited_once() + assert await mcp_per_user_token_cache.get_token("alice", "srv") is None + + +@pytest.mark.asyncio +async def test_unverified_legacy_cache_cannot_bypass_enforcement(monkeypatch): + from litellm.caching.caching import DualCache + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server import db as module + from litellm.proxy._experimental.mcp_server.oauth2_token_cache import mcp_per_user_token_cache + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="srv", + name="srv", + transport="http", + auth_type="oauth2", + oauth_identity_binding={"mode": "enforce", "issuer": "https://idp.example", "audiences": ["client"]}, + ) + monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache()) + monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) + monkeypatch.setattr(module, "get_user_oauth_credential", AsyncMock(return_value={"access_token": "bob"})) + await mcp_per_user_token_cache.set("alice", "srv", "bob", 60) + assert await module.resolve_user_oauth_access_token("alice", server) is None + assert await mcp_per_user_token_cache.get("alice", "srv") is None diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_identity_binding.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_identity_binding.py index 6f291865e20..0036035f448 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_identity_binding.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_identity_binding.py @@ -13,14 +13,14 @@ from pydantic import ValidationError from litellm.proxy._experimental.mcp_server.oauth_identity_binding import ( RefreshOwnershipProven, - VerifiedRefreshToken, RefreshTokenPresented, - current_binding_proof, + VerifiedRefreshToken, _discover_jwks_url, _fetch_issuer_jwks, _load_caller_principal, _load_stored_refresh_token, _select_signing_key, + current_binding_proof, enforce_oauth_identity_binding, ) from litellm.types.mcp import MCPAuth, MCPTransport @@ -471,19 +471,20 @@ async def test_refresh_with_mismatched_id_token_rejected(): @pytest.mark.asyncio -async def test_audit_mode_logs_but_does_not_reject(): +async def test_audit_mode_logs_but_does_not_reject(caplog): token: Final = _sign_id_token({"email": "mallory@example.com", "email_verified": True}) result: Final = await enforce_oauth_identity_binding( server=_server(mode="audit"), token_response={"access_token": "at", "id_token": token}, litellm_user_id="user-a", grant_type="authorization_code", - expected_nonce="test-nonce", refresh_ownership=None, jwks_fetcher=_jwks_fetcher, caller_principal_loader=_caller_loader("alice@example.com"), ) assert result is None + assert "oauth_principal_mismatch" in caplog.text + assert "nonce" not in caplog.text @pytest.mark.asyncio @@ -642,6 +643,25 @@ async def test_binding_proof_rejects_changed_user_or_policy(): def test_identity_binding_rejects_modes_without_gateway_credential_custody(auth_type): with pytest.raises(ValidationError, match="gateway-managed per-user"): MCPServer( - server_id="srv", name="srv", transport=MCPTransport.http, auth_type=auth_type, + server_id="srv", + name="srv", + transport=MCPTransport.http, + auth_type=auth_type, oauth_identity_binding=_server().oauth_identity_binding, ) + + +@pytest.mark.asyncio +async def test_audit_matching_login_without_nonce_does_not_report_failure(caplog): + token: Final = _sign_id_token({"email": "alice@example.com", "email_verified": True}) + result: Final = await enforce_oauth_identity_binding( + server=_server(mode="audit"), + token_response={"access_token": "at", "id_token": token}, + litellm_user_id="user-a", + grant_type="authorization_code", + refresh_ownership=None, + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + ) + assert result is None + assert "oauth_identity_binding audit" not in caplog.text From 93bd3d61b56a167b794c6b8cc2dcc6dcefaf4a47 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Thu, 10 Sep 2026 13:20:49 -0700 Subject: [PATCH 29/33] fix(mcp): preserve verified identity when warming OAuth cache --- .../mcp_server/discoverable_endpoints.py | 1 + .../mcp_server/test_discoverable_endpoints.py | 43 +++++++++++++++++++ 2 files changed, 44 insertions(+) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index c8e02be308b..920fe17e43e 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -634,6 +634,7 @@ async def _store_per_user_token_server_side( server_id=server.server_id, access_token=access_token, ttl=ttl, + identity_binding_proof=identity_binding_proof, ) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 42187cb1771..ff649075468 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -11103,3 +11103,46 @@ async def test_identity_bound_authorization_carries_nonce_and_caller_through_cal assert sealed.upstream_code == "upstream-code" assert sealed.oauth_nonce == query["nonce"][0] assert returned["state"] == ["client-state"] + + +@pytest.mark.asyncio +async def test_enforced_login_warms_verified_token_readable_without_database_lookup(monkeypatch): + from types import SimpleNamespace + from litellm.caching.caching import DualCache + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server import db, mcp_server_manager + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _store_per_user_token_server_side + from litellm.proxy._experimental.mcp_server.oauth2_token_cache import mcp_per_user_token_cache + from litellm.proxy._experimental.mcp_server.oauth_identity_binding import current_binding_proof + from litellm.proxy._experimental.mcp_server.outbound_credentials.dual_cache_token_backend import DualCacheTokenCacheBackend + from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import CachedOAuthTokenStore + from litellm.proxy._experimental.mcp_server.outbound_credentials.per_user_oauth_store import LazyPerUserOAuthTokenStore + from litellm.proxy._experimental.mcp_server.outbound_credentials.v2_token_store import V2PerUserTokenStore + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="srv", name="srv", transport="http", auth_type="oauth2", + oauth_identity_binding={"mode": "enforce", "issuer": "https://idp.example", "audiences": ["client"], + "caller_field": "user_id", "principal_claim": "sub"}, + ) + proof = await current_binding_proof(server.oauth_identity_binding, "alice", "srv") + cache = DualCache() + monkeypatch.setenv("LITELLM_SALT_KEY", "test-cache-warm-encryption-salt") + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) + monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", SimpleNamespace(invalidate_user_oauth_token_cache=AsyncMock())) + monkeypatch.setattr(db, "store_user_oauth_credential", AsyncMock()) + await _store_per_user_token_server_side( + server=server, user_id="alice", token_response={"access_token": "alice-token", "refresh_token": "private", "expires_in": 3600}, + identity_binding_proof=proof, + ) + read = AsyncMock(return_value=None) + cached = CachedOAuthTokenStore( + V2PerUserTokenStore(read), default_ttl_seconds=300, backend=DualCacheTokenCacheBackend(cache, mcp_per_user_token_cache._codec()), + ) + store = LazyPerUserOAuthTokenStore(lambda _: server, store_builder=lambda _: (cached, True), redis_available=lambda: True) + token = await store.fetch("alice", "srv") + assert token is not None and token.access_token == "alice-token" + assert token.identity_binding_proof == proof + assert token.refresh_token is None + read.assert_not_awaited() From 1987e6e290a8344515f595538e85c3ae00e69a96 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 10 Sep 2026 14:54:46 -0700 Subject: [PATCH 30/33] test(proxy): pass the request to get_marketplace in the archive marketplace test --- .../proxy/anthropic_endpoints/test_claude_code_marketplace.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_marketplace.py b/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_marketplace.py index a7c2bd7ba20..a585666743f 100644 --- a/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_marketplace.py +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_marketplace.py @@ -481,7 +481,7 @@ async def test_archive_source_registers_and_is_served_verbatim_in_marketplace(): assert response.action == "created" assert response.plugin.source == _ARCHIVE_SOURCE - marketplace = json.loads((await get_marketplace()).body) + marketplace = json.loads((await get_marketplace(request=MagicMock())).body) assert marketplace["plugins"] == [{"name": "s3-skill", "source": _ARCHIVE_SOURCE, "version": "1.0.0"}] From 6632e8b74fecaeb569abdd12be43f6870b45c3b5 Mon Sep 17 00:00:00 2001 From: Kishorekarthik P Date: Sat, 1 Aug 2026 12:57:22 +0530 Subject: [PATCH 31/33] fix(proxy): let internal users read request/response for their own spend logs The Logs drawer gets messages/response from GET /spend/logs/ui/{request_id}; the list endpoint omits those heavy columns for every caller, admins included. That detail route was missing from LiteLLMRoutes.spend_tracking_routes, and check_route_access anchors patterns, so /spend/logs/ui never matched it. Every internal_user got a 403 before the handler ran and the UI fell back to the "Request/Response Data Not Available" banner, even on their own requests Adds the route to spend_tracking_routes so internal_user, internal_user_view_only, admin_viewer and org_admin all inherit it, and drops the now-redundant explicit entry from admin_viewer_routes. The handler already authorizes non-admins per row via _assert_user_can_view_request_id, so no handler-side scoping change is needed That helper returned silently when no spend-log row existed, which the detail handler treats as authorized before asking every custom logger for the payload by raw request_id. With retention pruning the row can be gone while the payload is still in cold storage, so opening the route would have let a non-admin read another tenant's prompt out of S3/GCS. A missing row now falls through to the same 403 as a foreign row, which also removes the exists-but-not-yours oracle Fixes #34099 --- litellm/proxy/_types.py | 9 +- .../spend_management_endpoints.py | 12 +- .../proxy/auth/test_route_checks.py | 76 +++++++ .../test_spend_management_endpoints.py | 186 ++++++++++++++++++ 4 files changed, 273 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 23dc8237160..aa585120e5b 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -703,6 +703,7 @@ class LiteLLMRoutes(enum.Enum): "/spend/logs", "/spend/logs/v2", "/spend/logs/ui", + "/spend/logs/ui/{request_id}", "/spend/logs/session/ui", "/key/spend/report", "/user/spend/report", @@ -932,10 +933,10 @@ class LiteLLMRoutes(enum.Enum): # PROXY_ADMIN_VIEW_ONLY — the route gate must match). "/customer/list", "/customer/info", - # UI Logs page detail drawer (single + session) and the filter facets. - # The list endpoint `/spend/logs/ui` is covered via - # spend_tracking_routes below. - "/spend/logs/ui/{logId}", + # UI Logs page session detail drawer and the end-user filter facet. + # The list endpoint `/spend/logs/ui` and the single-log detail route + # `/spend/logs/ui/{request_id}` are covered via spend_tracking_routes + # below. "/spend/logs/session/ui", "/management/v1/spend_logs/end_users", "/management/v1/spend_logs/users", diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index f8831ca4152..88cdf37cabb 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -4633,16 +4633,16 @@ async def _assert_user_can_view_request_id( Verify the requesting non-admin user is allowed to view this spend-log row. Allowed when the log belongs to the user directly, or to one of their permitted teams (admin or ``/spend/logs`` permission). - Raises HTTP 403 if not. + Raises HTTP 403 if not, including when no spend-log row exists for the + request_id (e.g. it was pruned by retention), so a missing row can't be + used to read a payload out of cold storage via the detail endpoint. """ row: Final = await _find_spend_log_row(prisma_client, request_id) - if row is None: + + if row is not None and row.user is not None and row.user == user_api_key_dict.user_id: return - if row.user is not None and row.user == user_api_key_dict.user_id: - return - - if row.team_id: + if row is not None and row.team_id: can_view: Final = await _can_team_member_view_log( prisma_client=prisma_client, user_api_key_dict=user_api_key_dict, diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index 4f2ff1582dd..38a0e85c1ea 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -2033,6 +2033,82 @@ def test_proxy_admin_viewer_can_access_logs_page_endpoints(route): ) +@pytest.mark.parametrize( + "user_role", + [LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY], +) +def test_internal_user_can_access_logs_drawer_detail_route(user_role): + """ + The Logs drawer detail fetch (GET /spend/logs/ui/{request_id}) must pass + route_checks for plain internal users, not just admins — the handler + itself already self-authorizes row ownership via + _assert_user_can_view_request_id. + """ + route = "/spend/logs/ui/abc-request-id" + user_obj = LiteLLM_UserTable( + user_id="internal_user", + user_email="user@example.com", + user_role=user_role.value, + ) + valid_token = UserAPIKeyAuth( + user_id="internal_user", + user_role=user_role.value, + ) + request = MagicMock(spec=Request) + request.query_params = {} + + try: + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=user_obj, + _user_role=user_role.value, + route=route, + request=request, + valid_token=valid_token, + request_data={}, + ) + except Exception as e: + pytest.fail(f"{user_role.value} should be able to access {route}. Got error: {str(e)}") + + +@pytest.mark.parametrize( + "route_group_name", + [ + "spend_tracking_routes", + "internal_user_routes", + "internal_user_view_only_routes", + "admin_viewer_routes", + "org_admin_allowed_routes", + ], +) +def test_logs_drawer_detail_route_in_every_route_group(route_group_name): + """ + /spend/logs/ui/{request_id} must be reachable through + RouteChecks.check_route_access under each role's own route group, so a + partial revert (removing the route from `spend_tracking_routes` while + leaving `non_proxy_admin_allowed_routes_check` alone) is also caught. + """ + from litellm.proxy._types import LiteLLMRoutes + + allowed_routes = getattr(LiteLLMRoutes, route_group_name).value + assert RouteChecks.check_route_access( + route="/spend/logs/ui/req-34099", allowed_routes=allowed_routes + ) + + +def test_logs_drawer_detail_route_allowed_for_scoped_virtual_key(): + """ + A virtual key scoped to `allowed_routes=["spend_tracking_routes"]` must be + able to reach the Logs drawer detail route. + """ + valid_token = UserAPIKeyAuth( + user_id="scoped_key_user", + allowed_routes=["spend_tracking_routes"], + ) + assert RouteChecks.is_virtual_key_allowed_to_call_route( + route="/spend/logs/ui/req-34099", valid_token=valid_token + ) + + @pytest.mark.parametrize("route", ADMIN_VIEWER_LOGS_PAGE_ROUTES) def test_internal_user_blocked_from_admin_viewer_logs_routes(route): """ diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 329a33eb440..53accfd08ee 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -455,6 +455,34 @@ async def test_assert_user_can_view_request_id_rejects_both_users_none(): assert exc_info.value.status_code == 403 +@pytest.mark.asyncio +async def test_assert_user_can_view_request_id_rejects_missing_row(): + """ + A request_id with no spend-log row (e.g. pruned by retention) must not + authorize reading the payload from cold storage; a missing row is not + the same as an owned row. + """ + + class MockSpendLogs: + async def find_unique(self, where, include=None): + return None + + class MockDB: + def __init__(self): + self.litellm_spendlogs = MockSpendLogs() + + class MockPrisma: + def __init__(self): + self.db = MockDB() + + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_1") + with pytest.raises(HTTPException) as exc_info: + await spend_management_endpoints._assert_user_can_view_request_id( + MockPrisma(), auth, "req-missing-row" + ) + assert exc_info.value.status_code == 403 + + def test_ui_view_request_response_forbids_non_admin_without_db(client, monkeypatch): """ Without prisma, non-admins cannot be authorized to read request/response @@ -6777,3 +6805,161 @@ async def test_ui_view_spend_logs_search_returns_flat_rows_when_grouping_by_sess assert "next_session_cursor" not in data finally: app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +def _fake_prisma_with_owned_spend_log(owner_user_id, messages_json, response_json): + class _Row: + user = owner_user_id + team_id = None + + class _SpendLogs: + async def find_unique(self, where, include=None): + return _Row() + + class _DB: + def __init__(self): + self.litellm_spendlogs = _SpendLogs() + + async def query_raw(self, _sql, *_args): + return [ + { + "messages": messages_json, + "response": response_json, + "proxy_server_request": "{}", + "metadata": "{}", + } + ] + + class _Prisma: + def __init__(self): + self.db = _DB() + + return _Prisma() + + +def test_ui_view_request_response_internal_user_owner_gets_payload(client, monkeypatch): + """ + An internal_user who owns the spend-log row can fetch the Logs drawer + detail payload for their own request (regression for #34099, where the + route was blocked for INTERNAL_USER before reaching this ownership check). + """ + messages_json = json.dumps([{"role": "user", "content": "hi"}]) + response_json = json.dumps({"choices": [{"message": {"content": "hello"}}]}) + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + _fake_prisma_with_owned_spend_log("user_a", messages_json, response_json), + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_a" + ) + try: + response = client.get( + "/spend/logs/ui/req-owned-by-user-a", + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200 + body = response.json() + assert json.loads(body["messages"]) == [{"role": "user", "content": "hi"}] + assert json.loads(body["response"]) == { + "choices": [{"message": {"content": "hello"}}] + } + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +class _RecordingAdditionalLoggingUtils: + """Injectable custom logger that records every request_id it's asked for.""" + + def __init__(self, payload): + self._payload = payload + self.requested_ids = [] + + async def get_request_response_payload(self, request_id, start_time_utc, end_time_utc): + self.requested_ids.append(request_id) + return self._payload + + +def test_ui_view_request_response_internal_user_non_owner_forbidden(client, monkeypatch): + """ + A different internal_user requesting someone else's row is forbidden; + guards against _assert_user_can_view_request_id being skipped in the + detail-drawer handler. Also proves the handler stops before it ever asks + a custom logger or the DB for the payload. + """ + messages_json = json.dumps([{"role": "user", "content": "hi"}]) + response_json = json.dumps({"choices": [{"message": {"content": "hello"}}]}) + fake_prisma = _fake_prisma_with_owned_spend_log("user_a", messages_json, response_json) + original_query_raw = fake_prisma.db.query_raw + query_raw_calls = [] + + async def _spy_query_raw(*args, **kwargs): + query_raw_calls.append((args, kwargs)) + return await original_query_raw(*args, **kwargs) + + fake_prisma.db.query_raw = _spy_query_raw + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", fake_prisma) + + custom_logger = _RecordingAdditionalLoggingUtils({"messages": "should-not-be-returned"}) + monkeypatch.setattr( + litellm.logging_callback_manager, + "get_active_additional_logging_utils_from_custom_logger", + lambda: [custom_logger], + ) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_b" + ) + try: + response = client.get( + "/spend/logs/ui/req-owned-by-user-a", + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 403 + assert custom_logger.requested_ids == [] + assert query_raw_calls == [] + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +def test_ui_view_request_response_internal_user_missing_row_forbidden(client, monkeypatch): + """ + Regression for the fail-open in _assert_user_can_view_request_id: a + request_id with no spend-log row (e.g. pruned by retention) must be + denied before the handler ever consults a custom logger, otherwise a + non-admin who guesses/obtains a request_id could read another tenant's + payload out of cold storage. Fails if `if row is None: return` is + reintroduced. + """ + + class _SpendLogs: + async def find_unique(self, where, include=None): + return None + + class _DB: + def __init__(self): + self.litellm_spendlogs = _SpendLogs() + + from types import SimpleNamespace + + fake_prisma = SimpleNamespace(db=_DB()) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", fake_prisma) + + custom_logger = _RecordingAdditionalLoggingUtils({"messages": "should-not-be-returned"}) + monkeypatch.setattr( + litellm.logging_callback_manager, + "get_active_additional_logging_utils_from_custom_logger", + lambda: [custom_logger], + ) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_a" + ) + try: + response = client.get( + "/spend/logs/ui/req-pruned", + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 403 + assert custom_logger.requested_ids == [] + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) From d98522b6f662fec6f600ab94817855cb7ff7edce Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 10 Sep 2026 22:14:37 +0000 Subject: [PATCH 32/33] feat(proxy): share database connections across workers with an in-container pgbouncer (#39683) * feat(proxy): share database connections across workers with an in-container pgbouncer Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): parse pgbouncer options iteratively to satisfy the recursion gate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): refuse pgbouncer with token db auth and retry failed pooler restarts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): build pgbouncer 1.25.2 from a pinned source archive and verify pooler replacements The public Wolfi repository only carries pgbouncer 1.24.1-r3, which the image scan rejects (CVE-2026-6664, CVE-2026-6665, CVE-2026-6666, CVE-2025-12819). All three images now compile the checksummed 1.25.2 release in a builder stage. The supervisor now waits for a replacement pooler to listen before treating it as recovered, ends and retries one that never does, and takes the same lock for stop() and spawn so no replacement can be started after shutdown began. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): refuse to start pgbouncer on a loopback port another process already owns Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): count pgbouncer ready only once its own unix socket answers, not any listener on the port Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): refuse pgbouncer older than 1.19, whose unix socket cannot vouch for the tcp port Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(helm,terraform): expose the in-container pgbouncer pool for the componentized gateway Add database.connectionPool to helm/litellm and gateway_connection_pool_* to terraform/litellm/aws so the componentized gateway can receive the LITELLM_PGBOUNCER_* env the classic image already honours. Both reject the pool under IAM or Entra token auth at render/plan time: the pooler holds one static database password for the life of the pod or task. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(gateway): launch the componentized gateway image through a pgbouncer-aware supervisor (#40592) The componentized gateway image started uvicorn directly, so the in-container PgBouncer never ran for it: every worker opened its own Prisma pool to the database. It also passed no keep-alive timeout, so behind a load balancer with a 60s idle timeout uvicorn's 5s default closed idle connections first and the balancer returned 502s on scale-out gateway.launch assembles DATABASE_URL, starts PgBouncer once per pod when LITELLM_PGBOUNCER_ENABLED is set, hands the workers the loopback URL and then runs uvicorn on gateway.main:app with KEEPALIVE_TIMEOUT as --timeout-keep-alive. The image builds PgBouncer 1.25.2 from a checksummed tarball, copies the compiled Rust extension into the /app source tree it imports from (it was only in site-packages, which PYTHONPATH=/app shadows) and asserts the native bridge loads. The app user is added to stats_users so operators can read the PgBouncer console with the application credentials The supervisor returns the pooled URL instead of writing into the mapping it was handed, a database user whose name PgBouncer would split into several stats_users entries is refused before the config is written, and the launcher tests drive main() with an injected serve callable Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * docs(terraform): describe the gateway.launch pooler entrypoint in the aws module README Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): run pgbouncer exit hooks only in the parent and copy the CA into the runtime dir Gunicorn workers inherit the parent's atexit table, so a recycled worker (max_requests) stopped the shared pooler and removed its runtime dir, then hung in the inherited Popen lock. The hooks now no-op unless os.getpid() is the process that started PgBouncer A verified TLS upstream named the operator's CA bundle directly, which is often a 0600 root-owned file that nobody (the user PgBouncer drops to) cannot read, so every server connection failed with "failed to load CA". The bundle is copied into the runtime dir next to the ini and chowned with it Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(terraform): run the gateway through gateway.launch and add the gcp connection-pool variables Cloud Run and ECS overrode the image command with uvicorn gateway.main:app, which skips the supervisor that starts the in-container PgBouncer, so LITELLM_PGBOUNCER_ENABLED was inert on both stacks. Both now exec python -m gateway.launch (under ddtrace-run when USE_DDTRACE is set), and the gcp module gains gateway_connection_pool_enabled / gateway_pool_max_db_connections / gateway_pool_max_client_conn wired to the gateway service only The test_launch password_env fixture now restores DATABASE_URL even when it was unset: monkeypatch.delenv records nothing for an absent var, so main() left postgresql://...@db.internal in the xdist worker's environ and the key-rotation e2e test in the same proxy-infra shard stopped skipping and tried to reach db.internal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(pgbouncer): keep channel_binding and gssencmode off the loopback URL Prisma would demand TLS channel binding from a pooler that only speaks plain TCP on 127.0.0.1. Also pass the request the marketplace test started needing after #40518 landed on top of #40496 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/test-unit.yml | 1 + Dockerfile | 19 +- docker/Dockerfile.database | 19 +- docker/Dockerfile.non_root | 19 +- gateway/Dockerfile | 28 +- gateway/launch.py | 71 ++ helm/litellm-helm/templates/deployment.yaml | 8 + .../tests/connection_pool_tests.yaml | 61 ++ helm/litellm-helm/values.yaml | 14 + helm/litellm/templates/_helpers.tpl | 17 + .../litellm/templates/gateway/deployment.yaml | 3 + helm/litellm/tests/connection_pool_tests.yaml | 117 ++++ helm/litellm/values.yaml | 20 + litellm/proxy/db/pgbouncer.py | 525 +++++++++++++++ litellm/proxy/proxy_cli.py | 16 + terraform/litellm/aws/README.md | 33 + terraform/litellm/aws/ecs.tf | 14 +- .../aws/tests/connection_pool.tftest.hcl | 123 ++++ terraform/litellm/aws/variables.tf | 39 ++ terraform/litellm/gcp/README.md | 34 + terraform/litellm/gcp/cloudrun.tf | 10 +- .../gcp/tests/connection_pool.tftest.hcl | 135 ++++ terraform/litellm/gcp/variables.tf | 39 ++ tests/test_gateway/test_launch.py | 155 +++++ tests/test_litellm/proxy/db/test_pgbouncer.py | 624 ++++++++++++++++++ .../test_litellm/test_component_entrypoint.py | 72 +- 26 files changed, 2195 insertions(+), 21 deletions(-) create mode 100644 gateway/launch.py create mode 100644 helm/litellm-helm/tests/connection_pool_tests.yaml create mode 100644 helm/litellm/tests/connection_pool_tests.yaml create mode 100644 litellm/proxy/db/pgbouncer.py create mode 100644 terraform/litellm/aws/tests/connection_pool.tftest.hcl create mode 100644 terraform/litellm/gcp/tests/connection_pool.tftest.hcl create mode 100644 tests/test_gateway/test_launch.py create mode 100644 tests/test_litellm/proxy/db/test_pgbouncer.py diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index cc606339a20..f55c87c2ae5 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -200,6 +200,7 @@ jobs: tests/test_litellm/proxy/types_utils tests/test_litellm/proxy/logging_endpoints tests/test_litellm/proxy/test_*.py + tests/test_gateway workers: 4 reruns: 2 timeout-minutes: 20 diff --git a/Dockerfile b/Dockerfile index 0a92aa9a68c..759dac76795 100644 --- a/Dockerfile +++ b/Dockerfile @@ -8,9 +8,25 @@ ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a # Pinned by digest like the other base images; bump explicitly on Node upgrades. ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43 +# Checksum from https://www.pgbouncer.org/downloads/ (the Wolfi repo only carries 1.24.x) +ARG PGBOUNCER_VERSION=1.25.2 +ARG PGBOUNCER_SHA256=924ad35113fd0a71c8e2dbe85b5d03445532e2b7b37a9f8a48983beea238b332 FROM $UV_IMAGE AS uvbin +FROM $LITELLM_BUILD_IMAGE AS pgbouncer-builder +ARG PGBOUNCER_VERSION +ARG PGBOUNCER_SHA256 +USER root +RUN apk add --no-cache build-base pkgconf libevent-dev openssl-dev curl +WORKDIR /build +RUN curl -fsSL -o pgbouncer.tar.gz "https://www.pgbouncer.org/downloads/files/${PGBOUNCER_VERSION}/pgbouncer-${PGBOUNCER_VERSION}.tar.gz" && \ + echo "${PGBOUNCER_SHA256} pgbouncer.tar.gz" | sha256sum -c - && \ + tar xzf pgbouncer.tar.gz --strip-components=1 && \ + ./configure --prefix=/usr/local --with-openssl=/usr && \ + make -j"$(nproc)" pgbouncer && \ + install -m 0755 pgbouncer /usr/local/bin/pgbouncer + # Admin UI builder. Pinned to the build platform so the architecture-independent # Next.js static export compiles once natively even in a multi-arch build, # instead of once per target arch under QEMU. @@ -110,7 +126,8 @@ USER root RUN echo "https://packages.wolfi.dev/os" >> /etc/apk/repositories # node (without npm) is required by the prisma CLI at runtime -RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile +RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile libevent +COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer WORKDIR /app ENV PATH="/app/.venv/bin:${PATH}" \ diff --git a/docker/Dockerfile.database b/docker/Dockerfile.database index e9ad2849bb2..b0bf935c616 100644 --- a/docker/Dockerfile.database +++ b/docker/Dockerfile.database @@ -8,9 +8,25 @@ ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a # Pinned by digest like the other base images; bump explicitly on Node upgrades. ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43 +# Checksum from https://www.pgbouncer.org/downloads/ (the Wolfi repo only carries 1.24.x) +ARG PGBOUNCER_VERSION=1.25.2 +ARG PGBOUNCER_SHA256=924ad35113fd0a71c8e2dbe85b5d03445532e2b7b37a9f8a48983beea238b332 FROM $UV_IMAGE AS uvbin +FROM $LITELLM_BUILD_IMAGE AS pgbouncer-builder +ARG PGBOUNCER_VERSION +ARG PGBOUNCER_SHA256 +USER root +RUN apk add --no-cache build-base pkgconf libevent-dev openssl-dev curl +WORKDIR /build +RUN curl -fsSL -o pgbouncer.tar.gz "https://www.pgbouncer.org/downloads/files/${PGBOUNCER_VERSION}/pgbouncer-${PGBOUNCER_VERSION}.tar.gz" && \ + echo "${PGBOUNCER_SHA256} pgbouncer.tar.gz" | sha256sum -c - && \ + tar xzf pgbouncer.tar.gz --strip-components=1 && \ + ./configure --prefix=/usr/local --with-openssl=/usr && \ + make -j"$(nproc)" pgbouncer && \ + install -m 0755 pgbouncer /usr/local/bin/pgbouncer + # Admin UI builder. Pinned to the build platform so the architecture-independent # Next.js static export compiles once natively even in a multi-arch build, # instead of once per target arch under QEMU. @@ -101,7 +117,8 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime USER root # node (without npm) is required by the prisma CLI at runtime -RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile +RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile libevent +COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer WORKDIR /app ENV PATH="/app/.venv/bin:${PATH}" \ diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index edf20e8bbff..5d729046678 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -7,9 +7,25 @@ ARG PROXY_EXTRAS_SOURCE=published ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a # Pinned by digest like the other base images; bump explicitly on Node upgrades. ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43 +# Checksum from https://www.pgbouncer.org/downloads/ (the Wolfi repo only carries 1.24.x) +ARG PGBOUNCER_VERSION=1.25.2 +ARG PGBOUNCER_SHA256=924ad35113fd0a71c8e2dbe85b5d03445532e2b7b37a9f8a48983beea238b332 FROM $UV_IMAGE AS uvbin +FROM $LITELLM_BUILD_IMAGE AS pgbouncer-builder +ARG PGBOUNCER_VERSION +ARG PGBOUNCER_SHA256 +USER root +RUN apk add --no-cache build-base pkgconf libevent-dev openssl-dev curl +WORKDIR /build +RUN curl -fsSL -o pgbouncer.tar.gz "https://www.pgbouncer.org/downloads/files/${PGBOUNCER_VERSION}/pgbouncer-${PGBOUNCER_VERSION}.tar.gz" && \ + echo "${PGBOUNCER_SHA256} pgbouncer.tar.gz" | sha256sum -c - && \ + tar xzf pgbouncer.tar.gz --strip-components=1 && \ + ./configure --prefix=/usr/local --with-openssl=/usr && \ + make -j"$(nproc)" pgbouncer && \ + install -m 0755 pgbouncer /usr/local/bin/pgbouncer + # Admin UI builder. Pinned to the build platform so the architecture-independent # Next.js static export compiles once natively even in a multi-arch build, # instead of once per target arch under QEMU. @@ -128,8 +144,9 @@ RUN for i in 1 2 3; do \ apk upgrade --no-cache && break || sleep 5; \ done && \ for i in 1 2 3; do \ - apk add --no-cache python-3.13 bash openssl tzdata libsndfile nodejs && break || sleep 5; \ + apk add --no-cache python-3.13 bash openssl tzdata libsndfile nodejs libevent && break || sleep 5; \ done +COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer # Copy only what runtime needs. The application is installed inside the venv; # the rest of the builder's /app is source and build metadata that must not diff --git a/gateway/Dockerfile b/gateway/Dockerfile index 308d70a6b26..33d3791dbba 100644 --- a/gateway/Dockerfile +++ b/gateway/Dockerfile @@ -1,9 +1,25 @@ ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a +# Checksum from https://www.pgbouncer.org/downloads/ (the Wolfi repo only carries 1.24.x) +ARG PGBOUNCER_VERSION=1.25.2 +ARG PGBOUNCER_SHA256=924ad35113fd0a71c8e2dbe85b5d03445532e2b7b37a9f8a48983beea238b332 FROM $UV_IMAGE AS uvbin +FROM $LITELLM_BUILD_IMAGE AS pgbouncer-builder +ARG PGBOUNCER_VERSION +ARG PGBOUNCER_SHA256 +USER root +RUN apk add --no-cache build-base pkgconf libevent-dev openssl-dev curl +WORKDIR /build +RUN curl -fsSL -o pgbouncer.tar.gz "https://www.pgbouncer.org/downloads/files/${PGBOUNCER_VERSION}/pgbouncer-${PGBOUNCER_VERSION}.tar.gz" && \ + echo "${PGBOUNCER_SHA256} pgbouncer.tar.gz" | sha256sum -c - && \ + tar xzf pgbouncer.tar.gz --strip-components=1 && \ + ./configure --prefix=/usr/local --with-openssl=/usr && \ + make -j"$(nproc)" pgbouncer && \ + install -m 0755 pgbouncer /usr/local/bin/pgbouncer + # ---------- Builder ---------- FROM $LITELLM_BUILD_IMAGE AS builder @@ -61,6 +77,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \ --extra bedrock-realtime \ --python python3.13 +# PYTHONPATH=/app makes the source tree shadow the installed package, so the +# compiled Rust extension must live next to the source or it is never imported. +RUN cp "$(python -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')"/litellm/rust_bridge/_native*.so litellm/rust_bridge/ + RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ npm_config_cache=/root/.npm \ prisma generate --schema=./schema.prisma @@ -73,7 +93,7 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime USER root RUN for i in 1 2 3; do \ - apk add --no-cache bash openssl tzdata python-3.13 libsndfile libatomic && break; \ + apk add --no-cache bash openssl tzdata python-3.13 libsndfile libatomic libevent && break; \ [ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \ sleep 5; \ done @@ -90,15 +110,17 @@ ENV HOME=/home/nonroot \ COPY --from=builder --chown=nonroot:nonroot /app /app COPY --from=builder /opt/prisma /opt/prisma +COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \ find /app/.venv -type d -path "*/tornado/test" -delete && \ chmod -R a+rX /opt/prisma && \ - python -c "from prisma.client import BINARY_PATHS; paths = list(BINARY_PATHS.query_engine.values()); assert paths and all(p.startswith('/opt/prisma/') for p in paths), paths" + python -c "from prisma.client import BINARY_PATHS; paths = list(BINARY_PATHS.query_engine.values()); assert paths and all(p.startswith('/opt/prisma/') for p in paths), paths" && \ + python -c "import litellm; from litellm.rust_bridge.loader import native_bridge_available; assert litellm.__file__ == '/app/litellm/__init__.py', litellm.__file__; assert native_bridge_available()" USER nonroot EXPOSE 4000/tcp -ENTRYPOINT ["sh", "-c", "exec /app/docker/component_entrypoint.sh uvicorn gateway.main:app --workers \"${NUM_WORKERS:-1}\" \"$@\"", "--"] +ENTRYPOINT ["sh", "-c", "exec /app/docker/component_entrypoint.sh python -m gateway.launch --workers \"${NUM_WORKERS:-1}\" \"$@\"", "--"] CMD ["--host", "0.0.0.0", "--port", "4000"] diff --git a/gateway/launch.py b/gateway/launch.py new file mode 100644 index 00000000000..6b28fafdcf6 --- /dev/null +++ b/gateway/launch.py @@ -0,0 +1,71 @@ +"""Gateway supervisor: assemble DATABASE_URL, start the in-container PgBouncer, then run uvicorn. + +``gateway/main.py`` assembles ``DATABASE_URL`` inside every uvicorn worker, which +is fine for a plain Postgres URL but not for the pooler: PgBouncer must be +started exactly once per pod, before the workers fork, and the workers must be +handed the loopback URL it listens on. A pre-existing ``DATABASE_URL`` wins in +``DatabaseURLSettings.apply_to_env`` under password auth, so setting it here is +enough for every worker to pick the pooled URL up unchanged. + +Run with: + python -m gateway.launch --workers 4 --host 0.0.0.0 --port 4000 +""" + +import os +import sys +from collections.abc import Callable, Mapping, Sequence +from typing import Final + +from uvicorn.main import main as uvicorn_main + +from litellm.proxy.db.db_url_settings import DatabaseURLSettings +from litellm.proxy.db.pgbouncer import PgBouncerError, PgBouncerSettings, start_in_container_pgbouncer + +GATEWAY_APP: Final = "gateway.main:app" +KEEPALIVE_FLAG: Final = "--timeout-keep-alive" + + +def uvicorn_argv(argv: Sequence[str], environ: Mapping[str, str]) -> tuple[str, ...]: + """Honor ``KEEPALIVE_TIMEOUT`` like ``proxy_cli.py`` does, unless the flag was passed explicitly.""" + keepalive: Final = environ.get("KEEPALIVE_TIMEOUT") + if keepalive is None or any(arg == KEEPALIVE_FLAG or arg.startswith(f"{KEEPALIVE_FLAG}=") for arg in argv): + return (GATEWAY_APP, *argv) + return (GATEWAY_APP, *argv, KEEPALIVE_FLAG, keepalive) + + +def pool_database_url( + settings: DatabaseURLSettings, + pgbouncer: PgBouncerSettings, + environ: Mapping[str, str], +) -> str | PgBouncerError | None: + """Start the in-container PgBouncer and return its loopback URL, or None when ``pgbouncer.enabled`` is off. + + The upstream URL is whatever ``apply_to_env`` assembled from the discrete + ``DATABASE_*`` vars (or an operator-pinned ``DATABASE_URL``). Token auth is + rejected by the pooler itself, since it holds one password for its lifetime. + """ + if not pgbouncer.enabled: + return None + upstream_url: Final = environ.get("DATABASE_URL") + if upstream_url is None: + return PgBouncerError("LITELLM_PGBOUNCER_ENABLED is set but no DATABASE_URL could be assembled") + return start_in_container_pgbouncer(pgbouncer, upstream_url, token_auth_enabled=settings.token_auth() is not None) + + +def _serve(argv: Sequence[str]) -> None: + uvicorn_main(tuple(argv), prog_name="uvicorn") + + +def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) -> None: + settings: Final = DatabaseURLSettings.from_env() + settings.apply_to_env() + pooled_url: Final = pool_database_url(settings, PgBouncerSettings(), os.environ) + if isinstance(pooled_url, PgBouncerError): + sys.exit(f"LiteLLM gateway: in-container pgbouncer could not start: {pooled_url.reason}") + if pooled_url is not None: + os.environ["DATABASE_URL"] = pooled_url + serve(uvicorn_argv(argv, os.environ)) + + +if __name__ == "__main__": + main(sys.argv[1:]) diff --git a/helm/litellm-helm/templates/deployment.yaml b/helm/litellm-helm/templates/deployment.yaml index f7c918a6827..834071eb9b2 100644 --- a/helm/litellm-helm/templates/deployment.yaml +++ b/helm/litellm-helm/templates/deployment.yaml @@ -117,6 +117,14 @@ spec: - name: DATABASE_URL_READ_REPLICA value: {{ .Values.db.readReplicaUrl | quote }} {{- end }} + {{- if .Values.db.connectionPool.enabled }} + - name: LITELLM_PGBOUNCER_ENABLED + value: "true" + - name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS + value: {{ .Values.db.connectionPool.maxDbConnections | quote }} + - name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN + value: {{ .Values.db.connectionPool.maxClientConn | quote }} + {{- end }} - name: PROXY_MASTER_KEY valueFrom: secretKeyRef: diff --git a/helm/litellm-helm/tests/connection_pool_tests.yaml b/helm/litellm-helm/tests/connection_pool_tests.yaml new file mode 100644 index 00000000000..af23512dafc --- /dev/null +++ b/helm/litellm-helm/tests/connection_pool_tests.yaml @@ -0,0 +1,61 @@ +suite: test in-container connection pool +templates: + - deployment.yaml + - configmap-litellm.yaml +tests: + - it: should not emit pgbouncer env vars by default + template: deployment.yaml + asserts: + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_PGBOUNCER_ENABLED + value: "true" + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS + value: "20" + + - it: should enable the pool with the default sizing when connectionPool.enabled is set + template: deployment.yaml + set: + db.connectionPool.enabled: true + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_PGBOUNCER_ENABLED + value: "true" + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS + value: "20" + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN + value: "1000" + + - it: should pass custom sizing through as strings next to the worker count + template: deployment.yaml + set: + numWorkers: 4 + db.connectionPool.enabled: true + db.connectionPool.maxDbConnections: 8 + db.connectionPool.maxClientConn: 400 + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS + value: "8" + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN + value: "400" + - contains: + path: spec.template.spec.containers[0].args + content: "4" diff --git a/helm/litellm-helm/values.yaml b/helm/litellm-helm/values.yaml index 907f85de8fc..db596e2d68e 100644 --- a/helm/litellm-helm/values.yaml +++ b/helm/litellm-helm/values.yaml @@ -355,6 +355,20 @@ db: # only (e.g. when IAM_TOKEN_DB_AUTH supplies the token at runtime). readReplicaUrl: "" + # In-container connection pool (PgBouncer, transaction mode) shared by every + # worker in the pod. Without it each --num_workers worker opens its own + # connection_limit connections to Postgres, so a pod's footprint against the + # database's connection ceiling is workers x connection_limit and grows with + # every replica. With it, the pod holds at most maxDbConnections upstream + # connections no matter how many workers run; the workers connect to the pool + # over loopback, with no extra network hop. Migrations still go straight to + # Postgres. Starting profile for numWorkers: 4 is maxDbConnections: 20, so + # a database with a 5000-connection ceiling fits roughly 200 replicas. + connectionPool: + enabled: false + maxDbConnections: 20 + maxClientConn: 1000 + # Use the Stackgres Helm chart to deploy an instance of a Stackgres cluster. # The Stackgres Operator must already be installed within the target # Kubernetes cluster. diff --git a/helm/litellm/templates/_helpers.tpl b/helm/litellm/templates/_helpers.tpl index c459512c7b9..4ad3cfe0484 100644 --- a/helm/litellm/templates/_helpers.tpl +++ b/helm/litellm/templates/_helpers.tpl @@ -360,6 +360,23 @@ harmless no-op for the Job and authoritative for the app pods. {{- end }} {{- end -}} +{{/* +In-container PgBouncer env for the gateway container. Fails at render time under IAM or Entra auth: the pooler holds one static password for the life of the pod. +*/}} +{{- define "litellm.connectionPoolEnv" -}} +{{- if or .Values.database.writer.useIAMAuth .Values.database.writer.useAzureEntraAuth }} +{{- fail "database.connectionPool.enabled cannot be combined with database.writer.useIAMAuth or database.writer.useAzureEntraAuth: the in-container pgbouncer holds a static database password and cannot follow a rotating token. Disable the pool or use a static database password" }} +{{- end }} +{{- with .Values.database.connectionPool -}} +- name: LITELLM_PGBOUNCER_ENABLED + value: "true" +- name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS + value: {{ required "database.connectionPool.maxDbConnections is required when the pool is enabled" .maxDbConnections | quote }} +- name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN + value: {{ required "database.connectionPool.maxClientConn is required when the pool is enabled" .maxClientConn | quote }} +{{- end }} +{{- end -}} + {{/* PodDisruptionBudget shared by gateway, backend, and ui. diff --git a/helm/litellm/templates/gateway/deployment.yaml b/helm/litellm/templates/gateway/deployment.yaml index 9cb6b07e77b..1e9041f0a33 100644 --- a/helm/litellm/templates/gateway/deployment.yaml +++ b/helm/litellm/templates/gateway/deployment.yaml @@ -61,6 +61,9 @@ spec: - name: NUM_WORKERS value: {{ .Values.gateway.numWorkers | quote }} {{- end }} + {{- if .Values.database.connectionPool.enabled }} + {{- include "litellm.connectionPoolEnv" $ | nindent 12 }} + {{- end }} {{- if .Values.billingMetrics.enabled }} {{- include "litellm.billingMetricsEnv" . | nindent 12 }} {{- end }} diff --git a/helm/litellm/tests/connection_pool_tests.yaml b/helm/litellm/tests/connection_pool_tests.yaml new file mode 100644 index 00000000000..31be7575f3c --- /dev/null +++ b/helm/litellm/tests/connection_pool_tests.yaml @@ -0,0 +1,117 @@ +suite: test in-container connection pool env vars +templates: + - gateway/deployment.yaml + - gateway/configmap.yaml + - backend/deployment.yaml + - backend/configmap.yaml +values: + - ./values/required.yaml +tests: + - it: renders no pool env by default + template: gateway/deployment.yaml + asserts: + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_PGBOUNCER_ENABLED + any: true + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS + any: true + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN + any: true + + - it: enabled pool renders the three pgbouncer vars with the configured sizes + template: gateway/deployment.yaml + set: + gateway.numWorkers: 4 + database.connectionPool.enabled: true + database.connectionPool.maxDbConnections: 8 + database.connectionPool.maxClientConn: 250 + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_PGBOUNCER_ENABLED + value: "true" + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS + value: "8" + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN + value: "250" + - contains: + path: spec.template.spec.containers[0].env + content: + name: NUM_WORKERS + value: "4" + + - it: enabled pool uses the chart default sizes + template: gateway/deployment.yaml + set: + database.connectionPool.enabled: true + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS + value: "20" + - contains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN + value: "1000" + + - it: backend never gets the pool env + template: backend/deployment.yaml + set: + database.connectionPool.enabled: true + asserts: + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_PGBOUNCER_ENABLED + any: true + + - it: pool with IAM auth fails at render time + template: gateway/deployment.yaml + set: + database.connectionPool.enabled: true + database.writer.useIAMAuth: true + asserts: + - failedTemplate: + errorMessage: "database.connectionPool.enabled cannot be combined with database.writer.useIAMAuth or database.writer.useAzureEntraAuth: the in-container pgbouncer holds a static database password and cannot follow a rotating token. Disable the pool or use a static database password" + + - it: pool with Entra auth fails at render time + template: gateway/deployment.yaml + set: + database.connectionPool.enabled: true + database.writer.useAzureEntraAuth: true + asserts: + - failedTemplate: + errorMessage: "database.connectionPool.enabled cannot be combined with database.writer.useIAMAuth or database.writer.useAzureEntraAuth: the in-container pgbouncer holds a static database password and cannot follow a rotating token. Disable the pool or use a static database password" + + - it: IAM auth without the pool still renders + template: gateway/deployment.yaml + set: + database.writer.useIAMAuth: true + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: IAM_TOKEN_DB_AUTH + value: "true" + - notContains: + path: spec.template.spec.containers[0].env + content: + name: LITELLM_PGBOUNCER_ENABLED + any: true diff --git a/helm/litellm/values.yaml b/helm/litellm/values.yaml index be1e6d43987..45ab5229f38 100644 --- a/helm/litellm/values.yaml +++ b/helm/litellm/values.yaml @@ -225,6 +225,26 @@ database: usernameKey: username passwordKey: password + # In-container connection pool (PgBouncer, transaction mode) shared by every + # gateway worker in the pod. Without it each of the `gateway.numWorkers` + # workers opens its own Prisma pool straight to Postgres, so a pod's + # footprint against the database's connection ceiling is + # numWorkers x connection_limit and grows with every replica. With it, the + # pod holds at most maxDbConnections upstream connections no matter how many + # workers run; the workers connect to the pool over loopback, with no extra + # network hop. The chart emits LITELLM_PGBOUNCER_ENABLED / + # LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS / LITELLM_PGBOUNCER_MAX_CLIENT_CONN on + # the gateway container only: the backend runs a single worker and the + # migrations Job must keep a direct connection. The pool holds a static + # password, so it cannot be combined with `database.writer.useIAMAuth` or + # `useAzureEntraAuth` (rendering fails). Starting profile for + # `gateway.numWorkers: 4` is maxDbConnections: 20, so a database with a + # 5000-connection ceiling fits roughly 200 gateway replicas. + connectionPool: + enabled: false + maxDbConnections: 20 + maxClientConn: 1000 + # Optional Redis. Leave host empty to disable. # # This is the proxy's coordination store: cross-pod tpm/rpm rate limits, spend diff --git a/litellm/proxy/db/pgbouncer.py b/litellm/proxy/db/pgbouncer.py new file mode 100644 index 00000000000..7792867600c --- /dev/null +++ b/litellm/proxy/db/pgbouncer.py @@ -0,0 +1,525 @@ +"""In-container PgBouncer shared by every proxy worker. + +Each uvicorn worker owns a Prisma query engine with its own pool of +``connection_limit`` server connections, so the connections a pod holds open +against Postgres scale as ``workers * connection_limit`` and a database with a +fixed connection ceiling runs out of room as pods and workers are added. + +When ``LITELLM_PGBOUNCER_ENABLED`` is set, the supervisor process starts one +PgBouncer next to the workers (no extra network hop: it listens on loopback +inside the pod) in transaction pooling mode, points ``DATABASE_URL`` at it +with ``pgbouncer=true`` so Prisma stops using server-side prepared statements, +and keeps it running for the life of the proxy. Every worker's pool then +becomes cheap client connections to PgBouncer while the upstream connection +count is capped at ``LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS`` per pod, no matter +how many workers run. + +Migrations and the schema diff run in the supervisor before the pooler is +started, so they always go straight to Postgres. ``DATABASE_URL_READ_REPLICA`` +is left untouched. + +The pooler holds the database password from startup, so it cannot be combined +with ``IAM_TOKEN_DB_AUTH`` or ``AZURE_POSTGRESQL_AUTH``: those rotate the +password inside every worker on their own schedule, and PgBouncer would keep +authenticating upstream with the expired token. +""" + +from __future__ import annotations + +import atexit +import os +import re +import shlex +import shutil +import socket +import subprocess +import tempfile +import threading +import time +import urllib.parse +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final + +from pydantic import Field +from pydantic_settings import BaseSettings, SettingsConfigDict + +from litellm._logging import verbose_proxy_logger +from litellm.proxy.db.token_auth import AZURE_POSTGRESQL_AUTH_ENV_VAR, IAM_TOKEN_DB_AUTH_ENV_VAR + +PGBOUNCER_ENV_PREFIX: Final = "LITELLM_PGBOUNCER_" +PGBOUNCER_LISTEN_ADDR: Final = "127.0.0.1" +PGBOUNCER_INI_NAME: Final = "pgbouncer.ini" +PGBOUNCER_USERLIST_NAME: Final = "userlist.txt" +PGBOUNCER_CA_NAME: Final = "server-ca.pem" +PGBOUNCER_RESTART_DELAY_SECONDS: Final = 1.0 +PGBOUNCER_READY_TIMEOUT_SECONDS: Final = 15.0 +PGBOUNCER_STOP_GRACE_SECONDS: Final = 10.0 +PGBOUNCER_UNPRIVILEGED_USER: Final = "nobody" +PGBOUNCER_MIN_VERSION: Final = (1, 19) +PGBOUNCER_VERSION_PATTERN: Final = re.compile(r"PgBouncer (\d+)\.(\d+)") +PGBOUNCER_LIST_DELIMITER_PATTERN: Final = re.compile(r"[,\s]") +PGBOUNCER_TOKEN_AUTH_CONFLICT: Final = ( + f"the in-container pgbouncer cannot be combined with {IAM_TOKEN_DB_AUTH_ENV_VAR} or " + f"{AZURE_POSTGRESQL_AUTH_ENV_VAR}: each worker rotates the database password on its own schedule and the pooler " + "would keep using the expired token upstream. Disable the pooler or use a static database password" +) + +# Prisma's client-side TLS params describe the hop to Postgres, which becomes +# PgBouncer's server side. They move into ``server_tls_*`` and must not stay on +# the loopback URL: the listener speaks plain TCP and Prisma would refuse it +# under ``sslmode=require`` or ``channel_binding=require``. +PRISMA_TLS_PARAM_KEYS: Final[frozenset[str]] = frozenset( + {"sslmode", "sslcert", "sslaccept", "sslidentity", "sslpassword", "channel_binding", "gssencmode"} +) +POOLED_URL_DROPPED_KEYS: Final[frozenset[str]] = PRISMA_TLS_PARAM_KEYS | frozenset(("options", "pgbouncer")) +PGBOUNCER_SSLMODES: Final[frozenset[str]] = frozenset( + {"disable", "allow", "prefer", "require", "verify-ca", "verify-full"} +) + + +class PgBouncerSettings(BaseSettings): + """``LITELLM_PGBOUNCER_*`` env vars, read once in the supervisor.""" + + model_config = SettingsConfigDict( + env_prefix=PGBOUNCER_ENV_PREFIX, case_sensitive=False, extra="ignore", frozen=True + ) + + enabled: bool = False + port: int = Field(default=6432, ge=1, le=65535) + max_db_connections: int = Field(default=20, ge=1) + max_client_conn: int = Field(default=1000, ge=1) + binary: str = "pgbouncer" + + +@dataclass(frozen=True, slots=True) +class PgBouncerPlan: + ini: str + userlist: str + pooled_url: str + ca_source: str | None = None + + +@dataclass(frozen=True, slots=True) +class PgBouncerError: + reason: str + + +def _single_quoted(value: str) -> str: + """Quote for SQL and for PgBouncer's ``[databases]`` connection string: both double a literal ``'``.""" + return "'" + value.replace("'", "''") + "'" + + +def _userlist_quote(value: str) -> str: + return '"' + value.replace('"', '""') + '"' + + +def _option_settings(tokens: Sequence[str]) -> tuple[str, ...] | None: + """The ``name=value`` settings in a libpq ``options`` string, or None if it holds anything else. + + Accepts ``-c name=value``, ``-cname=value`` and ``--name=value``; a + detached ``-c`` is folded into the token that follows it first. + """ + folded: Final = tuple( + f"-c{tokens[index + 1]}" if token == "-c" and index + 1 < len(tokens) else token + for index, token in enumerate(tokens) + if index == 0 or tokens[index - 1] != "-c" + ) + settings: Final = tuple(token[2:] for token in folded if token.startswith(("-c", "--")) and "=" in token[2:]) + return settings if len(settings) == len(folded) else None + + +def _connect_query(options: str) -> str | PgBouncerError: + """Turn Prisma's ``options=-c name=value ...`` startup param into ``SET`` statements. + + PgBouncer rejects any ``-c`` setting in ``options`` that is not one of the + handful it tracks (``statement_timeout`` and ``lock_timeout`` are not), so + the settings are applied to each new server connection instead. Every + client shares them, which is what the single ``DATABASE_URL`` gave anyway. + """ + settings: Final = _option_settings(tuple(shlex.split(options))) + if settings is None: + return PgBouncerError(f"cannot translate the DATABASE_URL options {options!r} into PgBouncer settings") + return "; ".join( + f"SET {name.strip()} TO {_single_quoted(value.strip())}" + for name, value in (setting.split("=", 1) for setting in settings) + ) + + +def _server_tls_settings(sslmode: str, sslcert: str, sslaccept: str, ca_path: Path) -> tuple[str, ...] | PgBouncerError: + """``server_tls_*`` lines naming ``ca_path``, the runtime-dir copy of the bundle: the original (or the + 0600 root pinned by ``pin_bundle_root``) is often unreadable for the user PgBouncer drops to.""" + if sslmode not in PGBOUNCER_SSLMODES: + return PgBouncerError(f"unsupported sslmode {sslmode!r} on DATABASE_URL") + verify: Final = sslmode in ("verify-ca", "verify-full") or (sslmode == "require" and sslaccept == "strict") + if verify and not sslcert: + return PgBouncerError( + "DATABASE_URL asks for a verified TLS connection but names no CA bundle; " + "add sslcert= (or sslrootcert=) so the in-container PgBouncer can verify Postgres" + ) + mode: Final = "verify-full" if verify else sslmode + return (f"server_tls_sslmode = {mode}", *((f"server_tls_ca_file = {ca_path}",) if sslcert else ())) + + +def plan_pgbouncer( + upstream_url: str, + settings: PgBouncerSettings, + runtime_dir: Path, + run_as_user: str | None, +) -> PgBouncerPlan | PgBouncerError: + """Render the PgBouncer config for ``upstream_url`` and the loopback URL Prisma uses instead. + + Params describing Prisma's own pool (``connection_limit``, ``pool_timeout``, + ...) stay on the pooled URL; the TLS params and ``options`` describe the hop + to Postgres and move into the PgBouncer config. ``run_as_user`` is the + unprivileged user PgBouncer drops to when the proxy runs as root, which + PgBouncer itself refuses to do. + """ + parsed: Final = urllib.parse.urlsplit(upstream_url) + params: Final[Mapping[str, str]] = MappingProxyType( + dict(urllib.parse.parse_qsl(parsed.query, keep_blank_values=True)) + ) + dbname: Final = urllib.parse.unquote(parsed.path.lstrip("/")) + username: Final = urllib.parse.unquote(parsed.username or "") + password: Final = None if parsed.password is None else urllib.parse.unquote(parsed.password) + if not parsed.hostname or not username or password is None or not dbname: + return PgBouncerError( + "DATABASE_URL must carry a host, user, password and database name for the in-container PgBouncer" + ) + if PGBOUNCER_LIST_DELIMITER_PATTERN.search(username): + return PgBouncerError( + f"the database user {username!r} cannot be named in PgBouncer's stats_users list: " + "PgBouncer splits list settings on commas and whitespace and has no quoting for them" + ) + if "sslidentity" in params: + return PgBouncerError("client certificates (sslidentity) are not supported with the in-container PgBouncer") + tls: Final = _server_tls_settings( + params.get("sslmode", "prefer"), + params.get("sslcert", ""), + params.get("sslaccept", ""), + runtime_dir / PGBOUNCER_CA_NAME, + ) + if isinstance(tls, PgBouncerError): + return tls + connect_query: Final = _connect_query(params["options"]) if params.get("options") else "" + if isinstance(connect_query, PgBouncerError): + return connect_query + upstream: Final = " ".join( + ( + f"host={_single_quoted(parsed.hostname)}", + f"port={parsed.port or 5432}", + f"dbname={_single_quoted(dbname)}", + f"user={_single_quoted(username)}", + f"password={_single_quoted(password)}", + *((f"connect_query={_single_quoted(connect_query)}",) if connect_query else ()), + ) + ) + ini: Final = "\n".join( + ( + "[databases]", + f"{dbname} = {upstream}", + "", + "[pgbouncer]", + f"listen_addr = {PGBOUNCER_LISTEN_ADDR}", + f"listen_port = {settings.port}", + f"unix_socket_dir = {runtime_dir}", + f"auth_file = {runtime_dir / PGBOUNCER_USERLIST_NAME}", + "auth_type = scram-sha-256", + f"stats_users = {username}", + "pool_mode = transaction", + f"max_client_conn = {settings.max_client_conn}", + f"default_pool_size = {settings.max_db_connections}", + f"max_db_connections = {settings.max_db_connections}", + "ignore_startup_parameters = extra_float_digits", + *tls, + *((f"user = {run_as_user}",) if run_as_user else ()), + "", + ) + ) + userlist: Final = f"{_userlist_quote(username)} {_userlist_quote(password)}\n" + pooled_query: Final = urllib.parse.urlencode( + (*((key, value) for key, value in params.items() if key not in POOLED_URL_DROPPED_KEYS), ("pgbouncer", "true")) + ) + credentials: Final = f"{urllib.parse.quote(username, safe='')}:{urllib.parse.quote(password, safe='')}" + pooled_url: Final = urllib.parse.urlunsplit( + parsed._replace(netloc=f"{credentials}@{PGBOUNCER_LISTEN_ADDR}:{settings.port}", query=pooled_query) + ) + return PgBouncerPlan(ini=ini, userlist=userlist, pooled_url=pooled_url, ca_source=params.get("sslcert") or None) + + +def write_pgbouncer_files(plan: PgBouncerPlan, runtime_dir: Path, run_as_user: str | None) -> Path | PgBouncerError: + """Write the ini, userlist (both hold the password, so mode 0600) and CA copy, and return the ini path. + + ``run_as_user`` is the user PgBouncer drops to when started as root; it has + to own the files it re-reads on reload and the socket directory. + """ + ini_path: Final = runtime_dir / PGBOUNCER_INI_NAME + userlist_path: Final = runtime_dir / PGBOUNCER_USERLIST_NAME + ca_path: Final = runtime_dir / PGBOUNCER_CA_NAME + if plan.ca_source is not None: + try: + shutil.copyfile(plan.ca_source, ca_path) + except OSError as error: + return PgBouncerError(f"cannot read the CA bundle {plan.ca_source!r} named by sslcert: {error}") + for path, content in ((userlist_path, plan.userlist), (ini_path, plan.ini)): + path.touch(mode=0o600) + path.write_text(content, encoding="utf-8") + if run_as_user is not None: + runtime_dir.chmod(0o700) + for path in (runtime_dir, ini_path, userlist_path, *((ca_path,) if plan.ca_source is not None else ())): + shutil.chown(path, user=run_as_user) + return ini_path + + +def _port_open(port: int) -> bool: + try: + with socket.create_connection((PGBOUNCER_LISTEN_ADDR, port), timeout=0.5): + return True + except OSError: + return False + + +def _unix_socket_open(path: Path) -> bool: + with socket.socket(socket.AF_UNIX) as probe: + probe.settimeout(0.5) + try: + probe.connect(str(path)) + except OSError: + return False + return True + + +def unix_socket_path(runtime_dir: Path, port: int) -> Path: + return runtime_dir / f".s.PGSQL.{port}" + + +def pgbouncer_version(binary: str) -> tuple[int, int] | PgBouncerError: + """``(major, minor)`` from `` --version``. + + Readiness relies on PgBouncer exiting when it cannot bind its TCP port, + which it does from 1.19 on. Older releases log a warning and serve the unix + socket alone, so their socket would vouch for a port held by someone else. + """ + try: + output: Final = subprocess.run( + (binary, "--version"), capture_output=True, text=True, check=False, timeout=10 + ).stdout + except (OSError, subprocess.TimeoutExpired) as run_error: + return PgBouncerError(f"could not run {binary!r} --version: {run_error}") + found: Final = PGBOUNCER_VERSION_PATTERN.search(output) + if found is None: + return PgBouncerError(f"{binary!r} --version did not report a PgBouncer version: {output.strip()!r}") + return int(found[1]), int(found[2]) + + +def _end(process: subprocess.Popen[bytes]) -> None: + if process.poll() is not None: + return + process.terminate() + try: + process.wait(timeout=PGBOUNCER_STOP_GRACE_SECONDS) + except subprocess.TimeoutExpired: + process.kill() + process.wait() + + +class PgBouncerProcess: + """Runs ``argv`` as a foreground child and restarts it whenever it exits on its own. + + Prisma reconnects by itself after a failed query, so a PgBouncer crash + costs the requests in flight plus one failed query per idle pooled + connection the crash severed, and nothing else once the replacement is + listening again. A replacement that cannot be spawned, finds its port + taken, exits again or never starts listening is retried every + ``restart_delay_seconds`` until ``stop`` is called. + + A connect probe of ``port`` cannot tell the child from another process + that grabbed the port after the availability check, so readiness also + needs ``socket_path``: the unix socket PgBouncer creates in the private + runtime directory, which it only does once every TCP listener is bound + (PgBouncer 1.19 or newer, see ``pgbouncer_version``). + """ + + def __init__( + self, + argv: Sequence[str], + port: int, + socket_path: Path, + restart_delay_seconds: float = PGBOUNCER_RESTART_DELAY_SECONDS, + ready_timeout_seconds: float = PGBOUNCER_READY_TIMEOUT_SECONDS, + ) -> None: + self.argv: Final = tuple(argv) + self.port: Final = port + self.socket_path: Final = socket_path + self.restart_delay_seconds: Final = restart_delay_seconds + self.ready_timeout_seconds: Final = ready_timeout_seconds + self._stopping: Final = threading.Event() + self._lock: Final = threading.Lock() + self._process: subprocess.Popen[bytes] | None = None + + @property + def pid(self) -> int | None: + with self._lock: + return None if self._process is None else self._process.pid + + def _spawn(self) -> subprocess.Popen[bytes] | PgBouncerError | None: + """Start a child, or None once ``stop`` ran; both take the lock so no child can slip in after a stop. + + The port has to be free first: a listener that is already there would + pass the readiness check while the child fails to bind. + """ + with self._lock: + if self._stopping.is_set(): + return None + if _port_open(self.port): + return PgBouncerError(f"{PGBOUNCER_LISTEN_ADDR}:{self.port} is already in use by another process") + try: + process: Final = subprocess.Popen(self.argv) + except OSError as spawn_error: + return PgBouncerError(f"could not start {self.argv[0]!r}: {spawn_error}") + self._process = process + return process + + def _wait_ready(self, process: subprocess.Popen[bytes]) -> PgBouncerError | None: + deadline: Final = time.monotonic() + self.ready_timeout_seconds + while time.monotonic() < deadline: + if process.poll() is not None: + return PgBouncerError(f"pgbouncer exited with status {process.returncode} during startup") + if _port_open(self.port) and _unix_socket_open(self.socket_path): + return None + time.sleep(0.1) + if _port_open(self.port): + return PgBouncerError( + f"{PGBOUNCER_LISTEN_ADDR}:{self.port} is served by another process, not the pgbouncer that was started" + ) + return PgBouncerError( + f"pgbouncer did not start listening on {PGBOUNCER_LISTEN_ADDR}:{self.port} " + f"within {self.ready_timeout_seconds:.0f}s" + ) + + def start(self) -> PgBouncerError | None: + """Spawn PgBouncer, wait until it listens on port and unix socket, then supervise it from a daemon thread.""" + process: Final = self._spawn() + if process is None: + return PgBouncerError("pgbouncer was stopped before it started") + if isinstance(process, PgBouncerError): + return process + not_ready: Final = self._wait_ready(process) + if not_ready is not None: + self.stop() + return not_ready + self._watch(process) + return None + + def _watch(self, process: subprocess.Popen[bytes]) -> None: + threading.Thread( + target=self._supervise, args=(process,), daemon=True, name="litellm-pgbouncer-supervisor" + ).start() + + def _supervise(self, process: subprocess.Popen[bytes]) -> None: + status: Final = process.wait() + if self._stopping.is_set(): + return + verbose_proxy_logger.error( + "In-container pgbouncer (pid %s) exited with status %s; restarting in %.1fs.", + process.pid, + status, + self.restart_delay_seconds, + ) + self._restart_after_delay() + + def _restart_after_delay(self) -> None: + time.sleep(self.restart_delay_seconds) + process: Final = self._spawn() + if process is None: + return + if isinstance(process, PgBouncerError): + self._retry_restart(process.reason) + return + not_ready: Final = self._wait_ready(process) + if not_ready is None: + self._watch(process) + return + _end(process) + self._retry_restart(not_ready.reason) + + def _retry_restart(self, reason: str) -> None: + if self._stopping.is_set(): + return + verbose_proxy_logger.error( + "In-container pgbouncer could not be restarted (%s); retrying in %.1fs.", reason, self.restart_delay_seconds + ) + threading.Thread(target=self._restart_after_delay, daemon=True, name="litellm-pgbouncer-supervisor").start() + + def stop(self) -> None: + with self._lock: + self._stopping.set() + process: Final = self._process + if process is not None: + _end(process) + + +def _only_in_this_process(action: Callable[[], None]) -> Callable[[], None]: + """An exit hook that does nothing in a forked child, which inherits the parent's ``atexit`` table.""" + owner_pid: Final = os.getpid() + + def run() -> None: + if os.getpid() == owner_pid: + action() + + return run + + +def start_in_container_pgbouncer( + settings: PgBouncerSettings, + upstream_url: str, + token_auth_enabled: bool = False, + register_exit_hook: Callable[[Callable[[], None]], object] = atexit.register, +) -> str | PgBouncerError: + """Start the pooler for ``upstream_url`` and return the loopback URL the workers must use. + + The pooler lives as long as this process: it is stopped from the exit hooks + once the worker manager has returned, and only by the process that started + it (gunicorn forks its workers, so they carry the hooks too). PgBouncer + refuses to run as root, so a root proxy (the default image) has it drop to + ``nobody``. + """ + if token_auth_enabled: + return PgBouncerError(PGBOUNCER_TOKEN_AUTH_CONFLICT) + version: Final = pgbouncer_version(settings.binary) + if isinstance(version, PgBouncerError): + return version + if version < PGBOUNCER_MIN_VERSION: + return PgBouncerError( + f"PgBouncer {version[0]}.{version[1]} keeps running after failing to bind its TCP port, so the proxy " + f"cannot tell it apart from another listener; {PGBOUNCER_MIN_VERSION[0]}.{PGBOUNCER_MIN_VERSION[1]} " + "or newer is required" + ) + runtime_dir: Final = Path(tempfile.mkdtemp(prefix="litellm-pgbouncer-")) + register_exit_hook(_only_in_this_process(lambda: shutil.rmtree(runtime_dir, ignore_errors=True))) + run_as_user: Final = PGBOUNCER_UNPRIVILEGED_USER if os.geteuid() == 0 else None + plan: Final = plan_pgbouncer(upstream_url, settings, runtime_dir, run_as_user) + if isinstance(plan, PgBouncerError): + return plan + ini_path: Final = write_pgbouncer_files(plan, runtime_dir, run_as_user) + if isinstance(ini_path, PgBouncerError): + return ini_path + pooler: Final = PgBouncerProcess( + argv=(settings.binary, str(ini_path)), + port=settings.port, + socket_path=unix_socket_path(runtime_dir, settings.port), + ) + failed: Final = pooler.start() + if failed is not None: + return failed + register_exit_hook(_only_in_this_process(pooler.stop)) + verbose_proxy_logger.info( + "In-container pgbouncer (pid %s) listening on %s:%s; capping this pod at %s upstream database connections.", + pooler.pid, + PGBOUNCER_LISTEN_ADDR, + settings.port, + settings.max_db_connections, + ) + return plan.pooled_url diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index e245367b1b4..16d76ff0415 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -18,6 +18,7 @@ from pydantic import BaseModel, ConfigDict import litellm from litellm.constants import DEFAULT_NUM_WORKERS_LITELLM_PROXY +from litellm.proxy.db.pgbouncer import PgBouncerError, PgBouncerSettings, start_in_container_pgbouncer from litellm.proxy.db.query_engine_reaper import start_query_engine_reaper if TYPE_CHECKING: @@ -1377,6 +1378,21 @@ def run_server( print( f"Unable to connect to DB. DATABASE_URL found in environment, but prisma package not found." # noqa: F541 ) + pgbouncer_settings: Final = PgBouncerSettings() + upstream_database_url: Final = os.getenv("DATABASE_URL") + if pgbouncer_settings.enabled and upstream_database_url is not None: + pooled_database_url: Final = start_in_container_pgbouncer( + pgbouncer_settings, upstream_database_url, token_auth_enabled=wants_rds_iam or wants_azure_entra + ) + if isinstance(pooled_database_url, PgBouncerError): + print( + f"\033[1;31mLiteLLM Proxy: LITELLM_PGBOUNCER_ENABLED is set but the in-container pgbouncer " + f"could not start: {pooled_database_url.reason}\033[0m", + file=sys.stderr, + flush=True, + ) + sys.exit(1) + os.environ["DATABASE_URL"] = pooled_database_url if port == 4000 and ProxyInitializationHelpers._is_port_in_use(port): port = random.randint(1024, 49152) if prometheus_metrics_port == port: diff --git a/terraform/litellm/aws/README.md b/terraform/litellm/aws/README.md index 6d3e0269a15..986428151e0 100644 --- a/terraform/litellm/aws/README.md +++ b/terraform/litellm/aws/README.md @@ -258,6 +258,39 @@ gateway_metrics_port = 4001 gateway_metrics_scrape_cidrs = ["10.0.0.0/16"] ``` +### In-container connection pool + +Each of the `gateway_num_workers` uvicorn workers opens its own Prisma pool +straight to Postgres, so one task holds `workers x connection_limit` +connections and the fleet's footprint against the database ceiling grows with +every task. `gateway_connection_pool_enabled` runs a PgBouncer (transaction +mode, loopback) inside the gateway container that all workers share, capping +the task at `gateway_pool_max_db_connections` upstream connections however +many workers it runs. `gateway_pool_max_client_conn` bounds the worker-side +connections the pooler accepts. The module sets +`LITELLM_PGBOUNCER_ENABLED`, `LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS` and +`LITELLM_PGBOUNCER_MAX_CLIENT_CONN` on the gateway container only; the backend +and the migration task keep the direct connection. + +```hcl +create_database = false +database_url = "postgresql://litellm:@db.internal:5432/litellm" +gateway_num_workers = 4 +gateway_connection_pool_enabled = true +gateway_pool_max_db_connections = 20 +gateway_pool_max_client_conn = 1000 +``` + +The pool needs a static database password, so it is only valid with an +existing database via `database_url`. The module-created Aurora authenticates +with rotating IAM tokens (see [Aurora + IAM auth](#aurora--iam-auth)), which +the pooler cannot follow, and `terraform plan` rejects that combination. + +The componentized `gateway_image` starts through `python -m gateway.launch`, +which reads these variables, starts the pooler once per task and hands the +workers its loopback URL; the classic `litellm` image honours them the same +way. + ### Scaling the gateway on requests and tokens By default the gateway service target-tracks CPU (`gateway_cpu_target`) and diff --git a/terraform/litellm/aws/ecs.tf b/terraform/litellm/aws/ecs.tf index aa2c3d558e2..149842c14ff 100644 --- a/terraform/litellm/aws/ecs.tf +++ b/terraform/litellm/aws/ecs.tf @@ -213,6 +213,12 @@ locals { # otherwise we keep the image's ENTRYPOINT and only override `command`. gateway_uvicorn_args = "--host 0.0.0.0 --port 4000 --workers ${var.gateway_num_workers}" + gateway_pool_env = var.gateway_connection_pool_enabled ? [ + { name = "LITELLM_PGBOUNCER_ENABLED", value = "true" }, + { name = "LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS", value = tostring(var.gateway_pool_max_db_connections) }, + { name = "LITELLM_PGBOUNCER_MAX_CLIENT_CONN", value = tostring(var.gateway_pool_max_client_conn) }, + ] : [] + metrics_enabled = var.gateway_metrics_port != null metrics_multiproc_dir = "/tmp/litellm_prometheus_multiproc" metrics_volume = "prometheus-multiproc" @@ -253,7 +259,7 @@ locals { backend_uvicorn_args = "--host 0.0.0.0 --port 4001" - gateway_launch_cmd = "case \"$USE_DDTRACE\" in [Tt][Rr][Uu][Ee]) export DD_TRACE_OPENAI_ENABLED=\"False\"; exec ddtrace-run uvicorn gateway.main:app ${local.gateway_uvicorn_args};; *) exec uvicorn gateway.main:app ${local.gateway_uvicorn_args};; esac" + gateway_launch_cmd = "case \"$USE_DDTRACE\" in [Tt][Rr][Uu][Ee]) export DD_TRACE_OPENAI_ENABLED=\"False\"; exec ddtrace-run python -m gateway.launch ${local.gateway_uvicorn_args};; *) exec python -m gateway.launch ${local.gateway_uvicorn_args};; esac" backend_launch_cmd = "case \"$USE_DDTRACE\" in [Tt][Rr][Uu][Ee]) export DD_TRACE_OPENAI_ENABLED=\"False\"; exec ddtrace-run uvicorn backend.main:app ${local.backend_uvicorn_args};; *) exec uvicorn backend.main:app ${local.backend_uvicorn_args};; esac" gateway_proxy_overrides = local.proxy_config_enabled ? { @@ -298,6 +304,11 @@ resource "aws_ecs_task_definition" "gateway" { ) error_message = "billing_metrics_client_cert_pem and billing_metrics_client_key_pem are both required when billing_metrics_endpoint is set." } + + precondition { + condition = !var.gateway_connection_pool_enabled || local.byo_database + error_message = "gateway_connection_pool_enabled requires an existing database via database_url with create_database = false: the module-created Aurora authenticates with IAM tokens, which the in-container pgbouncer cannot follow because it holds a static database password." + } } family = "${local.name}-gateway" @@ -323,6 +334,7 @@ resource "aws_ecs_task_definition" "gateway" { local.gateway_extra_env_list, local.proxy_config_env, local.metrics_env, + local.gateway_pool_env, ) secrets = concat(local.shared_secrets, local.gateway_extra_secrets_list) mountPoints = local.metrics_mount_points diff --git a/terraform/litellm/aws/tests/connection_pool.tftest.hcl b/terraform/litellm/aws/tests/connection_pool.tftest.hcl new file mode 100644 index 00000000000..38aa7051134 --- /dev/null +++ b/terraform/litellm/aws/tests/connection_pool.tftest.hcl @@ -0,0 +1,123 @@ +# Plan-only coverage for the in-container PgBouncer knobs on the gateway task. +# `mock_provider` keeps this offline: no AWS credentials, no API calls, no +# resources. Run from terraform/litellm/aws with `terraform test`. + +mock_provider "aws" { + mock_data "aws_iam_policy_document" { + defaults = { + json = "{\"Version\":\"2012-10-17\",\"Statement\":[]}" + } + } +} +mock_provider "random" {} + +variables { + region = "us-east-1" + tenant = "acme" + env = "test" + azs = ["us-east-1a", "us-east-1b"] + allow_plaintext_alb = true +} + +run "pool_off_by_default" { + command = plan + + assert { + condition = length(local.gateway_pool_env) == 0 + error_message = "The gateway must get no LITELLM_PGBOUNCER_* env unless gateway_connection_pool_enabled is set." + } +} + +run "pool_enabled_renders_the_three_vars_with_configured_sizes" { + command = plan + + variables { + create_database = false + database_url = "postgresql://litellm:pw@db.internal:5432/litellm" + gateway_num_workers = 4 + gateway_connection_pool_enabled = true + gateway_pool_max_db_connections = 8 + gateway_pool_max_client_conn = 250 + } + + assert { + condition = alltrue([ + length(local.gateway_pool_env) == 3, + local.gateway_pool_env[0].name == "LITELLM_PGBOUNCER_ENABLED" && local.gateway_pool_env[0].value == "true", + local.gateway_pool_env[1].name == "LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS" && local.gateway_pool_env[1].value == "8", + local.gateway_pool_env[2].name == "LITELLM_PGBOUNCER_MAX_CLIENT_CONN" && local.gateway_pool_env[2].value == "250", + ]) + error_message = "The pool env must carry the enabled flag and the configured sizes as strings." + } +} + +run "pool_enabled_uses_the_module_default_sizes" { + command = plan + + variables { + create_database = false + database_url = "postgresql://litellm:pw@db.internal:5432/litellm" + gateway_connection_pool_enabled = true + } + + assert { + condition = alltrue([ + length(local.gateway_pool_env) == 3, + local.gateway_pool_env[1].name == "LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS" && local.gateway_pool_env[1].value == "20", + local.gateway_pool_env[2].name == "LITELLM_PGBOUNCER_MAX_CLIENT_CONN" && local.gateway_pool_env[2].value == "1000", + ]) + error_message = "The pool env must fall back to the module defaults of 20 upstream and 1000 client connections." + } +} + +run "gateway_starts_through_the_pool_aware_launcher" { + command = plan + + variables { + gateway_num_workers = 4 + } + + assert { + condition = alltrue([ + strcontains(local.gateway_launch_cmd, "exec python -m gateway.launch --host 0.0.0.0 --port 4000 --workers 4"), + strcontains(local.gateway_launch_cmd, "exec ddtrace-run python -m gateway.launch --host 0.0.0.0 --port 4000 --workers 4"), + !strcontains(local.gateway_launch_cmd, "uvicorn gateway.main:app"), + local.gateway_proxy_overrides.command[0] == local.gateway_launch_cmd, + ]) + error_message = "The gateway must start through gateway.launch (with and without ddtrace) so the pooler starts once before uvicorn forks the workers." + } +} + +run "pool_with_module_created_iam_aurora_fails_at_plan" { + command = plan + + variables { + gateway_connection_pool_enabled = true + } + + expect_failures = [ + aws_ecs_task_definition.gateway, + ] +} + +run "pool_without_any_database_fails_at_plan" { + command = plan + + variables { + create_database = false + gateway_connection_pool_enabled = true + } + + expect_failures = [ + aws_ecs_task_definition.gateway, + ] +} + +run "module_created_iam_aurora_without_the_pool_still_plans" { + command = plan + + assert { + condition = contains(local.managed_db_env, { name = "IAM_TOKEN_DB_AUTH", value = "true" }) + error_message = "Without the pool the module-created Aurora must keep IAM token auth." + } +} diff --git a/terraform/litellm/aws/variables.tf b/terraform/litellm/aws/variables.tf index 10a9b392673..ec8363dbbaf 100644 --- a/terraform/litellm/aws/variables.tf +++ b/terraform/litellm/aws/variables.tf @@ -200,6 +200,45 @@ variable "gateway_num_workers" { } } +variable "gateway_connection_pool_enabled" { + description = <<-EOT + Run an in-container PgBouncer (transaction mode, loopback) in each gateway + task, shared by every uvicorn worker. Without it each of the + `gateway_num_workers` workers opens its own Prisma pool straight to + Postgres, so a task's footprint against the database connection ceiling is + workers x connection_limit and grows with every task. Sets + LITELLM_PGBOUNCER_ENABLED / LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS / + LITELLM_PGBOUNCER_MAX_CLIENT_CONN on the gateway container only. Requires + an existing database via `database_url`: the module-created Aurora + authenticates with IAM tokens, which the pooler cannot follow because it + holds one static password for the life of the task. + EOT + type = bool + default = false +} + +variable "gateway_pool_max_db_connections" { + description = "Upstream Postgres connections one gateway task may hold when gateway_connection_pool_enabled is set, regardless of gateway_num_workers. 20 suits 4 workers; a 5000-connection database then fits roughly 200 tasks." + type = number + default = 20 + + validation { + condition = var.gateway_pool_max_db_connections >= 1 + error_message = "gateway_pool_max_db_connections must be >= 1." + } +} + +variable "gateway_pool_max_client_conn" { + description = "Client connections the in-container PgBouncer accepts from the gateway workers when gateway_connection_pool_enabled is set." + type = number + default = 1000 + + validation { + condition = var.gateway_pool_max_client_conn >= 1 + error_message = "gateway_pool_max_client_conn must be >= 1." + } +} + variable "backend_cpu" { description = "Fargate CPU units for the backend task (1024 = 1 vCPU)." type = number diff --git a/terraform/litellm/gcp/README.md b/terraform/litellm/gcp/README.md index 66095ec0108..9c71d3e15b7 100644 --- a/terraform/litellm/gcp/README.md +++ b/terraform/litellm/gcp/README.md @@ -287,6 +287,40 @@ the gateway on GKE with the Helm chart's `targetTokensPerSecond` (see "Dependencies only" below) rather than wiring the counter into Cloud Monitoring, which the autoscaler would ignore +### In-container connection pool + +Each of the `gateway_num_workers` uvicorn workers opens its own Prisma pool +straight to Cloud SQL, so one instance holds `workers x connection_limit` +connections and the fleet's footprint against the database ceiling grows with +every instance Cloud Run adds. `gateway_connection_pool_enabled` runs a +PgBouncer (transaction mode, loopback) inside the gateway container that all +workers share, capping the instance at `gateway_pool_max_db_connections` +upstream connections however many workers it runs. +`gateway_pool_max_client_conn` bounds the worker-side connections the pooler +accepts. The module sets `LITELLM_PGBOUNCER_ENABLED`, +`LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS` and `LITELLM_PGBOUNCER_MAX_CLIENT_CONN` +on the gateway service only; the backend service and the migrations job keep +the direct connection + +```hcl +gateway_num_workers = 4 +gateway_connection_pool_enabled = true +gateway_pool_max_db_connections = 20 +gateway_pool_max_client_conn = 1000 +``` + +The pooler holds one static database password for the life of the instance. +This stack authenticates to Cloud SQL with the Secret Manager password (see +[Database authentication](#database-authentication)), so nothing else is +needed; a Cloud SQL Auth Proxy sidecar with IAM auth would not work with the +pool + +The gateway container starts through `python -m gateway.launch` (the +componentized image's own entrypoint) rather than `uvicorn` directly. The +launcher reads these variables, starts the pooler once per instance before +uvicorn forks the workers and hands them its loopback `DATABASE_URL`. It also +honours `KEEPALIVE_TIMEOUT` from `gateway_extra_env` the way the image does + ## Tenant deployment Every resource the stack creates is named `${tenant}-litellm-${env}` (or diff --git a/terraform/litellm/gcp/cloudrun.tf b/terraform/litellm/gcp/cloudrun.tf index 405893b132b..78d4ffa2152 100644 --- a/terraform/litellm/gcp/cloudrun.tf +++ b/terraform/litellm/gcp/cloudrun.tf @@ -138,10 +138,16 @@ locals { "export DATABASE_URL_READ_REPLICA=\"postgresql://$${DATABASE_USER}:$${DATABASE_PASSWORD}@$${DATABASE_HOST_READ_REPLICA}:$${DATABASE_PORT_READ_REPLICA}/$${DATABASE_NAME}\"", ] + gateway_pool_env = var.gateway_connection_pool_enabled ? [ + { name = "LITELLM_PGBOUNCER_ENABLED", value = "true" }, + { name = "LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS", value = tostring(var.gateway_pool_max_db_connections) }, + { name = "LITELLM_PGBOUNCER_MAX_CLIENT_CONN", value = tostring(var.gateway_pool_max_client_conn) }, + ] : [] + gateway_uvicorn_args = "--host 0.0.0.0 --port 4000 --workers ${var.gateway_num_workers}" backend_uvicorn_args = "--host 0.0.0.0 --port 4001" - gateway_launch_cmd = "case \"$USE_DDTRACE\" in [Tt][Rr][Uu][Ee]) export DD_TRACE_OPENAI_ENABLED=\"False\"; exec ddtrace-run uvicorn gateway.main:app ${local.gateway_uvicorn_args};; *) exec uvicorn gateway.main:app ${local.gateway_uvicorn_args};; esac" + gateway_launch_cmd = "case \"$USE_DDTRACE\" in [Tt][Rr][Uu][Ee]) export DD_TRACE_OPENAI_ENABLED=\"False\"; exec ddtrace-run python -m gateway.launch ${local.gateway_uvicorn_args};; *) exec python -m gateway.launch ${local.gateway_uvicorn_args};; esac" backend_launch_cmd = "case \"$USE_DDTRACE\" in [Tt][Rr][Uu][Ee]) export DD_TRACE_OPENAI_ENABLED=\"False\"; exec ddtrace-run uvicorn backend.main:app ${local.backend_uvicorn_args};; *) exec uvicorn backend.main:app ${local.backend_uvicorn_args};; esac" gateway_args = join(" && ", concat( @@ -229,7 +235,7 @@ resource "google_cloud_run_v2_service" "gateway" { } dynamic "env" { - for_each = concat(local.shared_env_kv, local.gateway_otel_env_kv, local.billing_metrics_env_kv, local.gateway_extra_env_kv, local.proxy_config_env, local.metrics_env_kv) + for_each = concat(local.shared_env_kv, local.gateway_otel_env_kv, local.billing_metrics_env_kv, local.gateway_extra_env_kv, local.proxy_config_env, local.metrics_env_kv, local.gateway_pool_env) content { name = env.value.name value = env.value.value diff --git a/terraform/litellm/gcp/tests/connection_pool.tftest.hcl b/terraform/litellm/gcp/tests/connection_pool.tftest.hcl new file mode 100644 index 00000000000..439fe6b1d71 --- /dev/null +++ b/terraform/litellm/gcp/tests/connection_pool.tftest.hcl @@ -0,0 +1,135 @@ +# Plan-only coverage for the in-container PgBouncer knobs on the gateway +# service. `mock_provider` keeps this offline: no GCP credentials, no API +# calls, no resources. Run from terraform/litellm/gcp with `terraform test`. + +mock_provider "google" { + mock_resource "google_redis_instance" { + defaults = { + host = "10.0.0.4" + port = 6379 + server_ca_certs = [{ + cert = "-----BEGIN CERTIFICATE-----\nmock\n-----END CERTIFICATE-----" + }] + } + } +} + +mock_provider "google-beta" {} +mock_provider "random" {} + +variables { + project_id = "test-project" + tenant = "tenant" + env = "test" + allow_plaintext_lb = true + image_registry = "us-central1-docker.pkg.dev/test-project/litellm" +} + +run "pool_off_by_default" { + command = plan + + assert { + condition = length(local.gateway_pool_env) == 0 + error_message = "The gateway must get no LITELLM_PGBOUNCER_* env unless gateway_connection_pool_enabled is set." + } + + assert { + condition = !anytrue([ + for e in google_cloud_run_v2_service.gateway[0].template[0].containers[0].env : startswith(e.name, "LITELLM_PGBOUNCER_") + ]) + error_message = "The gateway service must carry no LITELLM_PGBOUNCER_* env by default." + } +} + +run "pool_enabled_renders_the_three_vars_with_configured_sizes" { + command = plan + + variables { + gateway_num_workers = 4 + gateway_connection_pool_enabled = true + gateway_pool_max_db_connections = 8 + gateway_pool_max_client_conn = 250 + } + + assert { + condition = alltrue([ + length(local.gateway_pool_env) == 3, + local.gateway_pool_env[0].name == "LITELLM_PGBOUNCER_ENABLED" && local.gateway_pool_env[0].value == "true", + local.gateway_pool_env[1].name == "LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS" && local.gateway_pool_env[1].value == "8", + local.gateway_pool_env[2].name == "LITELLM_PGBOUNCER_MAX_CLIENT_CONN" && local.gateway_pool_env[2].value == "250", + ]) + error_message = "The pool env must carry the enabled flag and the configured sizes as strings." + } + + assert { + condition = alltrue([ + contains([for e in google_cloud_run_v2_service.gateway[0].template[0].containers[0].env : e.name], "LITELLM_PGBOUNCER_ENABLED"), + contains([for e in google_cloud_run_v2_service.gateway[0].template[0].containers[0].env : e.name], "LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS"), + contains([for e in google_cloud_run_v2_service.gateway[0].template[0].containers[0].env : e.name], "LITELLM_PGBOUNCER_MAX_CLIENT_CONN"), + ]) + error_message = "The gateway service must receive all three LITELLM_PGBOUNCER_* env vars." + } + + assert { + condition = !anytrue(concat( + [for e in google_cloud_run_v2_service.backend[0].template[0].containers[0].env : startswith(e.name, "LITELLM_PGBOUNCER_")], + [for e in google_cloud_run_v2_job.migrations[0].template[0].template[0].containers[0].env : startswith(e.name, "LITELLM_PGBOUNCER_")], + )) + error_message = "The backend service and the migrations job must keep their direct database connection." + } +} + +run "pool_enabled_uses_the_module_default_sizes" { + command = plan + + variables { + gateway_connection_pool_enabled = true + } + + assert { + condition = alltrue([ + length(local.gateway_pool_env) == 3, + local.gateway_pool_env[1].name == "LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS" && local.gateway_pool_env[1].value == "20", + local.gateway_pool_env[2].name == "LITELLM_PGBOUNCER_MAX_CLIENT_CONN" && local.gateway_pool_env[2].value == "1000", + ]) + error_message = "The pool env must fall back to the module defaults of 20 upstream and 1000 client connections." + } +} + +run "gateway_starts_through_the_pool_aware_launcher" { + command = plan + + variables { + gateway_num_workers = 4 + } + + assert { + condition = alltrue([ + strcontains(local.gateway_launch_cmd, "exec python -m gateway.launch --host 0.0.0.0 --port 4000 --workers 4"), + strcontains(local.gateway_launch_cmd, "exec ddtrace-run python -m gateway.launch --host 0.0.0.0 --port 4000 --workers 4"), + !strcontains(local.gateway_launch_cmd, "uvicorn gateway.main:app"), + endswith(google_cloud_run_v2_service.gateway[0].template[0].containers[0].args[0], local.gateway_launch_cmd), + ]) + error_message = "The gateway must start through gateway.launch (with and without ddtrace) so the pooler starts once before uvicorn forks the workers." + } + + assert { + condition = strcontains(local.backend_launch_cmd, "uvicorn backend.main:app") + error_message = "The backend has no workers to share a pooler and keeps starting uvicorn directly." + } +} + +run "pool_sizes_below_one_fail_at_plan" { + command = plan + + variables { + gateway_connection_pool_enabled = true + gateway_pool_max_db_connections = 0 + gateway_pool_max_client_conn = 0 + } + + expect_failures = [ + var.gateway_pool_max_db_connections, + var.gateway_pool_max_client_conn, + ] +} diff --git a/terraform/litellm/gcp/variables.tf b/terraform/litellm/gcp/variables.tf index a753b4584a6..f298a0431c0 100644 --- a/terraform/litellm/gcp/variables.tf +++ b/terraform/litellm/gcp/variables.tf @@ -206,6 +206,45 @@ variable "gateway_num_workers" { } } +variable "gateway_connection_pool_enabled" { + description = <<-EOT + Run an in-container PgBouncer (transaction mode, loopback) in each gateway + instance, shared by every uvicorn worker. Without it each of the + `gateway_num_workers` workers opens its own Prisma pool straight to + Cloud SQL, so an instance's footprint against the database connection + ceiling is workers x connection_limit and grows with every instance. Sets + LITELLM_PGBOUNCER_ENABLED / LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS / + LITELLM_PGBOUNCER_MAX_CLIENT_CONN on the gateway container only. The + module's Cloud SQL authenticates with the static password in Secret + Manager, which is what the pooler needs. Mirrors the AWS stack's + gateway_connection_pool_enabled. + EOT + type = bool + default = false +} + +variable "gateway_pool_max_db_connections" { + description = "Upstream Cloud SQL connections one gateway instance may hold when gateway_connection_pool_enabled is set, regardless of gateway_num_workers. 20 suits 4 workers." + type = number + default = 20 + + validation { + condition = var.gateway_pool_max_db_connections >= 1 + error_message = "gateway_pool_max_db_connections must be >= 1." + } +} + +variable "gateway_pool_max_client_conn" { + description = "Client connections the in-container PgBouncer accepts from the gateway workers when gateway_connection_pool_enabled is set." + type = number + default = 1000 + + validation { + condition = var.gateway_pool_max_client_conn >= 1 + error_message = "gateway_pool_max_client_conn must be >= 1." + } +} + # Cloud Run autoscales out of the box (request-rate driven). The min/max # bounds mirror the HPA replica bounds in helm/litellm/values.yaml so each # stack scales over the same range. Cloud Run has no direct CPU-utilization diff --git a/tests/test_gateway/test_launch.py b/tests/test_gateway/test_launch.py new file mode 100644 index 00000000000..964a7021416 --- /dev/null +++ b/tests/test_gateway/test_launch.py @@ -0,0 +1,155 @@ +import os +import socket +import sys +import textwrap +import urllib.parse +from pathlib import Path +from typing import Final, cast + +import pytest +from uvicorn.importer import import_from_string +from uvicorn.main import main as uvicorn_main + +import gateway.main +from gateway.launch import GATEWAY_APP, main, pool_database_url, uvicorn_argv +from litellm.proxy.db.db_url_settings import DatabaseURLSettings +from litellm.proxy.db.pgbouncer import PgBouncerError, PgBouncerSettings + +DB_ENV: Final = { + "DATABASE_HOST": "db.internal", + "DATABASE_PORT": "5432", + "DATABASE_USER": "litellm_pool", + "DATABASE_NAME": "litellm", + "DATABASE_PASSWORD": "p@ss", +} + + +def _free_port() -> int: + with socket.socket() as probe: + probe.bind(("127.0.0.1", 0)) + return cast(tuple[str, int], probe.getsockname())[1] + + +def _fake_pooler(tmp_path: Path) -> Path: + script: Final = tmp_path / "fake-pgbouncer" + script.write_text( + textwrap.dedent( + f"""\ + #!{sys.executable} + import configparser, select, socket, sys + if sys.argv[1:] == ["--version"]: + print("PgBouncer 1.25.2") + sys.exit(0) + ini = configparser.ConfigParser() + ini.read(sys.argv[1]) + port = ini.getint("pgbouncer", "listen_port") + tcp = socket.socket() + tcp.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + tcp.bind(("127.0.0.1", port)) + tcp.listen() + unix = socket.socket(socket.AF_UNIX) + unix.bind(ini.get("pgbouncer", "unix_socket_dir") + f"/.s.PGSQL.{{port}}") + unix.listen() + while True: + for ready in select.select([tcp, unix], [], [])[0]: + ready.accept()[0].close() + """ + ) + ) + script.chmod(0o700) + return script + + +def _query(url: str) -> dict[str, str]: + return dict(urllib.parse.parse_qsl(urllib.parse.urlsplit(url).query)) + + +@pytest.fixture +def password_env(monkeypatch: pytest.MonkeyPatch) -> dict[str, str]: + for var in ("DATABASE_URL", "IAM_TOKEN_DB_AUTH", "AZURE_POSTGRESQL_AUTH", "DATABASE_HOST_READ_REPLICA"): + monkeypatch.setenv(var, "") + monkeypatch.delenv(var) + for var, value in DB_ENV.items(): + monkeypatch.setenv(var, value) + return dict(DB_ENV) + + +def _uvicorn_params(argv: tuple[str, ...]) -> dict[str, object]: + return uvicorn_main.make_context("uvicorn", list(argv)).params + + +class TestUvicornArgv: + def test_keepalive_env_reaches_uvicorn(self): + params: Final = _uvicorn_params(uvicorn_argv(("--workers", "4"), {"KEEPALIVE_TIMEOUT": "75"})) + assert params["app"] == GATEWAY_APP + assert params["workers"] == 4 + assert params["timeout_keep_alive"] == 75 + + def test_unset_env_keeps_the_uvicorn_default(self): + assert _uvicorn_params(uvicorn_argv(("--workers", "4"), {}))["timeout_keep_alive"] == 5 + + def test_an_explicit_flag_wins_over_the_env(self): + argv: Final = uvicorn_argv(("--timeout-keep-alive", "30"), {"KEEPALIVE_TIMEOUT": "75"}) + assert _uvicorn_params(argv)["timeout_keep_alive"] == 30 + + def test_the_app_uvicorn_is_told_to_serve_is_the_trimmed_gateway(self): + assert import_from_string(cast(str, _uvicorn_params(uvicorn_argv((), {}))["app"])) is gateway.main.app + + +class TestPoolDatabaseUrl: + def test_a_disabled_pooler_yields_no_url_to_install(self, password_env: dict[str, str]): + settings: Final = DatabaseURLSettings.from_env() + settings.apply_to_env() + environ: Final = {"DATABASE_URL": "postgresql://litellm_pool:p%40ss@db.internal:5432/litellm"} + assert pool_database_url(settings, PgBouncerSettings(enabled=False), environ) is None + + def test_a_missing_upstream_url_is_reported(self, password_env: dict[str, str]): + environ: Final[dict[str, str]] = {} + outcome: Final = pool_database_url(DatabaseURLSettings.from_env(), PgBouncerSettings(enabled=True), environ) + assert isinstance(outcome, PgBouncerError) + assert "DATABASE_URL" in outcome.reason + + def test_token_auth_is_refused(self, password_env: dict[str, str], monkeypatch: pytest.MonkeyPatch, tmp_path: Path): + monkeypatch.setenv("IAM_TOKEN_DB_AUTH", "true") + environ: Final = {"DATABASE_URL": "postgresql://litellm:token@db.internal:5432/litellm"} + outcome: Final = pool_database_url( + DatabaseURLSettings.from_env(), + PgBouncerSettings(enabled=True, port=_free_port(), binary=str(_fake_pooler(tmp_path))), + environ, + ) + assert isinstance(outcome, PgBouncerError) + assert "IAM_TOKEN_DB_AUTH" in outcome.reason + + +class TestMain: + def test_workers_inherit_the_loopback_url_the_supervisor_installed( + self, password_env: dict[str, str], monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ): + port: Final = _free_port() + monkeypatch.setenv("LITELLM_PGBOUNCER_ENABLED", "true") + monkeypatch.setenv("LITELLM_PGBOUNCER_PORT", str(port)) + monkeypatch.setenv("LITELLM_PGBOUNCER_BINARY", str(_fake_pooler(tmp_path))) + monkeypatch.setenv("KEEPALIVE_TIMEOUT", "75") + served: Final[list[tuple[str, ...]]] = [] + main(("--workers", "4"), serve=lambda argv: served.append(tuple(argv))) + + pooled: Final = os.environ["DATABASE_URL"] + assert urllib.parse.urlsplit(pooled).netloc == f"litellm_pool:p%40ss@127.0.0.1:{port}" + assert _query(pooled)["pgbouncer"] == "true" + assert _uvicorn_params(served[0])["timeout_keep_alive"] == 75 + + DatabaseURLSettings.from_env().apply_to_env() + worker_url: Final = os.environ["DATABASE_URL"] + assert urllib.parse.urlsplit(worker_url).netloc == f"litellm_pool:p%40ss@127.0.0.1:{port}" + assert _query(worker_url)["pgbouncer"] == "true" + + def test_a_pooler_that_cannot_start_stops_the_gateway_before_uvicorn( + self, password_env: dict[str, str], monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ): + monkeypatch.setenv("LITELLM_PGBOUNCER_ENABLED", "true") + monkeypatch.setenv("LITELLM_PGBOUNCER_BINARY", str(tmp_path / "missing-pgbouncer")) + served: Final[list[tuple[str, ...]]] = [] + with pytest.raises(SystemExit) as stopped: + main(("--workers", "4"), serve=lambda argv: served.append(tuple(argv))) + assert "missing-pgbouncer" in str(stopped.value) + assert served == [] diff --git a/tests/test_litellm/proxy/db/test_pgbouncer.py b/tests/test_litellm/proxy/db/test_pgbouncer.py new file mode 100644 index 00000000000..da1d4c91f09 --- /dev/null +++ b/tests/test_litellm/proxy/db/test_pgbouncer.py @@ -0,0 +1,624 @@ +import configparser +import logging +import os +import signal +import socket +import stat +import sys +import tempfile +import textwrap +import time +import urllib.parse +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from typing import Final, cast + +import pytest + +from litellm._logging import verbose_proxy_logger +from litellm.proxy.db.pgbouncer import ( + PgBouncerError, + PgBouncerPlan, + PgBouncerProcess, + PgBouncerSettings, + pgbouncer_version, + plan_pgbouncer, + start_in_container_pgbouncer, + unix_socket_path, + write_pgbouncer_files, +) + +UPSTREAM: Final = ( + "postgresql://app:p%40ss%27w@db.internal:5433/litellm" + "?schema=public&connection_limit=10&pool_timeout=20" + "&sslmode=require&sslaccept=strict&sslcert=/certs/ca.pem" + "&options=-c%20statement_timeout%3D7000%20-c%20lock_timeout%3D3000" +) +SETTINGS: Final = PgBouncerSettings(enabled=True, port=6543, max_db_connections=8, max_client_conn=400) + + +def _plan(url: str = UPSTREAM, run_as_user: str | None = None) -> PgBouncerPlan: + plan: Final = plan_pgbouncer(url, SETTINGS, Path("/run/pgb"), run_as_user) + assert isinstance(plan, PgBouncerPlan), plan + return plan + + +def _ini(plan: PgBouncerPlan) -> configparser.ConfigParser: + parser: Final = configparser.ConfigParser(interpolation=None) + parser.read_string(plan.ini) + return parser + + +def _query(url: str) -> dict[str, str]: + return dict(urllib.parse.parse_qsl(urllib.parse.urlsplit(url).query, keep_blank_values=True)) + + +class TestPlanPgBouncer: + def test_upstream_credentials_and_timeouts_move_into_the_pgbouncer_config(self): + ini: Final = _ini(_plan()) + assert ini["databases"]["litellm"] == ( + "host='db.internal' port=5433 dbname='litellm' user='app' password='p@ss''w' " + "connect_query='SET statement_timeout TO ''7000''; SET lock_timeout TO ''3000'''" + ) + assert _plan().userlist == '"app" "p@ss\'w"\n' + + def test_an_upstream_without_a_port_is_reached_on_the_postgres_default(self): + ini: Final = _ini(_plan("postgresql://app:pw@db/litellm")) + assert ini["databases"]["litellm"] == "host='db' port=5432 dbname='litellm' user='app' password='pw'" + + def test_pool_is_sized_from_settings_in_transaction_mode(self): + pgb: Final = _ini(_plan())["pgbouncer"] + assert pgb["pool_mode"] == "transaction" + assert pgb["max_db_connections"] == "8" + assert pgb["default_pool_size"] == "8" + assert pgb["max_client_conn"] == "400" + assert pgb["auth_type"] == "scram-sha-256" + assert pgb["listen_addr"] == "127.0.0.1" + assert pgb["listen_port"] == "6543" + assert pgb["auth_file"] == "/run/pgb/userlist.txt" + assert pgb["unix_socket_dir"] == "/run/pgb" + + def test_the_app_user_can_read_the_pgbouncer_console(self): + assert _ini(_plan())["pgbouncer"]["stats_users"] == "app" + + @pytest.mark.parametrize("user", ["app,admin", "app%20admin", "app%09admin"]) + def test_a_user_pgbouncer_would_split_into_several_console_users_is_refused(self, user: str): + outcome: Final = plan_pgbouncer(f"postgresql://{user}:pw@db/litellm", SETTINGS, Path("/run/pgb"), None) + assert isinstance(outcome, PgBouncerError) + assert "stats_users" in outcome.reason + + def test_pooled_url_points_prisma_at_loopback_without_prepared_statements(self): + pooled: Final = urllib.parse.urlsplit(_plan().pooled_url) + assert (pooled.hostname, pooled.port, pooled.path) == ("127.0.0.1", 6543, "/litellm") + assert (pooled.username, pooled.password) == ("app", "p%40ss%27w") + assert _query(_plan().pooled_url) == { + "schema": "public", + "connection_limit": "10", + "pool_timeout": "20", + "pgbouncer": "true", + } + + @pytest.mark.parametrize("hop_param", ["channel_binding=require", "gssencmode=require"]) + def test_transport_params_for_the_postgres_hop_stay_off_the_plain_tcp_loopback_url(self, hop_param: str): + pooled: Final = _plan(f"postgresql://app:pw@db/litellm?connection_limit=5&{hop_param}").pooled_url + assert _query(pooled) == {"connection_limit": "5", "pgbouncer": "true"} + + def test_verified_tls_becomes_server_side_verify_full_with_a_ca_copy_in_the_runtime_dir(self): + plan: Final = _plan() + pgb: Final = _ini(plan)["pgbouncer"] + assert pgb["server_tls_sslmode"] == "verify-full" + assert pgb["server_tls_ca_file"] == "/run/pgb/server-ca.pem" + assert plan.ca_source == "/certs/ca.pem" + + def test_unverified_require_stays_require_without_a_ca_file(self): + plan: Final = _plan("postgresql://app:pw@db/litellm?sslmode=require") + pgb: Final = _ini(plan)["pgbouncer"] + assert pgb["server_tls_sslmode"] == "require" + assert "server_tls_ca_file" not in pgb + assert plan.ca_source is None + + def test_no_tls_params_default_to_prefer(self): + assert _ini(_plan("postgresql://app:pw@db/litellm"))["pgbouncer"]["server_tls_sslmode"] == "prefer" + + def test_verification_without_a_ca_bundle_is_refused(self): + outcome: Final = plan_pgbouncer( + "postgresql://app:pw@db/litellm?sslmode=require&sslaccept=strict", SETTINGS, Path("/run/pgb"), None + ) + assert isinstance(outcome, PgBouncerError) + assert "sslcert" in outcome.reason + + def test_client_certificates_are_refused(self): + outcome: Final = plan_pgbouncer( + "postgresql://app:pw@db/litellm?sslidentity=/certs/client.p12", SETTINGS, Path("/run/pgb"), None + ) + assert isinstance(outcome, PgBouncerError) + assert "sslidentity" in outcome.reason + + @pytest.mark.parametrize( + "url", + [ + "postgresql://app@db/litellm", + "postgresql://app:pw@db", + "postgresql://:pw@db/litellm", + ], + ) + def test_urls_missing_forwardable_credentials_are_refused(self, url: str): + outcome: Final = plan_pgbouncer(url, SETTINGS, Path("/run/pgb"), None) + assert isinstance(outcome, PgBouncerError) + + def test_every_options_spelling_becomes_a_set_statement(self): + options: Final = urllib.parse.quote("-c a=1 -cb=2 --c=3") + ini: Final = _ini(_plan(f"postgresql://app:pw@db/litellm?options={options}")) + assert ini["databases"]["litellm"].endswith("connect_query='SET a TO ''1''; SET b TO ''2''; SET c TO ''3'''") + + def test_options_that_are_not_settings_are_refused(self): + outcome: Final = plan_pgbouncer( + "postgresql://app:pw@db/litellm?options=-c%20search_path", SETTINGS, Path("/run/pgb"), None + ) + assert isinstance(outcome, PgBouncerError) + assert "options" in outcome.reason + + def test_run_as_user_is_only_written_when_given(self): + assert _ini(_plan(run_as_user="nobody"))["pgbouncer"]["user"] == "nobody" + assert "user" not in _ini(_plan())["pgbouncer"] + + +class TestWritePgBouncerFiles: + def test_files_hold_the_plan_and_are_private_to_the_owner(self, tmp_path: Path): + plan: Final = _plan("postgresql://app:pw@db/litellm") + ini_path: Final = write_pgbouncer_files(plan, tmp_path, None) + assert isinstance(ini_path, Path), ini_path + assert ini_path == tmp_path / "pgbouncer.ini" + assert ini_path.read_text() == plan.ini + assert (tmp_path / "userlist.txt").read_text() == plan.userlist + for path in (ini_path, tmp_path / "userlist.txt"): + assert stat.S_IMODE(path.stat().st_mode) == 0o600 + assert not (tmp_path / "server-ca.pem").exists() + + def test_the_ca_bundle_is_copied_next_to_the_ini_pgbouncer_reads(self, tmp_path: Path): + bundle: Final = tmp_path / "rds-root.pem" + bundle.write_text("-----BEGIN CERTIFICATE-----\nMIIB\n-----END CERTIFICATE-----\n") + runtime_dir: Final = tmp_path / "run" + runtime_dir.mkdir() + plan: Final = plan_pgbouncer( + f"postgresql://app:pw@db/litellm?sslmode=verify-full&sslcert={bundle}", SETTINGS, runtime_dir, None + ) + assert isinstance(plan, PgBouncerPlan), plan + ini_path: Final = write_pgbouncer_files(plan, runtime_dir, None) + assert isinstance(ini_path, Path), ini_path + ca_file: Final = Path(_ini(plan)["pgbouncer"]["server_tls_ca_file"]) + assert ca_file.parent == runtime_dir + assert ca_file.read_text() == bundle.read_text() + + def test_an_unreadable_ca_bundle_is_reported(self, tmp_path: Path): + plan: Final = plan_pgbouncer( + f"postgresql://app:pw@db/litellm?sslmode=verify-full&sslcert={tmp_path / 'missing.pem'}", + SETTINGS, + tmp_path, + None, + ) + assert isinstance(plan, PgBouncerPlan), plan + outcome: Final = write_pgbouncer_files(plan, tmp_path, None) + assert isinstance(outcome, PgBouncerError) + assert "missing.pem" in outcome.reason + assert not (tmp_path / "pgbouncer.ini").exists() + + +def _bound_port(sock: socket.socket) -> int: + return cast(tuple[str, int], sock.getsockname())[1] + + +def _free_port() -> int: + with socket.socket() as probe: + probe.bind(("127.0.0.1", 0)) + return _bound_port(probe) + + +def _fake_pooler( + tmp_path: Path, + port: int, + exit_immediately: bool = False, + port_file: Path | None = None, + bind_delay_seconds: float = 0.0, + version_banner: str = "PgBouncer 1.25.2\nlibevent 2.1.13-stable", +) -> Path: + """An executable that listens like PgBouncer: on the TCP port first, then on ``.s.PGSQL.`` in the socket dir. + + Port and socket dir come from the ini it is given, else from ``port`` and + ``tmp_path``. With ``port_file`` each start reads the port from that file + instead. ``bind_delay_seconds`` holds the bind back, like a slow start. + ``--version`` prints ``version_banner``. + """ + script: Final = tmp_path / "fake-pgbouncer" + script.write_text( + textwrap.dedent( + f"""\ + #!{sys.executable} + import configparser, os, pathlib, select, socket, sys, time + if sys.argv[1:] == ["--version"]: + print({version_banner!r}) + sys.exit(0) + if {exit_immediately!r}: + sys.exit(3) + ini = configparser.ConfigParser() + ini.read(sys.argv[1:2]) + port = ini.getint("pgbouncer", "listen_port", fallback={port}) + if not {port_file is None!r}: + port = int(pathlib.Path({str(port_file)!r}).read_text()) + socket_dir = ini.get("pgbouncer", "unix_socket_dir", fallback={str(tmp_path)!r}) + socket_path = f"{{socket_dir}}/.s.PGSQL.{{port}}" + time.sleep({bind_delay_seconds!r}) + listener = socket.socket() + listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + listener.bind(("127.0.0.1", port)) + listener.listen() + if os.path.exists(socket_path): + os.unlink(socket_path) + unix_listener = socket.socket(socket.AF_UNIX) + unix_listener.bind(socket_path) + unix_listener.listen() + while True: + for ready in select.select([listener, unix_listener], [], [])[0]: + conn, _ = ready.accept() + conn.close() + """ + ) + ) + script.chmod(0o700) + return script + + +def _listening(port: int) -> bool: + try: + with socket.create_connection(("127.0.0.1", port), timeout=0.5): + return True + except OSError: + return False + + +def _wait_until(condition: Callable[[], bool], timeout_seconds: float = 5.0) -> bool: + deadline: Final = time.monotonic() + timeout_seconds + while time.monotonic() < deadline: + if condition(): + return True + time.sleep(0.05) + return False + + +class TestPgBouncerProcess: + def test_start_waits_for_the_listener_and_stop_ends_it(self, tmp_path: Path): + port: Final = _free_port() + pooler: Final = PgBouncerProcess( + argv=(str(_fake_pooler(tmp_path, port)),), port=port, socket_path=unix_socket_path(tmp_path, port) + ) + assert pooler.start() is None + assert _listening(port) + pid: Final = pooler.pid + assert pid is not None + pooler.stop() + assert _wait_until(lambda: not _listening(port)) + with pytest.raises(ProcessLookupError): + os.kill(pid, 0) + + def test_a_crashed_pooler_is_restarted_with_a_new_pid(self, tmp_path: Path): + port: Final = _free_port() + pooler: Final = PgBouncerProcess( + argv=(str(_fake_pooler(tmp_path, port)),), + port=port, + socket_path=unix_socket_path(tmp_path, port), + restart_delay_seconds=0.1, + ) + assert pooler.start() is None + first_pid: Final = pooler.pid + assert first_pid is not None + os.kill(first_pid, signal.SIGKILL) + assert _wait_until(lambda: pooler.pid not in (None, first_pid) and _listening(port)) + pooler.stop() + assert _wait_until(lambda: not _listening(port)) + + def test_a_failed_restart_is_retried_until_the_pooler_is_back( + self, tmp_path: Path, caplog: pytest.LogCaptureFixture + ): + port: Final = _free_port() + script: Final = _fake_pooler(tmp_path, port) + pooler: Final = PgBouncerProcess( + argv=(str(script),), port=port, socket_path=unix_socket_path(tmp_path, port), restart_delay_seconds=0.1 + ) + assert pooler.start() is None + first_pid: Final = pooler.pid + assert first_pid is not None + hidden: Final = script.rename(tmp_path / "hidden") + with caplog.at_level(logging.ERROR, logger=verbose_proxy_logger.name): + os.kill(first_pid, signal.SIGKILL) + assert _wait_until(lambda: any("could not be restarted" in record.message for record in caplog.records)) + assert not _listening(port) + hidden.rename(script) + assert _wait_until(lambda: pooler.pid not in (None, first_pid) and _listening(port)) + pooler.stop() + assert _wait_until(lambda: not _listening(port)) + + def test_a_replacement_that_never_listens_is_replaced_again(self, tmp_path: Path, caplog: pytest.LogCaptureFixture): + port: Final = _free_port() + port_file: Final = tmp_path / "port" + port_file.write_text(str(port)) + script: Final = _fake_pooler(tmp_path, port, port_file=port_file) + pooler: Final = PgBouncerProcess( + argv=(str(script),), + port=port, + socket_path=unix_socket_path(tmp_path, port), + restart_delay_seconds=0.1, + ready_timeout_seconds=0.3, + ) + assert pooler.start() is None + first_pid: Final = pooler.pid + assert first_pid is not None + wrong_port: Final = _free_port() + port_file.write_text(str(wrong_port)) + with caplog.at_level(logging.ERROR, logger=verbose_proxy_logger.name): + os.kill(first_pid, signal.SIGKILL) + assert _wait_until(lambda: _listening(wrong_port)) + assert _wait_until(lambda: any("did not start listening" in record.message for record in caplog.records)) + port_file.write_text(str(port)) + assert _wait_until(lambda: _listening(port)) + assert _wait_until(lambda: not _listening(wrong_port)) + pooler.stop() + assert _wait_until(lambda: not _listening(port)) + + def test_stopping_during_the_restart_delay_leaves_no_pooler_behind(self, tmp_path: Path): + port: Final = _free_port() + pooler: Final = PgBouncerProcess( + argv=(str(_fake_pooler(tmp_path, port)),), + port=port, + socket_path=unix_socket_path(tmp_path, port), + restart_delay_seconds=0.3, + ) + assert pooler.start() is None + first_pid: Final = pooler.pid + assert first_pid is not None + os.kill(first_pid, signal.SIGKILL) + assert _wait_until(lambda: not _listening(port)) + pooler.stop() + time.sleep(1.0) + assert not _listening(port) + assert pooler.pid == first_pid + + def test_a_stopped_pooler_is_not_restarted(self, tmp_path: Path, caplog: pytest.LogCaptureFixture): + port: Final = _free_port() + pooler: Final = PgBouncerProcess( + argv=(str(_fake_pooler(tmp_path, port)),), + port=port, + socket_path=unix_socket_path(tmp_path, port), + restart_delay_seconds=0.1, + ) + assert pooler.start() is None + with caplog.at_level(logging.ERROR, logger=verbose_proxy_logger.name): + pooler.stop() + time.sleep(0.5) + assert not _listening(port) + assert caplog.records == [] + + def test_a_pooler_that_exits_during_startup_is_reported(self, tmp_path: Path): + port: Final = _free_port() + pooler: Final = PgBouncerProcess( + argv=(str(_fake_pooler(tmp_path, port, exit_immediately=True)),), + port=port, + socket_path=unix_socket_path(tmp_path, port), + ) + outcome: Final = pooler.start() + assert isinstance(outcome, PgBouncerError) + assert "status 3" in outcome.reason + + def test_a_missing_binary_is_reported(self, tmp_path: Path): + outcome: Final = PgBouncerProcess( + argv=("/nonexistent/pgbouncer",), port=_free_port(), socket_path=tmp_path / "sock" + ).start() + assert isinstance(outcome, PgBouncerError) + assert "/nonexistent/pgbouncer" in outcome.reason + + def test_a_port_owned_by_someone_else_is_refused_before_spawning(self, tmp_path: Path): + with socket.socket() as squatter: + squatter.bind(("127.0.0.1", 0)) + squatter.listen() + port: Final = _bound_port(squatter) + pooler: Final = PgBouncerProcess( + argv=(str(_fake_pooler(tmp_path, port)),), port=port, socket_path=unix_socket_path(tmp_path, port) + ) + outcome: Final = pooler.start() + assert isinstance(outcome, PgBouncerError) + assert f"127.0.0.1:{port} is already in use" in outcome.reason + assert pooler.pid is None + + def test_a_replacement_waits_until_a_squatter_leaves_the_port( + self, tmp_path: Path, caplog: pytest.LogCaptureFixture + ): + port: Final = _free_port() + pooler: Final = PgBouncerProcess( + argv=(str(_fake_pooler(tmp_path, port)),), + port=port, + socket_path=unix_socket_path(tmp_path, port), + restart_delay_seconds=0.5, + ) + assert pooler.start() is None + first_pid: Final = pooler.pid + assert first_pid is not None + os.kill(first_pid, signal.SIGKILL) + assert _wait_until(lambda: not _listening(port)) + with socket.socket() as squatter, caplog.at_level(logging.ERROR, logger=verbose_proxy_logger.name): + squatter.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + squatter.bind(("127.0.0.1", port)) + squatter.listen() + assert _wait_until(lambda: any("already in use" in record.message for record in caplog.records)) + assert pooler.pid == first_pid + assert _wait_until(lambda: pooler.pid not in (None, first_pid) and _listening(port)) + pooler.stop() + assert _wait_until(lambda: not _listening(port)) + + def test_a_listener_that_grabs_the_port_after_the_spawn_is_not_taken_for_the_pooler(self, tmp_path: Path): + port: Final = _free_port() + pooler: Final = PgBouncerProcess( + argv=(str(_fake_pooler(tmp_path, port, bind_delay_seconds=0.5)),), + port=port, + socket_path=unix_socket_path(tmp_path, port), + ready_timeout_seconds=3.0, + ) + with socket.socket() as squatter, ThreadPoolExecutor(max_workers=1) as starter: + squatter.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + starting: Final = starter.submit(pooler.start) + assert _wait_until(lambda: pooler.pid is not None) + squatter.bind(("127.0.0.1", port)) + squatter.listen() + outcome: Final = starting.result() + assert isinstance(outcome, PgBouncerError) + assert "exited with status 1" in outcome.reason + + def test_a_port_served_by_a_stranger_while_the_pooler_is_still_starting_is_reported(self, tmp_path: Path): + port: Final = _free_port() + pooler: Final = PgBouncerProcess( + argv=(str(_fake_pooler(tmp_path, port, bind_delay_seconds=30.0)),), + port=port, + socket_path=unix_socket_path(tmp_path, port), + ready_timeout_seconds=0.5, + ) + with socket.socket() as squatter, ThreadPoolExecutor(max_workers=1) as starter: + squatter.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + starting: Final = starter.submit(pooler.start) + assert _wait_until(lambda: pooler.pid is not None) + squatter.bind(("127.0.0.1", port)) + squatter.listen() + outcome: Final = starting.result() + assert isinstance(outcome, PgBouncerError) + assert f"127.0.0.1:{port} is served by another process" in outcome.reason + pid: Final = pooler.pid + assert pid is not None + with pytest.raises(ProcessLookupError): + os.kill(pid, 0) + + def test_a_pooler_that_never_listens_times_out(self, tmp_path: Path): + port: Final = _free_port() + pooler: Final = PgBouncerProcess( + argv=(str(_fake_pooler(tmp_path, _free_port())),), + port=port, + socket_path=unix_socket_path(tmp_path, port), + ready_timeout_seconds=0.5, + ) + outcome: Final = pooler.start() + assert isinstance(outcome, PgBouncerError) + assert "did not start listening" in outcome.reason + pid: Final = pooler.pid + assert pid is not None + with pytest.raises(ProcessLookupError): + os.kill(pid, 0) + + +def _runtime_dir_listening_on(port: int) -> Path: + matches: Final = tuple( + ini.parent + for ini in Path(tempfile.gettempdir()).glob("litellm-pgbouncer-*/pgbouncer.ini") + if f"listen_port = {port}\n" in ini.read_text() + ) + assert len(matches) == 1, matches + return matches[0] + + +class TestStartInContainerPgBouncer: + def test_returns_the_loopback_url_once_the_pooler_listens(self, tmp_path: Path): + port: Final = _free_port() + settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(_fake_pooler(tmp_path, port))) + pooled: Final = start_in_container_pgbouncer(settings, "postgresql://app:pw@db/litellm?connection_limit=5") + assert pooled == f"postgresql://app:pw@127.0.0.1:{port}/litellm?connection_limit=5&pgbouncer=true" + assert _listening(port) + + @pytest.mark.filterwarnings("ignore:This process .* is multi-threaded:DeprecationWarning") + def test_a_forked_worker_exiting_leaves_the_pooler_and_its_files_to_the_parent(self, tmp_path: Path): + port: Final = _free_port() + settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(_fake_pooler(tmp_path, port))) + exit_hooks: Final[list[Callable[[], None]]] = [] + pooled: Final = start_in_container_pgbouncer( + settings, "postgresql://app:pw@db/litellm", register_exit_hook=exit_hooks.append + ) + assert isinstance(pooled, str), pooled + runtime_dir: Final = _runtime_dir_listening_on(port) + + worker: Final = os.fork() + if worker == 0: + try: + for hook in exit_hooks: + hook() + finally: + os._exit(0) + if not _wait_until(lambda: os.waitpid(worker, os.WNOHANG) != (0, 0)): + os.kill(worker, signal.SIGKILL) + pytest.fail("the forked worker did not exit: an exit hook blocked on state inherited from the parent") + assert _listening(port) + assert (runtime_dir / "pgbouncer.ini").exists() + + for hook in exit_hooks: + hook() + assert _wait_until(lambda: not _listening(port)) + assert not runtime_dir.exists() + + def test_a_bad_upstream_url_is_reported_without_starting_anything(self, tmp_path: Path): + port: Final = _free_port() + settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(_fake_pooler(tmp_path, port))) + outcome: Final = start_in_container_pgbouncer(settings, "postgresql://app@db/litellm") + assert isinstance(outcome, PgBouncerError) + assert not _listening(port) + + def test_token_auth_is_refused_without_starting_anything(self, tmp_path: Path): + port: Final = _free_port() + settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(_fake_pooler(tmp_path, port))) + outcome: Final = start_in_container_pgbouncer( + settings, "postgresql://app:pw@db/litellm", token_auth_enabled=True + ) + assert isinstance(outcome, PgBouncerError) + assert "IAM_TOKEN_DB_AUTH" in outcome.reason + assert "AZURE_POSTGRESQL_AUTH" in outcome.reason + assert not _listening(port) + + def test_a_pgbouncer_that_survives_a_failed_tcp_bind_is_refused_without_starting(self, tmp_path: Path): + port: Final = _free_port() + binary: Final = _fake_pooler(tmp_path, port, version_banner="PgBouncer 1.18.1\nlibevent 2.1.12-stable") + settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(binary)) + outcome: Final = start_in_container_pgbouncer(settings, "postgresql://app:pw@db/litellm") + assert isinstance(outcome, PgBouncerError) + assert "PgBouncer 1.18" in outcome.reason + assert "1.19" in outcome.reason + assert not _listening(port) + + def test_the_first_version_that_dies_on_a_failed_tcp_bind_is_accepted(self, tmp_path: Path): + port: Final = _free_port() + binary: Final = _fake_pooler(tmp_path, port, version_banner="PgBouncer 1.19.0") + settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(binary)) + assert start_in_container_pgbouncer(settings, "postgresql://app:pw@db/litellm") == ( + f"postgresql://app:pw@127.0.0.1:{port}/litellm?pgbouncer=true" + ) + assert _listening(port) + + +class TestPgBouncerVersion: + def test_reads_major_and_minor_from_the_banner(self, tmp_path: Path): + assert pgbouncer_version(str(_fake_pooler(tmp_path, _free_port()))) == (1, 25) + + def test_a_binary_that_cannot_run_is_reported(self, tmp_path: Path): + outcome: Final = pgbouncer_version(str(tmp_path / "missing-pgbouncer")) + assert isinstance(outcome, PgBouncerError) + assert "missing-pgbouncer" in outcome.reason + + def test_a_banner_without_a_version_is_reported(self, tmp_path: Path): + outcome: Final = pgbouncer_version(str(_fake_pooler(tmp_path, _free_port(), version_banner="something else"))) + assert isinstance(outcome, PgBouncerError) + assert "something else" in outcome.reason + + +class TestPgBouncerSettings: + def test_reads_the_litellm_pgbouncer_env_vars(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("LITELLM_PGBOUNCER_ENABLED", "true") + monkeypatch.setenv("LITELLM_PGBOUNCER_PORT", "7000") + monkeypatch.setenv("LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS", "12") + settings: Final = PgBouncerSettings() + assert (settings.enabled, settings.port, settings.max_db_connections) == (True, 7000, 12) + + def test_defaults_are_off(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("LITELLM_PGBOUNCER_ENABLED", raising=False) + assert PgBouncerSettings().enabled is False diff --git a/tests/test_litellm/test_component_entrypoint.py b/tests/test_litellm/test_component_entrypoint.py index 09837d2b233..0c2a533b8bc 100644 --- a/tests/test_litellm/test_component_entrypoint.py +++ b/tests/test_litellm/test_component_entrypoint.py @@ -38,12 +38,16 @@ _STUB_TEMPLATE = """#!/bin/sh _ENTRYPOINT_RE = re.compile(r"^ENTRYPOINT\s+(\[.*\])\s*$", re.MULTILINE) _CMD_RE = re.compile(r"^CMD\s+(\[.*\])\s*$", re.MULTILINE) _COPY_RE = re.compile(r"^COPY\s+(?!--from)(\S+)\s+(\S+)\s*$", re.MULTILINE) -_APP_TARGET_RE = re.compile(r"(?:gateway|backend)\.main:app") +_APP_TARGET_RE = re.compile(r"(?:gateway|backend)\.main:app|gateway\.launch") _TF_STRING_LOCAL_RE = re.compile(r'^\s*(\w+)\s*=\s*"((?:[^"\\]|\\.)*)"\s*$', re.MULTILINE) _TF_INTERPOLATION_RE = re.compile(r"\$\{(local|var)\.(\w+)\}") TERRAFORM_LAUNCH_SITES = {TERRAFORM_ECS: 2, TERRAFORM_CLOUDRUN: 2} TERRAFORM_VAR_STUBS = {"gateway_num_workers": "2"} +COMPONENT_LAUNCHERS = { + "gateway": ("python", "-m", "gateway.launch"), + "backend": ("uvicorn", "backend.main:app"), +} _MAX_INTERPOLATION_PASSES = 5 @@ -63,7 +67,7 @@ def _run_entrypoint( """Run `script` with stubbed executables on PATH and return the recorded lines.""" bin_dir = tmp_path / "bin" bin_dir.mkdir(parents=True) - _write_stubs(bin_dir, ("ddtrace-run", "uvicorn", "litellm")) + _write_stubs(bin_dir, ("ddtrace-run", "uvicorn", "python", "litellm")) record = tmp_path / "record.txt" env = { @@ -330,22 +334,63 @@ def test_entrypoint_script_has_no_carriage_returns() -> None: @pytest.mark.parametrize( - "dockerfile, app_target", + "dockerfile, launcher", [ - (GATEWAY_DOCKERFILE, "gateway.main:app"), - (BACKEND_DOCKERFILE, "backend.main:app"), + (GATEWAY_DOCKERFILE, "python -m gateway.launch"), + (BACKEND_DOCKERFILE, "uvicorn backend.main:app"), ], ) -def test_component_images_launch_uvicorn_through_the_entrypoint(dockerfile: Path, app_target: str) -> None: +def test_component_images_launch_uvicorn_through_the_entrypoint(dockerfile: Path, launcher: str) -> None: entrypoint = " ".join(_entrypoint_argv(dockerfile)) assert IMAGE_ENTRYPOINT_PATH in entrypoint, f"{dockerfile} bypasses the ddtrace-aware entrypoint" - assert app_target in entrypoint - assert entrypoint.index(IMAGE_ENTRYPOINT_PATH) < entrypoint.index("uvicorn"), ( + assert launcher in entrypoint + assert entrypoint.index(IMAGE_ENTRYPOINT_PATH) < entrypoint.index(launcher), ( f"{dockerfile} must invoke uvicorn through the entrypoint, not the other way around" ) +@pytest.mark.parametrize( + "use_ddtrace, num_workers, expected_exec, expected_args", + [ + (None, "4", "exec=python", "args=-m gateway.launch --workers 4 --host 0.0.0.0 --port 4000"), + (None, None, "exec=python", "args=-m gateway.launch --workers 1 --host 0.0.0.0 --port 4000"), + ("true", "4", "exec=ddtrace-run", "args=python -m gateway.launch --workers 4 --host 0.0.0.0 --port 4000"), + ], +) +def test_gateway_image_execs_the_supervisor_with_its_worker_count( + use_ddtrace: str | None, num_workers: str | None, expected_exec: str, expected_args: str, tmp_path: Path +) -> None: + """Run the gateway image's ENTRYPOINT + CMD and record what the container execs. + + The Dockerfile's `/app/...` script path is resolved to the checked-in script and `python` + is stubbed on PATH, so the assertion is on the argv `gateway.launch` receives, not on the + Dockerfile text. + """ + entrypoint = tuple( + part.replace(IMAGE_ENTRYPOINT_PATH, str(COMPONENT_ENTRYPOINT)) for part in _entrypoint_argv(GATEWAY_DOCKERFILE) + ) + bin_dir = tmp_path / "bin" + bin_dir.mkdir(parents=True) + _write_stubs(bin_dir, ("ddtrace-run", "python", "uvicorn")) + record = tmp_path / "record.txt" + overrides = {"USE_DDTRACE": use_ddtrace, "NUM_WORKERS": num_workers} + env = { + **{k: v for k, v in os.environ.items() if k not in ("DD_TRACE_OPENAI_ENABLED", *overrides)}, + **{k: v for k, v in overrides.items() if v is not None}, + "PATH": f"{bin_dir}{os.pathsep}{os.environ['PATH']}", + "RECORD": str(record), + "PYTHONPATH": PYTHONPATH_SENTINEL, + } + + result = subprocess.run( + [*entrypoint, *_cmd_argv(GATEWAY_DOCKERFILE)], env=env, capture_output=True, text=True, check=False + ) + + assert result.returncode == 0, f"stdout={result.stdout} stderr={result.stderr}" + assert tuple(record.read_text().splitlines())[:2] == (expected_exec, expected_args) + + @pytest.mark.parametrize("dockerfile", [GATEWAY_DOCKERFILE, BACKEND_DOCKERFILE]) def test_component_images_make_the_entrypoint_executable(dockerfile: Path) -> None: body = dockerfile.read_text() @@ -367,17 +412,18 @@ def test_terraform_launch_command_matches_the_script_contract( implementations under the same environment and asserts they agree on which binary is exec'd and on whether the openai integration is disabled. """ - app_target = f"{component}.main:app" + launcher = COMPONENT_LAUNCHERS[component] + app_target = " ".join(launcher[1:]) command = _resolve_tf_local(terraform_file, f"{component}_launch_cmd") bin_dir = tmp_path / "bin" bin_dir.mkdir(parents=True) - _write_stubs(bin_dir, ("ddtrace-run", "uvicorn")) + _write_stubs(bin_dir, ("ddtrace-run", "uvicorn", "python")) from_terraform = _run_shell_command(command, bin_dir, tmp_path / "terraform.txt", use_ddtrace) from_script = _run_entrypoint( COMPONENT_ENTRYPOINT, - ("uvicorn", app_target), + launcher, use_ddtrace=use_ddtrace, tmp_path=tmp_path / "script", ) @@ -387,12 +433,14 @@ def test_terraform_launch_command_matches_the_script_contract( ) assert from_terraform[2] == from_script[2], f"{terraform_file} disagrees with the script on the openai integration" assert app_target in from_terraform[1] + assert "gateway.main:app" not in from_terraform[1], f"{terraform_file} bypasses the gateway.launch supervisor" if use_ddtrace in TRUTHY_USE_DDTRACE: assert from_terraform[0] == "exec=ddtrace-run" + assert from_terraform[1].startswith(f"args={launcher[0]} ") assert from_terraform[2] == "DD_TRACE_OPENAI_ENABLED=False" else: - assert from_terraform[0] == "exec=uvicorn" + assert from_terraform[0] == f"exec={launcher[0]}" assert from_terraform[2] == "DD_TRACE_OPENAI_ENABLED=" From 032a9ebc4b370eb6f53ca00ce0755b069ddc1331 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 10 Sep 2026 15:46:01 -0700 Subject: [PATCH 33/33] chore(github): disable blank issues so filers must use a template (#40629) --- .github/ISSUE_TEMPLATE/config.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/ISSUE_TEMPLATE/config.yml b/.github/ISSUE_TEMPLATE/config.yml index cbf380bac01..fc86c229c81 100644 --- a/.github/ISSUE_TEMPLATE/config.yml +++ b/.github/ISSUE_TEMPLATE/config.yml @@ -1,4 +1,4 @@ -blank_issues_enabled: true +blank_issues_enabled: false contact_links: - name: Schedule Demo url: https://enterprise.litellm.ai/demo