From 6d7a80ac755bf2e3c14e7b76af910ca28cea13c2 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Mon, 13 Jul 2026 23:59:18 -0700 Subject: [PATCH 01/12] feat(mcp): aggregate gateway DCR discovery front door behind mcp_gateway_dcr --- .../mcp_server/auth/user_api_key_auth_mcp.py | 156 ++++-- .../mcp_server/discoverable_endpoints.py | 98 ++++ .../_experimental/mcp_server/oauth_utils.py | 23 + .../auth/test_user_api_key_auth_mcp.py | 133 ++++++ .../mcp_server/test_discoverable_endpoints.py | 449 +++++------------- 5 files changed, 489 insertions(+), 370 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index d2f3efbc54e..418e63468b4 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -10,6 +10,10 @@ from typing_extensions import assert_never import litellm from litellm._logging import verbose_logger +from litellm.proxy._experimental.mcp_server.oauth_utils import ( + get_request_base_url, + is_mcp_gateway_dcr_enabled, +) from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( BridgeEnvelopeAdmitted, BridgeEnvelopeInvalid, @@ -120,6 +124,96 @@ def _has_client_supplied_mcp_auth( return bool(mcp_auth_header) or bool(mcp_server_auth_headers) +def _is_aggregate_gateway_dcr_challenge_scope( + route: str, + mcp_servers: list[str] | None, + mcp_auth_header: str | None, + mcp_server_auth_headers: dict[str, dict[str, str]] | None, + exc: Exception, +) -> bool: + """True when an unauthenticated request to the aggregate ``/mcp`` endpoint + should receive the RFC 9728 401 challenge that advertises the gateway as + the authorization server (``mcp_gateway_dcr`` front door). + + Fires only for a genuine 401 on the aggregate scope: any named target + (path or ``x-mcp-servers``) belongs to the per-server challenge paths, and + client-supplied MCP auth headers mean the caller is not a cold-start DCR + client. Fails closed to the original admission error otherwise.""" + if not is_mcp_gateway_dcr_enabled(): + return False + if not _is_litellm_auth_admission_error(exc): + return False + if mcp_servers: + return False + if _has_client_supplied_mcp_auth(mcp_auth_header, mcp_server_auth_headers): + return False + return len(MCPRequestHandler._extract_target_server_names_from_path(route)) == 0 + + +def _aggregate_gateway_dcr_challenge(request: Request, invalid_token: bool) -> HTTPException: + """The RFC 9728 challenge for the aggregate endpoint: points the client at + the gateway's own protected-resource metadata so a DCR client discovers + the gateway as its authorization server and starts the sign-in flow. + + ``invalid_token`` adds the RFC 6750 error code for a request that DID + present a bearer that failed admission (expired or revoked), telling + spec-compliant clients to re-authorize rather than retry; a request with + no credentials at all gets the bare challenge per RFC 6750 section 3.1.""" + error_attr = 'error="invalid_token", ' if invalid_token else "" + resource_metadata_url = f"{get_request_base_url(request)}/.well-known/oauth-protected-resource/mcp" + return HTTPException( + status_code=401, + detail={ + "error": "authentication_required", + "message": "Authenticate with the gateway to use the MCP endpoint.", + }, + headers={"WWW-Authenticate": f'Bearer {error_attr}resource_metadata="{resource_metadata_url}"'}, + ) + + +def _admission_failure_fallback( + request: Request, + request_route: str, + mcp_servers: list[str] | None, + mcp_auth_header: str | None, + mcp_server_auth_headers: dict[str, dict[str, str]] | None, + exc: Exception, + bearer_presented: bool, +) -> UserAPIKeyAuth: + """Map a failed LiteLLM admission to its anonymous fallback or challenge. + + Two fallbacks exist, both gated on a genuine 401 with no client-supplied + MCP auth headers. The pass-through cold start (RFC 9728 / MCP + Authorization spec discovery return) admits anonymously so the route's + 401 emitter can produce the per-server challenge. The aggregate + gateway-DCR scope converts the failure into the gateway's own + resource_metadata challenge, with the RFC 6750 ``invalid_token`` error + code when the caller DID present a bearer (an expired gateway session + must re-authorize, not retry a dead token). Anything else re-raises the + original admission error unchanged.""" + mcp_servers_from_path = _parse_mcp_server_names_from_path(request_route, mcp_servers) + if ( + mcp_servers_from_path is not None + and not _has_client_supplied_mcp_auth(mcp_auth_header, mcp_server_auth_headers) + and _is_litellm_auth_admission_error(exc) + and _is_mcp_passthrough_cold_start( + mcp_servers_from_path, + client_ip=IPAddressUtils.get_mcp_client_ip(request), + ) + ): + verbose_logger.debug("MCP pass-through cold start: deferring admission to route 401 emitter") + return UserAPIKeyAuth() + if _is_aggregate_gateway_dcr_challenge_scope( + route=request_route, + mcp_servers=mcp_servers, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + exc=exc, + ): + raise _aggregate_gateway_dcr_challenge(request, invalid_token=bearer_presented) from exc + raise exc + + class MCPRequestHandler: """ Class to handle MCP request processing, including: @@ -271,56 +365,32 @@ class MCPRequestHandler: elif oauth2_headers: # Authorization on a non-delegated server: the bearer must be a real # LiteLLM credential, so a failed validation is a genuine 401/403 and - # propagates. The sole anonymous fallback is the auth_type=none - # pass-through cold-start (RFC 9728 discovery return), gated on a 401 - # so a recognized-but-forbidden key still fails closed. - client_ip = IPAddressUtils.get_mcp_client_ip(request) + # propagates unless a fallback in _admission_failure_fallback applies. try: validated_user_api_key_auth = await user_api_key_auth(api_key=litellm_api_key, request=request) except (HTTPException, ProxyException) as e: - # ProxyException.code is normalized to str (possibly "None"), so - # compare both int and str forms rather than coercing. - status = e.status_code if isinstance(e, HTTPException) else e.code - is_unauthenticated = status in (401, "401") - mcp_servers_from_path = _parse_mcp_server_names_from_path(request_route, mcp_servers) - if ( - is_unauthenticated - and mcp_servers_from_path is not None - and not _has_client_supplied_mcp_auth( - mcp_auth_header, - mcp_server_auth_headers, - ) - and _is_mcp_passthrough_cold_start(mcp_servers_from_path, client_ip=client_ip) - ): - verbose_logger.debug( - "MCP pass-through return: forwarding Authorization as upstream OAuth token for delegated auth" - ) - validated_user_api_key_auth = UserAPIKeyAuth() - else: - raise + validated_user_api_key_auth = _admission_failure_fallback( + request=request, + request_route=request_route, + mcp_servers=mcp_servers, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + exc=e, + bearer_presented=True, + ) else: try: validated_user_api_key_auth = await user_api_key_auth(api_key=litellm_api_key, request=request) except (HTTPException, ProxyException) as exc: - # Cold-start MCP OAuth discovery: RFC 9728 / MCP Authorization spec - # require unauthenticated requests to protected resources to receive - # 401 + WWW-Authenticate. Defer to _raise_preemptive_401_for_unauthenticated_servers - # for pass-through servers instead of surfacing a generic admission error. - mcp_servers_from_path = _parse_mcp_server_names_from_path(request_route, mcp_servers) - client_ip = IPAddressUtils.get_mcp_client_ip(request) - if ( - mcp_servers_from_path is not None - and not _has_client_supplied_mcp_auth( - mcp_auth_header, - mcp_server_auth_headers, - ) - and _is_litellm_auth_admission_error(exc) - and _is_mcp_passthrough_cold_start(mcp_servers_from_path, client_ip=client_ip) - ): - verbose_logger.debug("MCP pass-through cold start: deferring admission to route 401 emitter") - validated_user_api_key_auth = UserAPIKeyAuth() - else: - raise + validated_user_api_key_auth = _admission_failure_fallback( + request=request, + request_route=request_route, + mcp_servers=mcp_servers, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + exc=exc, + bearer_presented=False, + ) return ( validated_user_api_key_auth, diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 1af64749304..075ece42b04 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -42,6 +42,7 @@ from litellm.proxy._experimental.mcp_server.faults import ( from litellm.proxy._experimental.mcp_server.oauth_utils import ( TOKEN_NO_CACHE_HEADERS, get_request_base_url, + is_mcp_gateway_dcr_enabled, validate_trusted_redirect_uri, ) from litellm.proxy.auth.ip_address_utils import IPAddressUtils @@ -1708,6 +1709,12 @@ async def _build_oauth_protected_resource_response( global_mcp_server_manager, ) + # With the gateway-level DCR front door enabled, unnamed discovery + # describes the gateway itself as the authorization server for the + # aggregate /mcp resource instead of narrowing to one server. + if mcp_server_name is None and is_mcp_gateway_dcr_enabled(): + return _build_aggregate_protected_resource_response(request) + request_base_url = get_request_base_url(request) client_ip = IPAddressUtils.get_mcp_client_ip(request) @@ -1838,6 +1845,92 @@ def _jwt_auth_issuers() -> list: return issuers +def _build_aggregate_protected_resource_response(request: Request) -> dict: + """RFC 9728 metadata for the aggregate /mcp resource: the gateway itself is + the authorization server. No per-server names or scopes leak here; access + is resolved after sign-in from the authenticated user's grants. + + The advertised authorization server is ``{base}/mcp`` (not the bare + origin) so RFC 8414 path-insertion resolves its metadata at + ``/.well-known/oauth-authorization-server/mcp``, a route this module + owns. The bare-origin well-known is registered first by the BYOK OAuth + feature and describes the BYOK flow, so it must not be the aggregate + discovery entry point (same pattern as the per-server documents, which + advertise ``{base}/{server_name}``).""" + request_base_url = get_request_base_url(request) + return { + "authorization_servers": [f"{request_base_url}/mcp"], + "resource": f"{request_base_url}/mcp", + "scopes_supported": [], + } + + +def _build_aggregate_authorization_server_response(request: Request) -> dict: + """RFC 8414 metadata for the gateway as the aggregate authorization server. + + The issuer is ``{base}/mcp`` and must stay equal to the value the + aggregate protected-resource document advertises: spec clients verify the + issuer in the metadata matches the one that derived the well-known URL. + Advertises the root /authorize, /token, and /register endpoints and + ``token_endpoint_auth_methods_supported: ["none", ...]`` because DCR + clients (Claude Desktop, MCP Inspector) register as public clients; PKCE + S256 is mandatory in the gateway's authorize flow.""" + request_base_url = get_request_base_url(request) + return { + "issuer": f"{request_base_url}/mcp", + "authorization_endpoint": f"{request_base_url}/authorize", + "token_endpoint": f"{request_base_url}/token", + "registration_endpoint": f"{request_base_url}/register", + "response_types_supported": ["code"], + "scopes_supported": [], + "grant_types_supported": ["authorization_code", "refresh_token"], + "code_challenge_methods_supported": ["S256"], + "token_endpoint_auth_methods_supported": ["none", "client_secret_post"], + } + + +def _raise_404_unless_gateway_dcr_enabled() -> None: + """The aggregate well-known routes exist only under the gateway-level DCR + front door; flag-off they 404 exactly like the previously-absent routes so + discovery behavior is byte-identical for existing deployments.""" + if is_mcp_gateway_dcr_enabled(): + return + raise HTTPException(status_code=404, detail="Not Found") + + +# RFC 9728 path-appended discovery for the aggregate /mcp endpoint. A client +# pointed at {base}/mcp inserts the well-known segment before the resource +# path, so this exact route must exist for aggregate discovery to work at all. +# Declared before the parameterized well-known routes below: Starlette matches +# in registration order, and /.well-known/oauth-authorization-server/{name} +# would otherwise capture the "/mcp" suffix as a server name. +@router.get( + f"/.well-known/oauth-protected-resource{'' if get_server_root_path() == '/' else get_server_root_path()}/mcp" +) +async def oauth_protected_resource_aggregate(request: Request): + """ + OAuth protected resource discovery for the aggregate /mcp endpoint + (gateway-level DCR front door; 404 when the flag is off). + """ + _raise_404_unless_gateway_dcr_enabled() + return _build_aggregate_protected_resource_response(request) + + +@router.get( + f"/.well-known/oauth-authorization-server{'' if get_server_root_path() == '/' else get_server_root_path()}/mcp" +) +async def oauth_authorization_server_aggregate(request: Request): + """ + OAuth authorization server discovery for the aggregate /mcp endpoint, the + RFC 8414 path-inserted form for a client that treats {base}/mcp as its + authorization base URL (gateway-level DCR front door; 404 when the flag + is off, indistinguishable from an unknown server name on the + parameterized route below). + """ + _raise_404_unless_gateway_dcr_enabled() + return _build_aggregate_authorization_server_response(request) + + # Standard MCP pattern: /.well-known/oauth-protected-resource/mcp/{server_name} # This is the pattern expected by standard MCP clients (mcp-inspector, VSCode Copilot) @router.get( @@ -1897,6 +1990,11 @@ def _build_oauth_authorization_server_response( global_mcp_server_manager, ) + # With the gateway-level DCR front door enabled, unnamed discovery keeps + # advertising the gateway's own /authorize, /token, and /register. + if mcp_server_name is None and is_mcp_gateway_dcr_enabled(): + return _build_aggregate_authorization_server_response(request) + request_base_url = get_request_base_url(request) client_ip = IPAddressUtils.get_mcp_client_ip(request) diff --git a/litellm/proxy/_experimental/mcp_server/oauth_utils.py b/litellm/proxy/_experimental/mcp_server/oauth_utils.py index 6edb22dd858..74b56cf424b 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_utils.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_utils.py @@ -70,6 +70,29 @@ def _origin_label(scheme: str, netloc: str) -> str: return f"{scheme}://{netloc}" if netloc else f"{scheme}://" +MCP_GATEWAY_DCR_SETTING = "mcp_gateway_dcr" + + +def is_mcp_gateway_dcr_enabled() -> bool: + """True when ``general_settings.mcp_gateway_dcr`` opts this deployment into + the gateway-level DCR front door for the aggregate ``/mcp`` endpoint: root + OAuth discovery advertises the gateway itself as the authorization server + (instead of resolving the single configured oauth2 server), and the + anonymous aggregate 401 carries the RFC 9728 ``resource_metadata`` + challenge so DCR clients (Claude Desktop, MCP Inspector) can start the + sign-in flow. Off by default; flag-off behavior is unchanged.""" + from litellm.proxy.proxy_server import general_settings # noqa: PLC0415 # circular import at module load + + if not isinstance(general_settings, dict): + return False + raw = general_settings.get(MCP_GATEWAY_DCR_SETTING) + if isinstance(raw, bool): + return raw + if isinstance(raw, str): + return raw.strip().lower() == "true" + return False + + def _resolve_proxy_base_url_env() -> Optional[str]: global _warned_invalid_proxy_base_url configured = os.environ.get("PROXY_BASE_URL", "").strip() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 9375f7481c8..e1fd3bca7bb 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -6131,3 +6131,136 @@ class TestMCPDcrBridgeDelegateAdmission: route="/mcp/bridge_delegate_server", ) assert exc_info.value.status_code == 500 + + +@pytest.mark.asyncio +class TestAggregateGatewayDcrChallenge: + """The mcp_gateway_dcr front door: a 401 on the aggregate /mcp scope must + carry the RFC 9728 resource_metadata challenge pointing at the gateway's + own protected-resource metadata, and must NOT fire for named-server + targets, explicit litellm keys, non-401 failures, or with the flag off.""" + + _FLAG_PATCH_TARGET = "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.is_mcp_gateway_dcr_enabled" + _AUTH_PATCH_TARGET = "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth" + _EXPECTED_RESOURCE_METADATA = 'resource_metadata="http://testserver/.well-known/oauth-protected-resource/mcp"' + + def _scope(self, path="/mcp", extra_headers=()): + return { + "type": "http", + "method": "POST", + "path": path, + "headers": [(b"host", b"testserver"), *extra_headers], + } + + def _auth_401(self): + async def _raise(api_key, request): + raise ProxyException( + message="Authentication Error: Invalid API key", + type="auth_error", + param="api_key", + code=401, + ) + + return _raise + + async def test_challenge_on_anonymous_aggregate_mcp(self): + """Anonymous request to the aggregate /mcp with the flag on: 401 plus + the bare bearer challenge (no error attribute, RFC 6750 section 3.1).""" + with ( + patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), + patch(self._FLAG_PATCH_TARGET, return_value=True), + ): + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(self._scope()) + assert exc_info.value.status_code == 401 + www_authenticate = (exc_info.value.headers or {})["WWW-Authenticate"] + assert www_authenticate == f"Bearer {self._EXPECTED_RESOURCE_METADATA}" + + async def test_challenge_invalid_token_on_failed_bearer(self): + """A bearer that fails LiteLLM admission at aggregate scope (an expired + gateway session, a revoked key) re-challenges with error=invalid_token + so a spec client re-authorizes instead of retrying the dead token.""" + with ( + patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), + patch(self._FLAG_PATCH_TARGET, return_value=True), + ): + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request( + self._scope(extra_headers=((b"authorization", b"Bearer expired-session-token"),)) + ) + assert exc_info.value.status_code == 401 + www_authenticate = (exc_info.value.headers or {})["WWW-Authenticate"] + assert www_authenticate == f'Bearer error="invalid_token", {self._EXPECTED_RESOURCE_METADATA}' + + async def test_no_challenge_when_flag_off(self): + """Flag off: the original admission error propagates untouched, both + with and without a bearer.""" + for extra_headers in ((), ((b"authorization", b"Bearer some-token"),)): + with ( + patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), + patch(self._FLAG_PATCH_TARGET, return_value=False), + ): + with pytest.raises(ProxyException) as exc_info: + await MCPRequestHandler.process_mcp_request(self._scope(extra_headers=extra_headers)) + assert str(exc_info.value.code) == "401" + + async def test_no_challenge_for_explicit_litellm_key(self): + """An explicit x-litellm-api-key declares a litellm-key client; a typo + there must surface the real auth error, never a DCR challenge that + would send SDKs into a sign-in flow.""" + with ( + patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), + patch(self._FLAG_PATCH_TARGET, return_value=True), + ): + with pytest.raises(ProxyException): + await MCPRequestHandler.process_mcp_request( + self._scope(extra_headers=((b"x-litellm-api-key", b"sk-typo"),)) + ) + + async def test_no_challenge_for_named_servers_header(self): + """x-mcp-servers names explicit targets; the per-server challenge paths + own those, so the aggregate challenge must not fire.""" + with ( + patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), + patch(self._FLAG_PATCH_TARGET, return_value=True), + ): + with pytest.raises(ProxyException): + await MCPRequestHandler.process_mcp_request( + self._scope(extra_headers=((b"x-mcp-servers", b"github"),)) + ) + + async def test_no_challenge_for_path_named_server(self): + """/mcp/{server} targets one server; the aggregate challenge must not + fire even when that server does not resolve.""" + with ( + patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), + patch(self._FLAG_PATCH_TARGET, return_value=True), + ): + with pytest.raises(ProxyException): + await MCPRequestHandler.process_mcp_request(self._scope(path="/mcp/github")) + + async def test_no_challenge_for_client_supplied_mcp_auth(self): + """Per-server x-mcp-{alias}-authorization headers mean the caller is + not a cold-start DCR client; keep the original error.""" + with ( + patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), + patch(self._FLAG_PATCH_TARGET, return_value=True), + ): + with pytest.raises(ProxyException): + await MCPRequestHandler.process_mcp_request( + self._scope(extra_headers=((b"x-mcp-github-authorization", b"Bearer upstream"),)) + ) + + async def test_no_challenge_for_non_401_failure(self): + """Only genuine 401s convert to a challenge; a 500 stays a 500.""" + + async def _raise_500(api_key, request): + raise ProxyException(message="boom", type="server_error", param=None, code=500) + + with ( + patch(self._AUTH_PATCH_TARGET, side_effect=_raise_500), + patch(self._FLAG_PATCH_TARGET, return_value=True), + ): + with pytest.raises(ProxyException) as exc_info: + await MCPRequestHandler.process_mcp_request(self._scope()) + assert str(exc_info.value.code) == "500" 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 6f2f24df8fa..30a814c8ea4 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 @@ -7132,371 +7132,166 @@ async def test_token_exchange_unreadable_body_still_renders_oauth_fault(): assert body == {"error": "server_error", "error_description": "upstream token endpoint returned HTTP 400"} +def _patch_gateway_dcr_flag(enabled: bool): + return patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.is_mcp_gateway_dcr_enabled", + return_value=enabled, + ) + + @pytest.mark.asyncio -async def test_persist_dcr_client_for_config_server_uses_side_store(): - """A config.yaml-declared OAuth2 DCR server has no LiteLLM_MCPServerTable row, so - update_mcp_server returns None. The minted client must then persist to the server-scoped - OAuth-client store keyed by server_id (never a shadow server row), overlay onto the in-memory - server so refresh can authenticate this process, and never call update_server(None) (which - previously raised AttributeError on .approval_status, was swallowed, and reported a 200 that - persisted nothing).""" +async def test_gateway_dcr_root_discovery_describes_gateway_not_single_server(): + """Flag on: root discovery must keep describing the gateway as the + authorization server for the aggregate /mcp resource even when exactly one + OAuth2 server exists (flag off, resolution narrows to that server; that + behavior is pinned by test_discovery_root_includes_server_name_prefix).""" + from fastapi import Request + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - _persist_dcr_client_registration, + _build_oauth_authorization_server_response, + _build_oauth_protected_resource_response, ) from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) - from litellm.proxy._types import MCPTransport - from litellm.types.mcp_server.mcp_server_manager import MCPServer - config_server = MCPServer( - server_id="config_faros", - name="config_faros", - server_name="config_faros", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id=None, - client_secret=None, - authorization_url="https://provider.example/oauth/authorize", - token_url="https://provider.example/oauth/token", - registration_url="https://provider.example/oauth/register", - ) + global_mcp_server_manager.registry.clear() + oauth2_server = _create_oauth2_server() + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - mock_upsert = AsyncMock() - mock_update_server = AsyncMock() + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://llm.example.com/" + mock_request.headers = {} - with ( - patch.object(global_mcp_server_manager, "is_config_declared_server", return_value=True), - patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), - patch( - "litellm.proxy._experimental.mcp_server.db.update_mcp_server", - new=AsyncMock(return_value=None), - ), - patch( - "litellm.proxy._experimental.mcp_server.db.get_mcp_server", - new=AsyncMock(return_value=None), - ), - patch( - "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", - new=AsyncMock(return_value=None), - ), - patch( - "litellm.proxy._experimental.mcp_server.db.upsert_mcp_server_oauth_client_credentials", - new=mock_upsert, - ), - patch.object(global_mcp_server_manager, "update_server", new=mock_update_server), - ): - result = await _persist_dcr_client_registration( - mcp_server=config_server, - registration_response={ - "client_id": "minted-client", - "client_secret": "minted-secret", - "token_endpoint_auth_method": "client_secret_basic", - }, - current_redirect_uri="https://proxy.litellm.example/callback", - ) + try: + with _patch_gateway_dcr_flag(True): + authorization_response = _build_oauth_authorization_server_response( + request=mock_request, + mcp_server_name=None, + ) + resource_response = await _build_oauth_protected_resource_response( + request=mock_request, + mcp_server_name=None, + use_standard_pattern=True, + ) - assert result == "persisted" + assert authorization_response["issuer"] == "https://llm.example.com/mcp" + assert authorization_response["authorization_endpoint"] == "https://llm.example.com/authorize" + assert authorization_response["token_endpoint"] == "https://llm.example.com/token" + assert authorization_response["registration_endpoint"] == "https://llm.example.com/register" + assert "none" in authorization_response["token_endpoint_auth_methods_supported"] + assert authorization_response["code_challenge_methods_supported"] == ["S256"] + assert authorization_response["scopes_supported"] == [] - mock_upsert.assert_called_once() - assert mock_upsert.call_args.kwargs["server_id"] == "config_faros" - stored = mock_upsert.call_args.kwargs["credentials"] - assert stored["client_id"] == "minted-client" - assert stored["client_secret"] == "minted-secret" - assert stored["token_endpoint_auth_method"] == "client_secret_basic" - assert stored["redirect_uris"] == ["https://proxy.litellm.example/callback"] - - assert config_server.client_id == "minted-client" - assert config_server.client_secret == "minted-secret" - assert config_server.token_endpoint_auth_method == "client_secret_basic" - - mock_update_server.assert_not_called() + assert resource_response["resource"] == "https://llm.example.com/mcp" + assert resource_response["authorization_servers"] == ["https://llm.example.com/mcp"] + assert resource_response["scopes_supported"] == [] + finally: + global_mcp_server_manager.registry.clear() @pytest.mark.asyncio -async def test_hydrate_config_server_applies_stored_dcr_client(monkeypatch): - """On restart a config server's in-memory object has no client_id; hydration overlays the - persisted DCR client from the server-scoped store, decrypting the encrypted-at-rest blob, so the - refresh_token grant can authenticate as the registered client instead of re-authenticating.""" - import litellm.proxy.common_utils.encrypt_decrypt_utils as enc - from litellm.litellm_core_utils.safe_json_dumps import safe_dumps - from litellm.proxy._experimental.mcp_server.db import encrypt_credentials +async def test_gateway_dcr_named_discovery_unaffected_by_flag(): + """Flag on must not change named-server discovery: a named oauth2 server + still resolves to its own per-server document.""" + from fastapi import Request + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - hydrate_config_server_dcr_client, - ) - from litellm.proxy._types import MCPTransport - from litellm.types.mcp_server.mcp_server_manager import MCPServer - - server = MCPServer( - server_id="config_faros", - name="config_faros", - server_name="config_faros", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id=None, - ) - - monkeypatch.setattr(enc, "_get_salt_key", lambda: "salt-hydrate-key") - stored_blob = safe_dumps( - encrypt_credentials( - credentials={ - "client_id": "stored-client", - "client_secret": "stored-secret", - "token_endpoint_auth_method": "client_secret_basic", - "redirect_uris": ["https://proxy.litellm.example/callback"], - }, - encryption_key="salt-hydrate-key", - ) - ) - assert "stored-client" not in stored_blob and "stored-secret" not in stored_blob - - with ( - patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), - patch( - "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", - new=AsyncMock(return_value=stored_blob), - ), - ): - applied = await hydrate_config_server_dcr_client(server) - - assert applied is True - assert server.client_id == "stored-client" - assert server.client_secret == "stored-secret" - assert server.token_endpoint_auth_method == "client_secret_basic" - - -@pytest.mark.asyncio -async def test_reuse_config_server_reads_store_with_real_crypto(monkeypatch): - """A config-declared server (rowless) keeps its DCR client in the store, so the reuse read - resolves it from the store and decrypts the encrypted-at-rest client, mirroring the write path so - a re-authorize reuses the client instead of re-minting one.""" - import litellm.proxy.common_utils.encrypt_decrypt_utils as enc - from litellm.litellm_core_utils.safe_json_dumps import safe_dumps - from litellm.proxy._experimental.mcp_server.db import encrypt_credentials - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - _reuse_persisted_dcr_client_if_available, + _build_oauth_authorization_server_response, ) from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) - from litellm.proxy._types import MCPTransport - from litellm.types.mcp_server.mcp_server_manager import MCPServer - server = MCPServer( - server_id="config_faros", - name="config_faros", - server_name="config_faros", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id=None, - ) + global_mcp_server_manager.registry.clear() + oauth2_server = _create_oauth2_server() + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - monkeypatch.setattr(enc, "_get_salt_key", lambda: "salt-reuse-key") - blob = safe_dumps( - encrypt_credentials( - credentials={"client_id": "stored-client", "client_secret": "sec", "redirect_uris": ["https://x/callback"]}, - encryption_key="salt-reuse-key", - ) - ) - assert "stored-client" not in blob - store_lookup = AsyncMock(return_value=blob) - with ( - patch.object(global_mcp_server_manager, "is_config_declared_server", return_value=True), - patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), - patch("litellm.proxy._experimental.mcp_server.db.get_mcp_server", new=AsyncMock(return_value=None)), - patch( - "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", - new=store_lookup, - ), - ): - result = await _reuse_persisted_dcr_client_if_available(server, current_redirect_uri="https://x/callback") + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://llm.example.com/" + mock_request.headers = {} - assert result is True - assert server.client_id == "stored-client" - store_lookup.assert_awaited_once() + try: + with _patch_gateway_dcr_flag(True): + response = _build_oauth_authorization_server_response( + request=mock_request, + mcp_server_name="test_oauth", + ) + assert "/test_oauth/authorize" in response["authorization_endpoint"] + assert response["scopes_supported"] == ["read", "write"] + finally: + global_mcp_server_manager.registry.clear() -@pytest.mark.asyncio -async def test_temp_server_is_not_persisted_to_store(): - """A rowless server that is NOT config-declared (a throwaway /server/oauth/session server) must - not leave a permanent store row on persist, and the read must never consult the store for it. Its - minted client is overlaid in memory for the session only.""" - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - _persist_dcr_client_registration, - _reuse_persisted_dcr_client_if_available, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.proxy._types import MCPTransport - from litellm.types.mcp_server.mcp_server_manager import MCPServer +def test_aggregate_wellknown_routes_404_when_flag_off(): + """Flag off, the aggregate well-known routes answer 404 exactly like the + previously-absent routes: discovery behavior is byte-identical for + existing deployments.""" + from fastapi import FastAPI + from fastapi.testclient import TestClient - temp = MCPServer( - server_id="temp-uuid", - name="temp", - server_name="temp", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id=None, - authorization_url="https://p.example/authorize", - token_url="https://p.example/token", - registration_url="https://p.example/register", - ) + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import router - upsert = AsyncMock() - store_read = AsyncMock(return_value=None) - with ( - patch.object(global_mcp_server_manager, "is_config_declared_server", return_value=False), - patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), - patch("litellm.proxy._experimental.mcp_server.db.update_mcp_server", new=AsyncMock(return_value=None)), - patch("litellm.proxy._experimental.mcp_server.db.get_mcp_server", new=AsyncMock(return_value=None)), - patch("litellm.proxy._experimental.mcp_server.db.upsert_mcp_server_oauth_client_credentials", new=upsert), - patch( - "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", - new=store_read, - ), - patch.object(global_mcp_server_manager, "update_server", new=AsyncMock()), - ): - result = await _persist_dcr_client_registration( - temp, {"client_id": "temp-client", "client_secret": "s"}, "https://x/callback" - ) - reused = await _reuse_persisted_dcr_client_if_available( - MCPServer( - server_id="temp-uuid", - name="temp", - server_name="temp", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id=None, - ), - current_redirect_uri="https://x/callback", - ) + app = FastAPI() + app.include_router(router) + client = TestClient(app) - assert result == "persisted" - assert temp.client_id == "temp-client" - upsert.assert_not_called() - store_read.assert_not_called() - assert reused is False + with _patch_gateway_dcr_flag(False): + assert client.get("/.well-known/oauth-protected-resource/mcp").status_code == 404 + assert client.get("/.well-known/oauth-authorization-server/mcp").status_code == 404 -@pytest.mark.asyncio -async def test_hydrate_does_not_overwrite_explicit_config_client_id(): - """An explicit client_id set in config.yaml wins: hydration must not overwrite it with a stale - persisted store client, and must not even read the store when config already supplied a client.""" - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - hydrate_config_server_dcr_client, - ) - from litellm.proxy._types import MCPTransport - from litellm.types.mcp_server.mcp_server_manager import MCPServer +def test_aggregate_wellknown_routes_serve_gateway_metadata_when_flag_on(): + """Flag on, both path-appended aggregate routes serve the gateway + documents. Exercises real routing, so this also pins registration order: + /.well-known/oauth-authorization-server/{name} would otherwise capture + the /mcp suffix as a server name and 404.""" + from fastapi import FastAPI + from fastapi.testclient import TestClient - server = MCPServer( - server_id="config_static", - name="config_static", - server_name="config_static", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id="explicit-from-config", - ) - store_read = AsyncMock( - return_value={"client_id": "stale-store-client", "client_secret": "x", "redirect_uris": []} - ) - with ( - patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), - patch( - "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", - new=store_read, - ), - ): - applied = await hydrate_config_server_dcr_client(server) - - assert applied is False - assert server.client_id == "explicit-from-config" - store_read.assert_not_called() - - -@pytest.mark.asyncio -async def test_reuse_does_not_inherit_store_client_when_a_row_exists(): - """Security: a server that HAS a LiteLLM_MCPServerTable row reads its DCR client only from that - row, never from the server-scoped store. server_id is caller-settable on create, so a submitted - server whose id collides with a config-declared server must not be able to load that config - server's client from the store and send it to its own token endpoint. A row that exists but has - no client_id yields no reusable client and must not fall back to the store.""" - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - _reuse_persisted_dcr_client_if_available, - ) - from litellm.proxy._types import MCPTransport - from litellm.types.mcp_server.mcp_server_manager import MCPServer - - submitted = MCPServer( - server_id="collides_with_config", - name="submitted", - server_name="submitted", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id=None, - ) - - row_without_client = MagicMock() - row_without_client.credentials = None - row_without_client.server_id = "collides_with_config" - store_lookup = AsyncMock( - return_value={"client_id": "config-secret-client", "client_secret": "leak", "redirect_uris": []} - ) - with ( - patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), - patch( - "litellm.proxy._experimental.mcp_server.db.get_mcp_server", - new=AsyncMock(return_value=row_without_client), - ), - patch( - "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", - new=store_lookup, - ), - ): - result = await _reuse_persisted_dcr_client_if_available(submitted, current_redirect_uri="https://x/callback") - - assert result is False - assert submitted.client_id is None - store_lookup.assert_not_called() - - -@pytest.mark.asyncio -async def test_load_servers_from_config_hydrates_dcr_clients(): - """load_servers_from_config must invoke DCR-client hydration so config servers pick up their - persisted client on startup; deleting the call site leaves a restarted server with no client_id - and forces re-authentication on every token expiry.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import router from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) - hydrate_spy = AsyncMock() - with patch.object(global_mcp_server_manager, "_hydrate_config_servers_dcr_clients", new=hydrate_spy): - await global_mcp_server_manager.load_servers_from_config({}) + global_mcp_server_manager.registry.clear() + app = FastAPI() + app.include_router(router) + client = TestClient(app) - hydrate_spy.assert_awaited_once() + with _patch_gateway_dcr_flag(True): + prm = client.get("/.well-known/oauth-protected-resource/mcp") + asm = client.get("/.well-known/oauth-authorization-server/mcp") + + assert prm.status_code == 200 + assert prm.json()["resource"] == "http://testserver/mcp" + assert prm.json()["authorization_servers"] == ["http://testserver/mcp"] + + assert asm.status_code == 200 + assert asm.json()["issuer"] == "http://testserver/mcp" + assert asm.json()["authorization_endpoint"] == "http://testserver/authorize" -@pytest.mark.asyncio -async def test_reload_servers_from_database_hydrates_dcr_clients(): - """load_servers_from_config runs before the DB connects at startup, so its hydration no-ops; - reload_servers_from_database runs after the DB connects and must hydrate config servers' persisted - DCR clients too, or a fresh pod has no client_id for a config server and forces re-authentication - on the first token refresh.""" - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, +def test_is_mcp_gateway_dcr_enabled_reads_general_settings(): + """The flag reader accepts YAML booleans and env-interpolated strings, and + fails closed on anything else.""" + from litellm.proxy._experimental.mcp_server.oauth_utils import ( + is_mcp_gateway_dcr_enabled, ) + from litellm.proxy.proxy_server import general_settings - prisma = MagicMock() - prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[]) - - hydrate_spy = AsyncMock() - with ( - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", - return_value=prisma, - ), - patch.object(global_mcp_server_manager, "_hydrate_config_servers_dcr_clients", new=hydrate_spy), + for raw, expected in ( + (True, True), + (False, False), + ("true", True), + ("True", True), + ("false", False), + ("yes", False), + (1, False), + (None, False), ): - await global_mcp_server_manager.reload_servers_from_database() + with patch.dict(general_settings, {"mcp_gateway_dcr": raw}): + assert is_mcp_gateway_dcr_enabled() is expected, f"raw={raw!r}" - hydrate_spy.assert_awaited_once() + with patch.dict(general_settings, {}, clear=True): + assert is_mcp_gateway_dcr_enabled() is False From 14b1647cd66e1da9216939a128b17a3268eee765 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 14 Jul 2026 17:27:22 -0700 Subject: [PATCH 02/12] refactor(mcp): make the aggregate DCR front door always-on, remove the mcp_gateway_dcr flag The flag guarded no breaking change: the aggregate discovery lives at new /mcp-suffixed routes, the challenge only fires at aggregate scope, and the authorize/token/register/admission arms self-gate on the llm_dcrc_/llm_session_ prefixes. Bare-origin and per-server discovery are left exactly as they were, and a server literally named mcp keeps its own discovery via disambiguation, so turning it on for everyone changes nothing about existing flows. --- .../mcp_server/auth/user_api_key_auth_mcp.py | 5 +- .../mcp_server/discoverable_endpoints.py | 55 +++-- .../_experimental/mcp_server/oauth_utils.py | 23 -- .../auth/test_user_api_key_auth_mcp.py | 24 +- .../mcp_server/test_discoverable_endpoints.py | 212 ++++++------------ 5 files changed, 104 insertions(+), 215 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 418e63468b4..815be2229fc 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -12,7 +12,6 @@ import litellm from litellm._logging import verbose_logger from litellm.proxy._experimental.mcp_server.oauth_utils import ( get_request_base_url, - is_mcp_gateway_dcr_enabled, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( BridgeEnvelopeAdmitted, @@ -133,14 +132,12 @@ def _is_aggregate_gateway_dcr_challenge_scope( ) -> bool: """True when an unauthenticated request to the aggregate ``/mcp`` endpoint should receive the RFC 9728 401 challenge that advertises the gateway as - the authorization server (``mcp_gateway_dcr`` front door). + the authorization server. Fires only for a genuine 401 on the aggregate scope: any named target (path or ``x-mcp-servers``) belongs to the per-server challenge paths, and client-supplied MCP auth headers mean the caller is not a cold-start DCR client. Fails closed to the original admission error otherwise.""" - if not is_mcp_gateway_dcr_enabled(): - return False if not _is_litellm_auth_admission_error(exc): return False if mcp_servers: diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 075ece42b04..08827accf3c 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -42,7 +42,6 @@ from litellm.proxy._experimental.mcp_server.faults import ( from litellm.proxy._experimental.mcp_server.oauth_utils import ( TOKEN_NO_CACHE_HEADERS, get_request_base_url, - is_mcp_gateway_dcr_enabled, validate_trusted_redirect_uri, ) from litellm.proxy.auth.ip_address_utils import IPAddressUtils @@ -1709,12 +1708,6 @@ async def _build_oauth_protected_resource_response( global_mcp_server_manager, ) - # With the gateway-level DCR front door enabled, unnamed discovery - # describes the gateway itself as the authorization server for the - # aggregate /mcp resource instead of narrowing to one server. - if mcp_server_name is None and is_mcp_gateway_dcr_enabled(): - return _build_aggregate_protected_resource_response(request) - request_base_url = get_request_base_url(request) client_ip = IPAddressUtils.get_mcp_client_ip(request) @@ -1889,13 +1882,20 @@ def _build_aggregate_authorization_server_response(request: Request) -> dict: } -def _raise_404_unless_gateway_dcr_enabled() -> None: - """The aggregate well-known routes exist only under the gateway-level DCR - front door; flag-off they 404 exactly like the previously-absent routes so - discovery behavior is byte-identical for existing deployments.""" - if is_mcp_gateway_dcr_enabled(): - return - raise HTTPException(status_code=404, detail="Not Found") +def _mcp_named_server_exists(request: Request) -> bool: + """True when a server literally named ``mcp`` is configured and visible to this caller. + + Its per-server authorization-server document is served at + ``/.well-known/oauth-authorization-server/mcp``, a single segment that collides with the + aggregate path. When such a server exists the real server wins the route, so that + deployment keeps its per-server discovery regardless of whether the aggregate front door + is on.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # circular import with mcp_server_manager at module load + global_mcp_server_manager, + ) + + client_ip = IPAddressUtils.get_mcp_client_ip(request) + return global_mcp_server_manager.get_mcp_server_by_name("mcp", client_ip=client_ip) is not None # RFC 9728 path-appended discovery for the aggregate /mcp endpoint. A client @@ -1909,10 +1909,12 @@ def _raise_404_unless_gateway_dcr_enabled() -> None: ) async def oauth_protected_resource_aggregate(request: Request): """ - OAuth protected resource discovery for the aggregate /mcp endpoint - (gateway-level DCR front door; 404 when the flag is off). + OAuth protected resource discovery for the aggregate /mcp endpoint. + + The single-segment ``/mcp`` path does not collide with any per-server PRM pattern + (those are two-segment: ``/mcp/{server}`` or ``/{server}/mcp``), so this unambiguously + describes the aggregate resource. """ - _raise_404_unless_gateway_dcr_enabled() return _build_aggregate_protected_resource_response(request) @@ -1921,13 +1923,15 @@ async def oauth_protected_resource_aggregate(request: Request): ) async def oauth_authorization_server_aggregate(request: Request): """ - OAuth authorization server discovery for the aggregate /mcp endpoint, the - RFC 8414 path-inserted form for a client that treats {base}/mcp as its - authorization base URL (gateway-level DCR front door; 404 when the flag - is off, indistinguishable from an unknown server name on the - parameterized route below). + OAuth authorization server discovery for the aggregate /mcp endpoint, the RFC 8414 + path-inserted form for a client that treats {base}/mcp as its authorization base URL. + + This single-segment path collides with the parameterized ``/{mcp_server_name}`` route + below, so a server literally named ``mcp`` wins it and keeps its per-server discovery; + only when no such server exists is the aggregate document served. """ - _raise_404_unless_gateway_dcr_enabled() + if _mcp_named_server_exists(request): + return _build_oauth_authorization_server_response(request=request, mcp_server_name="mcp") return _build_aggregate_authorization_server_response(request) @@ -1990,11 +1994,6 @@ def _build_oauth_authorization_server_response( global_mcp_server_manager, ) - # With the gateway-level DCR front door enabled, unnamed discovery keeps - # advertising the gateway's own /authorize, /token, and /register. - if mcp_server_name is None and is_mcp_gateway_dcr_enabled(): - return _build_aggregate_authorization_server_response(request) - request_base_url = get_request_base_url(request) client_ip = IPAddressUtils.get_mcp_client_ip(request) diff --git a/litellm/proxy/_experimental/mcp_server/oauth_utils.py b/litellm/proxy/_experimental/mcp_server/oauth_utils.py index 74b56cf424b..6edb22dd858 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_utils.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_utils.py @@ -70,29 +70,6 @@ def _origin_label(scheme: str, netloc: str) -> str: return f"{scheme}://{netloc}" if netloc else f"{scheme}://" -MCP_GATEWAY_DCR_SETTING = "mcp_gateway_dcr" - - -def is_mcp_gateway_dcr_enabled() -> bool: - """True when ``general_settings.mcp_gateway_dcr`` opts this deployment into - the gateway-level DCR front door for the aggregate ``/mcp`` endpoint: root - OAuth discovery advertises the gateway itself as the authorization server - (instead of resolving the single configured oauth2 server), and the - anonymous aggregate 401 carries the RFC 9728 ``resource_metadata`` - challenge so DCR clients (Claude Desktop, MCP Inspector) can start the - sign-in flow. Off by default; flag-off behavior is unchanged.""" - from litellm.proxy.proxy_server import general_settings # noqa: PLC0415 # circular import at module load - - if not isinstance(general_settings, dict): - return False - raw = general_settings.get(MCP_GATEWAY_DCR_SETTING) - if isinstance(raw, bool): - return raw - if isinstance(raw, str): - return raw.strip().lower() == "true" - return False - - def _resolve_proxy_base_url_env() -> Optional[str]: global _warned_invalid_proxy_base_url configured = os.environ.get("PROXY_BASE_URL", "").strip() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index e1fd3bca7bb..568081e0673 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -6138,9 +6138,8 @@ class TestAggregateGatewayDcrChallenge: """The mcp_gateway_dcr front door: a 401 on the aggregate /mcp scope must carry the RFC 9728 resource_metadata challenge pointing at the gateway's own protected-resource metadata, and must NOT fire for named-server - targets, explicit litellm keys, non-401 failures, or with the flag off.""" + targets, explicit litellm keys, or non-401 failures.""" - _FLAG_PATCH_TARGET = "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.is_mcp_gateway_dcr_enabled" _AUTH_PATCH_TARGET = "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth" _EXPECTED_RESOURCE_METADATA = 'resource_metadata="http://testserver/.well-known/oauth-protected-resource/mcp"' @@ -6164,11 +6163,10 @@ class TestAggregateGatewayDcrChallenge: return _raise async def test_challenge_on_anonymous_aggregate_mcp(self): - """Anonymous request to the aggregate /mcp with the flag on: 401 plus + """Anonymous request to the aggregate /mcp: 401 plus the bare bearer challenge (no error attribute, RFC 6750 section 3.1).""" with ( patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), - patch(self._FLAG_PATCH_TARGET, return_value=True), ): with pytest.raises(HTTPException) as exc_info: await MCPRequestHandler.process_mcp_request(self._scope()) @@ -6182,7 +6180,6 @@ class TestAggregateGatewayDcrChallenge: so a spec client re-authorizes instead of retrying the dead token.""" with ( patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), - patch(self._FLAG_PATCH_TARGET, return_value=True), ): with pytest.raises(HTTPException) as exc_info: await MCPRequestHandler.process_mcp_request( @@ -6192,25 +6189,12 @@ class TestAggregateGatewayDcrChallenge: www_authenticate = (exc_info.value.headers or {})["WWW-Authenticate"] assert www_authenticate == f'Bearer error="invalid_token", {self._EXPECTED_RESOURCE_METADATA}' - async def test_no_challenge_when_flag_off(self): - """Flag off: the original admission error propagates untouched, both - with and without a bearer.""" - for extra_headers in ((), ((b"authorization", b"Bearer some-token"),)): - with ( - patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), - patch(self._FLAG_PATCH_TARGET, return_value=False), - ): - with pytest.raises(ProxyException) as exc_info: - await MCPRequestHandler.process_mcp_request(self._scope(extra_headers=extra_headers)) - assert str(exc_info.value.code) == "401" - async def test_no_challenge_for_explicit_litellm_key(self): """An explicit x-litellm-api-key declares a litellm-key client; a typo there must surface the real auth error, never a DCR challenge that would send SDKs into a sign-in flow.""" with ( patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), - patch(self._FLAG_PATCH_TARGET, return_value=True), ): with pytest.raises(ProxyException): await MCPRequestHandler.process_mcp_request( @@ -6222,7 +6206,6 @@ class TestAggregateGatewayDcrChallenge: own those, so the aggregate challenge must not fire.""" with ( patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), - patch(self._FLAG_PATCH_TARGET, return_value=True), ): with pytest.raises(ProxyException): await MCPRequestHandler.process_mcp_request( @@ -6234,7 +6217,6 @@ class TestAggregateGatewayDcrChallenge: fire even when that server does not resolve.""" with ( patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), - patch(self._FLAG_PATCH_TARGET, return_value=True), ): with pytest.raises(ProxyException): await MCPRequestHandler.process_mcp_request(self._scope(path="/mcp/github")) @@ -6244,7 +6226,6 @@ class TestAggregateGatewayDcrChallenge: not a cold-start DCR client; keep the original error.""" with ( patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), - patch(self._FLAG_PATCH_TARGET, return_value=True), ): with pytest.raises(ProxyException): await MCPRequestHandler.process_mcp_request( @@ -6259,7 +6240,6 @@ class TestAggregateGatewayDcrChallenge: with ( patch(self._AUTH_PATCH_TARGET, side_effect=_raise_500), - patch(self._FLAG_PATCH_TARGET, return_value=True), ): with pytest.raises(ProxyException) as exc_info: await MCPRequestHandler.process_mcp_request(self._scope()) 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 30a814c8ea4..af4b2caca72 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 @@ -7132,19 +7132,74 @@ async def test_token_exchange_unreadable_body_still_renders_oauth_fault(): assert body == {"error": "server_error", "error_description": "upstream token endpoint returned HTTP 400"} -def _patch_gateway_dcr_flag(enabled: bool): - return patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.is_mcp_gateway_dcr_enabled", - return_value=enabled, +def test_aggregate_wellknown_routes_serve_gateway_metadata(): + """Both path-appended aggregate routes serve the gateway documents. Exercises real + routing, so this also pins registration order: the parameterized + /.well-known/oauth-authorization-server/{name} route would otherwise capture the /mcp + suffix as a server name.""" + from fastapi import FastAPI + from fastapi.testclient import TestClient + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import router + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, ) + global_mcp_server_manager.registry.clear() + app = FastAPI() + app.include_router(router) + client = TestClient(app) + + prm = client.get("/.well-known/oauth-protected-resource/mcp") + asm = client.get("/.well-known/oauth-authorization-server/mcp") + + assert prm.status_code == 200 + assert prm.json()["resource"] == "http://testserver/mcp" + assert prm.json()["authorization_servers"] == ["http://testserver/mcp"] + + assert asm.status_code == 200 + assert asm.json()["issuer"] == "http://testserver/mcp" + assert asm.json()["authorization_endpoint"] == "http://testserver/authorize" + assert "none" in asm.json()["token_endpoint_auth_methods_supported"] + + +def test_as_aggregate_route_prefers_a_real_server_named_mcp(): + """A server literally named ``mcp`` wins the single-segment + /.well-known/oauth-authorization-server/mcp route (it collides with the parameterized + /{server_name} route) and keeps its per-server discovery; the aggregate document is + served only when no such server exists.""" + from fastapi import FastAPI + from fastapi.testclient import TestClient + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import router + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + global_mcp_server_manager.registry.clear() + server_named_mcp = _create_oauth2_server(server_id="mcp_srv", name="mcp", server_name="mcp", alias="mcp") + global_mcp_server_manager.registry[server_named_mcp.server_id] = server_named_mcp + app = FastAPI() + app.include_router(router) + client = TestClient(app) + + try: + asm = client.get("/.well-known/oauth-authorization-server/mcp") + assert asm.status_code == 200 + # the real server's own document (issuer is the bare origin, endpoint is /mcp/authorize), + # not the aggregate one (whose issuer would be {base}/mcp) + assert asm.json()["issuer"] == "http://testserver" + assert "/mcp/authorize" in asm.json()["authorization_endpoint"] + finally: + global_mcp_server_manager.registry.clear() + @pytest.mark.asyncio -async def test_gateway_dcr_root_discovery_describes_gateway_not_single_server(): - """Flag on: root discovery must keep describing the gateway as the - authorization server for the aggregate /mcp resource even when exactly one - OAuth2 server exists (flag off, resolution narrows to that server; that - behavior is pinned by test_discovery_root_includes_server_name_prefix).""" +async def test_bare_origin_discovery_resolves_single_server_not_aggregate(): + """The always-on aggregate front door must not change bare-origin discovery: with one + oauth2 server configured, the no-suffix /.well-known/oauth-{authorization-server, + protected-resource} still resolves THAT server, so an existing single-server deployment's + discovery is unchanged. The aggregate document lives only at the /mcp-suffixed routes.""" from fastapi import Request from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( @@ -7164,134 +7219,15 @@ async def test_gateway_dcr_root_discovery_describes_gateway_not_single_server(): mock_request.headers = {} try: - with _patch_gateway_dcr_flag(True): - authorization_response = _build_oauth_authorization_server_response( - request=mock_request, - mcp_server_name=None, - ) - resource_response = await _build_oauth_protected_resource_response( - request=mock_request, - mcp_server_name=None, - use_standard_pattern=True, - ) - - assert authorization_response["issuer"] == "https://llm.example.com/mcp" - assert authorization_response["authorization_endpoint"] == "https://llm.example.com/authorize" - assert authorization_response["token_endpoint"] == "https://llm.example.com/token" - assert authorization_response["registration_endpoint"] == "https://llm.example.com/register" - assert "none" in authorization_response["token_endpoint_auth_methods_supported"] - assert authorization_response["code_challenge_methods_supported"] == ["S256"] - assert authorization_response["scopes_supported"] == [] - - assert resource_response["resource"] == "https://llm.example.com/mcp" - assert resource_response["authorization_servers"] == ["https://llm.example.com/mcp"] - assert resource_response["scopes_supported"] == [] + authorization_response = _build_oauth_authorization_server_response( + request=mock_request, mcp_server_name=None + ) + resource_response = await _build_oauth_protected_resource_response( + request=mock_request, mcp_server_name=None, use_standard_pattern=True + ) + # per-server, not aggregate: the single server's name is in the endpoints + assert "/test_oauth/authorize" in authorization_response["authorization_endpoint"] + assert authorization_response["issuer"] == "https://llm.example.com" + assert resource_response["authorization_servers"] == ["https://llm.example.com/test_oauth"] finally: global_mcp_server_manager.registry.clear() - - -@pytest.mark.asyncio -async def test_gateway_dcr_named_discovery_unaffected_by_flag(): - """Flag on must not change named-server discovery: a named oauth2 server - still resolves to its own per-server document.""" - from fastapi import Request - - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - _build_oauth_authorization_server_response, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - - global_mcp_server_manager.registry.clear() - oauth2_server = _create_oauth2_server() - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - mock_request = MagicMock(spec=Request) - mock_request.base_url = "https://llm.example.com/" - mock_request.headers = {} - - try: - with _patch_gateway_dcr_flag(True): - response = _build_oauth_authorization_server_response( - request=mock_request, - mcp_server_name="test_oauth", - ) - assert "/test_oauth/authorize" in response["authorization_endpoint"] - assert response["scopes_supported"] == ["read", "write"] - finally: - global_mcp_server_manager.registry.clear() - - -def test_aggregate_wellknown_routes_404_when_flag_off(): - """Flag off, the aggregate well-known routes answer 404 exactly like the - previously-absent routes: discovery behavior is byte-identical for - existing deployments.""" - from fastapi import FastAPI - from fastapi.testclient import TestClient - - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import router - - app = FastAPI() - app.include_router(router) - client = TestClient(app) - - with _patch_gateway_dcr_flag(False): - assert client.get("/.well-known/oauth-protected-resource/mcp").status_code == 404 - assert client.get("/.well-known/oauth-authorization-server/mcp").status_code == 404 - - -def test_aggregate_wellknown_routes_serve_gateway_metadata_when_flag_on(): - """Flag on, both path-appended aggregate routes serve the gateway - documents. Exercises real routing, so this also pins registration order: - /.well-known/oauth-authorization-server/{name} would otherwise capture - the /mcp suffix as a server name and 404.""" - from fastapi import FastAPI - from fastapi.testclient import TestClient - - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import router - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - - global_mcp_server_manager.registry.clear() - app = FastAPI() - app.include_router(router) - client = TestClient(app) - - with _patch_gateway_dcr_flag(True): - prm = client.get("/.well-known/oauth-protected-resource/mcp") - asm = client.get("/.well-known/oauth-authorization-server/mcp") - - assert prm.status_code == 200 - assert prm.json()["resource"] == "http://testserver/mcp" - assert prm.json()["authorization_servers"] == ["http://testserver/mcp"] - - assert asm.status_code == 200 - assert asm.json()["issuer"] == "http://testserver/mcp" - assert asm.json()["authorization_endpoint"] == "http://testserver/authorize" - - -def test_is_mcp_gateway_dcr_enabled_reads_general_settings(): - """The flag reader accepts YAML booleans and env-interpolated strings, and - fails closed on anything else.""" - from litellm.proxy._experimental.mcp_server.oauth_utils import ( - is_mcp_gateway_dcr_enabled, - ) - from litellm.proxy.proxy_server import general_settings - - for raw, expected in ( - (True, True), - (False, False), - ("true", True), - ("True", True), - ("false", False), - ("yes", False), - (1, False), - (None, False), - ): - with patch.dict(general_settings, {"mcp_gateway_dcr": raw}): - assert is_mcp_gateway_dcr_enabled() is expected, f"raw={raw!r}" - - with patch.dict(general_settings, {}, clear=True): - assert is_mcp_gateway_dcr_enabled() is False From 5e1050709dec103f021cd6246050fc7d10668012 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 14 Jul 2026 22:37:35 -0700 Subject: [PATCH 03/12] fix(mcp): reserve mcp for the aggregate AS and root-path the discovery challenges Two RFC 9728 / 8414 discovery fixes on the aggregate front door, both raised by Bugbot on this PR The aggregate authorization-server document at /.well-known/oauth-authorization-server/mcp used to defer to a per-server row literally named "mcp", serving issuer {base} while the aggregate protected-resource document advertises {base}/mcp as its authorization server. A spec client following that chain fails the RFC 8414 issuer check and cannot sign in. The single segment /mcp is now reserved for the aggregate so the issuer stays {base}/mcp and matches the protected-resource document; a server named "mcp" keeps its standard two-segment discovery at /.well-known/oauth-authorization-server/mcp/mcp The 401 challenges built the resource_metadata URL as {base}/.well-known/oauth-protected-resource/mcp with no SERVER_ROOT_PATH segment, but the routes are registered with the path-inserted root segment, so a proxy mounted under a sub-path pointed DCR clients at a URL that 404s. Both the aggregate challenge and the pre-existing per-server pass-through challenge now derive the path from one well_known_root_suffix helper that the route registrations also use, so the advertised URL cannot drift from the served route --- .../mcp_server/auth/user_api_key_auth_mcp.py | 5 +- .../mcp_server/discoverable_endpoints.py | 54 +++++-------------- .../_experimental/mcp_server/oauth_utils.py | 12 +++++ .../proxy/_experimental/mcp_server/server.py | 6 ++- .../auth/test_user_api_key_auth_mcp.py | 17 ++++++ .../mcp_server/test_discoverable_endpoints.py | 45 ++++++++++++---- 6 files changed, 87 insertions(+), 52 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 815be2229fc..f1fcc95c532 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -12,6 +12,7 @@ import litellm from litellm._logging import verbose_logger from litellm.proxy._experimental.mcp_server.oauth_utils import ( get_request_base_url, + well_known_root_suffix, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( BridgeEnvelopeAdmitted, @@ -157,7 +158,9 @@ def _aggregate_gateway_dcr_challenge(request: Request, invalid_token: bool) -> H spec-compliant clients to re-authorize rather than retry; a request with no credentials at all gets the bare challenge per RFC 6750 section 3.1.""" error_attr = 'error="invalid_token", ' if invalid_token else "" - resource_metadata_url = f"{get_request_base_url(request)}/.well-known/oauth-protected-resource/mcp" + resource_metadata_url = ( + f"{get_request_base_url(request)}/.well-known/oauth-protected-resource{well_known_root_suffix()}/mcp" + ) return HTTPException( status_code=401, detail={ diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 08827accf3c..9ea452b9aa8 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -43,6 +43,7 @@ from litellm.proxy._experimental.mcp_server.oauth_utils import ( TOKEN_NO_CACHE_HEADERS, get_request_base_url, validate_trusted_redirect_uri, + well_known_root_suffix, ) from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.proxy.common_utils.encrypt_decrypt_utils import ( @@ -50,7 +51,6 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( encrypt_value_helper, ) from litellm.proxy.common_utils.http_parsing_utils import _read_request_body -from litellm.proxy.utils import get_server_root_path from litellm.types.mcp import MCPAuth, MCPCredentials from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -1882,31 +1882,13 @@ def _build_aggregate_authorization_server_response(request: Request) -> dict: } -def _mcp_named_server_exists(request: Request) -> bool: - """True when a server literally named ``mcp`` is configured and visible to this caller. - - Its per-server authorization-server document is served at - ``/.well-known/oauth-authorization-server/mcp``, a single segment that collides with the - aggregate path. When such a server exists the real server wins the route, so that - deployment keeps its per-server discovery regardless of whether the aggregate front door - is on.""" - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # circular import with mcp_server_manager at module load - global_mcp_server_manager, - ) - - client_ip = IPAddressUtils.get_mcp_client_ip(request) - return global_mcp_server_manager.get_mcp_server_by_name("mcp", client_ip=client_ip) is not None - - # RFC 9728 path-appended discovery for the aggregate /mcp endpoint. A client # pointed at {base}/mcp inserts the well-known segment before the resource # path, so this exact route must exist for aggregate discovery to work at all. # Declared before the parameterized well-known routes below: Starlette matches # in registration order, and /.well-known/oauth-authorization-server/{name} # would otherwise capture the "/mcp" suffix as a server name. -@router.get( - f"/.well-known/oauth-protected-resource{'' if get_server_root_path() == '/' else get_server_root_path()}/mcp" -) +@router.get(f"/.well-known/oauth-protected-resource{well_known_root_suffix()}/mcp") async def oauth_protected_resource_aggregate(request: Request): """ OAuth protected resource discovery for the aggregate /mcp endpoint. @@ -1918,28 +1900,26 @@ async def oauth_protected_resource_aggregate(request: Request): return _build_aggregate_protected_resource_response(request) -@router.get( - f"/.well-known/oauth-authorization-server{'' if get_server_root_path() == '/' else get_server_root_path()}/mcp" -) +@router.get(f"/.well-known/oauth-authorization-server{well_known_root_suffix()}/mcp") async def oauth_authorization_server_aggregate(request: Request): """ OAuth authorization server discovery for the aggregate /mcp endpoint, the RFC 8414 path-inserted form for a client that treats {base}/mcp as its authorization base URL. - This single-segment path collides with the parameterized ``/{mcp_server_name}`` route - below, so a server literally named ``mcp`` wins it and keeps its per-server discovery; - only when no such server exists is the aggregate document served. + The single-segment /mcp is reserved for the aggregate so the discovery chain stays + consistent: the aggregate protected-resource document advertises {base}/mcp as its + authorization server, so the document served here must have issuer {base}/mcp. A server + literally named ``mcp`` therefore does not take this route; it keeps its standard + two-segment discovery at /.well-known/oauth-authorization-server/mcp/mcp. Letting the + per-server row win here instead would serve an issuer of {base} against a resource that + advertised {base}/mcp, which fails the RFC 8414 issuer check and breaks the front door. """ - if _mcp_named_server_exists(request): - return _build_oauth_authorization_server_response(request=request, mcp_server_name="mcp") return _build_aggregate_authorization_server_response(request) # Standard MCP pattern: /.well-known/oauth-protected-resource/mcp/{server_name} # This is the pattern expected by standard MCP clients (mcp-inspector, VSCode Copilot) -@router.get( - f"/.well-known/oauth-protected-resource{'' if get_server_root_path() == '/' else get_server_root_path()}/mcp/{{mcp_server_name}}" -) +@router.get(f"/.well-known/oauth-protected-resource{well_known_root_suffix()}/mcp/{{mcp_server_name}}") async def oauth_protected_resource_mcp_standard(request: Request, mcp_server_name: str): """ OAuth protected resource discovery endpoint using standard MCP URL pattern. @@ -1959,9 +1939,7 @@ async def oauth_protected_resource_mcp_standard(request: Request, mcp_server_nam # LiteLLM legacy pattern: /.well-known/oauth-protected-resource/{server_name}/mcp # Kept for backward compatibility with existing deployments -@router.get( - f"/.well-known/oauth-protected-resource{'' if get_server_root_path() == '/' else get_server_root_path()}/{{mcp_server_name}}/mcp" -) +@router.get(f"/.well-known/oauth-protected-resource{well_known_root_suffix()}/{{mcp_server_name}}/mcp") @router.get("/.well-known/oauth-protected-resource") async def oauth_protected_resource_mcp(request: Request, mcp_server_name: Optional[str] = None): """ @@ -2031,9 +2009,7 @@ def _build_oauth_authorization_server_response( # Standard MCP pattern: /.well-known/oauth-authorization-server/mcp/{server_name} -@router.get( - f"/.well-known/oauth-authorization-server{'' if get_server_root_path() == '/' else get_server_root_path()}/mcp/{{mcp_server_name}}" -) +@router.get(f"/.well-known/oauth-authorization-server{well_known_root_suffix()}/mcp/{{mcp_server_name}}") async def oauth_authorization_server_mcp_standard(request: Request, mcp_server_name: str): """ OAuth authorization server discovery endpoint using standard MCP URL pattern. @@ -2048,9 +2024,7 @@ async def oauth_authorization_server_mcp_standard(request: Request, mcp_server_n # LiteLLM legacy pattern and root endpoint -@router.get( - f"/.well-known/oauth-authorization-server{'' if get_server_root_path() == '/' else get_server_root_path()}/{{mcp_server_name}}" -) +@router.get(f"/.well-known/oauth-authorization-server{well_known_root_suffix()}/{{mcp_server_name}}") @router.get("/.well-known/oauth-authorization-server") async def oauth_authorization_server_mcp(request: Request, mcp_server_name: Optional[str] = None): """ diff --git a/litellm/proxy/_experimental/mcp_server/oauth_utils.py b/litellm/proxy/_experimental/mcp_server/oauth_utils.py index 6edb22dd858..ccee3fc8ac0 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_utils.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_utils.py @@ -132,6 +132,18 @@ def get_request_base_url(request: Request) -> str: return urlunparse((scheme, _strip_default_port(scheme, netloc), parsed.path, "", "", "")) +def well_known_root_suffix() -> str: + """The ``SERVER_ROOT_PATH`` segment inserted into a ``.well-known`` path (RFC 8414 / 9728 + path insertion), empty for a root-mounted proxy or an explicit ``/``. + + The discovery route registrations and the 401 challenges that advertise those routes both + derive their path from this one function, so the ``resource_metadata`` URL a client is told + to fetch cannot drift from the route that actually serves it. + """ + root = os.getenv("SERVER_ROOT_PATH", "") + return "" if root == "/" else root + + def validate_loopback_redirect_uri(redirect_uri: str) -> None: """Require a loopback ``redirect_uri`` (OAuth 2.1 §4.1.2.1 + RFC 8252 §7.3 native-app pattern). MCP clients are native apps that listen on diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index a8ab0937124..a9840bdc02f 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -48,6 +48,7 @@ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( ) from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( get_request_base_url, + well_known_root_suffix, ) from litellm.proxy._experimental.mcp_server.exceptions import ( MCPToolResultError, @@ -3525,9 +3526,10 @@ if MCP_AVAILABLE: base_url = get_request_base_url(request) _path = scope.get("_original_path") or scope.get("path", "") or "" + suffix = well_known_root_suffix() if _path.startswith(f"/{server_name}/mcp"): - return f"{base_url}/.well-known/oauth-protected-resource/{server_name}/mcp" - return f"{base_url}/.well-known/oauth-protected-resource/mcp/{server_name}" + return f"{base_url}/.well-known/oauth-protected-resource{suffix}/{server_name}/mcp" + return f"{base_url}/.well-known/oauth-protected-resource{suffix}/mcp/{server_name}" def _get_passthrough_www_authenticate( scope: Scope, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 568081e0673..7b05b8c9dd0 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -6189,6 +6189,23 @@ class TestAggregateGatewayDcrChallenge: www_authenticate = (exc_info.value.headers or {})["WWW-Authenticate"] assert www_authenticate == f'Bearer error="invalid_token", {self._EXPECTED_RESOURCE_METADATA}' + async def test_challenge_inserts_server_root_path(self): + """With SERVER_ROOT_PATH set the resource_metadata URL must carry the same path-inserted + root segment the aggregate PRM route is registered with (both derive it from + well_known_root_suffix), so a DCR client behind a sub-path is pointed at a route that + exists instead of a 404. Regression: the challenge used to hard-code /mcp and omit the + root path the route inserts.""" + import os + + with ( + patch.dict(os.environ, {"SERVER_ROOT_PATH": "/litellm"}), + patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()), + ): + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(self._scope()) + www_authenticate = (exc_info.value.headers or {})["WWW-Authenticate"] + assert 'resource_metadata="http://testserver/.well-known/oauth-protected-resource/litellm/mcp"' in www_authenticate + async def test_no_challenge_for_explicit_litellm_key(self): """An explicit x-litellm-api-key declares a litellm-key client; a typo there must surface the real auth error, never a DCR challenge that 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 af4b2caca72..7e9ff4692b5 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 @@ -7163,11 +7163,13 @@ def test_aggregate_wellknown_routes_serve_gateway_metadata(): assert "none" in asm.json()["token_endpoint_auth_methods_supported"] -def test_as_aggregate_route_prefers_a_real_server_named_mcp(): - """A server literally named ``mcp`` wins the single-segment - /.well-known/oauth-authorization-server/mcp route (it collides with the parameterized - /{server_name} route) and keeps its per-server discovery; the aggregate document is - served only when no such server exists.""" +def test_as_aggregate_route_reserves_mcp_for_the_aggregate(): + """The single-segment /.well-known/oauth-authorization-server/mcp is reserved for the + aggregate even when a server is literally named ``mcp``. The aggregate protected-resource + document advertises {base}/mcp as its authorization server, so the document served here + must carry issuer {base}/mcp for the RFC 8414 issuer check to pass. Letting the per-server + row win (issuer {base}) breaks that chain, so the aggregate wins and the mcp-named server + keeps its standard two-segment discovery at /.well-known/oauth-authorization-server/mcp/mcp.""" from fastapi import FastAPI from fastapi.testclient import TestClient @@ -7186,14 +7188,39 @@ def test_as_aggregate_route_prefers_a_real_server_named_mcp(): try: asm = client.get("/.well-known/oauth-authorization-server/mcp") assert asm.status_code == 200 - # the real server's own document (issuer is the bare origin, endpoint is /mcp/authorize), - # not the aggregate one (whose issuer would be {base}/mcp) - assert asm.json()["issuer"] == "http://testserver" - assert "/mcp/authorize" in asm.json()["authorization_endpoint"] + # the aggregate document, whose issuer matches what the aggregate PRM advertises + assert asm.json()["issuer"] == "http://testserver/mcp" + + prm = client.get("/.well-known/oauth-protected-resource/mcp") + assert prm.status_code == 200 + assert prm.json()["authorization_servers"] == [asm.json()["issuer"]] + + # the mcp-named server keeps its own document on the standard two-segment route + per_server = client.get("/.well-known/oauth-authorization-server/mcp/mcp") + assert per_server.status_code == 200 + assert "/mcp/authorize" in per_server.json()["authorization_endpoint"] finally: global_mcp_server_manager.registry.clear() +def test_well_known_root_suffix_reflects_server_root_path(): + """The single path segment both the discovery routes and the 401 challenges insert for RFC + 8414/9728 path insertion: empty for a root-mounted proxy or an explicit ``/``, the configured + path otherwise. Sharing this one function is what keeps the advertised resource_metadata URL + equal to the route that serves it.""" + import os + from unittest.mock import patch + + from litellm.proxy._experimental.mcp_server.oauth_utils import well_known_root_suffix + + with patch.dict(os.environ, {"SERVER_ROOT_PATH": ""}): + assert well_known_root_suffix() == "" + with patch.dict(os.environ, {"SERVER_ROOT_PATH": "/"}): + assert well_known_root_suffix() == "" + with patch.dict(os.environ, {"SERVER_ROOT_PATH": "/litellm"}): + assert well_known_root_suffix() == "/litellm" + + @pytest.mark.asyncio async def test_bare_origin_discovery_resolves_single_server_not_aggregate(): """The always-on aggregate front door must not change bare-origin discovery: with one From 70bc9523ba94615167b3728efeff6331306f4937 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Wed, 15 Jul 2026 00:14:10 -0700 Subject: [PATCH 04/12] test(mcp): isolate MCP discovery tests from a leaked SERVER_ROOT_PATH tests/test_litellm/proxy/test_custom_proxy.py sets SERVER_ROOT_PATH at import time (its app mounts under a custom path) and never restores it, so in a shared shard the value leaks into the process. The discovery routes and the 401 challenges now read SERVER_ROOT_PATH to path-insert it where they previously ignored it, so a leaked value rewrites every resource_metadata URL and the exact-URL assertions in the delegate, pass-through, and aggregate challenge tests fail depending on shard order An autouse fixture clears SERVER_ROOT_PATH for the MCP discovery tests so they deterministically exercise the default root-mounted deployment; the tests that assert a sub-path deployment set the value explicitly within their own body. No assertion changed; the leak was invisible before only because the code ignored the variable --- .../_experimental/mcp_server/conftest.py | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/conftest.py diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py b/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py new file mode 100644 index 00000000000..b477bf3f406 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/conftest.py @@ -0,0 +1,22 @@ +import os + +import pytest + + +@pytest.fixture(autouse=True) +def _hermetic_server_root_path(): + """Isolate MCP discovery tests from a leaked ``SERVER_ROOT_PATH``. + + ``tests/test_litellm/proxy/test_custom_proxy.py`` sets ``SERVER_ROOT_PATH`` at import time + (its app mounts under a custom path) and never restores it, so in a shared shard the value + leaks into this process. The discovery routes and the 401 challenges read it, so a leaked + value would silently rewrite every ``resource_metadata`` URL and make these tests depend on + shard ordering. Clearing it here pins the default (root-mounted) deployment; a test that + exercises a sub-path deployment sets the value explicitly within its own body. + """ + saved = os.environ.pop("SERVER_ROOT_PATH", None) + try: + yield + finally: + if saved is not None: + os.environ["SERVER_ROOT_PATH"] = saved From 8a57067d4a0e0f17605eb20911a4ba8b84598d35 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Thu, 16 Jul 2026 17:04:00 -0700 Subject: [PATCH 05/12] refactor(mcp): scope the root-path helper to the aggregate front door only The SERVER_ROOT_PATH fix for the per-server pass-through challenge belongs with its sibling in exceptions.py (both fabricate a per-server resource_metadata URL and both omit the root segment), and both are pre-existing paths unrelated to the aggregate discovery this PR adds. Reverting the server.py change keeps this PR to the aggregate front door and avoids leaving the two per-server challenge builders inconsistent; the per-server root-path fix lands as its own change covering both sites. --- .../proxy/_experimental/mcp_server/server.py | 6 +- .../mcp_server/test_discoverable_endpoints.py | 370 ++++++++++++++++++ 2 files changed, 372 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index a9840bdc02f..a8ab0937124 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -48,7 +48,6 @@ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( ) from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( get_request_base_url, - well_known_root_suffix, ) from litellm.proxy._experimental.mcp_server.exceptions import ( MCPToolResultError, @@ -3526,10 +3525,9 @@ if MCP_AVAILABLE: base_url = get_request_base_url(request) _path = scope.get("_original_path") or scope.get("path", "") or "" - suffix = well_known_root_suffix() if _path.startswith(f"/{server_name}/mcp"): - return f"{base_url}/.well-known/oauth-protected-resource{suffix}/{server_name}/mcp" - return f"{base_url}/.well-known/oauth-protected-resource{suffix}/mcp/{server_name}" + return f"{base_url}/.well-known/oauth-protected-resource/{server_name}/mcp" + return f"{base_url}/.well-known/oauth-protected-resource/mcp/{server_name}" def _get_passthrough_www_authenticate( scope: Scope, 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 7e9ff4692b5..eb8b4a89721 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 @@ -7132,6 +7132,376 @@ async def test_token_exchange_unreadable_body_still_renders_oauth_fault(): assert body == {"error": "server_error", "error_description": "upstream token endpoint returned HTTP 400"} +@pytest.mark.asyncio +async def test_persist_dcr_client_for_config_server_uses_side_store(): + """A config.yaml-declared OAuth2 DCR server has no LiteLLM_MCPServerTable row, so + update_mcp_server returns None. The minted client must then persist to the server-scoped + OAuth-client store keyed by server_id (never a shadow server row), overlay onto the in-memory + server so refresh can authenticate this process, and never call update_server(None) (which + previously raised AttributeError on .approval_status, was swallowed, and reported a 200 that + persisted nothing).""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _persist_dcr_client_registration, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + config_server = MCPServer( + server_id="config_faros", + name="config_faros", + server_name="config_faros", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id=None, + client_secret=None, + authorization_url="https://provider.example/oauth/authorize", + token_url="https://provider.example/oauth/token", + registration_url="https://provider.example/oauth/register", + ) + + mock_upsert = AsyncMock() + mock_update_server = AsyncMock() + + with ( + patch.object(global_mcp_server_manager, "is_config_declared_server", return_value=True), + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch( + "litellm.proxy._experimental.mcp_server.db.update_mcp_server", + new=AsyncMock(return_value=None), + ), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server", + new=AsyncMock(return_value=None), + ), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", + new=AsyncMock(return_value=None), + ), + patch( + "litellm.proxy._experimental.mcp_server.db.upsert_mcp_server_oauth_client_credentials", + new=mock_upsert, + ), + patch.object(global_mcp_server_manager, "update_server", new=mock_update_server), + ): + result = await _persist_dcr_client_registration( + mcp_server=config_server, + registration_response={ + "client_id": "minted-client", + "client_secret": "minted-secret", + "token_endpoint_auth_method": "client_secret_basic", + }, + current_redirect_uri="https://proxy.litellm.example/callback", + ) + + assert result == "persisted" + + mock_upsert.assert_called_once() + assert mock_upsert.call_args.kwargs["server_id"] == "config_faros" + stored = mock_upsert.call_args.kwargs["credentials"] + assert stored["client_id"] == "minted-client" + assert stored["client_secret"] == "minted-secret" + assert stored["token_endpoint_auth_method"] == "client_secret_basic" + assert stored["redirect_uris"] == ["https://proxy.litellm.example/callback"] + + assert config_server.client_id == "minted-client" + assert config_server.client_secret == "minted-secret" + assert config_server.token_endpoint_auth_method == "client_secret_basic" + + mock_update_server.assert_not_called() + + +@pytest.mark.asyncio +async def test_hydrate_config_server_applies_stored_dcr_client(monkeypatch): + """On restart a config server's in-memory object has no client_id; hydration overlays the + persisted DCR client from the server-scoped store, decrypting the encrypted-at-rest blob, so the + refresh_token grant can authenticate as the registered client instead of re-authenticating.""" + import litellm.proxy.common_utils.encrypt_decrypt_utils as enc + from litellm.litellm_core_utils.safe_json_dumps import safe_dumps + from litellm.proxy._experimental.mcp_server.db import encrypt_credentials + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + hydrate_config_server_dcr_client, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="config_faros", + name="config_faros", + server_name="config_faros", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id=None, + ) + + monkeypatch.setattr(enc, "_get_salt_key", lambda: "salt-hydrate-key") + stored_blob = safe_dumps( + encrypt_credentials( + credentials={ + "client_id": "stored-client", + "client_secret": "stored-secret", + "token_endpoint_auth_method": "client_secret_basic", + "redirect_uris": ["https://proxy.litellm.example/callback"], + }, + encryption_key="salt-hydrate-key", + ) + ) + assert "stored-client" not in stored_blob and "stored-secret" not in stored_blob + + with ( + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", + new=AsyncMock(return_value=stored_blob), + ), + ): + applied = await hydrate_config_server_dcr_client(server) + + assert applied is True + assert server.client_id == "stored-client" + assert server.client_secret == "stored-secret" + assert server.token_endpoint_auth_method == "client_secret_basic" + + +@pytest.mark.asyncio +async def test_reuse_config_server_reads_store_with_real_crypto(monkeypatch): + """A config-declared server (rowless) keeps its DCR client in the store, so the reuse read + resolves it from the store and decrypts the encrypted-at-rest client, mirroring the write path so + a re-authorize reuses the client instead of re-minting one.""" + import litellm.proxy.common_utils.encrypt_decrypt_utils as enc + from litellm.litellm_core_utils.safe_json_dumps import safe_dumps + from litellm.proxy._experimental.mcp_server.db import encrypt_credentials + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _reuse_persisted_dcr_client_if_available, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="config_faros", + name="config_faros", + server_name="config_faros", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id=None, + ) + + monkeypatch.setattr(enc, "_get_salt_key", lambda: "salt-reuse-key") + blob = safe_dumps( + encrypt_credentials( + credentials={"client_id": "stored-client", "client_secret": "sec", "redirect_uris": ["https://x/callback"]}, + encryption_key="salt-reuse-key", + ) + ) + assert "stored-client" not in blob + store_lookup = AsyncMock(return_value=blob) + with ( + patch.object(global_mcp_server_manager, "is_config_declared_server", return_value=True), + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch("litellm.proxy._experimental.mcp_server.db.get_mcp_server", new=AsyncMock(return_value=None)), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", + new=store_lookup, + ), + ): + result = await _reuse_persisted_dcr_client_if_available(server, current_redirect_uri="https://x/callback") + + assert result is True + assert server.client_id == "stored-client" + store_lookup.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_temp_server_is_not_persisted_to_store(): + """A rowless server that is NOT config-declared (a throwaway /server/oauth/session server) must + not leave a permanent store row on persist, and the read must never consult the store for it. Its + minted client is overlaid in memory for the session only.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _persist_dcr_client_registration, + _reuse_persisted_dcr_client_if_available, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + temp = MCPServer( + server_id="temp-uuid", + name="temp", + server_name="temp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id=None, + authorization_url="https://p.example/authorize", + token_url="https://p.example/token", + registration_url="https://p.example/register", + ) + + upsert = AsyncMock() + store_read = AsyncMock(return_value=None) + with ( + patch.object(global_mcp_server_manager, "is_config_declared_server", return_value=False), + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch("litellm.proxy._experimental.mcp_server.db.update_mcp_server", new=AsyncMock(return_value=None)), + patch("litellm.proxy._experimental.mcp_server.db.get_mcp_server", new=AsyncMock(return_value=None)), + patch("litellm.proxy._experimental.mcp_server.db.upsert_mcp_server_oauth_client_credentials", new=upsert), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", + new=store_read, + ), + patch.object(global_mcp_server_manager, "update_server", new=AsyncMock()), + ): + result = await _persist_dcr_client_registration( + temp, {"client_id": "temp-client", "client_secret": "s"}, "https://x/callback" + ) + reused = await _reuse_persisted_dcr_client_if_available( + MCPServer( + server_id="temp-uuid", + name="temp", + server_name="temp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id=None, + ), + current_redirect_uri="https://x/callback", + ) + + assert result == "persisted" + assert temp.client_id == "temp-client" + upsert.assert_not_called() + store_read.assert_not_called() + assert reused is False + + +@pytest.mark.asyncio +async def test_hydrate_does_not_overwrite_explicit_config_client_id(): + """An explicit client_id set in config.yaml wins: hydration must not overwrite it with a stale + persisted store client, and must not even read the store when config already supplied a client.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + hydrate_config_server_dcr_client, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="config_static", + name="config_static", + server_name="config_static", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="explicit-from-config", + ) + store_read = AsyncMock( + return_value={"client_id": "stale-store-client", "client_secret": "x", "redirect_uris": []} + ) + with ( + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", + new=store_read, + ), + ): + applied = await hydrate_config_server_dcr_client(server) + + assert applied is False + assert server.client_id == "explicit-from-config" + store_read.assert_not_called() + + +@pytest.mark.asyncio +async def test_reuse_does_not_inherit_store_client_when_a_row_exists(): + """Security: a server that HAS a LiteLLM_MCPServerTable row reads its DCR client only from that + row, never from the server-scoped store. server_id is caller-settable on create, so a submitted + server whose id collides with a config-declared server must not be able to load that config + server's client from the store and send it to its own token endpoint. A row that exists but has + no client_id yields no reusable client and must not fall back to the store.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _reuse_persisted_dcr_client_if_available, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + submitted = MCPServer( + server_id="collides_with_config", + name="submitted", + server_name="submitted", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id=None, + ) + + row_without_client = MagicMock() + row_without_client.credentials = None + row_without_client.server_id = "collides_with_config" + store_lookup = AsyncMock( + return_value={"client_id": "config-secret-client", "client_secret": "leak", "redirect_uris": []} + ) + with ( + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server", + new=AsyncMock(return_value=row_without_client), + ), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", + new=store_lookup, + ), + ): + result = await _reuse_persisted_dcr_client_if_available(submitted, current_redirect_uri="https://x/callback") + + assert result is False + assert submitted.client_id is None + store_lookup.assert_not_called() + + +@pytest.mark.asyncio +async def test_load_servers_from_config_hydrates_dcr_clients(): + """load_servers_from_config must invoke DCR-client hydration so config servers pick up their + persisted client on startup; deleting the call site leaves a restarted server with no client_id + and forces re-authentication on every token expiry.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + hydrate_spy = AsyncMock() + with patch.object(global_mcp_server_manager, "_hydrate_config_servers_dcr_clients", new=hydrate_spy): + await global_mcp_server_manager.load_servers_from_config({}) + + hydrate_spy.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_reload_servers_from_database_hydrates_dcr_clients(): + """load_servers_from_config runs before the DB connects at startup, so its hydration no-ops; + reload_servers_from_database runs after the DB connects and must hydrate config servers' persisted + DCR clients too, or a fresh pod has no client_id for a config server and forces re-authentication + on the first token refresh.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + prisma = MagicMock() + prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[]) + + hydrate_spy = AsyncMock() + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=prisma, + ), + patch.object(global_mcp_server_manager, "_hydrate_config_servers_dcr_clients", new=hydrate_spy), + ): + await global_mcp_server_manager.reload_servers_from_database() + + hydrate_spy.assert_awaited_once() + + def test_aggregate_wellknown_routes_serve_gateway_metadata(): """Both path-appended aggregate routes serve the gateway documents. Exercises real routing, so this also pins registration order: the parameterized From cc45d18e9c41b19d9eabe077288f7fcc6a080b11 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 18 Jul 2026 19:24:56 -0700 Subject: [PATCH 06/12] feat(complexity-router): add return_raw_model_name toggle for response model field (#33875) * feat(complexity-router): optionally return raw model name Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): restore asyncio import Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(tests): preserve staging asyncio import Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): drop unused local asyncio import Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(dashboard): add complexity router raw model toggle Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(complexity-router): move metadata key constant to constants.py Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(proxy-tests): preserve module spacing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Krrish Dholakia Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 1 + litellm/proxy/common_request_processing.py | 12 +++++++- litellm/proxy/proxy_server.py | 4 +++ .../complexity_router/complexity_router.py | 7 +++++ .../complexity_router/config.py | 8 +++++ .../proxy_server/test_streaming_helpers.py | 16 ++++++++++ .../proxy/test_common_request_processing.py | 29 +++++++++++++++++-- .../router_strategy/test_complexity_router.py | 24 +++++++++++++++ .../add_model/ComplexityRouterConfig.test.tsx | 14 +++++++++ .../add_model/ComplexityRouterConfig.tsx | 25 +++++++++++++++- .../add_model/add_auto_router_tab.tsx | 2 ++ .../build_complexity_router_config.test.ts | 11 +++++++ .../build_complexity_router_config.ts | 4 +++ .../edit_auto_router_modal.test.ts | 10 +++++++ .../edit_auto_router_modal.tsx | 3 ++ 15 files changed, 166 insertions(+), 4 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 6432e2176c7..05944c81ea2 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1292,6 +1292,7 @@ MAXIMUM_TRACEBACK_LINES_TO_LOG = int(os.getenv("MAXIMUM_TRACEBACK_LINES_TO_LOG", X_LITELLM_DISABLE_CALLBACKS = "x-litellm-disable-callbacks" LITELLM_METADATA_FIELD = "litellm_metadata" OLD_LITELLM_METADATA_FIELD = "metadata" +RETURN_RAW_MODEL_NAME_METADATA_KEY = "_complexity_router_return_raw_model_name" LITELLM_TRUNCATED_PAYLOAD_FIELD = "litellm_truncated" LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE = ( "Truncation is a DB storage safeguard. " diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 1dc0ee3f947..3f9929f81da 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -33,6 +33,7 @@ from litellm.constants import ( LITELLM_DETAILED_TIMING, LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED, MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG, + RETURN_RAW_MODEL_NAME_METADATA_KEY, STREAM_SSE_DATA_PREFIX, ) from litellm.integrations.custom_guardrail import CustomGuardrail @@ -91,6 +92,13 @@ _CLIENT_DISCONNECTED_ERROR_INFORMATION: StandardLoggingPayloadErrorInformation = } +def _should_return_raw_model_name(request_data: dict[str, object]) -> bool: + return any( + isinstance(metadata, dict) and metadata.get(RETURN_RAW_MODEL_NAME_METADATA_KEY) is True + for metadata in (request_data.get("metadata"), request_data.get("litellm_metadata")) + ) + + def _apply_client_disconnect_metadata(target_metadata: Optional[dict[str, object]]) -> None: if target_metadata is None: return @@ -672,6 +680,7 @@ def _override_openai_response_model( response_obj: Any, requested_model: str, log_context: str, + return_raw_model_name: bool = False, ) -> None: """ Force the OpenAI-compatible `model` field in the response to match what the client requested. @@ -695,7 +704,7 @@ def _override_openai_response_model( 3. If this was a fastest_response batch completion, use the winning model's model group name instead of the comma-separated list the client sent. """ - if not requested_model: + if return_raw_model_name or not requested_model: return hidden_params = get_hidden_params_dict(response_obj) @@ -1938,6 +1947,7 @@ class ProxyBaseLLMRequestProcessing: response_obj=response, requested_model=requested_model_from_client, log_context=f"litellm_call_id={logging_obj.litellm_call_id}", + return_raw_model_name=_should_return_raw_model_name(self.data), ) hidden_params = get_hidden_params_dict(response) # get any updated response headers diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index aed345c5db4..3b40abed19e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -291,6 +291,7 @@ from litellm.proxy.caching_routes import router as caching_router from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, _is_azure_model_router_request, + _should_return_raw_model_name, create_response, ) from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy @@ -7076,6 +7077,9 @@ def _restamp_streaming_chunk_model( fallback_was_attempted: bool = False, fallback_model_from_metadata: str | None = None, ) -> tuple[Any, bool]: + if _should_return_raw_model_name(request_data): + return chunk, model_mismatch_logged + target_model = fallback_model_from_metadata if fallback_was_attempted else requested_model_from_client # Always return the client-requested model name (not provider-prefixed internal identifiers) # on streaming chunks. diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 695d8b8aeaa..e5268b5107b 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -23,6 +23,7 @@ from typing import TYPE_CHECKING, Any, Literal, Union, cast from pydantic import BaseModel from litellm._logging import verbose_router_logger +from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger from litellm.types.utils import ModelResponse @@ -956,6 +957,12 @@ class ComplexityRouter(CustomLogger): """ from litellm.types.router import PreRoutingHookResponse + if self.config.return_raw_model_name: + metadata_key = "litellm_metadata" if "litellm_metadata" in request_kwargs else "metadata" + metadata = request_kwargs.setdefault(metadata_key, {}) + if isinstance(metadata, dict): + metadata[RETURN_RAW_MODEL_NAME_METADATA_KEY] = True + use_session_affinity = self.config.session_affinity and not self.config.plugins session_id = self._get_session_id_from_request_kwargs(request_kwargs) if use_session_affinity else None cache_key = self._get_session_affinity_cache_key(session_id, request_kwargs) if session_id is not None else None diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index 17c2c287dde..7437138fbb7 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -311,6 +311,14 @@ class ComplexityRouterConfig(BaseModel): description="Default model to use if tier cannot be determined", ) + return_raw_model_name: bool = Field( + default=False, + description=( + "Return the resolved raw model name in the response model field instead of " + "the client-requested complexity-router alias" + ), + ) + # Classifier strategy classifier_type: Literal["heuristic", "llm"] = Field( default="heuristic", diff --git a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py index 699606b5277..f7e2d276a2e 100644 --- a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py @@ -19,6 +19,7 @@ import json import pytest +from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY import litellm.proxy.proxy_server as ps from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.proxy_server import ( @@ -272,6 +273,21 @@ def test_restamp_streaming_chunk_model_overrides_model_on_basemodel(): assert snapshot == {"model": "gpt-4", "logged": True, "same_object": True} +@pytest.mark.parametrize("return_raw_model_name", [False, True]) +def test_restamp_streaming_chunk_model_respects_raw_model_name_toggle(return_raw_model_name): + chunk = _simple_chunk(model="gpt-4o-mini") + new_chunk, logged = _restamp_streaming_chunk_model( + chunk=chunk, + requested_model_from_client="auto_router/complexity_router", + request_data={"metadata": {RETURN_RAW_MODEL_NAME_METADATA_KEY: return_raw_model_name}}, + model_mismatch_logged=False, + ) + + expected_model = "gpt-4o-mini" if return_raw_model_name else "auto_router/complexity_router" + assert new_chunk.model == expected_model + assert logged is (not return_raw_model_name) + + def test_restamp_streaming_chunk_model_overrides_model_on_dict(): chunk = {"model": "internal", "choices": []} new_chunk, logged = _restamp_streaming_chunk_model( diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index ebfbb46053d..58f81cdad35 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -11,6 +11,7 @@ from fastapi.responses import JSONResponse, StreamingResponse import litellm from litellm._uuid import uuid +from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.opentelemetry import UserAPIKeyAuth from litellm.proxy.common_request_processing import ( @@ -27,6 +28,7 @@ from litellm.proxy.common_request_processing import ( _is_azure_model_router_request, _override_openai_response_model, _parse_event_data_for_error, + _should_return_raw_model_name, _UpstreamClosingStreamingResponse, create_response, ) @@ -1675,6 +1677,31 @@ class TestExtractErrorFromSSEChunk: class TestOverrideOpenAIResponseModel: """Tests for _override_openai_response_model function""" + @pytest.mark.parametrize("return_raw_model_name", [False, True]) + def test_raw_model_name_toggle(self, return_raw_model_name): + response_obj = {"model": "gpt-4o-mini"} + + _override_openai_response_model( + response_obj=response_obj, + requested_model="auto_router/complexity_router", + log_context="test_context", + return_raw_model_name=return_raw_model_name, + ) + + expected_model = "gpt-4o-mini" if return_raw_model_name else "auto_router/complexity_router" + assert response_obj["model"] == expected_model + + @pytest.mark.parametrize( + "request_data, expected", + [ + ({"metadata": {}}, False), + ({"metadata": {RETURN_RAW_MODEL_NAME_METADATA_KEY: True}}, True), + ({"litellm_metadata": {RETURN_RAW_MODEL_NAME_METADATA_KEY: True}}, True), + ], + ) + def test_raw_model_name_toggle_metadata(self, request_data, expected): + assert _should_return_raw_model_name(request_data) is expected + def test_override_model_preserves_fallback_model_when_fallback_occurred_object( self, ): @@ -3203,8 +3230,6 @@ class TestDisconnectGatherCleanup: async def test_base_process_llm_request_preserves_llm_error_after_gather( self, monkeypatch ): - import asyncio - import litellm.proxy.common_request_processing as cpr from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 280a0fe072a..ef70687bd97 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -20,6 +20,7 @@ import litellm from litellm import Router from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache +from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY from litellm.router_strategy.complexity_router.complexity_router import ( ComplexityRouter, DimensionScore, @@ -125,6 +126,29 @@ class TestComplexityRouterInit: ) assert router.config.default_model == "fallback-model" + @pytest.mark.asyncio + @pytest.mark.parametrize("return_raw_model_name", [False, True]) + async def test_pre_routing_hook_propagates_raw_model_response_setting( + self, mock_router_instance, basic_config, return_raw_model_name + ): + config = {**basic_config, "return_raw_model_name": return_raw_model_name} + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=config, + ) + request_kwargs = {} + + result = await router.async_pre_routing_hook( + model="test-router", + request_kwargs=request_kwargs, + messages=[{"role": "user", "content": "Hello"}], + ) + + assert result is not None + metadata = request_kwargs.get("metadata", {}) + assert metadata.get(RETURN_RAW_MODEL_NAME_METADATA_KEY, False) is return_raw_model_name + class TestTokenScoring: """Test token count scoring.""" diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx index a2e2ca21d00..6b3c1961468 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx @@ -77,6 +77,20 @@ describe("ComplexityRouterConfig", () => { expect(screen.queryByText("Classifier Model")).not.toBeInTheDocument(); }); + it("should toggle returning the raw model name", async () => { + const user = userEvent.setup(); + const onChange = vi.fn(); + renderWithProviders(); + + await user.click(screen.getByText("Advanced: Response Format")); + await user.click(screen.getByRole("switch")); + + expect(onChange).toHaveBeenCalledWith({ + ...defaultValue, + return_raw_model_name: true, + }); + }); + it("should reveal classifier model and timeout fields when llm is selected", () => { const onChange = vi.fn(); renderWithProviders(); diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index 8008012a95c..1f2edf697a9 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -1,5 +1,5 @@ import { InfoCircleOutlined } from "@ant-design/icons"; -import { Select as AntdSelect, Card, Collapse, Divider, Space, Tooltip, Typography } from "antd"; +import { Select as AntdSelect, Card, Collapse, Divider, Space, Switch, Tooltip, Typography } from "antd"; import React from "react"; import { ModelGroup } from "@/components/llm_calls/fetch_models"; import AdaptiveRoutingConfig from "./AdaptiveRoutingConfig"; @@ -44,6 +44,7 @@ export interface ComplexityRouterConfigValue { adaptive_weights?: AdaptiveRouterWeights; tier_distance_penalty?: number; adaptive_eligible?: AdaptiveEligible; + return_raw_model_name?: boolean; } interface ComplexityRouterConfigProps { @@ -218,6 +219,28 @@ const ComplexityRouterConfig: React.FC = ({ ), children: , }, + { + key: "response", + label: ( + + Advanced: Response Format + + ), + children: ( + <> +
+ onChange({ ...value, return_raw_model_name: returnRawModelName })} + /> + Return raw model name +
+ + Return the resolved underlying model name in responses instead of the autorouter alias. + + + ), + }, ...(onEscalationKeywordsChange ? [ { diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx index 6e7bc49afce..e7826e09dce 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx @@ -100,6 +100,7 @@ const AddAutoRouterTab: React.FC = ({ form, handleOk, acc adaptive_weights: adaptiveWeights = DEFAULT_ADAPTIVE_WEIGHTS, tier_distance_penalty: tierDistancePenalty = DEFAULT_TIER_DISTANCE_PENALTY, adaptive_eligible: adaptiveEligible = "all", + return_raw_model_name: returnRawModelName = false, } = complexityRouterConfig; const missingTiersError = getMissingTiersError(tiers); @@ -148,6 +149,7 @@ const AddAutoRouterTab: React.FC = ({ form, handleOk, acc adaptiveWeights, tierDistancePenalty, adaptiveEligible, + returnRawModelName, }; const submitValues = { diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts index 0c9c19d1286..b5973bf7101 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts @@ -26,6 +26,7 @@ const baseParams: BuildComplexityRouterConfigParams = { adaptiveWeights: { quality: 0.3, cost: 0.7 }, tierDistancePenalty: 0.5, adaptiveEligible: "all", + returnRawModelName: false, }; describe("buildComplexityRouterConfig", () => { @@ -164,6 +165,16 @@ describe("buildComplexityRouterConfig", () => { expect(config.adaptive_eligible).toBeUndefined(); }); + it("omits return_raw_model_name when disabled", () => { + const config = buildComplexityRouterConfig({ ...baseParams, returnRawModelName: false }); + expect(config.return_raw_model_name).toBeUndefined(); + }); + + it("includes return_raw_model_name when enabled", () => { + const config = buildComplexityRouterConfig({ ...baseParams, returnRawModelName: true }); + expect(config.return_raw_model_name).toBe(true); + }); + it("includes tier_distance_penalty when adaptive is enabled with eligible='all'", () => { const config = buildComplexityRouterConfig({ ...baseParams, diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts index 0b92dc1b02d..3b41b916611 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts @@ -21,6 +21,7 @@ export interface BuildComplexityRouterConfigParams { adaptiveWeights: AdaptiveRouterWeights; tierDistancePenalty: number; adaptiveEligible: AdaptiveEligible; + returnRawModelName: boolean; } export interface ComplexityRouterConfigPayload { @@ -37,6 +38,7 @@ export interface ComplexityRouterConfigPayload { adaptive_weights?: AdaptiveRouterWeights; tier_distance_penalty?: number; adaptive_eligible?: AdaptiveEligible; + return_raw_model_name?: boolean; } const TIER_KEYS: Array = ["SIMPLE", "MEDIUM", "COMPLEX", "REASONING"]; @@ -76,6 +78,7 @@ export const buildComplexityRouterConfig = ({ adaptiveWeights, tierDistancePenalty, adaptiveEligible, + returnRawModelName, }: BuildComplexityRouterConfigParams): ComplexityRouterConfigPayload => { const cleanedEscalationKeywords = escalationKeywords.map((keyword) => keyword.trim()).filter(Boolean); // Trim keywords and drop empty ones; drop any rule left with no keywords. Clicking @@ -104,5 +107,6 @@ export const buildComplexityRouterConfig = ({ ...(adaptiveEligible === "all" && { tier_distance_penalty: tierDistancePenalty }), adaptive_eligible: adaptiveEligible, }), + ...(returnRawModelName && { return_raw_model_name: true }), }; }; diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts index cd8093928d5..17fa810b529 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts @@ -18,6 +18,7 @@ const storedConfigValue = { adaptive_weights: { quality: 0.3, cost: 0.7 }, tier_distance_penalty: 0.8, adaptive_eligible: "all", + return_raw_model_name: true, }; const storedConfig = JSON.stringify(storedConfigValue); @@ -80,6 +81,15 @@ describe("buildUpdatedComplexityRouterConfig", () => { expect(updatedConfig).toEqual(expectedAdaptiveDisabledConfig); }); + it("includes return_raw_model_name only when enabled", () => { + const updatedConfig = buildUpdatedComplexityRouterConfig(storedConfig, { + ...classifiedTierValue, + return_raw_model_name: true, + }); + + expect(updatedConfig.return_raw_model_name).toBe(true); + }); + it("updates custom technical keywords when they are edited", () => { const updatedConfig = buildUpdatedComplexityRouterConfig(storedConfig, classifiedTierValue, ["postgres"]); diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx index f85cd16486a..46d7d41d9b3 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx @@ -38,6 +38,7 @@ const MANAGED_COMPLEXITY_ROUTER_KEYS = new Set([ "adaptive_weights", "tier_distance_penalty", "adaptive_eligible", + "return_raw_model_name", ]); const toRecord = (value: unknown): Record => { @@ -78,6 +79,7 @@ export const buildUpdatedComplexityRouterConfig = ( }), adaptive_eligible: adaptiveEligible, }), + ...(value.return_raw_model_name && { return_raw_model_name: true }), }; }; @@ -158,6 +160,7 @@ const EditAutoRouterModal: React.FC = ({ adaptive_weights: parsedConfig.adaptive_weights, tier_distance_penalty: parsedConfig.tier_distance_penalty, adaptive_eligible: parsedConfig.adaptive_eligible || "all", + return_raw_model_name: parsedConfig.return_raw_model_name || false, }); setCustomTechnicalKeywords( Array.isArray(parsedConfig.custom_technical_keywords) ? parsedConfig.custom_technical_keywords : [], From bd44c9e305b89526d4c5d773ee39ca935561b9c8 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Sat, 18 Jul 2026 20:36:51 -0700 Subject: [PATCH 07/12] fix(langfuse): send v4 ingestion header for otel callback (#33907) * fix(langfuse): send v4 ingestion header for otel callback * refactor(langfuse): inline otel ingestion header literals * test(langfuse): assert v4 ingestion header on dynamic key config paths * style: apply ruff format to langfuse otel header changes * chore(langfuse): drop stale development annotation on json import --------- Co-authored-by: Hassieb Pakzad <68423100+hassiebp@users.noreply.github.com> --- .../integrations/langfuse/langfuse_otel.py | 32 +++++++++++++-- .../integrations/test_langfuse_otel.py | 12 ++++-- tests/test_service_logger_otel.py | 39 +++++++++++++++++++ 3 files changed, 76 insertions(+), 7 deletions(-) diff --git a/litellm/integrations/langfuse/langfuse_otel.py b/litellm/integrations/langfuse/langfuse_otel.py index 449457bd123..d464d55453d 100644 --- a/litellm/integrations/langfuse/langfuse_otel.py +++ b/litellm/integrations/langfuse/langfuse_otel.py @@ -1,8 +1,8 @@ import base64 -import json # <--- NEW +import json import os from datetime import datetime -from typing import TYPE_CHECKING, Any, Optional, Union +from typing import TYPE_CHECKING, Any, Dict, Optional, Union from litellm._logging import verbose_logger from litellm.integrations.arize import _utils @@ -25,6 +25,8 @@ else: LANGFUSE_CLOUD_EU_ENDPOINT = "https://cloud.langfuse.com/api/public/otel" LANGFUSE_CLOUD_US_ENDPOINT = "https://us.cloud.langfuse.com/api/public/otel" +LANGFUSE_INGESTION_VERSION_HEADER = "x-langfuse-ingestion-version" +LANGFUSE_INGESTION_VERSION = "4" class LangfuseOtelLogger(OpenTelemetry): @@ -326,7 +328,9 @@ class LangfuseOtelLogger(OpenTelemetry): return OpenTelemetryConfig( exporter="otlp_http", endpoint=endpoint, - headers=f"Authorization={auth_header}", + headers=LangfuseOtelLogger._format_otel_headers( + LangfuseOtelLogger._build_langfuse_otel_headers(auth_header) + ), ) @staticmethod @@ -338,6 +342,26 @@ class LangfuseOtelLogger(OpenTelemetry): auth_header = base64.b64encode(auth_string.encode()).decode() return f"Basic {auth_header}" + @staticmethod + def _build_langfuse_otel_headers(auth_header: str) -> Dict[str, str]: + """ + Build the OTLP header set Langfuse expects. + + `x-langfuse-ingestion-version: 4` selects Langfuse's v4 ingestion path; + without it spans fall back to the older transformation path. + """ + return { + "Authorization": auth_header, + LANGFUSE_INGESTION_VERSION_HEADER: LANGFUSE_INGESTION_VERSION, + } + + @staticmethod + def _format_otel_headers(headers: Dict[str, str]) -> str: + """ + Serialize a header mapping into the comma-separated OTLP header string + """ + return ",".join(f"{key}={value}" for key, value in headers.items()) + def construct_dynamic_otel_headers( self, standard_callback_dynamic_params: StandardCallbackDynamicParams ) -> Optional[dict]: @@ -358,7 +382,7 @@ class LangfuseOtelLogger(OpenTelemetry): public_key=dynamic_langfuse_public_key, secret_key=dynamic_langfuse_secret_key, ) - dynamic_headers["Authorization"] = auth_header + dynamic_headers.update(LangfuseOtelLogger._build_langfuse_otel_headers(auth_header)) return dynamic_headers diff --git a/tests/test_litellm/integrations/test_langfuse_otel.py b/tests/test_litellm/integrations/test_langfuse_otel.py index 2f2675ca790..28f138c7acd 100644 --- a/tests/test_litellm/integrations/test_langfuse_otel.py +++ b/tests/test_litellm/integrations/test_langfuse_otel.py @@ -456,7 +456,7 @@ class TestLangfuseOtelKeyDynamicConfig: import base64 expected_auth = base64.b64encode(b"key_public:key_secret").decode() - assert config.headers == f"Authorization=Basic {expected_auth}" + assert config.headers == f"Authorization=Basic {expected_auth},x-langfuse-ingestion-version=4" def test_construct_dynamic_otel_config_host_without_protocol(self): with self._clean_env(): @@ -521,7 +521,10 @@ class TestLangfuseOtelKeyDynamicConfig: import base64 expected_auth = base64.b64encode(b"key_public:key_secret").decode() - assert exporter._headers == {"Authorization": f"Basic {expected_auth}"} + assert exporter._headers == { + "Authorization": f"Basic {expected_auth}", + "x-langfuse-ingestion-version": "4", + } def test_key_dynamic_params_reuse_cached_provider(self): with self._clean_env(): @@ -574,7 +577,10 @@ class TestLangfuseOtelKeyDynamicConfig: provider = next(iter(logger._tracer_provider_cache.values())) exporter = provider._active_span_processor._span_processors[0].span_exporter assert isinstance(exporter, OTLPSpanExporter) - assert exporter._headers == {"Authorization": f"Basic {secret}"} + assert exporter._headers == { + "Authorization": f"Basic {secret}", + "x-langfuse-ingestion-version": "4", + } class TestLangfuseOtelResponsesAPI: diff --git a/tests/test_service_logger_otel.py b/tests/test_service_logger_otel.py index 35070d55546..044d37d6781 100644 --- a/tests/test_service_logger_otel.py +++ b/tests/test_service_logger_otel.py @@ -12,6 +12,7 @@ from litellm.integrations.langfuse.langfuse_otel import LangfuseOtelLogger from litellm.integrations.opentelemetry import OpenTelemetry from litellm.types.services import ServiceTypes from litellm._service_logger import ServiceLogging +from litellm.types.utils import StandardCallbackDynamicParams class TestServiceLoggerOTEL(unittest.IsolatedAsyncioTestCase): @@ -108,6 +109,44 @@ class TestServiceLoggerOTEL(unittest.IsolatedAsyncioTestCase): "Generic OTEL logger should have received the log exactly once.", ) + @patch("litellm.integrations.opentelemetry.OpenTelemetry._init_tracing") + @patch("litellm.integrations.opentelemetry.OpenTelemetry._init_metrics") + @patch("litellm.integrations.opentelemetry.OpenTelemetry._init_logs") + async def test_langfuse_otel_env_config_includes_v4_ingestion_header( + self, mock_logs, mock_metrics, mock_tracing + ): + logger = LangfuseOtelLogger() + + headers = OpenTelemetry._get_headers_dictionary(logger.config.headers) + + self.assertEqual( + headers["x-langfuse-ingestion-version"], + "4", + ) + self.assertTrue(headers["Authorization"].startswith("Basic ")) + + @patch("litellm.integrations.opentelemetry.OpenTelemetry._init_tracing") + @patch("litellm.integrations.opentelemetry.OpenTelemetry._init_metrics") + @patch("litellm.integrations.opentelemetry.OpenTelemetry._init_logs") + async def test_langfuse_otel_dynamic_headers_include_v4_ingestion_header( + self, mock_logs, mock_metrics, mock_tracing + ): + logger = LangfuseOtelLogger() + + headers = logger.construct_dynamic_otel_headers( + StandardCallbackDynamicParams( + langfuse_public_key="pk-lf-dynamic", + langfuse_secret_key="sk-lf-dynamic", + ) + ) + + self.assertIsNotNone(headers) + self.assertEqual( + headers["x-langfuse-ingestion-version"], + "4", + ) + self.assertTrue(headers["Authorization"].startswith("Basic ")) + if __name__ == "__main__": unittest.main() From 34561482ed092d78c296cab7999486022af5a938 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 20 Jul 2026 07:30:49 -0700 Subject: [PATCH 08/12] feat(ui): add configuration tabs to the Cost Optimization page (#33899) * feat(ui): add configuration tabs to Cost Optimization page Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(ui): reuse AutoRouter v2 and Router Settings prompt-caching panel in Cost Optimization; clarify Headroom compression Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(ui): add experimental dashboard banner with feedback discussion link to Cost Optimization Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(ui): add savings methodology note and per-key/team compression enterprise callout to Cost Optimization Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(ui): assert active tab state in Cost Optimization tab-switch test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Krrish Dholakia Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../_components/AutorouterTab.tsx | 28 +++ .../_components/CostOptimizationView.test.tsx | 113 ++--------- .../_components/CostOptimizationView.tsx | 187 +++++------------- .../_components/PromptCachingTab.tsx | 52 +++++ .../_components/PromptCompressionTab.tsx | 176 +++++++++++++++++ .../_components/UsageTab.test.tsx | 108 ++++++++++ .../_components/UsageTab.tsx | 184 +++++++++++++++++ .../_components/helpers.test.ts | 37 ++++ .../cost-optimization/_components/helpers.ts | 39 ++++ .../_components/general_settings.tsx | 4 +- 10 files changed, 699 insertions(+), 229 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutorouterTab.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingTab.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCompressionTab.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/helpers.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/helpers.ts diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutorouterTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutorouterTab.tsx new file mode 100644 index 00000000000..3474036528a --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutorouterTab.tsx @@ -0,0 +1,28 @@ +"use client"; + +import React from "react"; +import { Form } from "antd"; + +import AddAutoRouterTab from "@/components/add_model/add_auto_router_tab"; + +interface AutorouterTabProps { + accessToken: string | null; + userId: string | null; + userRole: string; +} + +const AutorouterTab: React.FC = ({ accessToken, userRole }) => { + const [form] = Form.useForm(); + + if (!accessToken) { + return null; + } + + return ( +
+ form.resetFields()} accessToken={accessToken} userRole={userRole} /> +
+ ); +}; + +export default AutorouterTab; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.test.tsx index 92426d3ce04..46aa23fcfc0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.test.tsx @@ -1,109 +1,34 @@ -import { render } from "@testing-library/react"; +import { fireEvent, render } from "@testing-library/react"; import { describe, expect, it, vi } from "vitest"; -import type { DailyData, SpendMetrics } from "@/components/UsagePage/types"; - -const mockUsePaginatedDailyActivity = vi.fn(); - -vi.mock("@/app/(dashboard)/usage/_components/hooks/usePaginatedDailyActivity", () => ({ - usePaginatedDailyActivity: (args: unknown) => mockUsePaginatedDailyActivity(args), -})); - -vi.mock("@/components/networking", () => ({ - userDailyActivityCall: vi.fn(), -})); - -vi.mock("@/components/shared/advanced_date_picker", () => ({ - __esModule: true, - default: () =>
, -})); - -vi.mock("@/components/shared/charts", () => ({ - AreaChart: ({ data, categories }: { data: unknown; categories: string[] }) => ( -
- ), - DonutChart: ({ data, label }: { data: unknown; label: string }) => ( -
- ), -})); +vi.mock("./UsageTab", () => ({ __esModule: true, default: () =>
})); +vi.mock("./PromptCompressionTab", () => ({ __esModule: true, default: () =>
})); +vi.mock("./AutorouterTab", () => ({ __esModule: true, default: () =>
})); +vi.mock("./PromptCachingTab", () => ({ __esModule: true, default: () =>
})); import CostOptimizationView from "./CostOptimizationView"; -const baseMetrics = (overrides: Partial): SpendMetrics => ({ - spend: 0, - prompt_tokens: 0, - completion_tokens: 0, - total_tokens: 0, - api_requests: 0, - successful_requests: 0, - failed_requests: 0, - cache_read_input_tokens: 0, - cache_creation_input_tokens: 0, - ...overrides, -}); - -const day = (date: string, metrics: Partial): DailyData => ({ - date, - metrics: baseMetrics(metrics), - breakdown: { - models: {}, - model_groups: {}, - mcp_servers: {}, - providers: {}, - api_keys: {}, - entities: {}, - }, -}); - -const renderWith = (results: DailyData[]) => { - mockUsePaginatedDailyActivity.mockReturnValue({ data: { results }, loading: false, isFetchingMore: false }); - return render(); -}; +const renderView = () => render(); describe("CostOptimizationView", () => { - it("sums compression and caching dollars across days into the summary cards", () => { - const { getByText } = renderWith([ - day("2026-07-12", { - compression_savings_spend: 0.04, - prompt_caching_savings_spend: 0.006, - compression_saved_tokens: 40000, - }), - day("2026-07-13", { - compression_savings_spend: 0.1, - prompt_caching_savings_spend: 0.01, - compression_saved_tokens: 100000, - }), - ]); + it("renders all four cost-optimization tabs", () => { + const { getByText } = renderView(); - // compression 0.14 + caching 0.016 = 0.156 - expect(getByText("$0.1560")).toBeInTheDocument(); - expect(getByText("$0.1400")).toBeInTheDocument(); - expect(getByText("$0.0160")).toBeInTheDocument(); - expect(getByText("140,000 tokens compressed")).toBeInTheDocument(); + expect(getByText("Usage")).toBeInTheDocument(); + expect(getByText("Prompt Compression")).toBeInTheDocument(); + expect(getByText("Autorouter")).toBeInTheDocument(); + expect(getByText("Prompt Caching")).toBeInTheDocument(); }); - it("builds a per-day time series and per-driver donut from the daily rows", () => { - const { getByTestId } = renderWith([ - day("2026-07-12", { compression_savings_spend: 0.04, prompt_caching_savings_spend: 0.006 }), - day("2026-07-13", { compression_savings_spend: 0.1, prompt_caching_savings_spend: 0.01 }), - ]); + it("defaults to the Usage tab and switches the active tab on click", () => { + const { getByRole } = renderView(); - const series = JSON.parse(getByTestId("area-chart").getAttribute("data-series") ?? "[]"); - expect(series).toHaveLength(2); - expect(series[0]).toMatchObject({ Compression: 0.04, "Prompt caching": 0.006 }); - expect(series[1]).toMatchObject({ Compression: 0.1, "Prompt caching": 0.01 }); + expect(getByRole("tab", { name: "Usage" })).toHaveAttribute("aria-selected", "true"); + expect(getByRole("tab", { name: "Prompt Compression" })).toHaveAttribute("aria-selected", "false"); - const slices = JSON.parse(getByTestId("donut-chart").getAttribute("data-slices") ?? "[]"); - expect(slices).toEqual([ - { driver: "Compression", usd: expect.closeTo(0.14, 5) }, - { driver: "Prompt caching", usd: expect.closeTo(0.016, 5) }, - ]); - }); + fireEvent.click(getByRole("tab", { name: "Prompt Compression" })); - it("omits a driver slice when that driver has no savings", () => { - const { getByTestId } = renderWith([day("2026-07-12", { compression_savings_spend: 0.04 })]); - - const slices = JSON.parse(getByTestId("donut-chart").getAttribute("data-slices") ?? "[]"); - expect(slices).toEqual([{ driver: "Compression", usd: expect.closeTo(0.04, 5) }]); + expect(getByRole("tab", { name: "Usage" })).toHaveAttribute("aria-selected", "false"); + expect(getByRole("tab", { name: "Prompt Compression" })).toHaveAttribute("aria-selected", "true"); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.tsx index e45e3e23f1d..6e6830b8451 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.tsx @@ -1,16 +1,13 @@ "use client"; -import React, { useMemo, useState } from "react"; +import React from "react"; import { PiggyBank } from "lucide-react"; +import { Alert, Tabs } from "antd"; -import { AreaChart, DonutChart } from "@/components/shared/charts"; -import AdvancedDatePicker from "@/components/shared/advanced_date_picker"; -import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; -import { userDailyActivityCall } from "@/components/networking"; -import { DailyData, SpendMetrics } from "@/components/UsagePage/types"; -import { formatNumberWithCommas } from "@/utils/dataUtils"; -import { all_admin_roles } from "@/utils/roles"; -import { usePaginatedDailyActivity } from "@/app/(dashboard)/usage/_components/hooks/usePaginatedDailyActivity"; +import UsageTab from "./UsageTab"; +import PromptCompressionTab from "./PromptCompressionTab"; +import AutorouterTab from "./AutorouterTab"; +import PromptCachingTab from "./PromptCachingTab"; interface CostOptimizationViewProps { accessToken: string | null; @@ -18,138 +15,62 @@ interface CostOptimizationViewProps { userRole: string; } -type DateRange = { from?: Date; to?: Date }; - -const THIRTY_DAYS_MS = 30 * 24 * 60 * 60 * 1000; - -const usd = (value: number): string => { - const decimals = value > 0 && value < 1 ? 4 : 2; - return `$${formatNumberWithCommas(value, decimals)}`; -}; - -const shortDate = (iso: string): string => - new Date(`${iso}T00:00:00`).toLocaleDateString("en-US", { month: "short", day: "numeric" }); - -const compressionOf = (m: SpendMetrics): number => m.compression_savings_spend ?? 0; -const cachingOf = (m: SpendMetrics): number => m.prompt_caching_savings_spend ?? 0; -const savedTokensOf = (m: SpendMetrics): number => m.compression_saved_tokens ?? 0; - -const SummaryCard = ({ label, value, hint }: { label: string; value: string; hint?: string }) => ( - - - {label} - - -

{value}

- {hint &&

{hint}

} - - -); - const CostOptimizationView: React.FC = ({ accessToken, userId, userRole }) => { - const initialFrom = useMemo(() => new Date(new Date().getTime() - THIRTY_DAYS_MS), []); - const initialTo = useMemo(() => new Date(), []); - const [dateValue, setDateValue] = useState({ from: initialFrom, to: initialTo }); - - const startTime = dateValue.from ?? null; - const endTime = dateValue.to ?? null; - const isAdmin = all_admin_roles.includes(userRole); - const effectiveUserId = isAdmin ? null : userId; - - const { data, loading, isFetchingMore } = usePaginatedDailyActivity({ - fetchFn: userDailyActivityCall, - args: [accessToken, startTime, endTime, effectiveUserId], - enabled: !!accessToken && !!startTime && !!endTime, - }); - - const results = data.results as DailyData[]; - - const compressionTotal = useMemo(() => results.reduce((sum, d) => sum + compressionOf(d.metrics), 0), [results]); - const cachingTotal = useMemo(() => results.reduce((sum, d) => sum + cachingOf(d.metrics), 0), [results]); - const savedTokensTotal = useMemo(() => results.reduce((sum, d) => sum + savedTokensOf(d.metrics), 0), [results]); - const totalSaved = compressionTotal + cachingTotal; - - const overTime = useMemo( - () => - results.map((d) => ({ - date: shortDate(d.date), - Compression: compressionOf(d.metrics), - "Prompt caching": cachingOf(d.metrics), - })), - [results], - ); - - const byDriver = useMemo( - () => - [ - { driver: "Compression", usd: compressionTotal }, - { driver: "Prompt caching", usd: cachingTotal }, - ].filter((d) => d.usd > 0), - [compressionTotal, cachingTotal], - ); + const items = [ + { + key: "usage", + label: "Usage", + children: , + }, + { + key: "compression", + label: "Prompt Compression", + children: , + }, + { + key: "autorouter", + label: "Autorouter", + children: , + }, + { + key: "caching", + label: "Prompt Caching", + children: , + }, + ]; return (
-
-
-
- -

Cost Optimization

-
-

- Money saved by prompt compression and prompt caching across your requests -

+
+
+ +

Cost Optimization

- setDateValue(v)} /> +

+ Track and configure the mechanisms that save you money: prompt compression, prompt caching, and auto routing +

-
- - - -
+ + Have feedback? Join the discussion{" "} + + here + + + } + /> -
- - - Savings over time - - - - - - - - Savings by driver - - - - - -
+
); }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingTab.tsx new file mode 100644 index 00000000000..e6f73088824 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCachingTab.tsx @@ -0,0 +1,52 @@ +"use client"; + +import React, { useCallback, useEffect, useState } from "react"; + +import { getGeneralSettingsCall } from "@/components/networking"; +import NotificationsManager from "@/components/molecules/notifications_manager"; +import { + PromptCachingPanel, + generalSettingsItem, +} from "@/app/(dashboard)/router-settings/_components/general_settings"; + +interface PromptCachingTabProps { + accessToken: string | null; +} + +const PromptCachingTab: React.FC = ({ accessToken }) => { + const [settings, setSettings] = useState([]); + + const loadSettings = useCallback(() => { + if (!accessToken) { + return; + } + getGeneralSettingsCall(accessToken) + .then((data: generalSettingsItem[]) => setSettings(data)) + .catch((error) => { + console.error("Failed to load prompt caching settings:", error); + NotificationsManager.fromBackend("Failed to load prompt caching settings"); + }); + }, [accessToken]); + + useEffect(() => { + loadSettings(); + }, [loadSettings]); + + const handleChange = (fieldName: string, newValue: unknown) => { + setSettings((prev) => + prev.map((setting) => (setting.field_name === fieldName ? { ...setting, field_value: newValue } : setting)), + ); + }; + + if (!accessToken) { + return null; + } + + return ( +
+ +
+ ); +}; + +export default PromptCachingTab; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCompressionTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCompressionTab.tsx new file mode 100644 index 00000000000..071ad1d6bd5 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCompressionTab.tsx @@ -0,0 +1,176 @@ +"use client"; + +import React, { useCallback, useEffect, useState } from "react"; +import { Button, Form, Input, Switch } from "antd"; + +import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; +import { createGuardrailCall, getGuardrailsList } from "@/components/networking"; +import NotificationsManager from "@/components/molecules/notifications_manager"; +import { + buildCompressionGuardrailPayload, + compressionGuardrailsOf, + GuardrailListItem, + GuardrailListResponse, +} from "./helpers"; + +interface PromptCompressionTabProps { + accessToken: string | null; +} + +interface CompressionFormValues { + name: string; + apiBase: string; + defaultOn: boolean; +} + +const PromptCompressionTab: React.FC = ({ accessToken }) => { + const [form] = Form.useForm(); + const [guardrails, setGuardrails] = useState([]); + const [isLoading, setIsLoading] = useState(true); + const [isSaving, setIsSaving] = useState(false); + + const loadGuardrails = useCallback(() => { + if (!accessToken) { + return; + } + getGuardrailsList(accessToken) + .then((response) => setGuardrails(compressionGuardrailsOf(response as GuardrailListResponse))) + .catch((error) => { + console.error("Failed to load compression guardrails:", error); + NotificationsManager.fromBackend("Failed to load compression guardrails"); + }) + .finally(() => setIsLoading(false)); + }, [accessToken]); + + useEffect(() => { + loadGuardrails(); + }, [loadGuardrails]); + + const handleAdd = async (values: CompressionFormValues) => { + if (!accessToken) { + return; + } + setIsSaving(true); + try { + await createGuardrailCall( + accessToken, + buildCompressionGuardrailPayload({ + name: values.name, + apiBase: values.apiBase, + defaultOn: values.defaultOn ?? true, + }), + ); + NotificationsManager.success("Compression guardrail created"); + form.resetFields(); + await loadGuardrails(); + } catch (error) { + console.error("Failed to create compression guardrail:", error); + NotificationsManager.fromBackend("Failed to create compression guardrail"); + } finally { + setIsSaving(false); + } + }; + + return ( +
+ + + Headroom prompt compression + + +

+ Headroom is a native LiteLLM guardrail that compresses your prompts before they reach the model, so you pay + for fewer input tokens. The tokens it removes are priced and shown on the Usage tab as compression savings.{" "} + + Headroom setup docs + +

+ {isLoading &&

Loading...

} + {!isLoading && guardrails.length === 0 && ( +

+ No prompt compression guardrails configured yet. Add one below to start saving on input tokens +

+ )} + {!isLoading && guardrails.length > 0 && ( +
    + {guardrails.map((guardrail) => ( +
  • +
    +

    {guardrail.guardrail_name}

    +

    {guardrail.litellm_params?.api_base ?? ""}

    +
    + + {guardrail.litellm_params?.default_on ? "Always on" : "Opt-in"} + +
    • + ))} +
+ )} + + + + + + Add Headroom compression guardrail + + +
+ + + + + + + + + +
+

+ Applying compression to all requests is available to all users. Enabling it selectively per key or team + is a LiteLLM Enterprise feature. Get a trial key{" "} + + here + +

+
+
+ +
+
+ + +
+ ); +}; + +export default PromptCompressionTab; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx new file mode 100644 index 00000000000..048d9a34d31 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx @@ -0,0 +1,108 @@ +import { render } from "@testing-library/react"; +import { describe, expect, it, vi } from "vitest"; + +import type { DailyData, SpendMetrics } from "@/components/UsagePage/types"; + +const mockUsePaginatedDailyActivity = vi.fn(); + +vi.mock("@/app/(dashboard)/usage/_components/hooks/usePaginatedDailyActivity", () => ({ + usePaginatedDailyActivity: (args: unknown) => mockUsePaginatedDailyActivity(args), +})); + +vi.mock("@/components/networking", () => ({ + userDailyActivityCall: vi.fn(), +})); + +vi.mock("@/components/shared/advanced_date_picker", () => ({ + __esModule: true, + default: () =>
, +})); + +vi.mock("@/components/shared/charts", () => ({ + AreaChart: ({ data, categories }: { data: unknown; categories: string[] }) => ( +
+ ), + DonutChart: ({ data, label }: { data: unknown; label: string }) => ( +
+ ), +})); + +import UsageTab from "./UsageTab"; + +const baseMetrics = (overrides: Partial): SpendMetrics => ({ + spend: 0, + prompt_tokens: 0, + completion_tokens: 0, + total_tokens: 0, + api_requests: 0, + successful_requests: 0, + failed_requests: 0, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + ...overrides, +}); + +const day = (date: string, metrics: Partial): DailyData => ({ + date, + metrics: baseMetrics(metrics), + breakdown: { + models: {}, + model_groups: {}, + mcp_servers: {}, + providers: {}, + api_keys: {}, + entities: {}, + }, +}); + +const renderWith = (results: DailyData[]) => { + mockUsePaginatedDailyActivity.mockReturnValue({ data: { results }, loading: false, isFetchingMore: false }); + return render(); +}; + +describe("UsageTab", () => { + it("sums compression and caching dollars across days into the summary cards", () => { + const { getByText } = renderWith([ + day("2026-07-12", { + compression_savings_spend: 0.04, + prompt_caching_savings_spend: 0.006, + compression_saved_tokens: 40000, + }), + day("2026-07-13", { + compression_savings_spend: 0.1, + prompt_caching_savings_spend: 0.01, + compression_saved_tokens: 100000, + }), + ]); + + expect(getByText("$0.1560")).toBeInTheDocument(); + expect(getByText("$0.1400")).toBeInTheDocument(); + expect(getByText("$0.0160")).toBeInTheDocument(); + expect(getByText("140,000 tokens compressed")).toBeInTheDocument(); + }); + + it("builds a per-day time series and per-driver donut from the daily rows", () => { + const { getByTestId } = renderWith([ + day("2026-07-12", { compression_savings_spend: 0.04, prompt_caching_savings_spend: 0.006 }), + day("2026-07-13", { compression_savings_spend: 0.1, prompt_caching_savings_spend: 0.01 }), + ]); + + const series = JSON.parse(getByTestId("area-chart").getAttribute("data-series") ?? "[]"); + expect(series).toHaveLength(2); + expect(series[0]).toMatchObject({ Compression: 0.04, "Prompt caching": 0.006 }); + expect(series[1]).toMatchObject({ Compression: 0.1, "Prompt caching": 0.01 }); + + const slices = JSON.parse(getByTestId("donut-chart").getAttribute("data-slices") ?? "[]"); + expect(slices).toEqual([ + { driver: "Compression", usd: expect.closeTo(0.14, 5) }, + { driver: "Prompt caching", usd: expect.closeTo(0.016, 5) }, + ]); + }); + + it("omits a driver slice when that driver has no savings", () => { + const { getByTestId } = renderWith([day("2026-07-12", { compression_savings_spend: 0.04 })]); + + const slices = JSON.parse(getByTestId("donut-chart").getAttribute("data-slices") ?? "[]"); + expect(slices).toEqual([{ driver: "Compression", usd: expect.closeTo(0.04, 5) }]); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx new file mode 100644 index 00000000000..8e6fc40b5ad --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx @@ -0,0 +1,184 @@ +"use client"; + +import React, { useMemo, useState } from "react"; +import { Collapse } from "antd"; + +import { AreaChart, DonutChart } from "@/components/shared/charts"; +import AdvancedDatePicker from "@/components/shared/advanced_date_picker"; +import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; +import { userDailyActivityCall } from "@/components/networking"; +import { DailyData, SpendMetrics } from "@/components/UsagePage/types"; +import { formatNumberWithCommas } from "@/utils/dataUtils"; +import { all_admin_roles } from "@/utils/roles"; +import { usePaginatedDailyActivity } from "@/app/(dashboard)/usage/_components/hooks/usePaginatedDailyActivity"; + +interface UsageTabProps { + accessToken: string | null; + userId: string | null; + userRole: string; +} + +type DateRange = { from?: Date; to?: Date }; + +const THIRTY_DAYS_MS = 30 * 24 * 60 * 60 * 1000; + +const usd = (value: number): string => { + const decimals = value > 0 && value < 1 ? 4 : 2; + return `$${formatNumberWithCommas(value, decimals)}`; +}; + +const shortDate = (iso: string): string => + new Date(`${iso}T00:00:00`).toLocaleDateString("en-US", { month: "short", day: "numeric" }); + +const compressionOf = (m: SpendMetrics): number => m.compression_savings_spend ?? 0; +const cachingOf = (m: SpendMetrics): number => m.prompt_caching_savings_spend ?? 0; +const savedTokensOf = (m: SpendMetrics): number => m.compression_saved_tokens ?? 0; + +const MethodologyNote = () => ( + How savings are calculated, + children: ( +
+

+ Savings are computed for each request when it is logged, using the provider's reported usage and the + model's pricing, then summed into a daily rollup. Totals below are read from that rollup over the + selected date range, so the numbers never require a scan of raw request logs. +

+

+ Compression savings are the tokens Headroom removed before the call, priced at the model's input + rate: compression_saved_tokens * input_cost_per_token +

+

+ Prompt caching savings are the tokens the provider served from cache (Anthropic{" "} + cache_read_input_tokens, or OpenAI-style prompt_tokens_details.cached_tokens), + priced at the discount between the normal input rate and the cache-read rate:{" "} + cache_read_input_tokens * max(input_cost_per_token - cache_read_input_token_cost, 0) +

+

+ Total saved is the sum of both drivers. Models without a separate cache-read price in the pricing map + contribute zero caching savings rather than erroring. +

+
+ ), + }, + ]} + /> +); + +const SummaryCard = ({ label, value, hint }: { label: string; value: string; hint?: string }) => ( + + + {label} + + +

{value}

+ {hint &&

{hint}

} + + +); + +const UsageTab: React.FC = ({ accessToken, userId, userRole }) => { + const initialFrom = useMemo(() => new Date(new Date().getTime() - THIRTY_DAYS_MS), []); + const initialTo = useMemo(() => new Date(), []); + const [dateValue, setDateValue] = useState({ from: initialFrom, to: initialTo }); + + const startTime = dateValue.from ?? null; + const endTime = dateValue.to ?? null; + const isAdmin = all_admin_roles.includes(userRole); + const effectiveUserId = isAdmin ? null : userId; + + const { data, loading, isFetchingMore } = usePaginatedDailyActivity({ + fetchFn: userDailyActivityCall, + args: [accessToken, startTime, endTime, effectiveUserId], + enabled: !!accessToken && !!startTime && !!endTime, + }); + + const results = data.results as DailyData[]; + + const compressionTotal = useMemo(() => results.reduce((sum, d) => sum + compressionOf(d.metrics), 0), [results]); + const cachingTotal = useMemo(() => results.reduce((sum, d) => sum + cachingOf(d.metrics), 0), [results]); + const savedTokensTotal = useMemo(() => results.reduce((sum, d) => sum + savedTokensOf(d.metrics), 0), [results]); + const totalSaved = compressionTotal + cachingTotal; + + const overTime = useMemo( + () => + results.map((d) => ({ + date: shortDate(d.date), + Compression: compressionOf(d.metrics), + "Prompt caching": cachingOf(d.metrics), + })), + [results], + ); + + const byDriver = useMemo( + () => + [ + { driver: "Compression", usd: compressionTotal }, + { driver: "Prompt caching", usd: cachingTotal }, + ].filter((d) => d.usd > 0), + [compressionTotal, cachingTotal], + ); + + return ( +
+
+ + setDateValue(v)} /> +
+ +
+ + + +
+ +
+ + + Savings over time + + + + + + + + Savings by driver + + + + + +
+
+ ); +}; + +export default UsageTab; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/helpers.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/helpers.test.ts new file mode 100644 index 00000000000..239cded0815 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/helpers.test.ts @@ -0,0 +1,37 @@ +import { describe, expect, it } from "vitest"; + +import { buildCompressionGuardrailPayload, compressionGuardrailsOf } from "./helpers"; + +describe("compressionGuardrailsOf", () => { + it("keeps only headroom-provider guardrails and drops others", () => { + const filtered = compressionGuardrailsOf({ + guardrails: [ + { guardrail_id: "1", guardrail_name: "headroom-compression", litellm_params: { guardrail: "headroom" } }, + { guardrail_id: "2", guardrail_name: "pii-masker", litellm_params: { guardrail: "presidio" } }, + { guardrail_id: "3", guardrail_name: "no-params", litellm_params: null }, + ], + }); + + expect(filtered.map((g) => g.guardrail_id)).toEqual(["1"]); + }); +}); + +describe("buildCompressionGuardrailPayload", () => { + it("builds a headroom guardrail payload with trimmed fields", () => { + const payload = buildCompressionGuardrailPayload({ + name: " headroom-compression ", + apiBase: " https://compress ", + defaultOn: false, + }); + + expect(payload).toEqual({ + guardrail_name: "headroom-compression", + litellm_params: { + guardrail: "headroom", + mode: "pre_call", + api_base: "https://compress", + default_on: false, + }, + }); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/helpers.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/helpers.ts new file mode 100644 index 00000000000..7c8c92f8890 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/helpers.ts @@ -0,0 +1,39 @@ +export interface GuardrailLitellmParams { + guardrail?: string | null; + api_base?: string | null; + default_on?: boolean | null; +} + +export interface GuardrailListItem { + guardrail_id: string; + guardrail_name: string | null; + litellm_params?: GuardrailLitellmParams | null; +} + +export interface GuardrailListResponse { + guardrails?: GuardrailListItem[]; +} + +export const COMPRESSION_GUARDRAIL_PROVIDER = "headroom"; + +export const isCompressionGuardrail = (guardrail: GuardrailListItem): boolean => + (guardrail.litellm_params?.guardrail ?? "").toLowerCase() === COMPRESSION_GUARDRAIL_PROVIDER; + +export const compressionGuardrailsOf = (response: GuardrailListResponse): GuardrailListItem[] => + (response.guardrails ?? []).filter(isCompressionGuardrail); + +export interface CompressionGuardrailInput { + name: string; + apiBase: string; + defaultOn: boolean; +} + +export const buildCompressionGuardrailPayload = (input: CompressionGuardrailInput): Record => ({ + guardrail_name: input.name.trim(), + litellm_params: { + guardrail: COMPRESSION_GUARDRAIL_PROVIDER, + mode: "pre_call", + api_base: input.apiBase.trim(), + default_on: input.defaultOn, + }, +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx index 1e8658d5104..dda7a23a8d4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx @@ -33,7 +33,7 @@ interface GeneralSettingsPageProps { userID: string | null; } -interface generalSettingsItem { +export interface generalSettingsItem { field_name: string; field_type: string; field_value: any; @@ -90,7 +90,7 @@ const SettingValueEditor: React.FC<{ return null; }; -const PromptCachingPanel: React.FC<{ +export const PromptCachingPanel: React.FC<{ accessToken: string; settings: generalSettingsItem[]; onChange: (fieldName: string, newValue: any) => void; From 8d96e959db0e251d6fbda9d74360617d18aff806 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Mon, 20 Jul 2026 08:47:39 -0700 Subject: [PATCH 09/12] test(e2e): guard destructive spend-log truncate behind an explicit opt-in (#33751) tests/e2e/conftest.py's pytest_sessionfinish truncated LiteLLM_SpendLogs against whatever DATABASE_URL resolved to, gated only by "an e2e test body ran". Pointed at a shared or staging DB, a routine local run wiped real spend data. It also reached the truncate helper through a sys.path.insert into quota_management/spend_tracking/spend_e2e_client.py, a cross-suite import-by-path hack it then unwound in a finally. The cleanup now routes through a new run_spend_log_cleanup in a top-level tests/e2e/e2e_db.py, which fires the destructive truncate only when the operator set E2E_RESET_SPEND_LOGS=1 and an e2e test actually ran. Any other value (unset, 0, true, empty) leaves the DB untouched, so presence of the variable alone or a test run alone never arms the truncate. The decision plus the injectable truncate callable live in that pure helper, and conftest is a thin adapter that supplies os.environ.get(...), the session stash, and reset_spend_logs. reset_spend_logs itself moved from spend_e2e_client.py into e2e_db.py (implementation unchanged), sitting next to e2e_config and lifecycle so both conftest and any suite import it by name; the sys.path munging is gone. Nothing else imported reset_spend_logs, so spend_e2e_client.py drops the definition, its __all__ entry, and the now-unused os import. --- tests/e2e/conftest.py | 35 +++++------- tests/e2e/e2e_db.py | 56 +++++++++++++++++++ .../spend_tracking/spend_e2e_client.py | 19 ------- 3 files changed, 69 insertions(+), 41 deletions(-) create mode 100644 tests/e2e/e2e_db.py diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 5347fffca4d..609da6a9b07 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -14,14 +14,14 @@ shared fixtures build on it. """ import functools -import sys +import os from collections.abc import Iterator -from pathlib import Path import pytest import requests from e2e_config import CONTROL_PLANE_BASE_URL, PROXY_BASE_URL +from e2e_db import RESET_OPT_IN_ENV, reset_spend_logs, run_spend_log_cleanup from junit_properties import attach_result_properties from lifecycle import ProxyClientProvider, ResourceManager from proxy_client import ProxyClient, build_proxy_client @@ -107,26 +107,17 @@ def pytest_runtest_call(item: pytest.Item) -> None: def pytest_sessionfinish(session: pytest.Session, exitstatus: int) -> None: - """Once the whole e2e session is done (all suites), truncate the spend logs so - the DB doesn't accumulate test rows. Sessions where no e2e test body ran leave - the DB alone so a `DATABASE_URL` pointing at a shared instance is never wiped - without an e2e run. Best-effort: a cleanup failure (no DB reachable) must not - fail the run. The spend_tracking dir goes on sys.path only for this import and - is removed after, so a broader `pytest tests/` run is not left with a mutated - path.""" - if not session.stash.get(_E2E_TEST_RAN, False): - return - spend_dir = str(Path(__file__).parent / "quota_management" / "spend_tracking") - sys.path.insert(0, spend_dir) - try: - from spend_e2e_client import reset_spend_logs # pyright: ignore - - reset_spend_logs() - except Exception as exc: # noqa: BLE001 - cleanup is best-effort - print(f"spend-log cleanup best-effort failed: {exc}") - finally: - if spend_dir in sys.path: - sys.path.remove(spend_dir) + """Once the whole e2e session is done (all suites), optionally truncate the + spend logs so the DB doesn't accumulate test rows. The truncate is destructive + and irreversible, so it runs only when the operator explicitly opts in + (`E2E_RESET_SPEND_LOGS=1`) and an e2e test body actually ran; otherwise a + `DATABASE_URL` pointing at a shared or staging instance is left untouched. + Best-effort: a cleanup failure (no DB reachable) must not fail the run.""" + run_spend_log_cleanup( + opt_in=os.environ.get(RESET_OPT_IN_ENV), + e2e_test_ran=session.stash.get(_E2E_TEST_RAN, False), + truncate=reset_spend_logs, + ) @pytest.fixture(scope="session") diff --git a/tests/e2e/e2e_db.py b/tests/e2e/e2e_db.py new file mode 100644 index 00000000000..439dd519c03 --- /dev/null +++ b/tests/e2e/e2e_db.py @@ -0,0 +1,56 @@ +"""Shared, destructive DB helpers for the e2e harness. + +Kept at the top level next to e2e_config and lifecycle so every suite imports it +by name (`from e2e_db import ...`); no suite reaches into another's directory by +mutating sys.path. + +reset_spend_logs truncates LiteLLM_SpendLogs and cannot be undone, so the +session-finish cleanup routes through run_spend_log_cleanup, which fires the +truncate only on an explicit operator opt-in. "An e2e test ran" is necessary but +never sufficient: a DATABASE_URL pointing at a shared or staging instance must +not be wiped by a routine local run that merely exercised a test. +""" + +import os +from collections.abc import Callable + +RESET_OPT_IN_ENV = "E2E_RESET_SPEND_LOGS" + + +def run_spend_log_cleanup( + *, opt_in: str | None, e2e_test_ran: bool, truncate: Callable[[], None] +) -> bool: + """Invoke `truncate` iff the destructive spend-log reset is both opted into + and warranted, returning whether the truncate was attempted. + + The truncate fires only when the opt-in value is exactly "1" AND an e2e test + body actually ran. Any other opt-in value (unset, "0", "true", "") leaves the + DB untouched, so the destructive path is never armed by the env var's mere + presence or by a test run on its own. Best-effort: a truncate failure is + swallowed so cleanup never fails the session, so the returned bool reports + that the reset was attempted, not that the DB call succeeded. + """ + if opt_in != "1" or not e2e_test_ran: + return False + try: + truncate() + except Exception as exc: # noqa: BLE001 - cleanup is best-effort + print(f"spend-log cleanup best-effort failed: {exc}") + return True + + +def reset_spend_logs() -> None: + """Truncate LiteLLM_SpendLogs for a clean slate. No proxy endpoint deletes + spend logs (/global/spend/reset keeps them), so go to the DB directly. Uses + DATABASE_URL (default: the local docker postgres on its mapped host port; the + in-container `@db` host isn't resolvable from the host, so default to + localhost). + """ + import psycopg + + url = os.environ.get( + "DATABASE_URL", + "postgresql://llmproxy:dbpassword9090@localhost:5432/litellm", + ) + with psycopg.connect(url) as conn: + _ = conn.execute('TRUNCATE TABLE "LiteLLM_SpendLogs"') diff --git a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py index 0d49869aa91..26860212fa3 100644 --- a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py +++ b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py @@ -11,7 +11,6 @@ helpers from one place. from __future__ import annotations -import os import time from collections.abc import Callable from dataclasses import dataclass @@ -50,7 +49,6 @@ from models import ( __all__ = [ "SpendClient", "build_client", - "reset_spend_logs", "unique_marker", "unwrap", "is_ok", @@ -59,23 +57,6 @@ __all__ = [ ] -def reset_spend_logs() -> None: - """Truncate LiteLLM_SpendLogs for a clean slate. No proxy endpoint deletes - spend logs (/global/spend/reset keeps them), so go to the DB directly. Uses - DATABASE_URL (default: the local docker postgres on its mapped host port; note - the in-container `@db` host isn't resolvable from the host, so default to - localhost). - """ - import psycopg - - url = os.environ.get( - "DATABASE_URL", - "postgresql://llmproxy:dbpassword9090@localhost:5432/litellm", - ) - with psycopg.connect(url) as conn: - _ = conn.execute('TRUNCATE TABLE "LiteLLM_SpendLogs"') - - def _chat_body( model: str, content: str, From 067c9bbc96fc1fd50c8af1afde0b2f5f021f4801 Mon Sep 17 00:00:00 2001 From: Vineet Puranik <40868710+vineetpuranik@users.noreply.github.com> Date: Mon, 20 Jul 2026 08:52:06 -0700 Subject: [PATCH 10/12] chore(rust): migrate the litellm-rust workspace (core, ai-gateway, python-bridge) from Rust edition 2021 to edition 2024 (#33940) * chore(deps): update cargo.lock file after cargo update * chore(rust): migrate workspace crates to edition 2024 * chore(rust): migrate workspace crates to edition 2024 + fix clippy warnings after 2024 update * chore(rust): add rust version to cargo workspace file * chore(rust): fix clippy collapsible if warning --- litellm-rust/Cargo.lock | 454 +++++++----------- litellm-rust/Cargo.toml | 3 +- .../crates/ai-gateway/src/auth/mod.rs | 2 +- .../crates/ai-gateway/src/io/messages.rs | 2 +- litellm-rust/crates/ai-gateway/src/io/ocr.rs | 2 +- .../crates/ai-gateway/src/io/realtime.rs | 12 +- .../crates/ai-gateway/src/io/realtime_pool.rs | 24 +- .../crates/ai-gateway/src/io/responses_ws.rs | 23 +- litellm-rust/crates/ai-gateway/src/main.rs | 2 +- .../ai-gateway/src/messages/common_utils.rs | 4 +- .../crates/ai-gateway/src/messages/handler.rs | 2 +- .../crates/ai-gateway/src/messages/prepare.rs | 4 +- .../crates/ai-gateway/src/messages/tests.rs | 4 +- .../crates/ai-gateway/src/ocr/common_utils.rs | 4 +- .../crates/ai-gateway/src/ocr/handler.rs | 2 +- .../crates/ai-gateway/src/ocr/hooks.rs | 6 +- litellm-rust/crates/ai-gateway/src/ocr/mod.rs | 4 +- .../crates/ai-gateway/src/ocr/prepare.rs | 2 +- .../crates/ai-gateway/src/ocr/tests.rs | 22 +- .../crates/ai-gateway/src/python/config.rs | 2 +- .../crates/ai-gateway/src/routes/health.rs | 2 +- .../ai-gateway/src/routes/messages/mod.rs | 10 +- .../ai-gateway/src/routes/messages/service.rs | 2 +- .../ai-gateway/src/routes/realtime/mod.rs | 4 +- .../ai-gateway/src/routes/realtime/service.rs | 4 +- .../ai-gateway/src/routes/responses/mod.rs | 4 +- .../core/src/caching/in_memory_cache.rs | 2 +- .../azure_ai/messages/transformation.rs | 2 +- .../providers/azure_ai/ocr/transformation.rs | 16 +- .../core/src/providers/bedrock/aws_base.rs | 16 +- .../providers/mistral/ocr/transformation.rs | 2 +- .../openai/realtime/transformation.rs | 2 +- .../openai/responses/transformation.rs | 4 +- .../providers/vertex_ai/ocr/transformation.rs | 6 +- .../core/src/realtime/transformation.rs | 2 +- .../crates/core/src/responses/websocket.rs | 2 +- litellm-rust/crates/python-bridge/src/lib.rs | 4 +- 37 files changed, 293 insertions(+), 371 deletions(-) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 402d16715a3..ce28f737334 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -13,13 +13,13 @@ dependencies = [ [[package]] name = "async-trait" -version = "0.1.89" +version = "0.1.91" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9035ad2d096bed7955a320ee7e2230574d28fd3c3a0f186cbea1ff3c7eed5dbb" +checksum = "ae36dc4177970ef04fde5178d3e2429882def40e57a451f919c098f72baa6cec" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 3.0.0", ] [[package]] @@ -113,7 +113,7 @@ dependencies = [ "bytes-utils", "fastrand", "http 1.4.2", - "http-body 1.0.1", + "http-body 1.1.0", "percent-encoding", "pin-project-lite", "tracing", @@ -193,7 +193,7 @@ dependencies = [ "futures-core", "futures-util", "http 1.4.2", - "http-body 1.0.1", + "http-body 1.1.0", "http-body-util", "percent-encoding", "pin-project-lite", @@ -222,7 +222,7 @@ dependencies = [ "hyper-util", "pin-project-lite", "rustls 0.21.12", - "rustls 0.23.41", + "rustls 0.23.42", "rustls-native-certs", "rustls-pki-types", "tokio", @@ -279,7 +279,7 @@ dependencies = [ "http 0.2.12", "http 1.4.2", "http-body 0.4.6", - "http-body 1.0.1", + "http-body 1.1.0", "http-body-util", "pin-project-lite", "pin-utils", @@ -313,7 +313,7 @@ checksum = "221eaa237ddf1ca79b60d1372aad77e47f9c0ea5b3ce5099da8c61d027dc77b3" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -340,7 +340,7 @@ dependencies = [ "http 0.2.12", "http 1.4.2", "http-body 0.4.6", - "http-body 1.0.1", + "http-body 1.1.0", "http-body-util", "itoa", "num-integer", @@ -392,7 +392,7 @@ dependencies = [ "bytes", "futures-util", "http 1.4.2", - "http-body 1.0.1", + "http-body 1.1.0", "http-body-util", "hyper 1.10.1", "hyper-util", @@ -427,7 +427,7 @@ dependencies = [ "bytes", "futures-util", "http 1.4.2", - "http-body 1.0.1", + "http-body 1.1.0", "http-body-util", "mime", "pin-project-lite", @@ -456,9 +456,9 @@ dependencies = [ [[package]] name = "bitflags" -version = "2.13.0" +version = "2.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" +checksum = "b588b76d00fde79687d7646a9b5bdf3cc0f655e0bbd080335a95d7e96f3587da" [[package]] name = "block-buffer" @@ -492,9 +492,9 @@ checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" [[package]] name = "bytes" -version = "1.12.0" +version = "1.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8ae3f5d315924270530207e2a68396c3cc547f6dca3fbdca317cfb1a51edb593" +checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04" [[package]] name = "bytes-utils" @@ -508,9 +508,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.2.65" +version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e228eec9be7c17ccb640b59b36a5cd805ea2a564a4c5e162c2f659fea30d3b96" +checksum = "c89588d05638b5b4594a3348a2d6c20277e43a7f5c5202b05cc56888475a47b8" dependencies = [ "find-msvc-tools", "jobserver", @@ -526,9 +526,20 @@ checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" [[package]] name = "cfg_aliases" -version = "0.2.1" +version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" +checksum = "f079e83a288787bcd14a6aea84cee5c87a67c5a3e660c30f557a3d24761b3527" + +[[package]] +name = "chacha20" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d524456ba66e72eb8b115ff89e01e497f8e6d11d78b70b1aa13c0fbd97540a81" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "rand_core 0.10.1", +] [[package]] name = "cmake" @@ -655,7 +666,7 @@ checksum = "1ac70aa55017e108007fbaf5aa0f54b021c98f92ff8af59d42eda9da96e3dd4f" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -678,9 +689,9 @@ checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" [[package]] name = "fastrand" -version = "2.4.1" +version = "2.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9f1f227452a390804cdb637b74a86990f2a7d7ba4b7d5693aac9b4dd6defd8d6" +checksum = "da7c62ceae207dd37ea5b845da6a0696c799f85e97da1ab5b7910be3c1c80223" [[package]] name = "find-msvc-tools" @@ -711,9 +722,9 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" [[package]] name = "futures-channel" -version = "0.3.32" +version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "07bbe89c50d7a535e539b8c17bc0b49bdb77747034daa8087407d655f3f7cc1d" +checksum = "262590f4fe6afeb0bc83be1daa64e52657fe185690a958af7f3ad0e92085c5ae" dependencies = [ "futures-core", "futures-sink", @@ -721,44 +732,44 @@ dependencies = [ [[package]] name = "futures-core" -version = "0.3.32" +version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" +checksum = "2cd50c473c80f6d7c3670a752354b8e569b1a7cbfdc0419ec88e5edad85e0dc7" [[package]] name = "futures-io" -version = "0.3.32" +version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cecba35d7ad927e23624b22ad55235f2239cfa44fd10428eecbeba6d6a717718" +checksum = "4577ecaa3c4f96589d473f679a71b596316f6641bc350038b962a5daf0085d7a" [[package]] name = "futures-macro" -version = "0.3.32" +version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b" +checksum = "2d6d3cde68c518367be28956066ddfef33813991b77a55005a69dae04bf3b10b" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] name = "futures-sink" -version = "0.3.32" +version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c39754e157331b013978ec91992bde1ac089843443c49cbc7f46150b0fad0893" +checksum = "e34418ac499d6305c2fb5ad0ed2f6ac998c5f8ca209b4510f7f94242c647e307" [[package]] name = "futures-task" -version = "0.3.32" +version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "037711b3d59c33004d3856fbdc83b99d4ff37a24768fa1be9ce3538a1cde4393" +checksum = "b231ed28831efb4a61a08580c4bc233ec56bc009f4cd8f52da2c3cb97df0c109" [[package]] name = "futures-util" -version = "0.3.32" +version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" +checksum = "a77a90a256fce34da66415271e30f94ee91c57b04b8a2c042d9cf3220179deaa" dependencies = [ "futures-core", "futures-io", @@ -793,20 +804,6 @@ dependencies = [ "wasm-bindgen", ] -[[package]] -name = "getrandom" -version = "0.3.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" -dependencies = [ - "cfg-if", - "js-sys", - "libc", - "r-efi 5.3.0", - "wasip2", - "wasm-bindgen", -] - [[package]] name = "getrandom" version = "0.4.3" @@ -814,8 +811,11 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" dependencies = [ "cfg-if", + "js-sys", "libc", - "r-efi 6.0.0", + "r-efi", + "rand_core 0.10.1", + "wasm-bindgen", ] [[package]] @@ -917,9 +917,9 @@ dependencies = [ [[package]] name = "http-body" -version = "1.0.1" +version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1efedce1fb8e6913f23e0c92de8e62cd5b772a67e7b3946df930a62566c93184" +checksum = "ca2a8f2913ee65f60facd6a5905613afaa448497a0230cc41ce022d93290bc2c" dependencies = [ "bytes", "http 1.4.2", @@ -927,14 +927,14 @@ dependencies = [ [[package]] name = "http-body-util" -version = "0.1.3" +version = "0.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b021d93e26becf5dc7e1b75b1bed1fd93124b374ceb73f43d4d4eafec896a64a" +checksum = "e9f41fd6a08e4d4ec69df65976da761afd5ad5e58a9d4acb46bd1c953a9e3ff2" dependencies = [ "bytes", "futures-core", "http 1.4.2", - "http-body 1.0.1", + "http-body 1.1.0", "pin-project-lite", ] @@ -995,7 +995,7 @@ dependencies = [ "futures-core", "h2 0.4.15", "http 1.4.2", - "http-body 1.0.1", + "http-body 1.1.0", "httparse", "httpdate", "itoa", @@ -1029,7 +1029,7 @@ dependencies = [ "http 1.4.2", "hyper 1.10.1", "hyper-util", - "rustls 0.23.41", + "rustls 0.23.42", "rustls-native-certs", "tokio", "tokio-rustls 0.26.4", @@ -1048,13 +1048,13 @@ dependencies = [ "futures-channel", "futures-util", "http 1.4.2", - "http-body 1.0.1", + "http-body 1.1.0", "hyper 1.10.1", "ipnet", "libc", "percent-encoding", "pin-project-lite", - "socket2 0.6.4", + "socket2 0.6.5", "tokio", "tower-service", "tracing", @@ -1242,12 +1242,12 @@ dependencies = [ "aws-sigv4", "aws-smithy-runtime-api", "aws-types", - "rand 0.8.6", + "rand 0.8.7", "reqwest", "serde", "serde_json", "sha2 0.10.9", - "thiserror 2.0.18", + "thiserror 2.0.19", "tokio", ] @@ -1289,9 +1289,9 @@ checksum = "0e7465ac9959cc2b1404e8e2367b43684a6d13790fe23056cc8c6c5a6b7bcb94" [[package]] name = "memchr" -version = "2.8.2" +version = "2.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "88904434abc2901f197fe8cc55f0445e7ded921dba5911dad2e2b39b48e663c4" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" [[package]] name = "mime" @@ -1301,9 +1301,9 @@ checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a" [[package]] name = "mio" -version = "1.2.1" +version = "1.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "02bd0af71c67b473010cbbc60715ee815645a4dc942899111f494b4b737d6fda" +checksum = "30d65c71f1ce40ab09135ce117d742b9f8a19ff91a41a8b57ed50bc2de59c427" dependencies = [ "libc", "wasi", @@ -1378,9 +1378,9 @@ checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" [[package]] name = "portable-atomic" -version = "1.13.1" +version = "1.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" +checksum = "3d20d5497ef88037a52ff98267d066e7f11fcc5e99bbfbd58a42336193aacec3" [[package]] name = "potential_utf" @@ -1408,9 +1408,9 @@ dependencies = [ [[package]] name = "proc-macro2" -version = "1.0.106" +version = "1.0.107" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" dependencies = [ "unicode-ident", ] @@ -1471,7 +1471,7 @@ dependencies = [ "proc-macro2", "pyo3-macros-backend", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -1483,7 +1483,7 @@ dependencies = [ "heck", "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -1498,9 +1498,9 @@ dependencies = [ "quinn-proto", "quinn-udp", "rustc-hash", - "rustls 0.23.41", - "socket2 0.6.4", - "thiserror 2.0.18", + "rustls 0.23.42", + "socket2 0.6.5", + "thiserror 2.0.19", "tokio", "tracing", "web-time", @@ -1508,20 +1508,21 @@ dependencies = [ [[package]] name = "quinn-proto" -version = "0.11.15" +version = "0.11.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4fcb935c5bec503c2f0e306bdd3e58bb9029dcb14fa8d9ac76e3a5256ac0763e" +checksum = "2f4bfc015262b9df63c8845072ce59068853ff5872180c2ce2f13038b970e560" dependencies = [ "bytes", - "getrandom 0.3.4", + "getrandom 0.4.3", "lru-slab", - "rand 0.9.4", + "rand 0.10.2", + "rand_pcg", "ring", "rustc-hash", - "rustls 0.23.41", + "rustls 0.23.42", "rustls-pki-types", "slab", - "thiserror 2.0.18", + "thiserror 2.0.19", "tinyvec", "tracing", "web-time", @@ -1529,33 +1530,27 @@ dependencies = [ [[package]] name = "quinn-udp" -version = "0.5.14" +version = "0.5.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "addec6a0dcad8a8d96a771f815f0eaf55f9d1805756410b39f5fa81332574cbd" +checksum = "35a133f956daabe89a61a685c2649f13d82d5aa4bd5d12d1277e1072a21c0694" dependencies = [ "cfg_aliases", "libc", "once_cell", - "socket2 0.6.4", + "socket2 0.6.5", "tracing", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] name = "quote" -version = "1.0.46" +version = "1.0.47" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dfbc457d0c7a0759a614551b11a6409e5951f6c7537be1f1b7682b9ae9230368" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" dependencies = [ "proc-macro2", ] -[[package]] -name = "r-efi" -version = "5.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" - [[package]] name = "r-efi" version = "6.0.0" @@ -1564,23 +1559,24 @@ checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" [[package]] name = "rand" -version = "0.8.6" +version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5ca0ecfa931c29007047d1bc58e623ab12e5590e8c7cc53200d5202b69266d8a" +checksum = "22f6172bdec972074665ed81ed53b71da00bfc44b65a753cfde883ec4c702a1a" dependencies = [ "libc", - "rand_chacha 0.3.1", + "rand_chacha", "rand_core 0.6.4", ] [[package]] name = "rand" -version = "0.9.4" +version = "0.10.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "44c5af06bb1b7d3216d91932aed5265164bf384dc89cd6ba05cf59a35f5f76ea" +checksum = "c7f5fa3a058cd35567ef9bfa5e75732bee0f9e4c55fa90477bef2dfcdbc4be80" dependencies = [ - "rand_chacha 0.9.0", - "rand_core 0.9.5", + "chacha20", + "getrandom 0.4.3", + "rand_core 0.10.1", ] [[package]] @@ -1593,16 +1589,6 @@ dependencies = [ "rand_core 0.6.4", ] -[[package]] -name = "rand_chacha" -version = "0.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" -dependencies = [ - "ppv-lite86", - "rand_core 0.9.5", -] - [[package]] name = "rand_core" version = "0.6.4" @@ -1614,11 +1600,17 @@ dependencies = [ [[package]] name = "rand_core" -version = "0.9.5" +version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" + +[[package]] +name = "rand_pcg" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "caa0f4137e1c0a72f4c651489402276c8e8e1cf081f3b0ba156d2cbeef09e86a" dependencies = [ - "getrandom 0.3.4", + "rand_core 0.10.1", ] [[package]] @@ -1640,7 +1632,7 @@ dependencies = [ "futures-util", "h2 0.4.15", "http 1.4.2", - "http-body 1.0.1", + "http-body 1.1.0", "http-body-util", "hyper 1.10.1", "hyper-rustls 0.27.9", @@ -1650,7 +1642,7 @@ dependencies = [ "percent-encoding", "pin-project-lite", "quinn", - "rustls 0.23.41", + "rustls 0.23.42", "rustls-pki-types", "serde", "serde_json", @@ -1686,9 +1678,9 @@ dependencies = [ [[package]] name = "rustc-hash" -version = "2.1.2" +version = "2.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "94300abf3f1ae2e2b8ffb7b58043de3d399c73fa6f4b73826402a5c457614dbe" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" [[package]] name = "rustc_version" @@ -1713,9 +1705,9 @@ dependencies = [ [[package]] name = "rustls" -version = "0.23.41" +version = "0.23.42" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6b92b125634d9b795e7beca796cc790df15a7fb38323bf3196fda83292d06b1f" +checksum = "3c54fcab019b409d04215d3a17cb438fd7fbf192ee61461f20f4fe18704bc138" dependencies = [ "aws-lc-rs", "once_cell", @@ -1740,9 +1732,9 @@ dependencies = [ [[package]] name = "rustls-pki-types" -version = "1.14.1" +version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "30a7197ae7eb376e574fe940d068c30fe0462554a3ddbe4eca7838e049c937a9" +checksum = "764899a24af3980067ee14bc143654f297b22eaebfe3c7b6b211920a5a59b046" dependencies = [ "web-time", "zeroize", @@ -1772,9 +1764,9 @@ dependencies = [ [[package]] name = "rustversion" -version = "1.0.22" +version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" +checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" [[package]] name = "ryu" @@ -1832,9 +1824,9 @@ checksum = "8a7852d02fc848982e0c167ef163aaff9cd91dc640ba85e263cb1ce46fae51cd" [[package]] name = "serde" -version = "1.0.228" +version = "1.0.229" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" dependencies = [ "serde_core", "serde_derive", @@ -1842,22 +1834,22 @@ dependencies = [ [[package]] name = "serde_core" -version = "1.0.228" +version = "1.0.229" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" dependencies = [ "serde_derive", ] [[package]] name = "serde_derive" -version = "1.0.228" +version = "1.0.229" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 3.0.0", ] [[package]] @@ -1898,9 +1890,9 @@ dependencies = [ [[package]] name = "sha1" -version = "0.10.6" +version = "0.10.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba" +checksum = "a978451301f4db1d02937a4ab3ccce137717b81826e79b7d49ffe3244a13c3b8" dependencies = [ "cfg-if", "cpufeatures 0.2.17", @@ -1959,9 +1951,9 @@ dependencies = [ [[package]] name = "socket2" -version = "0.6.4" +version = "0.6.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "52d1cfed4120b4d927bf7c0f86d2087a4a7d6027c906d9f9d525a80573b9be51" +checksum = "c3d1e2c7f27f8d4cb10542a02c49005dbd6e93095799d6f3be745fae9f8fedd4" dependencies = [ "libc", "windows-sys 0.61.2", @@ -1981,9 +1973,20 @@ checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" [[package]] name = "syn" -version = "2.0.118" +version = "2.0.119" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1b9ae57f904213ebb649ce6895b8a66c66f0203b9319718f69a5612a065b1422" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2fac314a64dc9a36e61a9eb4261a5e9bbfbc922b27e518af97bc32b926cf967" dependencies = [ "proc-macro2", "quote", @@ -2007,7 +2010,7 @@ checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -2027,11 +2030,11 @@ dependencies = [ [[package]] name = "thiserror" -version = "2.0.18" +version = "2.0.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +checksum = "09a43598840e33d5b0331f38c5e30d13bb11c11210a4b58f0d9b18a5a5eefcd9" dependencies = [ - "thiserror-impl 2.0.18", + "thiserror-impl 2.0.19", ] [[package]] @@ -2042,18 +2045,18 @@ checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] name = "thiserror-impl" -version = "2.0.18" +version = "2.0.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +checksum = "43cbfe0cf76104d42a574802844187e84a305e531ed54455f11fbde0f10541cd" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 3.0.0", ] [[package]] @@ -2098,9 +2101,9 @@ dependencies = [ [[package]] name = "tinyvec" -version = "1.11.0" +version = "1.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3e61e67053d25a4e82c844e8424039d9745781b3fc4f32b8d55ed50f5f667ef3" +checksum = "bb4ebadaa0af04fab11ae01eb5f9fdb5f9c5b875506e210e71c07873528baa7f" dependencies = [ "tinyvec_macros", ] @@ -2113,28 +2116,28 @@ checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" [[package]] name = "tokio" -version = "1.52.3" +version = "1.53.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8fc7f01b389ac15039e4dc9531aa973a135d7a4135281b12d7c1bc79fd57fffe" +checksum = "d988bcd52dbe076d3d46903332f58c912b87a2c49b1428419a5845154762ffee" dependencies = [ "bytes", "libc", "mio", "pin-project-lite", - "socket2 0.6.4", + "socket2 0.6.5", "tokio-macros", "windows-sys 0.61.2", ] [[package]] name = "tokio-macros" -version = "2.7.0" +version = "2.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "385a6cb71ab9ab790c5fe8d67f1645e6c450a7ce006a33de03daa956cf70a496" +checksum = "6328af13490e73a9b4694030fafd93f8c8c6a9dede33e821c3fc63eddf8042ba" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -2153,7 +2156,7 @@ version = "0.26.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1729aa945f29d91ba541258c8df89027d5792d85a8841fb65e8bf0f4ede4ef61" dependencies = [ - "rustls 0.23.41", + "rustls 0.23.42", "tokio", ] @@ -2165,7 +2168,7 @@ checksum = "edc5f74e248dc973e0dbb7b74c7e0d6fcc301c694ff50049504004ef4d0cdcd9" dependencies = [ "futures-util", "log", - "rustls 0.23.41", + "rustls 0.23.42", "rustls-native-certs", "rustls-pki-types", "tokio", @@ -2212,7 +2215,7 @@ dependencies = [ "bytes", "futures-util", "http 1.4.2", - "http-body 1.0.1", + "http-body 1.1.0", "pin-project-lite", "tower", "tower-layer", @@ -2252,7 +2255,7 @@ checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -2282,8 +2285,8 @@ dependencies = [ "http 1.4.2", "httparse", "log", - "rand 0.8.6", - "rustls 0.23.41", + "rand 0.8.7", + "rustls 0.23.42", "rustls-pki-types", "sha1", "thiserror 1.0.69", @@ -2375,15 +2378,6 @@ version = "0.11.1+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" -[[package]] -name = "wasip2" -version = "1.0.4+wasi-0.2.12" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b67efb37e106e55ce722a510d6b5f9c17f083e5fc79afc2badeb12cc313d9487" -dependencies = [ - "wit-bindgen", -] - [[package]] name = "wasm-bindgen" version = "0.2.126" @@ -2426,7 +2420,7 @@ dependencies = [ "bumpalo", "proc-macro2", "quote", - "syn", + "syn 2.0.119", "wasm-bindgen-shared", ] @@ -2474,9 +2468,9 @@ dependencies = [ [[package]] name = "webpki-roots" -version = "1.0.8" +version = "1.0.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bf85cb06032201fa7c6f829d7db5a7e5aa45bcc0655327713065f6f0576731bf" +checksum = "7dcd9d09a39985f5344844e66b0c530a33843579125f23e21e9f0f220850f22a" dependencies = [ "rustls-pki-types", ] @@ -2493,16 +2487,7 @@ version = "0.52.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" dependencies = [ - "windows-targets 0.52.6", -] - -[[package]] -name = "windows-sys" -version = "0.60.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f2f500e4d28234f72040990ec9d39e3a6b950f9f22d3dba18416c35882612bcb" -dependencies = [ - "windows-targets 0.53.5", + "windows-targets", ] [[package]] @@ -2520,31 +2505,14 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" dependencies = [ - "windows_aarch64_gnullvm 0.52.6", - "windows_aarch64_msvc 0.52.6", - "windows_i686_gnu 0.52.6", - "windows_i686_gnullvm 0.52.6", - "windows_i686_msvc 0.52.6", - "windows_x86_64_gnu 0.52.6", - "windows_x86_64_gnullvm 0.52.6", - "windows_x86_64_msvc 0.52.6", -] - -[[package]] -name = "windows-targets" -version = "0.53.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4945f9f551b88e0d65f3db0bc25c33b8acea4d9e41163edf90dcd0b19f9069f3" -dependencies = [ - "windows-link", - "windows_aarch64_gnullvm 0.53.1", - "windows_aarch64_msvc 0.53.1", - "windows_i686_gnu 0.53.1", - "windows_i686_gnullvm 0.53.1", - "windows_i686_msvc 0.53.1", - "windows_x86_64_gnu 0.53.1", - "windows_x86_64_gnullvm 0.53.1", - "windows_x86_64_msvc 0.53.1", + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", ] [[package]] @@ -2553,102 +2521,48 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" -[[package]] -name = "windows_aarch64_gnullvm" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a9d8416fa8b42f5c947f8482c43e7d89e73a173cead56d044f6a56104a6d1b53" - [[package]] name = "windows_aarch64_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" -[[package]] -name = "windows_aarch64_msvc" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b9d782e804c2f632e395708e99a94275910eb9100b2114651e04744e9b125006" - [[package]] name = "windows_i686_gnu" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" -[[package]] -name = "windows_i686_gnu" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "960e6da069d81e09becb0ca57a65220ddff016ff2d6af6a223cf372a506593a3" - [[package]] name = "windows_i686_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" -[[package]] -name = "windows_i686_gnullvm" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fa7359d10048f68ab8b09fa71c3daccfb0e9b559aed648a8f95469c27057180c" - [[package]] name = "windows_i686_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" -[[package]] -name = "windows_i686_msvc" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e7ac75179f18232fe9c285163565a57ef8d3c89254a30685b57d83a38d326c2" - [[package]] name = "windows_x86_64_gnu" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" -[[package]] -name = "windows_x86_64_gnu" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9c3842cdd74a865a8066ab39c8a7a473c0778a3f29370b5fd6b4b9aa7df4a499" - [[package]] name = "windows_x86_64_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" -[[package]] -name = "windows_x86_64_gnullvm" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ffa179e2d07eee8ad8f57493436566c7cc30ac536a3379fdf008f47f6bb7ae1" - [[package]] name = "windows_x86_64_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" -[[package]] -name = "windows_x86_64_msvc" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650" - -[[package]] -name = "wit-bindgen" -version = "0.57.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" - [[package]] name = "writeable" version = "0.6.3" @@ -2680,28 +2594,28 @@ checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", "synstructure", ] [[package]] name = "zerocopy" -version = "0.8.52" +version = "0.8.54" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ce1022995ff5ff5d841ad7d994facc23098cd40152f2c1d11cd607c6f530653f" +checksum = "b7cbbc0a705a0fd05cc3676525980d2bf5a9bc4adac6d6475209a7887cf59d19" dependencies = [ "zerocopy-derive", ] [[package]] name = "zerocopy-derive" -version = "0.8.52" +version = "0.8.54" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1ae7f38b72ec2a254e2b87ef277cf2cd4fb97cbebf944faa6f33354da0867930" +checksum = "e2e817b7b52d0c7358d3246da9d69935ebb18116b2b102b4230dac079b4862f5" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] @@ -2721,7 +2635,7 @@ checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", "synstructure", ] @@ -2761,11 +2675,11 @@ checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.119", ] [[package]] name = "zmij" -version = "1.0.21" +version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index a3baa33e6cf..6d63be05d00 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -7,7 +7,8 @@ members = [ resolver = "2" [workspace.package] -edition = "2021" +edition = "2024" +rust-version = "1.88" license = "MIT" repository = "https://github.com/BerriAI/litellm" diff --git a/litellm-rust/crates/ai-gateway/src/auth/mod.rs b/litellm-rust/crates/ai-gateway/src/auth/mod.rs index 438a0513057..b09d8285c3a 100644 --- a/litellm-rust/crates/ai-gateway/src/auth/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/auth/mod.rs @@ -9,9 +9,9 @@ //! runs during extraction, before the handler body. Routes never re-implement it. use axum::extract::FromRequestParts; +use axum::http::StatusCode; use axum::http::header::AUTHORIZATION; use axum::http::request::Parts; -use axum::http::StatusCode; use sha2::{Digest, Sha256}; use subtle::ConstantTimeEq; diff --git a/litellm-rust/crates/ai-gateway/src/io/messages.rs b/litellm-rust/crates/ai-gateway/src/io/messages.rs index b784d2b62a1..86170e45678 100644 --- a/litellm-rust/crates/ai-gateway/src/io/messages.rs +++ b/litellm-rust/crates/ai-gateway/src/io/messages.rs @@ -1 +1 @@ -pub use crate::messages::{messages, MessagesRequest}; +pub use crate::messages::{MessagesRequest, messages}; diff --git a/litellm-rust/crates/ai-gateway/src/io/ocr.rs b/litellm-rust/crates/ai-gateway/src/io/ocr.rs index 55e02839c4e..2fc82f0b61f 100644 --- a/litellm-rust/crates/ai-gateway/src/io/ocr.rs +++ b/litellm-rust/crates/ai-gateway/src/io/ocr.rs @@ -1 +1 @@ -pub use crate::ocr::{ocr, OcrRequest}; +pub use crate::ocr::{OcrRequest, ocr}; diff --git a/litellm-rust/crates/ai-gateway/src/io/realtime.rs b/litellm-rust/crates/ai-gateway/src/io/realtime.rs index 40a38c1579a..845e7bf9527 100644 --- a/litellm-rust/crates/ai-gateway/src/io/realtime.rs +++ b/litellm-rust/crates/ai-gateway/src/io/realtime.rs @@ -15,16 +15,16 @@ use std::time::Duration; use futures_util::stream::{SplitSink, SplitStream}; use futures_util::{Sink, SinkExt, Stream, StreamExt}; +use litellm_core::CoreResult; use litellm_core::error::CoreError; use litellm_core::realtime::transformation::RealtimeProviderConfig; use litellm_core::realtime::types::RealtimeEvent; -use litellm_core::CoreResult; use tokio::net::TcpStream; -use tokio_tungstenite::tungstenite::client::IntoClientRequest; -use tokio_tungstenite::tungstenite::http::header::AUTHORIZATION; -use tokio_tungstenite::tungstenite::http::HeaderValue; use tokio_tungstenite::tungstenite::Message; -use tokio_tungstenite::{connect_async, MaybeTlsStream, WebSocketStream}; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; +use tokio_tungstenite::tungstenite::http::HeaderValue; +use tokio_tungstenite::tungstenite::http::header::AUTHORIZATION; +use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async}; use litellm_core::providers::openai::realtime::transformation::OPENAI_REALTIME_CONFIG; @@ -113,7 +113,7 @@ pub(crate) async fn read_event(upstream_rx: &mut UpstreamRx) -> CoreResult { return Err(CoreError::Network( "upstream closed before first event".to_string(), - )) + )); } _ => continue, } diff --git a/litellm-rust/crates/ai-gateway/src/io/realtime_pool.rs b/litellm-rust/crates/ai-gateway/src/io/realtime_pool.rs index bf8041f31d7..4a1a3cd1166 100644 --- a/litellm-rust/crates/ai-gateway/src/io/realtime_pool.rs +++ b/litellm-rust/crates/ai-gateway/src/io/realtime_pool.rs @@ -28,11 +28,11 @@ use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; use futures_util::StreamExt; -use litellm_core::realtime::types::RealtimeEvent; use litellm_core::CoreResult; +use litellm_core::realtime::types::RealtimeEvent; use crate::io::realtime::{ - dial_upstream, read_event, resolve_api_key, UpstreamRx, UpstreamTx, UpstreamWs, + UpstreamRx, UpstreamTx, UpstreamWs, dial_upstream, read_event, resolve_api_key, }; /// Default target warm sockets per key when pooling is enabled. @@ -473,8 +473,8 @@ pub fn upstream_key( /// unhealthy — we'd rather discard and fresh-dial than hand over a socket in an /// unexpected state. `Pending` (the healthy case) returns `false`. fn is_dead(rx: &mut UpstreamRx) -> bool { - use futures_util::task::noop_waker_ref; use futures_util::Stream; + use futures_util::task::noop_waker_ref; use std::pin::Pin; use std::task::{Context, Poll}; @@ -523,15 +523,15 @@ mod tests { )) .await; while let Some(Ok(msg)) = ws.next().await { - if let Message::Text(text) = msg { - if text.contains("response.create") { - for frame in [ - r#"{"type":"response.created"}"#, - r#"{"type":"response.output_audio.delta","delta":"AAAA"}"#, - r#"{"type":"response.done"}"#, - ] { - let _ = ws.send(Message::Text(frame.to_string())).await; - } + if let Message::Text(text) = msg + && text.contains("response.create") + { + for frame in [ + r#"{"type":"response.created"}"#, + r#"{"type":"response.output_audio.delta","delta":"AAAA"}"#, + r#"{"type":"response.done"}"#, + ] { + let _ = ws.send(Message::Text(frame.to_string())).await; } } } diff --git a/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs b/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs index ae6ad150bcf..9b51019f4bc 100644 --- a/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs +++ b/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs @@ -10,19 +10,18 @@ use litellm_core::responses::websocket::ResponsesWebSocketProviderConfig; use litellm_core::{CoreError, CoreResult}; use tokio::net::TcpStream; use tokio::sync::Mutex; -use tokio_tungstenite::tungstenite::client::IntoClientRequest; -use tokio_tungstenite::tungstenite::http::header::{HeaderName, AUTHORIZATION}; -use tokio_tungstenite::tungstenite::http::HeaderValue; use tokio_tungstenite::tungstenite::Message; -use tokio_tungstenite::{connect_async, MaybeTlsStream, WebSocketStream}; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; +use tokio_tungstenite::tungstenite::http::HeaderValue; +use tokio_tungstenite::tungstenite::http::header::{AUTHORIZATION, HeaderName}; +use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async}; use crate::constants::{ DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS, DEFAULT_RESPONSES_WS_IDLE_TIMEOUT_SECS, }; const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY"; -const MISSING_KEY_MESSAGE: &str = - "Missing OpenAI API Key - a Responses WebSocket call is being made but no key was passed via params or the OPENAI_API_KEY environment variable"; +const MISSING_KEY_MESSAGE: &str = "Missing OpenAI API Key - a Responses WebSocket call is being made but no key was passed via params or the OPENAI_API_KEY environment variable"; pub type ResponsesUpstreamWs = WebSocketStream>; type UpstreamTx = SplitSink; @@ -83,8 +82,8 @@ impl ResponsesWebSocketConnection { } pub async fn recv_text(&self) -> CoreResult> { - let mut socket = self.socket.lock().await; - let Some(socket) = socket.as_mut() else { + let mut socket_guard = self.socket.lock().await; + let Some(socket) = socket_guard.as_mut() else { return Ok(None); }; match socket.next().await { @@ -456,9 +455,11 @@ mod tests { assert_eq!(fourth.event_type, ResponsesWsEventType::ResponseCompleted); let observed: Vec<_> = observed_rx.collect().await; assert_eq!(observed.len(), 4); - assert!(observed - .iter() - .all(|event| event.event_type != ResponsesWsEventType::ResponseCreate)); + assert!( + observed + .iter() + .all(|event| event.event_type != ResponsesWsEventType::ResponseCreate) + ); } #[tokio::test] diff --git a/litellm-rust/crates/ai-gateway/src/main.rs b/litellm-rust/crates/ai-gateway/src/main.rs index f9ce97801d3..da3a486d4ee 100644 --- a/litellm-rust/crates/ai-gateway/src/main.rs +++ b/litellm-rust/crates/ai-gateway/src/main.rs @@ -11,7 +11,7 @@ use std::sync::Arc; -use litellm_ai_gateway::io::realtime_pool::{upstream_key, PoolConfig, RealtimePool}; +use litellm_ai_gateway::io::realtime_pool::{PoolConfig, RealtimePool, upstream_key}; use litellm_ai_gateway::routes; use litellm_ai_gateway::state::AppState; use litellm_core::router::{Deployment, LiteLLMParams, Router}; diff --git a/litellm-rust/crates/ai-gateway/src/messages/common_utils.rs b/litellm-rust/crates/ai-gateway/src/messages/common_utils.rs index 33894d0ee64..4b906155665 100644 --- a/litellm-rust/crates/ai-gateway/src/messages/common_utils.rs +++ b/litellm-rust/crates/ai-gateway/src/messages/common_utils.rs @@ -1,8 +1,8 @@ -use litellm_core::error::{json_type_name, CoreError}; +use litellm_core::CoreResult; +use litellm_core::error::{CoreError, json_type_name}; use litellm_core::messages::transformation::AnthropicMessagesProviderConfig; use litellm_core::providers::anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG; use litellm_core::providers::azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG; -use litellm_core::CoreResult; use serde_json::{Map, Value}; use crate::constants::MESSAGES_ERROR_BODY_MAX_CHARS; diff --git a/litellm-rust/crates/ai-gateway/src/messages/handler.rs b/litellm-rust/crates/ai-gateway/src/messages/handler.rs index d3b9d3b3fba..90c12367f50 100644 --- a/litellm-rust/crates/ai-gateway/src/messages/handler.rs +++ b/litellm-rust/crates/ai-gateway/src/messages/handler.rs @@ -1,5 +1,5 @@ -use litellm_core::error::CoreError; use litellm_core::CoreResult; +use litellm_core::error::CoreError; use serde_json::Value; use super::client::http_client; diff --git a/litellm-rust/crates/ai-gateway/src/messages/prepare.rs b/litellm-rust/crates/ai-gateway/src/messages/prepare.rs index 6176f9cb67f..624c3598fb0 100644 --- a/litellm-rust/crates/ai-gateway/src/messages/prepare.rs +++ b/litellm-rust/crates/ai-gateway/src/messages/prepare.rs @@ -1,7 +1,7 @@ -use litellm_core::messages::transformation::MessagesAuthStrategy; -use litellm_core::routing_utils::provider::{get_custom_llm_provider, CustomLlmProvider}; use litellm_core::CoreError; use litellm_core::CoreResult; +use litellm_core::messages::transformation::MessagesAuthStrategy; +use litellm_core::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider}; use super::common_utils::{has_header, messages_provider_config, string_headers}; use super::types::{MessagesRequest, ProviderMessagesRequest}; diff --git a/litellm-rust/crates/ai-gateway/src/messages/tests.rs b/litellm-rust/crates/ai-gateway/src/messages/tests.rs index 9b1cc45aacb..a2d0f6fae23 100644 --- a/litellm-rust/crates/ai-gateway/src/messages/tests.rs +++ b/litellm-rust/crates/ai-gateway/src/messages/tests.rs @@ -1,14 +1,14 @@ use std::time::Duration; use litellm_core::error::CoreError; -use serde_json::{json, Map, Value}; +use serde_json::{Map, Value, json}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; use super::common_utils::{ has_header, messages_provider_config, string_headers, truncate_error_body, }; -use super::{messages, MessagesRequest}; +use super::{MessagesRequest, messages}; async fn read_http_request(socket: &mut TcpStream) -> String { let mut request = Vec::new(); diff --git a/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs b/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs index d4b4d9338e7..7d164a80137 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs @@ -1,11 +1,11 @@ use std::net::IpAddr; use std::time::{Duration, Instant}; -use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; use base64::Engine; +use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; +use litellm_core::CoreResult; use litellm_core::error::CoreError; use litellm_core::ocr::transformation::OcrProviderConfig; -use litellm_core::CoreResult; use reqwest::Url; use serde_json::{Map, Value}; diff --git a/litellm-rust/crates/ai-gateway/src/ocr/handler.rs b/litellm-rust/crates/ai-gateway/src/ocr/handler.rs index 4d93c2a25db..381d22e9cea 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/handler.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/handler.rs @@ -1,6 +1,6 @@ +use litellm_core::CoreResult; use litellm_core::error::CoreError; use litellm_core::ocr::transformation::OcrResponseHandling; -use litellm_core::CoreResult; use serde_json::Value; use super::client::http_client; diff --git a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs index 6be74ed2714..ffe2e0122c0 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs @@ -1,11 +1,11 @@ use std::future::Future; use std::pin::Pin; +use litellm_core::CoreResult; use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming}; use litellm_core::error::CoreError; use litellm_core::ocr::transformation::OcrAuthStrategy; -use litellm_core::CoreResult; -use serde_json::{json, Map, Value}; +use serde_json::{Map, Value, json}; use super::common_utils::{ convert_document_url_to_data_uri, has_header, ocr_provider_config, string_headers, @@ -292,7 +292,7 @@ fn parse_ocr_pre_call_guardrail_request( Some(_) => { return Err(CoreError::InvalidRequest( "OCR pre_call guardrail optional_params must be an object".to_string(), - )) + )); } None => Map::new(), }; diff --git a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs index b54ee39b21d..ad346bc0c64 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs @@ -1,5 +1,5 @@ -use litellm_core::call_lifecycle::CallLifecycle; use litellm_core::CoreResult; +use litellm_core::call_lifecycle::CallLifecycle; use serde_json::Value; mod client; @@ -12,7 +12,7 @@ mod types; pub use types::OcrRequest; use handler::execute_ocr_provider_call; -use prepare::{prepare_ocr_call, PreparedOcrCall}; +use prepare::{PreparedOcrCall, prepare_ocr_call}; pub async fn ocr(request: OcrRequest<'_>) -> CoreResult { let PreparedOcrCall { request, hooks } = prepare_ocr_call(request); diff --git a/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs b/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs index 5a4b350a4c4..6231393c889 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs @@ -1,7 +1,7 @@ use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{SystemTime, UNIX_EPOCH}; -use litellm_core::routing_utils::provider::{get_custom_llm_provider, CustomLlmProvider}; +use litellm_core::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider}; use super::hooks::OcrLifecycleHooks; use super::types::{OcrRequest, PreparedOcrRequest}; diff --git a/litellm-rust/crates/ai-gateway/src/ocr/tests.rs b/litellm-rust/crates/ai-gateway/src/ocr/tests.rs index 35747dc6985..bb2a6b06501 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/tests.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/tests.rs @@ -3,12 +3,12 @@ use std::time::Duration; use litellm_core::error::CoreError; use litellm_core::ocr::transformation::OcrResponseHandling; -use serde_json::{json, Map, Value}; +use serde_json::{Map, Value, json}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; use super::common_utils::{has_header, ocr_provider_config, string_headers, truncate_error_body}; -use super::{ocr, OcrRequest}; +use super::{OcrRequest, ocr}; use crate::integrations::custom_guardrail::{ CustomGuardrail, GuardrailContext, GuardrailDecision, GuardrailError, GuardrailEventHook, GuardrailFuture, GuardrailRequest, @@ -228,19 +228,23 @@ fn truncate_error_body_does_not_split_multibyte_chars() { #[test] fn ocr_dispatch_supports_migrated_providers() { assert!(ocr_provider_config("mistral", "mistral-ocr-latest").is_some()); - assert!(ocr_provider_config("azure_ai", "pixtral-12b-2409") - .expect("azure ai config resolves") - .requires_data_uri_document()); + assert!( + ocr_provider_config("azure_ai", "pixtral-12b-2409") + .expect("azure ai config resolves") + .requires_data_uri_document() + ); assert_eq!( ocr_provider_config("azure_ai", "doc-intelligence/prebuilt-read") .expect("document intelligence config resolves") .response_handling(), OcrResponseHandling::AzureDocumentIntelligencePoll ); - assert!(ocr_provider_config("vertex_ai", "deepseek-ocr-maas") - .expect("vertex deepseek config resolves") - .supported_ocr_params() - .contains(&"temperature")); + assert!( + ocr_provider_config("vertex_ai", "deepseek-ocr-maas") + .expect("vertex deepseek config resolves") + .supported_ocr_params() + .contains(&"temperature") + ); assert!(ocr_provider_config("openai", "gpt-4o").is_none()); } diff --git a/litellm-rust/crates/ai-gateway/src/python/config.rs b/litellm-rust/crates/ai-gateway/src/python/config.rs index 54b7a53bafa..c028d3d6b51 100644 --- a/litellm-rust/crates/ai-gateway/src/python/config.rs +++ b/litellm-rust/crates/ai-gateway/src/python/config.rs @@ -7,9 +7,9 @@ //! //! Compiled only under the `python-config` feature. +use litellm_core::CoreResult; use litellm_core::error::CoreError; use litellm_core::router::{Deployment, Router}; -use litellm_core::CoreResult; use pyo3::prelude::*; use crate::gil; diff --git a/litellm-rust/crates/ai-gateway/src/routes/health.rs b/litellm-rust/crates/ai-gateway/src/routes/health.rs index 15c67fea325..c64ca3a7199 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/health.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/health.rs @@ -1,8 +1,8 @@ //! Health probes. Simple-route template: a `router()` plus its handlers, in one file. +use axum::Router; use axum::http::StatusCode; use axum::routing::get; -use axum::Router; use crate::state::AppState; diff --git a/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs b/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs index 933386282fa..a34b2edd7b8 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs @@ -2,13 +2,13 @@ mod service; +use axum::Router; use axum::body::Body; use axum::extract::{Json, State}; -use axum::http::header::{HeaderMap, HeaderValue, CACHE_CONTROL, CONTENT_TYPE}; use axum::http::StatusCode; +use axum::http::header::{CACHE_CONTROL, CONTENT_TYPE, HeaderMap, HeaderValue}; use axum::response::{IntoResponse, Response}; use axum::routing::post; -use axum::Router; use litellm_core::CoreError; use serde_json::{Map, Value}; @@ -125,9 +125,9 @@ mod tests { use std::sync::Arc; use axum::body::Body; - use axum::http::header::{CACHE_CONTROL, CONTENT_TYPE}; use axum::http::Request; use axum::http::StatusCode; + use axum::http::header::{CACHE_CONTROL, CONTENT_TYPE}; use litellm_core::router::{Deployment, LiteLLMParams, Router as ModelRouter}; use serde_json::json; use tokio::io::{AsyncReadExt, AsyncWriteExt}; @@ -439,8 +439,8 @@ mod tests { .await .expect("response body reads"); assert_eq!( - serde_json::from_slice::(&response_body).expect("error is json") - ["error"]["message"], + serde_json::from_slice::(&response_body).expect("error is json")["error"] + ["message"], "messages provider request failed" ); server.await.expect("upstream task completes"); diff --git a/litellm-rust/crates/ai-gateway/src/routes/messages/service.rs b/litellm-rust/crates/ai-gateway/src/routes/messages/service.rs index 7f00123ca39..75ed26e5be8 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/messages/service.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/messages/service.rs @@ -5,7 +5,7 @@ use litellm_core::{CoreError, CoreResult}; use serde_json::{Map, Value}; use crate::constants::ANTHROPIC_MESSAGES_PROVIDER; -use crate::messages::{execute_messages, MessagesRequest}; +use crate::messages::{MessagesRequest, execute_messages}; pub(crate) enum MessagesResponse { Json(Value), diff --git a/litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs b/litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs index c3f929f5f0b..f9144ad1fdb 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs @@ -6,17 +6,17 @@ mod service; -use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{SystemTime, UNIX_EPOCH}; use crate::io::realtime_pool::RealtimePool; +use axum::Router; use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade}; use axum::extract::{Query, State}; use axum::http::StatusCode; use axum::response::Response; use axum::routing::get; -use axum::Router; use futures_util::{SinkExt, StreamExt}; use litellm_core::realtime::types::RealtimeEvent; use litellm_core::router::Router as ModelRouter; diff --git a/litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs b/litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs index d6c31edd454..4ae8cfe7379 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs @@ -9,12 +9,12 @@ use std::time::Duration; -use crate::io::realtime_pool::{upstream_key, RealtimePool}; +use crate::io::realtime_pool::{RealtimePool, upstream_key}; use futures_util::{Sink, Stream}; +use litellm_core::CoreResult; use litellm_core::error::CoreError; use litellm_core::realtime::types::RealtimeEvent; use litellm_core::router::Router; -use litellm_core::CoreResult; /// Select a deployment for `model` and splice the client stream to the provider. /// diff --git a/litellm-rust/crates/ai-gateway/src/routes/responses/mod.rs b/litellm-rust/crates/ai-gateway/src/routes/responses/mod.rs index bdaffc97afb..a94853e106d 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/responses/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/responses/mod.rs @@ -1,15 +1,15 @@ mod service; -use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{SystemTime, UNIX_EPOCH}; +use axum::Router; use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade}; use axum::extract::{Query, State}; use axum::http::StatusCode; use axum::response::Response; use axum::routing::get; -use axum::Router; use futures_util::{Sink, SinkExt, StreamExt}; use litellm_core::responses::types::{ResponsesErrorFrame, ResponsesWsEvent, ResponsesWsEventType}; use litellm_core::router::Router as ModelRouter; diff --git a/litellm-rust/crates/core/src/caching/in_memory_cache.rs b/litellm-rust/crates/core/src/caching/in_memory_cache.rs index 0ceeedb8b71..45d4bd69b79 100644 --- a/litellm-rust/crates/core/src/caching/in_memory_cache.rs +++ b/litellm-rust/crates/core/src/caching/in_memory_cache.rs @@ -134,8 +134,8 @@ impl InMemoryCache { #[cfg(test)] mod tests { use std::sync::{ - atomic::{AtomicU64, Ordering}, Arc, + atomic::{AtomicU64, Ordering}, }; use super::InMemoryCache; diff --git a/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs b/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs index 13e79b087c7..6935bb4604b 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs +++ b/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs @@ -5,7 +5,7 @@ use crate::messages::types::{ MessageContent, SystemPrompt, }; use crate::providers::anthropic::messages::transformation::{ - non_empty, AnthropicMessagesConfig, ANTHROPIC_MESSAGES_CONFIG, + ANTHROPIC_MESSAGES_CONFIG, AnthropicMessagesConfig, non_empty, }; use serde_json::{Map, Value}; diff --git a/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs index 060073acd47..eabd15677cc 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs @@ -1,9 +1,9 @@ use std::collections::BTreeSet; -use crate::error::{json_type_name, CoreError, CoreResult}; +use crate::error::{CoreError, CoreResult, json_type_name}; use crate::ocr::transformation::{OcrAuthStrategy, OcrProviderConfig, OcrResponseHandling}; use crate::ocr::types::{OcrRequestData, OcrResponseData}; -use serde_json::{json, Map, Value}; +use serde_json::{Map, Value, json}; use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG; @@ -206,11 +206,11 @@ pub fn complete_document_intelligence_url( AZURE_DOCUMENT_INTELLIGENCE_API_VERSION ); - if let Some(pages) = optional_params.get("pages") { - if let Some(normalized) = normalize_pages_param(pages)? { - url.push_str("&pages="); - url.push_str(&normalized); - } + if let Some(pages) = optional_params.get("pages") + && let Some(normalized) = normalize_pages_param(pages)? + { + url.push_str("&pages="); + url.push_str(&normalized); } Ok(url) @@ -231,7 +231,7 @@ fn document_url_from_mistral_document(document: &Value) -> CoreResult<&str> { other => { return Err(CoreError::InvalidRequest(format!( "Invalid document type: {other}. Must be 'document_url' or 'image_url'" - ))) + ))); } }; object diff --git a/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs b/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs index 82d1e8fdf91..c5995732e41 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs @@ -5,10 +5,10 @@ use std::time::{SystemTime, UNIX_EPOCH}; use crate::caching::in_memory_cache::InMemoryCache; use crate::error::{CoreError, CoreResult}; -use aws_credential_types::provider::ProvideCredentials; use aws_credential_types::Credentials; +use aws_credential_types::provider::ProvideCredentials; use aws_sigv4::http_request::{ - sign, SignableBody, SignableRequest, SigningParams, SigningSettings, + SignableBody, SignableRequest, SigningParams, SigningSettings, sign, }; use aws_sigv4::sign::v4; use aws_smithy_runtime_api::client::identity::Identity; @@ -368,11 +368,11 @@ async fn is_already_running_as_role(role: &str, config: &AwsAuthConfig) -> CoreR if let (Ok(current_role), Ok(token_file)) = ( std::env::var(AWS_ROLE_ARN), std::env::var(AWS_WEB_IDENTITY_TOKEN_FILE), - ) { - if !token_file.is_empty() { - return Ok(same_role_arns(role, ¤t_role)); - } + ) && !token_file.is_empty() + { + return Ok(same_role_arns(role, ¤t_role)); } + let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest()); if let Some(region) = config.region_name.clone() { loader = loader.region(aws_types::region::Region::new(region)); @@ -639,7 +639,9 @@ mod tests { ); assert_eq!( signed.get("Authorization").map(String::as_str), - Some("AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20240102/us-east-1/bedrock/aws4_request, SignedHeaders=content-type;host;x-amz-date;x-amz-security-token, Signature=55c027ef47527d3ad63f1735f9d099efdbc99f296ff914bd94e727e24ec0e464") + Some( + "AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20240102/us-east-1/bedrock/aws4_request, SignedHeaders=content-type;host;x-amz-date;x-amz-security-token, Signature=55c027ef47527d3ad63f1735f9d099efdbc99f296ff914bd94e727e24ec0e464" + ) ); } diff --git a/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs index 1a33bc1e951..dc720cc4244 100644 --- a/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs @@ -1,4 +1,4 @@ -use crate::error::{json_type_name, CoreError, CoreResult}; +use crate::error::{CoreError, CoreResult, json_type_name}; use crate::ocr::transformation::OcrProviderConfig; use crate::ocr::types::{OcrRequestData, OcrResponseData}; use serde_json::{Map, Value}; diff --git a/litellm-rust/crates/core/src/providers/openai/realtime/transformation.rs b/litellm-rust/crates/core/src/providers/openai/realtime/transformation.rs index 626e4014ff9..b3f6b03b28a 100644 --- a/litellm-rust/crates/core/src/providers/openai/realtime/transformation.rs +++ b/litellm-rust/crates/core/src/providers/openai/realtime/transformation.rs @@ -1,6 +1,6 @@ +use crate::CoreResult; use crate::realtime::transformation::RealtimeProviderConfig; use crate::realtime::types::{RealtimeEvent, RealtimeTransformResult}; -use crate::CoreResult; /// Default OpenAI API base, used when the caller does not override `api_base`. pub const OPENAI_REALTIME_DEFAULT_API_BASE: &str = "https://api.openai.com"; diff --git a/litellm-rust/crates/core/src/providers/openai/responses/transformation.rs b/litellm-rust/crates/core/src/providers/openai/responses/transformation.rs index ece10971806..e15197c468c 100644 --- a/litellm-rust/crates/core/src/providers/openai/responses/transformation.rs +++ b/litellm-rust/crates/core/src/providers/openai/responses/transformation.rs @@ -1,6 +1,6 @@ -use crate::responses::types::{ResponsesWsEvent, ResponsesWsTransformResult}; -use crate::responses::websocket::{enforce_model, ResponsesWebSocketProviderConfig}; use crate::CoreResult; +use crate::responses::types::{ResponsesWsEvent, ResponsesWsTransformResult}; +use crate::responses::websocket::{ResponsesWebSocketProviderConfig, enforce_model}; pub struct OpenAIResponsesWsConfig; diff --git a/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs index 8639926c435..6300149c237 100644 --- a/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs @@ -1,7 +1,7 @@ -use crate::error::{json_type_name, CoreError, CoreResult}; +use crate::error::{CoreError, CoreResult, json_type_name}; use crate::ocr::transformation::OcrProviderConfig; use crate::ocr::types::{OcrRequestData, OcrResponseData}; -use serde_json::{json, Map, Value}; +use serde_json::{Map, Value, json}; use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG; @@ -140,7 +140,7 @@ fn document_content_item(document: &Value) -> CoreResult { other => { return Err(CoreError::InvalidRequest(format!( "Unsupported document type: {other}. Expected 'image_url' or 'document_url'" - ))) + ))); } }; let url = object diff --git a/litellm-rust/crates/core/src/realtime/transformation.rs b/litellm-rust/crates/core/src/realtime/transformation.rs index a4baa27a6c2..69b88687000 100644 --- a/litellm-rust/crates/core/src/realtime/transformation.rs +++ b/litellm-rust/crates/core/src/realtime/transformation.rs @@ -1,5 +1,5 @@ -use crate::realtime::types::{RealtimeEvent, RealtimeTransformResult}; use crate::CoreResult; +use crate::realtime::types::{RealtimeEvent, RealtimeTransformResult}; pub trait RealtimeProviderConfig { /// Build the upstream WebSocket URL (e.g. `wss://api.openai.com/v1/realtime?model=…`). diff --git a/litellm-rust/crates/core/src/responses/websocket.rs b/litellm-rust/crates/core/src/responses/websocket.rs index 1edffd44985..92dc19627a0 100644 --- a/litellm-rust/crates/core/src/responses/websocket.rs +++ b/litellm-rust/crates/core/src/responses/websocket.rs @@ -1,6 +1,6 @@ +use crate::CoreResult; use crate::constants::{OPENAI_RESPONSES_DEFAULT_API_BASE, OPENAI_RESPONSES_PATH}; use crate::responses::types::{ResponsesWsEvent, ResponsesWsEventType, ResponsesWsTransformResult}; -use crate::CoreResult; pub trait ResponsesWebSocketProviderConfig: Sync { fn supports_native_websocket(&self) -> bool { diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 1decb789a22..07429667644 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -1,8 +1,8 @@ use std::collections::HashMap; use std::time::Duration; -use litellm_ai_gateway::io::messages::{messages as run_messages, MessagesRequest}; -use litellm_ai_gateway::io::ocr::{ocr as run_ocr, OcrRequest}; +use litellm_ai_gateway::io::messages::{MessagesRequest, messages as run_messages}; +use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr}; use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection as RustResponsesWebSocketConnection; use litellm_core::error::CoreError; use pyo3::exceptions::{PyRuntimeError, PyValueError}; From 3fcd19d7ad5b06f839312a8b909e3d18dd7f2f80 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 20 Jul 2026 09:52:35 -0700 Subject: [PATCH 11/12] fix(fireworks_ai): restore Content-Type application/json header (fixes 415) (#33929) * fix(fireworks_ai): set Content-Type application/json in validate_environment Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(fireworks_ai): delegate chat validate_environment to OpenAIGPTConfig Instead of re-adding the JSON Content-Type default inside FireworksAIMixin, FireworksAIConfig now delegates header construction to OpenAIGPTConfig and only layers the Fireworks-specific x-session-affinity header on top, so the Content-Type default can no longer drift away from the OpenAI base and reintroduce the 415. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(fireworks_ai): cover missing api key error path in chat validate_environment Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Krrish Dholakia Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/fireworks_ai/chat/transformation.py | 26 ++++++ litellm/llms/fireworks_ai/common_utils.py | 19 +++-- .../test_fireworks_ai_chat_transformation.py | 84 +++++++++++++++++++ 3 files changed, 123 insertions(+), 6 deletions(-) diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index 319f03fea89..eeae8c76888 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -133,6 +133,32 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): def get_config(cls): return super().get_config() + def validate_environment( + self, + headers: dict, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + api_key: str | None = None, + api_base: str | None = None, + ) -> dict: + api_key = self._get_api_key(api_key) + if api_key is None: + raise ValueError("FIREWORKS_API_KEY is not set") + + validated_headers = OpenAIGPTConfig.validate_environment( + self, + headers=headers, + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + api_key=api_key, + api_base=api_base, + ) + return self._add_session_affinity_header(validated_headers, litellm_params) + def get_supported_openai_params(self, model: str): # Base parameters supported by all models supported_params = [ diff --git a/litellm/llms/fireworks_ai/common_utils.py b/litellm/llms/fireworks_ai/common_utils.py index 4e22445bcc0..51ed8afbbd2 100644 --- a/litellm/llms/fireworks_ai/common_utils.py +++ b/litellm/llms/fireworks_ai/common_utils.py @@ -64,9 +64,16 @@ class FireworksAIMixin: if api_key is None: raise ValueError("FIREWORKS_API_KEY is not set") - validated_headers = {"Authorization": "Bearer {}".format(api_key), **headers} - if not any(key.lower() == "x-session-affinity" for key in validated_headers): - session_id = get_fireworks_session_id(litellm_params) - if session_id: - validated_headers["x-session-affinity"] = session_id - return validated_headers + auth_headers = {"Authorization": "Bearer {}".format(api_key), **headers} + content_type_header = ( + {} if any(key.lower() == "content-type" for key in auth_headers) else {"Content-Type": "application/json"} + ) + return self._add_session_affinity_header({**auth_headers, **content_type_header}, litellm_params) + + def _add_session_affinity_header(self, headers: dict, litellm_params: dict) -> dict: + if any(key.lower() == "x-session-affinity" for key in headers): + return headers + session_id = get_fireworks_session_id(litellm_params) + if not session_id: + return headers + return {**headers, "x-session-affinity": session_id} diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index 6809799d34f..94945ed4bfb 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -123,6 +123,90 @@ def test_validate_environment_preserves_explicit_session_affinity_header(): assert headers["x-session-affinity"] == "explicit-session" +def test_validate_environment_sets_json_content_type(): + config = FireworksAIConfig() + + headers = config.validate_environment( + headers={}, + model="accounts/fireworks/models/test-model", + messages=[], + optional_params={}, + litellm_params={}, + api_key="test-key", + ) + + assert headers["Content-Type"] == "application/json" + + +def test_validate_environment_preserves_explicit_content_type(): + config = FireworksAIConfig() + + headers = config.validate_environment( + headers={"content-type": "multipart/form-data"}, + model="accounts/fireworks/models/test-model", + messages=[], + optional_params={}, + litellm_params={}, + api_key="test-key", + ) + + assert headers["content-type"] == "multipart/form-data" + assert "Content-Type" not in headers + + +def test_validate_environment_sets_json_content_type_with_session_affinity(): + config = FireworksAIConfig() + + headers = config.validate_environment( + headers={}, + model="accounts/fireworks/models/test-model", + messages=[], + optional_params={}, + litellm_params={"litellm_session_id": "session-123"}, + api_key="test-key", + ) + + assert headers["Content-Type"] == "application/json" + assert headers["Authorization"] == "Bearer test-key" + assert headers["x-session-affinity"] == "session-123" + + +def test_validate_environment_resolves_api_key_from_env_and_sets_content_type(monkeypatch): + monkeypatch.setenv("FIREWORKS_API_KEY", "fw-env-key") + config = FireworksAIConfig() + + headers = config.validate_environment( + headers={}, + model="accounts/fireworks/models/test-model", + messages=[], + optional_params={}, + litellm_params={}, + ) + + assert headers["Authorization"] == "Bearer fw-env-key" + assert headers["Content-Type"] == "application/json" + + +def test_validate_environment_raises_without_api_key(monkeypatch): + for env_var in ( + "FIREWORKS_API_KEY", + "FIREWORKS_AI_API_KEY", + "FIREWORKSAI_API_KEY", + "FIREWORKS_AI_TOKEN", + ): + monkeypatch.delenv(env_var, raising=False) + config = FireworksAIConfig() + + with pytest.raises(ValueError, match="FIREWORKS_API_KEY is not set"): + config.validate_environment( + headers={}, + model="accounts/fireworks/models/test-model", + messages=[], + optional_params={}, + litellm_params={}, + ) + + def test_get_fireworks_session_id_prefers_litellm_session_id_over_trace_id(): assert ( get_fireworks_session_id( From 479e997eed08e4393ceb7b0c4d39c896c4b19da3 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 20 Jul 2026 09:55:38 -0700 Subject: [PATCH 12/12] feat(spend): raise /spend/logs/v2 page_size cap to 1000 Clients exporting large spend-log ranges were forced into 100-row pages, which meant a bounded COUNT plus an increasingly deep OFFSET scan per request. Larger pages reduce both the request count and the cumulative OFFSET cost for the same result set. The handler already excludes the heavy JSON columns (messages, response, proxy_server_request) from the paginated SELECT and bounds the COUNT via SPEND_LOGS_PAGINATION_COUNT_CAP, so per-row cost does not grow with page size. 1000 matches the ceiling already used by the user and user-agent analytics list endpoints. --- .../spend_management_endpoints.py | 2 +- .../test_spend_management_endpoints.py | 53 +++++++++++++++++++ 2 files changed, 54 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 5a0b94d1524..55b50e7d9ff 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -1647,7 +1647,7 @@ async def ui_view_spend_logs( description="Time till which to view key spend", ), page: int = fastapi.Query(default=1, description="Page number for pagination", ge=1), - page_size: int = fastapi.Query(default=50, description="Number of items per page", ge=1, le=100), + page_size: int = fastapi.Query(default=50, description="Number of items per page", ge=1, le=1000), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), status_filter: str | None = fastapi.Query( default=None, description="Filter logs by status (e.g., success, failure)" 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 579f46c8c77..db72a7fb38c 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 @@ -1467,6 +1467,59 @@ async def test_ui_view_spend_logs_pagination(client, monkeypatch): app.dependency_overrides.pop(ps.user_api_key_auth, None) +@pytest.mark.parametrize( + "page_size, expected_status, expected_rows", + [ + (1000, 200, 1000), + (1001, 422, None), + ], +) +@pytest.mark.asyncio +async def test_ui_view_spend_logs_page_size_upper_bound( + client, monkeypatch, page_size, expected_status, expected_rows +): + mock_spend_logs = [ + { + "id": f"log{i}", + "request_id": f"req{i}", + "api_key": "sk-test-key", + "startTime": datetime.datetime.now(timezone.utc).isoformat(), + "model": "gpt-4", + } + for i in range(1200) + ] + + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + make_ui_spend_logs_mock_prisma(mock_spend_logs, lambda where: mock_spend_logs), + ) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + + try: + start_date, end_date = _default_date_range() + response = client.get( + "/spend/logs/v2", + params={ + "page": 1, + "page_size": page_size, + "start_date": start_date, + "end_date": end_date, + }, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert response.status_code == expected_status + if expected_status == 200: + data = response.json() + assert data["page_size"] == page_size + assert len(data["data"]) == expected_rows + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + @pytest.mark.asyncio async def test_ui_view_session_spend_logs_pagination(client, monkeypatch): mock_spend_logs = [