From 94e4dd725fec6904c21102e1b115d91e1767babd Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Sat, 12 Sep 2026 16:25:57 -0700 Subject: [PATCH] fix(mcp): require admission for delegated OAuth --- .../proxy/_experimental/mcp_server/AGENTS.md | 5 +- .../proxy/_experimental/mcp_server/CLAUDE.md | 2 +- .../mcp_server/auth/user_api_key_auth_mcp.py | 119 +++---- .../mcp_server/mcp_server_manager.py | 38 +-- .../proxy/_experimental/mcp_server/server.py | 88 +---- .../proxy/hooks/proxy_track_cost_callback.py | 1 + .../types/mcp_server/mcp_server_manager.py | 8 +- tests/e2e/coverage_registry/mcp.yaml | 2 +- .../auth/test_user_api_key_auth_mcp.py | 226 ++++--------- .../mcp_server/test_mcp_server.py | 299 +++++------------ .../mcp_server/test_mcp_server_manager.py | 309 +++++++++++------- .../hooks/test_proxy_track_cost_callback.py | 38 ++- 12 files changed, 433 insertions(+), 702 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/AGENTS.md b/litellm/proxy/_experimental/mcp_server/AGENTS.md index 6e1d121c3be..a35e5d1ce10 100644 --- a/litellm/proxy/_experimental/mcp_server/AGENTS.md +++ b/litellm/proxy/_experimental/mcp_server/AGENTS.md @@ -67,9 +67,8 @@ module materially harder to understand. auth, SSE, streamable HTTP, and stdio as separate flows. Do not collapse them behind a single generic branch unless tests prove every mode still behaves correctly. -- Be especially careful with `available_on_public_internet: false` combined with - `delegate_auth_to_upstream: true`. The local `CLAUDE.md` explains the anonymous - upstream PKCE path that must remain intentional. +- Be especially careful with legacy `delegate_auth_to_upstream: true`. The local + `CLAUDE.md` explains its admitted replacement and public discovery contract. - Keep database-backed fields in sync across migrations, typed models under `litellm/types/mcp.py` or `litellm/types/mcp_server/`, config loading, this package, and dashboard state when the field is user-visible. diff --git a/litellm/proxy/_experimental/mcp_server/CLAUDE.md b/litellm/proxy/_experimental/mcp_server/CLAUDE.md index 0ba8f73315f..7f8d06b4570 100644 --- a/litellm/proxy/_experimental/mcp_server/CLAUDE.md +++ b/litellm/proxy/_experimental/mcp_server/CLAUDE.md @@ -1 +1 @@ -MCP note: **`available_on_public_internet: false` with `delegate_auth_to_upstream: true` (oauth2, interactive - not `client_credentials`)** - LiteLLM still allows the anonymous upstream PKCE path (no proxy API key for `/authorize` and matching MCP routes). The internal-only flag mainly affects other surfaces (e.g. IP-based discovery). Rely on the upstream IdP and network policy; the dashboard shows a warning when both are set, and the proxy logs a warning when the server is loaded from config or the database +MCP note: **`auth_type: oauth2` with `delegate_auth_to_upstream: true` is deprecated** - LiteLLM admission is required for matching MCP routes. Use `auth_type: oauth_delegate` for client-forwarded OAuth. OAuth discovery endpoints stay public so clients can start the RFC 9728 flow 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 cd6739777dd..b0d57cb6228 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 @@ -129,10 +129,9 @@ def _is_mcp_passthrough_cold_start(mcp_servers: list[str] | None, client_ip: str spec-compliant WWW-Authenticate challenge instead of surfacing a generic admission error. - Uses "all" semantics (mirrors - :meth:`MCPRequestHandler._target_servers_delegate_auth_to_upstream`): one - non-passthrough target in a co-targeted set must not flip the bypass open - for the others. Fails closed when any target cannot be resolved.""" + Uses "all" semantics: one non-passthrough target in a co-targeted set must + not flip the bypass open for the others. Fails closed when any target + cannot be resolved.""" if not mcp_servers: return False from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( @@ -146,6 +145,27 @@ def _is_mcp_passthrough_cold_start(mcp_servers: list[str] | None, client_ip: str return True +def _is_legacy_delegate_cold_start(mcp_servers: list[str] | None, client_ip: str | None) -> bool: + """Allow only credential-free legacy delegates to reach the route's OAuth challenge.""" + if not mcp_servers: + return False + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + MCPServerManager, + global_mcp_server_manager, + ) + from litellm.types.mcp import MCPAuth + + for name in mcp_servers: + server = global_mcp_server_manager.get_mcp_server_by_name(name, client_ip=client_ip) + if server is None or server.auth_type != MCPAuth.oauth2: + return False + if server.delegate_auth_to_upstream is not True: + return False + if MCPServerManager.effective_oauth2_flow(server) == "client_credentials": + return False + return True + + def _is_litellm_auth_admission_error(exc: Exception) -> bool: if isinstance(exc, HTTPException): return exc.status_code == 401 @@ -277,9 +297,18 @@ def _admission_failure_fallback( 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), + and ( + _is_mcp_passthrough_cold_start( + mcp_servers_from_path, + client_ip=IPAddressUtils.get_mcp_client_ip(request), + ) + or ( + not bearer_presented + and _is_legacy_delegate_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") @@ -434,22 +463,6 @@ class MCPRequestHandler: api_key=f"Bearer {_get_bearer_token_or_received_api_key(litellm_api_key)}", request=request, ) - elif MCPRequestHandler._target_servers_delegate_auth_to_upstream( - path=request_route, - mcp_servers=mcp_servers, - client_ip=IPAddressUtils.get_mcp_client_ip(request), - ): - # Operator opted this oauth2 server into upstream-delegated auth: the - # client authenticates directly with the upstream MCP server, so any - # Authorization bearer is an upstream token, never a LiteLLM key. Skip - # LiteLLM validation entirely — covering both the no-credential - # discovery request and the authenticated call carrying the upstream - # bearer — so a tool call that succeeds never carries a phantom 401 - # auth span; the bearer is forwarded upstream unchanged. Gated by - # _target_servers_delegate_auth_to_upstream, which returns True only - # when EVERY target is auth_type=oauth2 with delegate_auth_to_upstream - # set; fails closed otherwise. - validated_user_api_key_auth = UserAPIKeyAuth() elif MCPRequestHandler._target_servers_are_true_passthrough( path=request_route, mcp_servers=mcp_servers, @@ -660,64 +673,6 @@ class MCPRequestHandler: return [single_server_match.group(1)] return [servers_and_path] - @staticmethod - def _target_servers_delegate_auth_to_upstream( - path: str, mcp_servers: list[str] | None, client_ip: str | None - ) -> bool: - """ - True only when EVERY MCP server the request targets is configured for - ``auth_type == oauth2`` AND has ``delegate_auth_to_upstream=True``. - Fails closed when any target does not opt in or cannot be resolved. - - Used by :meth:`process_mcp_request` to skip LiteLLM API-key/SSO auth - entirely (PKCE passthrough) so the client authenticates directly with - the upstream MCP server. Mixed-target requests (e.g. one delegated + - one non-delegated server) fall back to normal LiteLLM auth. - """ - # Inline imports avoid a circular dependency: mcp_server_manager imports - # from this module. - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - MCPServerManager, - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - - # Must mirror the downstream header-vs-path override - # (``extract_mcp_auth_context``) or an attacker could set - # ``x-mcp-servers`` to a delegate-enabled server while the URL path - # targets a non-delegate server, skipping LiteLLM auth for it. - target_names: Final = MCPRequestHandler._resolve_target_server_names(path=path, mcp_servers_header=mcp_servers) - if not target_names: - return False - - for name in target_names: - server = global_mcp_server_manager.get_mcp_server_by_name(name, client_ip=client_ip) - if server is None or server.auth_type != MCPAuth.oauth2: - return False - # `is True` is intentional: opt-in must be an explicit boolean - # True. A MagicMock attribute (in tests) or any other truthy - # non-bool must not silently enable the bypass. - if getattr(server, "delegate_auth_to_upstream", False) is not True: - return False - # Never delegate for M2M (client_credentials) servers: LiteLLM - # fetches the upstream token automatically using stored credentials, - # so allowing anonymous bypass would let any external caller invoke - # tools authenticated as LiteLLM's service account. - # - # Resolve the flow rather than reading has_client_credentials directly: - # this is a security gate, and a legacy row whose oauth2_flow was never - # stamped still carries the M2M credential shape (client_id/secret + - # token_url, no authorization_url). Treating an unstamped-but-M2M-shaped - # row as non-M2M here would reopen the anonymous bypass the explicit - # column no longer closes on its own. Shares the one resolution helper - # with the egress backstop and the anonymous-delegate allowlist; all fail - # closed on the ambiguous shape and are removed together once no null rows - # remain. A pure-PKCE delegate server (no stored credentials) resolves to a - # non-M2M flow and keeps its bypass. - if MCPServerManager.effective_oauth2_flow(server) == "client_credentials": - return False - return True - @staticmethod def _target_servers_are_true_passthrough(path: str, mcp_servers: list[str] | None, client_ip: str | None) -> bool: """ @@ -726,7 +681,7 @@ class MCPRequestHandler: Used by :meth:`process_mcp_request` to skip LiteLLM admission auth entirely: the gateway is a transparent proxy and the caller's ``Authorization`` is an upstream token, never a LiteLLM key. - Mirrors :meth:`_target_servers_delegate_auth_to_upstream`; a mixed-target request keeps normal auth. + A mixed-target request keeps normal auth. """ from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 3291d5effc2..24e3bc6efc4 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1453,22 +1453,19 @@ def _warn_on_server_name_fields( _warn("server_name", server_name) -def _warn_internal_delegate_pkce_if_applicable(server: MCPServer, *, source: str) -> None: - """Surface internal + upstream PKCE delegate in logs for operators.""" +def _warn_legacy_delegate_auth_if_applicable(server: MCPServer, *, source: str) -> None: + """Direct legacy delegated OAuth configurations to the admitted replacement.""" if server.auth_type != MCPAuth.oauth2: return if getattr(server, "delegate_auth_to_upstream", False) is not True: return - if getattr(server, "available_on_public_internet", True): - return if server.has_client_credentials: return label: Final = get_server_prefix(server) verbose_logger.warning( - "MCP server %r (id=%s, source=%s): internal-only (available_on_public_internet=false) " - "with delegate_auth_to_upstream=true. Anonymous callers can reach the upstream OAuth2 " - "/authorize flow and complete PKCE without a LiteLLM API key session; ensure the " - "upstream IdP and network enforce your access policy.", + "MCP server %r (id=%s, source=%s) uses deprecated auth_type=oauth2 with " + "delegate_auth_to_upstream=true. LiteLLM admission is now required; migrate to " + "auth_type=oauth_delegate for client-forwarded OAuth.", label, server.server_id, source, @@ -2640,7 +2637,7 @@ class MCPServerManager: oauth_identity_binding=server_config.get("oauth_identity_binding", None), ) self._assign_unique_short_prefix(new_server) - _warn_internal_delegate_pkce_if_applicable(new_server, source="config") + _warn_legacy_delegate_auth_if_applicable(new_server, source="config") _warn_config_id_jag_server_outruns_sso(new_server) self._invalidate_discovery_lists(server_id) self.config_mcp_servers[server_id] = new_server @@ -3185,7 +3182,7 @@ class MCPServerManager: timeout=getattr(mcp_server, "timeout", None), max_concurrent_requests=getattr(mcp_server, "max_concurrent_requests", None), ) - _warn_internal_delegate_pkce_if_applicable(new_server, source="database") + _warn_legacy_delegate_auth_if_applicable(new_server, source="database") self._set_oauth_discovery_deferred( new_server.server_id, _requires_oauth_discovery(server_url, use_issuer_anchor, new_server), @@ -3479,10 +3476,6 @@ class MCPServerManager: ) ) - # For anonymous callers (no user_id, no role), also surface any - # servers the operator has opted into upstream-delegated auth. - # These servers handle their own auth at the upstream level, so - # LiteLLM granting access here does not bypass any security gate. is_anonymous: Final = not ( user_api_key_auth and ( @@ -3492,23 +3485,12 @@ class MCPServerManager: ) ) if is_anonymous: - delegate_server_ids: Final = [ + passthrough_server_ids: Final = [ server.server_id for server in self.get_registry().values() - if ( - getattr(server, "auth_type", None) == MCPAuth.oauth2 - and getattr(server, "delegate_auth_to_upstream", False) is True - # M2M servers must not be exposed anonymously: an - # unauthenticated caller would get LiteLLM to proxy tool - # calls using its stored client_credentials. Resolve the flow - # rather than reading has_client_credentials so an unstamped - # M2M-shape row (null column, verbatim-read as non-M2M) still - # fails closed here, matching the anonymous-delegate auth gate. - and MCPServerManager.effective_oauth2_flow(server) != "client_credentials" - ) - or getattr(server, "auth_type", None) == MCPAuth.true_passthrough + if getattr(server, "auth_type", None) == MCPAuth.true_passthrough ] - combined_servers.update(delegate_server_ids) + combined_servers.update(passthrough_server_ids) restrict_allow_all: Final = ( resolved_general_settings.get("mcp_allow_all_keys_respects_mcp_scope", False) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index fc87db69e16..7feb1fd468d 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -4257,20 +4257,6 @@ if MCP_AVAILABLE: return None return _get_authorization_header_from_scope(scope) - def _is_delegate_upstream_probe_target(server: MCPServer) -> bool: - """Whether ``server`` is an interactive delegate-auth server whose client-supplied - token should be preflighted upstream. - - Mirrors the anonymous-delegate gate in ``get_allowed_mcp_servers``: the flow is - resolved via ``effective_oauth2_flow`` so an unstamped M2M-shape row fails closed - (its stored client credentials drive egress; the caller's bearer is irrelevant). - """ - return ( - server.auth_type == MCPAuth.oauth2 - and server.delegate_auth_to_upstream is True - and MCPServerManager.effective_oauth2_flow(server) != "client_credentials" - ) - async def _probe_upstream_auth( url: str, auth_header: str, @@ -4331,7 +4317,7 @@ if MCP_AVAILABLE: mcp_servers: list[str] | None, client_ip: str | None, ) -> None: - """Probe pass-through and delegate-auth upstream servers in parallel before the MCP session starts. + """Probe pass-through upstream servers in parallel before the MCP session starts. Only servers the caller's key is already authorized to reach are probed — the list is derived from _get_allowed_mcp_servers so that a user cannot @@ -4343,38 +4329,9 @@ if MCP_AVAILABLE: if the upstream accepts it but forbids the caller. Fails-open: network errors are logged and the request is allowed through. - Delegate-auth servers (``auth_type=oauth2`` + ``delegate_auth_to_upstream``) - are probed with the caller's bare ``Authorization`` bearer. That bearer is only - an upstream token (never a LiteLLM key) when admission took the delegate bypass, - so the delegate target is resolved through ``get_mcp_server_by_name`` -- the same - resolver admission used -- rather than the wider allowed-server prefix/access-group - matching. A name that only reaches a delegate server via server_id or an access - group would have been admitted as a real LiteLLM key, so probing it would leak that - key upstream; requiring the admission-resolver match closes that gap. Without the - probe a rejected token is absorbed by the tools/list handler and masked as an empty - tool list. Gated to single-server routes so one rejected token cannot 401 a - multi-server aggregate connect, matching the OBO preflight gating; the challenge - echoes the requested name so aliased routes get the same resource_metadata URL as - the tokenless preemptive challenge. """ forwarded_auth: Final = _get_forwarded_auth_from_scope(scope) - requested_single_target: Final = mcp_servers[0] if mcp_servers is not None and len(mcp_servers) == 1 else None - # The bare Authorization header (no x-litellm-api-key) is a valid upstream token - # only when admission classified it as one, i.e. the single requested name resolves - # to a delegate server under admission's own resolver. Resolve it the same way here - # so a server_id- or access-group-named delegate (which admission would have treated - # as a LiteLLM key) is never probed with that key. - delegate_server: Final = ( - global_mcp_server_manager.get_mcp_server_by_name(requested_single_target, client_ip=client_ip) - if requested_single_target - else None - ) - delegate_auth: Final = ( - _get_authorization_header_from_scope(scope) - if delegate_server is not None and _is_delegate_upstream_probe_target(delegate_server) - else None - ) - if not forwarded_auth and not delegate_auth: + if not forwarded_auth: return # Use the authorized server set, not the raw user-supplied names, so that @@ -4384,35 +4341,20 @@ if MCP_AVAILABLE: mcp_servers=mcp_servers, client_ip=client_ip, ) - passthrough_targets: Final[tuple[tuple[MCPServer, str, str], ...]] = ( - tuple( - (srv, forwarded_auth, srv.name) - for srv in allowed_servers - # Restrict to genuine OAuth pass-through servers (auth_type none + - # Authorization in extra_headers). Gateway-managed OAuth2 servers - # must not receive the ``resource_metadata=`` challenge emitted - # below — they require ``authorization_uri=`` pointing at the - # gateway AS metadata. ``is_oauth_passthrough`` already requires - # ``auth_type in (None, MCPAuth.none)``, which is mutually - # exclusive with ``has_client_credentials`` (oauth2 + M2M flow), - # so M2M servers are implicitly excluded here. - if srv.is_oauth_passthrough - ) - if forwarded_auth - else () + passthrough_targets: Final[tuple[tuple[MCPServer, str, str], ...]] = tuple( + (srv, forwarded_auth, srv.name) + for srv in allowed_servers + # Restrict to genuine OAuth pass-through servers (auth_type none + + # Authorization in extra_headers). Gateway-managed OAuth2 servers + # must not receive the ``resource_metadata=`` challenge emitted + # below — they require ``authorization_uri=`` pointing at the + # gateway AS metadata. ``is_oauth_passthrough`` already requires + # ``auth_type in (None, MCPAuth.none)``, which is mutually + # exclusive with ``has_client_credentials`` (oauth2 + M2M flow), + # so M2M servers are implicitly excluded here. + if srv.is_oauth_passthrough ) - # Probe the admission-resolved delegate server only when the caller is actually - # authorized for it (present in the IP-filtered allowed set), keyed by server_id. - delegate_targets: Final[tuple[tuple[MCPServer, str, str], ...]] = ( - tuple( - (srv, delegate_auth, requested_single_target) - for srv in allowed_servers - if delegate_server is not None and srv.server_id == delegate_server.server_id - ) - if delegate_auth and requested_single_target - else () - ) - probe_targets: Final = passthrough_targets + delegate_targets + probe_targets: Final = passthrough_targets if not probe_targets: return diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 64da2ad00f6..3a61f773001 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -64,6 +64,7 @@ _UNATTRIBUTED_TRACKABLE_CALL_TYPES: Final[frozenset[str]] = frozenset( CallTypes.pass_through.value, CallTypes.llm_passthrough_route.value, CallTypes.allm_passthrough_route.value, + CallTypes.call_mcp_tool.value, # CheckBatchCost's synthetic logging_obj for a completed managed batch carries # whatever LiteLLM_ManagedObjectTable stored at create time, and all of it is # None for a batch created before those columns were persisted, or by the master diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 58f9dadc1ce..985d31af997 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -155,11 +155,9 @@ class MCPServer(BaseModel): access_groups: list[str] | None = None allow_all_keys: bool = False available_on_public_internet: bool = True - # Explicit opt-in to upstream-delegated authentication for ``oauth2`` - # servers. When ``auth_type == oauth2`` and this is ``True``, MCP requests - # bypass LiteLLM API-key/SSO auth (and the pre-emptive 401) so the client - # completes PKCE directly with the upstream MCP server. See - # ``MCPRequestHandler._target_servers_delegate_auth_to_upstream``. + # Legacy opt-in to upstream-delegated authentication for ``oauth2`` + # servers. LiteLLM admission still applies; use ``oauth_delegate`` for the + # supported client-forwarded OAuth flow. # # Honored only for ``auth_type == oauth2``; ignored for any other # ``auth_type``. OAuth pass-through for non-oauth2 servers diff --git a/tests/e2e/coverage_registry/mcp.yaml b/tests/e2e/coverage_registry/mcp.yaml index f853e9ff8d6..a7d4135d550 100644 --- a/tests/e2e/coverage_registry/mcp.yaml +++ b/tests/e2e/coverage_registry/mcp.yaml @@ -78,7 +78,7 @@ auth_family: none assertions: [succeeds] source: "mcp_server_manager.py:1485-1492" - rationale: Public/anonymous servers; delegate_auth_to_upstream + rationale: Explicitly anonymous true_passthrough servers - id: mcp.call_tool.none.succeeds module: mcp tier: P1 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 c2f4f7163a0..27a8c7992fe 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 @@ -1483,12 +1483,10 @@ class TestMCPOAuth2AuthFlow: as LiteLLM API keys, causing auth failures and empty tool listings. """ - async def test_oauth2_token_in_authorization_header_fallback(self): + async def test_oauth2_token_in_authorization_header_requires_litellm_admission(self): """ - When only the Authorization header is present with a non-LiteLLM OAuth2 - token AND the target server delegates auth to upstream, LiteLLM skips its - own validation entirely (so the upstream token is never mistaken for a - virtual key) and forwards the bearer upstream. + A bare Authorization token on the legacy delegated mode must establish + a LiteLLM principal rather than entering anonymously. """ from litellm.types.mcp import MCPAuth @@ -1510,6 +1508,7 @@ class TestMCPOAuth2AuthFlow: patch( "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", new_callable=AsyncMock, + return_value=UserAPIKeyAuth(user_id="admitted-user"), ) as mock_auth, patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr, ): @@ -1524,9 +1523,8 @@ class TestMCPOAuth2AuthFlow: ) = await MCPRequestHandler.process_mcp_request(scope) assert isinstance(auth_result, UserAPIKeyAuth) - # The upstream token is never validated as a LiteLLM key ... - mock_auth.assert_not_called() - # ... and is preserved for upstream forwarding. + assert auth_result.user_id == "admitted-user" + mock_auth.assert_awaited_once() assert oauth2_headers.get("Authorization") == "Bearer atlassian-oauth2-access-token-xyz" async def test_explicit_litellm_key_with_oauth2_authorization(self): @@ -2367,11 +2365,9 @@ class TestMCPDelegateAuthToUpstream: """ Tests for the ``delegate_auth_to_upstream`` per-server flag. - When set on an ``auth_type=oauth2`` MCP server, LiteLLM must skip its own - API-key/SSO check entirely so the client completes PKCE directly with the - upstream MCP server. The gate must fail closed for any non-oauth2 server, - any mixed-target request, and any request where the target cannot be - resolved. + The legacy flag no longer bypasses LiteLLM admission. OAuth discovery may + still use the anonymous cold-start challenge, but a presented bearer must + authenticate to LiteLLM unless a separate admission credential is supplied. """ @staticmethod @@ -2386,6 +2382,13 @@ class TestMCPDelegateAuthToUpstream: delegate_auth_to_upstream=delegate_auth_to_upstream, ) + def test_legacy_delegate_cold_start_fails_closed_without_targets(self): + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( + _is_legacy_delegate_cold_start, + ) + + assert _is_legacy_delegate_cold_start(None, client_ip=None) is False + def test_build_mcp_server_table_preserves_delegate_auth_to_upstream(self): """Registry → API list rows must expose delegate_auth_to_upstream for the UI.""" from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( @@ -2439,11 +2442,10 @@ class TestMCPDelegateAuthToUpstream: not_passthrough = passthrough.model_copy(update={"oauth_passthrough": False}) assert manager._build_mcp_server_table(not_passthrough).oauth_passthrough is False - async def test_delegate_skips_litellm_auth_with_no_authorization(self): + async def test_delegate_without_authorization_attempts_litellm_auth_before_cold_start(self): """ - oauth2 + delegate_auth_to_upstream=True, no Authorization header at - all → anonymous UserAPIKeyAuth and ``user_api_key_auth`` is never - called. + A credential-free discovery request attempts LiteLLM admission before + the route emits its RFC 9728 challenge. """ from litellm.types.mcp import MCPAuth @@ -2457,6 +2459,8 @@ class TestMCPDelegateAuthToUpstream: with ( patch( "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + new_callable=AsyncMock, + side_effect=HTTPException(status_code=401, detail="No key provided"), ) as mock_auth, patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr, ): @@ -2466,17 +2470,12 @@ class TestMCPDelegateAuthToUpstream: ) auth_result, *_rest = await MCPRequestHandler.process_mcp_request(scope) assert isinstance(auth_result, UserAPIKeyAuth) - mock_auth.assert_not_called() + mock_auth.assert_awaited_once() - async def test_delegate_with_upstream_token_in_authorization_skips_litellm_auth( - self, - ): + async def test_delegate_with_only_upstream_token_requires_litellm_auth(self): """ - oauth2 + delegate_auth_to_upstream=True with an upstream OAuth token in - ``Authorization``: the delegate gate fires before any LiteLLM validation, - so ``user_api_key_auth`` is never called and the bearer is forwarded - upstream untouched. Skipping the doomed validation is what keeps a tool - call that actually succeeds from carrying a phantom 401 auth span. + An upstream token cannot establish a LiteLLM principal and must not + reopen anonymous admission. """ from litellm.types.mcp import MCPAuth @@ -2491,6 +2490,7 @@ class TestMCPDelegateAuthToUpstream: patch( "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", new_callable=AsyncMock, + side_effect=HTTPException(status_code=401, detail="Invalid API key"), ) as mock_auth, patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr, ): @@ -2498,17 +2498,11 @@ class TestMCPDelegateAuthToUpstream: auth_type=MCPAuth.oauth2, delegate_auth_to_upstream=True, ) - ( - auth_result, - _, - _, - _, - oauth2_headers, - _, - ) = await MCPRequestHandler.process_mcp_request(scope) - assert isinstance(auth_result, UserAPIKeyAuth) - assert oauth2_headers.get("Authorization") == "Bearer upstream-pkce-token" - mock_auth.assert_not_called() + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(scope) + + assert exc_info.value.status_code == 401 + mock_auth.assert_awaited_once() async def test_delegate_off_still_requires_litellm_auth(self): """ @@ -2687,15 +2681,11 @@ class TestMCPDelegateAuthToUpstream: assert auth_result.user_id == "real-user" mock_auth.assert_called_once() - async def test_authorization_bearer_on_delegate_server_treated_as_upstream(self): + async def test_authorization_bearer_on_delegate_server_establishes_litellm_principal(self): """ - On a delegate server the ``Authorization`` header is, by contract, an - upstream token rather than a LiteLLM key — even when it is sk-shaped. It - is forwarded upstream without LiteLLM validation, so ``user_api_key_auth`` - is not called and no LiteLLM identity is resolved. Callers who need - LiteLLM identity / spend tracking on a delegate server must supply - ``x-litellm-api-key`` (see - test_explicit_litellm_key_takes_precedence_over_delegate). + A bare Authorization bearer now follows normal LiteLLM admission. A + separate x-litellm-api-key is required when Authorization is intended + for the upstream server. """ from litellm.types.mcp import MCPAuth @@ -2727,9 +2717,9 @@ class TestMCPDelegateAuthToUpstream: _, ) = await MCPRequestHandler.process_mcp_request(scope) assert isinstance(auth_result, UserAPIKeyAuth) - assert auth_result.user_id is None + assert auth_result.user_id == "real-user" assert oauth2_headers.get("Authorization") == "Bearer sk-1234" - mock_auth.assert_not_called() + mock_auth.assert_awaited_once() async def test_delegate_ignored_for_client_credentials_server(self): """ @@ -2828,13 +2818,10 @@ class TestMCPDelegateAuthToUpstream: assert exc_info.value.status_code == 401 mock_auth.assert_called_once() - async def test_delegate_bypass_for_pure_pkce_server(self): + async def test_delegate_pkce_cold_start_attempts_litellm_auth(self): """ - oauth2 + delegate + oauth2_flow=None and NO stored client credentials - (pure PKCE, the common delegate case) → bypass must still fire. The - shape resolves to a non-M2M flow, so the security gate leaves it alone; - the fail-closed rule targets the M2M shape specifically, not every - unstamped row. + A pure PKCE server may defer a credential-free request to the route's + challenge, but normal LiteLLM admission still runs first. """ from litellm.types.mcp import MCPAuth from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -2869,13 +2856,12 @@ class TestMCPDelegateAuthToUpstream: ): mock_mgr.get_mcp_server_by_name.return_value = pkce_server auth, *_rest = await MCPRequestHandler.process_mcp_request(scope) - mock_auth.assert_not_called() + mock_auth.assert_awaited_once() assert auth.api_key is None - async def test_delegate_bypass_for_internal_server(self): + async def test_internal_delegate_cold_start_attempts_litellm_auth(self): """ - Delegate + oauth2 interactive servers bypass LiteLLM auth even when - ``available_on_public_internet`` is False (internal MCPs). + Internal delegated servers follow the same admission contract. """ from litellm.types.mcp import MCPAuth from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -2910,14 +2896,11 @@ class TestMCPDelegateAuthToUpstream: ): mock_mgr.get_mcp_server_by_name.return_value = internal_server auth, *_rest = await MCPRequestHandler.process_mcp_request(scope) - mock_auth.assert_not_called() + mock_auth.assert_awaited_once() assert auth.api_key is None - async def test_get_allowed_servers_excludes_client_credentials_delegate(self): - """ - get_allowed_mcp_servers must not surface M2M (client_credentials) delegate - servers to anonymous callers even if delegate_auth_to_upstream=True. - """ + async def test_get_allowed_servers_excludes_legacy_delegates(self): + """Legacy delegated servers are never added to anonymous access.""" from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( MCPServerManager, ) @@ -2955,102 +2938,7 @@ class TestMCPDelegateAuthToUpstream: ): result = await manager.get_allowed_mcp_servers(None) - assert "pkce-server" in result - assert "m2m-server" not in result - - async def test_get_allowed_servers_excludes_unstamped_m2m_shape_delegate(self): - """ - The anonymous allow-list must also exclude an M2M-shape delegate server whose - oauth2_flow was never stamped (null column, verbatim-read as non-M2M). Reading - the bare has_client_credentials here would surface it to anonymous callers; the - resolved-flow check fails closed on the shape, matching the auth gate. - """ - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - MCPServerManager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - - manager = MCPServerManager() - pkce_server = MCPServer( - server_id="pkce-server", - name="pkce_server", - transport="http", - auth_type=MCPAuth.oauth2, - delegate_auth_to_upstream=True, - available_on_public_internet=True, - ) - unstamped_m2m = MCPServer( - server_id="unstamped-m2m", - name="unstamped_m2m", - transport="http", - auth_type=MCPAuth.oauth2, - delegate_auth_to_upstream=True, - oauth2_flow=None, - client_id="cid", - client_secret="csecret", - token_url="https://idp.example.com/token", - ) - assert unstamped_m2m.has_client_credentials is False - manager.registry = { - pkce_server.server_id: pkce_server, - unstamped_m2m.server_id: unstamped_m2m, - } - - with patch.object( - MCPRequestHandler, - "get_allowed_mcp_servers", - new_callable=AsyncMock, - return_value=[], - ): - result = await manager.get_allowed_mcp_servers(None) - - assert "pkce-server" in result - assert "unstamped-m2m" not in result - - async def test_get_allowed_servers_includes_internal_delegate(self): - """ - Internal-only (available_on_public_internet=False) delegate servers - appear in the anonymous allow-list like public delegate servers. - """ - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - MCPServerManager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - - manager = MCPServerManager() - public_server = MCPServer( - server_id="public-server", - name="public_server", - transport="http", - auth_type=MCPAuth.oauth2, - delegate_auth_to_upstream=True, - available_on_public_internet=True, - ) - internal_server = MCPServer( - server_id="internal-server", - name="internal_server", - transport="http", - auth_type=MCPAuth.oauth2, - delegate_auth_to_upstream=True, - available_on_public_internet=False, - ) - manager.registry = { - public_server.server_id: public_server, - internal_server.server_id: internal_server, - } - - with patch.object( - MCPRequestHandler, - "get_allowed_mcp_servers", - new_callable=AsyncMock, - return_value=[], - ): - result = await manager.get_allowed_mcp_servers(None) - - assert "public-server" in result - assert "internal-server" in result + assert result == [] async def test_true_passthrough_skips_litellm_auth_anonymously(self): """auth_type=true_passthrough performs no admission auth: the caller's Authorization is an @@ -4453,7 +4341,9 @@ class TestAgentMCPPermissions: stack.enter_context(patcher) stack.enter_context( patch.object( # test-quality-ok: key resolution has its own tests; pin its grants here - MCPRequestHandler, "_get_allowed_mcp_servers_for_key", AsyncMock(return_value=["server-a", "server-b"]) + MCPRequestHandler, + "_get_allowed_mcp_servers_for_key", + AsyncMock(return_value=["server-a", "server-b"]), ) ) stack.enter_context( @@ -4479,7 +4369,9 @@ class TestAgentMCPPermissions: await MCPRequestHandler._get_allowed_mcp_servers_for_agent(user_api_key_auth) stack.enter_context( patch.object( # test-quality-ok: key resolution has its own tests; pin its grants here - MCPRequestHandler, "_get_allowed_mcp_servers_for_key", AsyncMock(return_value=["server-a", "server-b"]) + MCPRequestHandler, + "_get_allowed_mcp_servers_for_key", + AsyncMock(return_value=["server-a", "server-b"]), ) ) stack.enter_context( @@ -4503,9 +4395,15 @@ class TestAgentMCPPermissions: with contextlib.ExitStack() as stack: for patcher in self._agent_toolset_patches(agent_object_permission, mock_manager): stack.enter_context(patcher) - server_a_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server("server-a", user_api_key_auth) - server_b_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server("server-b", user_api_key_auth) - server_c_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server("server-c", user_api_key_auth) + server_a_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server( + "server-a", user_api_key_auth + ) + server_b_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server( + "server-b", user_api_key_auth + ) + server_c_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server( + "server-c", user_api_key_auth + ) assert sorted(server_a_tools) == ["tool_direct", "tool_via_toolset"] assert server_b_tools == ["tool_b"] diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 7eb4396e60d..38cfbd20df4 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -6157,15 +6157,8 @@ def _patch_delegate_resolver(server: MCPServer, *resolvable_names: str): @pytest.mark.asyncio -async def test_delegate_bad_token_gets_connect_time_401(): - """Regression (LIT-4194): a rejected upstream token on a delegate-auth server - must fail the connect with 401 + ``error="invalid_token"``, not be absorbed - into HTTP 200 + an empty tool list by the tools/list handler. - - Delegate-mode clients send only ``Authorization`` (no ``x-litellm-api-key``), - so ``_get_forwarded_auth_from_scope`` returns None and, before the fix, the - preflight returned early without probing. - """ +async def test_legacy_delegate_bare_token_is_not_probed_upstream(): # test-quality-ok: this removed security-sensitive egress has no return value; non-invocation is the contract + """A bare bearer is an admission credential and must never reach upstream.""" from litellm.proxy._experimental.mcp_server.server import ( _check_passthrough_upstream_auth, ) @@ -6184,48 +6177,6 @@ async def test_delegate_bad_token_gets_connect_time_401(): "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", new=AsyncMock(return_value=(401, 'Bearer realm="upstream", error="invalid_token"')), ) as probe, - ): - with pytest.raises(HTTPException) as exc_info: - await _check_passthrough_upstream_auth( - scope=scope, - user_api_key_auth=UserAPIKeyAuth(), - mcp_servers=["delegate_test"], - client_ip=None, - ) - - assert exc_info.value.status_code == 401 - challenge = exc_info.value.headers["www-authenticate"] - assert 'error="invalid_token"' in challenge - assert ( - 'resource_metadata="http://localhost:4000/.well-known/oauth-protected-resource/mcp/delegate_test"' in challenge - ) - probe.assert_awaited_once() - probe_url, probe_auth = probe.call_args.args - assert probe_url == "http://upstream:9401/mcp" - assert probe_auth == "Bearer bogus-token" - - -@pytest.mark.asyncio -async def test_delegate_valid_token_passes_preflight(): - """An upstream-accepted token must not be blocked by the delegate preflight.""" - from litellm.proxy._experimental.mcp_server.server import ( - _check_passthrough_upstream_auth, - ) - from litellm.proxy._types import UserAPIKeyAuth - - server = _delegate_auth_mcp_server() - scope = _delegate_scope([(b"authorization", b"Bearer good-token")]) - - with ( - _patch_delegate_resolver(server, "delegate_test"), - patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", - new=AsyncMock(return_value=[server]), - ), - patch( - "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", - new=AsyncMock(return_value=(200, None)), - ) as probe, ): await _check_passthrough_upstream_auth( scope=scope, @@ -6234,43 +6185,103 @@ async def test_delegate_valid_token_passes_preflight(): client_ip=None, ) - probe.assert_awaited_once() + probe.assert_not_awaited() @pytest.mark.asyncio -async def test_delegate_valid_token_forbidden_returns_403(): - """An upstream that accepts the token but forbids the caller (403) must surface - as a bare 403 with no ``WWW-Authenticate`` re-auth hint (a fresh token with the - same scopes would loop), not as an invalid_token challenge.""" +async def test_legacy_delegate_dual_credentials_are_not_probed_upstream(): # test-quality-ok: this removed security-sensitive egress has no return value; non-invocation is the contract + """A separate upstream bearer never triggers the removed legacy probe.""" from litellm.proxy._experimental.mcp_server.server import ( _check_passthrough_upstream_auth, ) from litellm.proxy._types import UserAPIKeyAuth server = _delegate_auth_mcp_server() - scope = _delegate_scope([(b"authorization", b"Bearer scoped-out-token")]) + scope = _delegate_scope( + [ + (b"x-litellm-api-key", b"sk-litellm-proxy-key"), + (b"authorization", b"Bearer upstream-token"), + ] + ) with ( - _patch_delegate_resolver(server, "delegate_test"), - patch( + patch( # test-quality-ok: isolate authorized-server resolution so this test targets the preflight boundary "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", new=AsyncMock(return_value=[server]), ), - patch( + patch( # test-quality-ok: the removed probe call is the security regression under test "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", - new=AsyncMock(return_value=(403, None)), - ), + new=AsyncMock(), + ) as probe, ): - with pytest.raises(HTTPException) as exc_info: - await _check_passthrough_upstream_auth( + await _check_passthrough_upstream_auth( + scope=scope, + user_api_key_auth=UserAPIKeyAuth(user_id="admitted-user"), + mcp_servers=["delegate_test"], + client_ip=None, + ) + + probe.assert_not_awaited() + + +@pytest.mark.parametrize( + "probe_status, expected_status", + [(200, None), (401, 401), (403, 403)], +) +@pytest.mark.asyncio +async def test_oauth_passthrough_preflight_preserves_status_contract(probe_status, expected_status): + from litellm.proxy._experimental.mcp_server.server import ( + _check_passthrough_upstream_auth, + ) + from litellm.proxy._types import UserAPIKeyAuth + + server = MCPServer( + server_id="passthrough-id", + name="passthrough_server", + url="https://upstream.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + oauth_passthrough=True, + extra_headers=["Authorization"], + ) + scope = _delegate_scope( + [ + (b"x-litellm-api-key", b"sk-litellm-proxy-key"), + (b"authorization", b"Bearer upstream-token"), + ] + ) + + with ( + patch( # test-quality-ok: isolate authorized-server resolution so this test exercises the preflight contract + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(return_value=[server]), + ), + patch( # test-quality-ok: the upstream transport boundary is the behavior being mapped to an HTTP response + "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", + new=AsyncMock(return_value=(probe_status, None)), + ) as probe, + ): + if expected_status is None: + result = await _check_passthrough_upstream_auth( scope=scope, - user_api_key_auth=UserAPIKeyAuth(), - mcp_servers=["delegate_test"], + user_api_key_auth=UserAPIKeyAuth(user_id="admitted-user"), + mcp_servers=["passthrough_server"], client_ip=None, ) + assert result is None + else: + with pytest.raises(HTTPException) as exc_info: + await _check_passthrough_upstream_auth( + scope=scope, + user_api_key_auth=UserAPIKeyAuth(user_id="admitted-user"), + mcp_servers=["passthrough_server"], + client_ip=None, + ) + assert exc_info.value.status_code == expected_status + if expected_status == 401: + assert "passthrough_server" in exc_info.value.headers["www-authenticate"] - assert exc_info.value.status_code == 403 - assert not (exc_info.value.headers or {}) + probe.assert_awaited_once_with("https://upstream.example.com/mcp", "Bearer upstream-token") @pytest.mark.asyncio @@ -6428,123 +6439,6 @@ async def test_delegate_not_probed_when_named_only_via_server_id(): probe.assert_not_awaited() -@pytest.mark.asyncio -async def test_delegate_preflight_with_unpatched_probe(): - """Integration across the preflight and the unpatched ``_probe_upstream_auth``, - mocked only at the httpx-client boundary (tests/test_litellm is mocked-only; the - real-network proof lives in the PR's live-proxy evidence). The mock honors the - ``AsyncHTTPHandler.post`` contract by raising ``httpx.HTTPStatusError`` on the - upstream 401, so the production ``except httpx.HTTPStatusError`` branch is the one - exercised. A rejected token surfaces as the connect-time 401 challenge; an - accepted token passes untouched, and the caller's bearer reaches the delegate URL.""" - import httpx - - from litellm.proxy._experimental.mcp_server.server import ( - _check_passthrough_upstream_auth, - ) - from litellm.proxy._types import UserAPIKeyAuth - - accepted = MagicMock() - accepted.status_code = 200 - accepted.headers = {} - rejected = MagicMock() - rejected.status_code = 401 - rejected.headers = {"www-authenticate": 'Bearer realm="stub-upstream", error="invalid_token"'} - - async def respond_by_token(url=None, headers=None, json=None, timeout=None, **kwargs): - if headers.get("Authorization") == "Bearer good-token": - return accepted - raise httpx.HTTPStatusError( - "401 Unauthorized", - request=httpx.Request("POST", url), - response=rejected, - ) - - mock_client = MagicMock() - mock_client.post = AsyncMock(side_effect=respond_by_token) - - server = _delegate_auth_mcp_server() - - with ( - _patch_delegate_resolver(server, "delegate_test"), - patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", - new=AsyncMock(return_value=[server]), - ), - patch( - "litellm.proxy._experimental.mcp_server.server.get_async_httpx_client", - return_value=mock_client, - ), - ): - with pytest.raises(HTTPException) as exc_info: - await _check_passthrough_upstream_auth( - scope=_delegate_scope([(b"authorization", b"Bearer bogus-token")]), - user_api_key_auth=UserAPIKeyAuth(), - mcp_servers=["delegate_test"], - client_ip=None, - ) - - await _check_passthrough_upstream_auth( - scope=_delegate_scope([(b"authorization", b"Bearer good-token")]), - user_api_key_auth=UserAPIKeyAuth(), - mcp_servers=["delegate_test"], - client_ip=None, - ) - - assert exc_info.value.status_code == 401 - challenge = exc_info.value.headers["www-authenticate"] - assert 'error="invalid_token"' in challenge - assert ( - 'resource_metadata="http://localhost:4000/.well-known/oauth-protected-resource/mcp/delegate_test"' in challenge - ) - probed_urls = [call.kwargs["url"] for call in mock_client.post.await_args_list] - assert probed_urls == ["http://upstream:9401/mcp", "http://upstream:9401/mcp"] - - -@pytest.mark.asyncio -async def test_delegate_challenge_echoes_requested_alias(): - """An alias-routed delegate request must be probed, and the challenge must echo - the requested alias (not the canonical server name) so the resource_metadata - URL matches what the tokenless preemptive challenge emits for the same route.""" - from litellm.proxy._experimental.mcp_server.server import ( - _check_passthrough_upstream_auth, - ) - from litellm.proxy._types import UserAPIKeyAuth - - server = _delegate_auth_mcp_server().model_copy(update={"alias": "dt-alias"}) - scope = { - "type": "http", - "method": "POST", - "path": "/mcp/dt-alias", - "scheme": "http", - "server": ("localhost", 4000), - "headers": [(b"authorization", b"Bearer bogus-token")], - } - - with ( - _patch_delegate_resolver(server, "dt-alias"), - patch( - "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", - new=AsyncMock(return_value=[server]), - ), - patch( - "litellm.proxy._experimental.mcp_server.server._probe_upstream_auth", - new=AsyncMock(return_value=(401, 'Bearer error="invalid_token"')), - ), - ): - with pytest.raises(HTTPException) as exc_info: - await _check_passthrough_upstream_auth( - scope=scope, - user_api_key_auth=UserAPIKeyAuth(), - mcp_servers=["dt-alias"], - client_ip=None, - ) - - challenge = exc_info.value.headers["www-authenticate"] - assert 'error="invalid_token"' in challenge - assert 'resource_metadata="http://localhost:4000/.well-known/oauth-protected-resource/mcp/dt-alias"' in challenge - - @pytest.mark.asyncio async def test_delegate_probe_not_fanned_out_to_access_group_members(): """A single access-group name passes the one-target route gate but must not fan @@ -6578,41 +6472,6 @@ async def test_delegate_probe_not_fanned_out_to_access_group_members(): probe.assert_not_awaited() -def test_is_delegate_upstream_probe_target_fails_closed_on_m2m_shape(): - """An unstamped M2M-shape row (null ``oauth2_flow`` + client credentials) - resolves to ``client_credentials`` and must not be probed with the caller's - bearer; its stored client credentials drive egress instead.""" - from litellm.proxy._experimental.mcp_server.server import ( - _is_delegate_upstream_probe_target, - ) - - assert _is_delegate_upstream_probe_target(_delegate_auth_mcp_server()) is True - - m2m_shape = MCPServer( - server_id="delegate-m2m", - name="delegate_m2m", - url="http://upstream:9401/mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - delegate_auth_to_upstream=True, - oauth2_flow=None, - token_url="http://idp:9000/token", - client_id="client", - client_secret="secret", - ) - assert _is_delegate_upstream_probe_target(m2m_shape) is False - - non_delegate = MCPServer( - server_id="oauth2-plain", - name="oauth2_plain", - url="http://upstream:9401/mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - oauth2_flow="authorization_code", - ) - assert _is_delegate_upstream_probe_target(non_delegate) is False - - @pytest.mark.asyncio async def test_create_mcp_client_sampling_disabled_by_default(): """Sampling callback must be None when allow_sampling is not set (default False).""" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index d2987c5112e..b4f1aeeef10 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -4469,7 +4469,9 @@ class TestMCPServerManager: @pytest.mark.parametrize("auth_type", [MCPAuth.none, MCPAuth.bearer_token, MCPAuth.api_key, MCPAuth.oauth2]) @pytest.mark.parametrize("is_byok", [False, True]) @pytest.mark.parametrize("scheme", ["http", "https"]) - async def test_openapi_health_loads_spec_without_mcp_handshake(self, respx_mock, monkeypatch, auth_type, is_byok, scheme): + async def test_openapi_health_loads_spec_without_mcp_handshake( + self, respx_mock, monkeypatch, auth_type, is_byok, scheme + ): monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") manager = MCPServerManager() server = MCPServer( @@ -4519,14 +4521,28 @@ class TestMCPServerManager: @pytest.mark.parametrize( ("failure", "expected_status", "expected_error"), [ - (httpx.Response(401, text="secret response content"), "unhealthy", "OpenAPI specification request failed (HTTP 401)"), + ( + httpx.Response(401, text="secret response content"), + "unhealthy", + "OpenAPI specification request failed (HTTP 401)", + ), (httpx.Response(404), "unhealthy", "OpenAPI specification request failed (HTTP 404)"), (httpx.Response(500), "unhealthy", "OpenAPI specification request failed (HTTP 500)"), - (httpx.ConnectError("secret network details"), "unhealthy", "OpenAPI specification could not be loaded (ConnectError)"), - (httpx.Response(200, text="secret invalid JSON body"), "unhealthy", "OpenAPI specification could not be loaded (JSONDecodeError)"), + ( + httpx.ConnectError("secret network details"), + "unhealthy", + "OpenAPI specification could not be loaded (ConnectError)", + ), + ( + httpx.Response(200, text="secret invalid JSON body"), + "unhealthy", + "OpenAPI specification could not be loaded (JSONDecodeError)", + ), ], ) - async def test_openapi_health_reports_safe_failures(self, respx_mock, monkeypatch, failure, expected_status, expected_error): + async def test_openapi_health_reports_safe_failures( + self, respx_mock, monkeypatch, failure, expected_status, expected_error + ): monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") manager = MCPServerManager() server = MCPServer( @@ -6916,51 +6932,6 @@ class TestMCPServerManager: == expected_server_ids ) - @pytest.mark.asyncio - async def test_get_allowed_mcp_servers_anonymous_delegate_requires_oauth2(self): - """Anonymous delegated auth listing should only include oauth2 servers.""" - from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( - MCPRequestHandler, - ) - - manager = MCPServerManager() - oauth_delegate_server = MCPServer( - server_id="oauth-delegate", - name="oauth_delegate", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - delegate_auth_to_upstream=True, - ) - api_key_delegate_server = MCPServer( - server_id="api-key-delegate", - name="api_key_delegate", - transport=MCPTransport.http, - auth_type=MCPAuth.api_key, - delegate_auth_to_upstream=True, - ) - oauth_non_delegate_server = MCPServer( - server_id="oauth-non-delegate", - name="oauth_non_delegate", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - delegate_auth_to_upstream=False, - ) - manager.registry = { - oauth_delegate_server.server_id: oauth_delegate_server, - api_key_delegate_server.server_id: api_key_delegate_server, - oauth_non_delegate_server.server_id: oauth_non_delegate_server, - } - - with patch.object( - MCPRequestHandler, - "get_allowed_mcp_servers", - new_callable=AsyncMock, - return_value=[], - ): - result = await manager.get_allowed_mcp_servers(None) - - assert set(result) == {"oauth-delegate"} - def test_get_mcp_server_from_tool_name_uses_server_name_not_name(self): """ Test that _get_mcp_server_from_tool_name uses server.server_name instead of server.name @@ -8088,9 +8059,9 @@ class TestMCPServerTokenExchangeColumns: assert rebuilt_table.token_exchange_profile == "entra_obo" -class TestInternalDelegatePkceWarningLog: +class TestLegacyDelegateAuthWarningLog: @pytest.mark.asyncio - async def test_build_mcp_server_logs_on_internal_delegate_interactive(self, caplog): + async def test_build_mcp_server_logs_deprecation_for_internal_delegate(self, caplog): caplog.set_level(logging.WARNING, logger="LiteLLM") manager = MCPServerManager() table_record = LiteLLM_MCPServerTable( @@ -8106,11 +8077,11 @@ class TestInternalDelegatePkceWarningLog: ) await manager.build_mcp_server_from_table(table_record) combined = " ".join(r.getMessage() for r in caplog.records) - assert "internal-only" in combined + assert "deprecated auth_type=oauth2" in combined assert "delegate_auth_to_upstream=true" in combined @pytest.mark.asyncio - async def test_build_mcp_server_no_internal_delegate_log_when_public(self, caplog): + async def test_build_mcp_server_logs_deprecation_for_public_delegate(self, caplog): caplog.set_level(logging.WARNING, logger="LiteLLM") manager = MCPServerManager() table_record = LiteLLM_MCPServerTable( @@ -8126,12 +8097,13 @@ class TestInternalDelegatePkceWarningLog: ) await manager.build_mcp_server_from_table(table_record) combined = " ".join(r.getMessage() for r in caplog.records) - assert "internal-only" not in combined + assert "deprecated auth_type=oauth2" in combined + assert "auth_type=oauth_delegate" in combined def test_warn_skipped_for_client_credentials(self, caplog): caplog.set_level(logging.WARNING, logger="LiteLLM") from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - _warn_internal_delegate_pkce_if_applicable, + _warn_legacy_delegate_auth_if_applicable, ) server = MCPServer( @@ -8144,9 +8116,9 @@ class TestInternalDelegatePkceWarningLog: available_on_public_internet=False, delegate_auth_to_upstream=True, ) - _warn_internal_delegate_pkce_if_applicable(server, source="test") + _warn_legacy_delegate_auth_if_applicable(server, source="test") combined = " ".join(r.getMessage() for r in caplog.records) - assert "internal-only" not in combined + assert "deprecated auth_type=oauth2" not in combined class TestHasClientCredentialsOAuth2Flow: @@ -12697,7 +12669,12 @@ async def test_debug_resolution_matches_final_header_conflict_winner( from litellm.proxy._experimental.mcp_server.auth.litellm_auth_handler import MCPAuthenticatedUser from litellm.proxy._experimental.mcp_server.mcp_debug import MCP_AUTH_DIAGNOSTICS_SCOPE_KEY, MCPAuthDiagnostics from litellm.proxy._experimental.mcp_server.outbound_credentials import ( - ApiKeyConfig, AuthorizationCodeConfig, NoneConfig, ServerSpec, SharedKey, UpstreamCredentialProvider, + ApiKeyConfig, + AuthorizationCodeConfig, + NoneConfig, + ServerSpec, + SharedKey, + UpstreamCredentialProvider, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import OAuthToken from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -12713,10 +12690,15 @@ async def test_debug_resolution_matches_final_header_conflict_winner( store = Store() context = MCPAuthenticatedUser(UserAPIKeyAuth(user_id="alice")) diagnostics = MCPAuthDiagnostics() - token = request_ctx.set(RequestContext( - request_id=1, meta=None, session=MagicMock(), lifespan_context=None, - request=Request({"type": "http", MCP_AUTH_DIAGNOSTICS_SCOPE_KEY: diagnostics}), - )) + token = request_ctx.set( + RequestContext( + request_id=1, + meta=None, + session=MagicMock(), + lifespan_context=None, + request=Request({"type": "http", MCP_AUTH_DIAGNOSTICS_SCOPE_KEY: diagnostics}), + ) + ) selected = { "stored": AuthorizationCodeConfig(), "static": ApiKeyConfig(key_source=SharedKey(value=SecretStr("static-token"))), @@ -12725,7 +12707,10 @@ async def test_debug_resolution_matches_final_header_conflict_winner( try: auth, remaining = await MCPServerManager()._resolve_v2_auth( server=MCPServer( - server_id="s", name="s", transport="http", url="https://up.example/mcp", + server_id="s", + name="s", + transport="http", + url="https://up.example/mcp", static_headers={"Authorization": "Bearer configured"}, ), spec=ServerSpec(server_id="s", resource="https://up.example/mcp", config=selected), @@ -12755,17 +12740,28 @@ async def test_debug_reports_legacy_signing_and_non_http_transport(transport: Li from litellm.types.mcp_server.mcp_server_manager import MCPServer diagnostics = MCPAuthDiagnostics() - token = request_ctx.set(RequestContext( - request_id=1, meta=None, session=MagicMock(), lifespan_context=None, - request=Request({"type": "http", MCP_AUTH_DIAGNOSTICS_SCOPE_KEY: diagnostics}), - )) + token = request_ctx.set( + RequestContext( + request_id=1, + meta=None, + session=MagicMock(), + lifespan_context=None, + request=Request({"type": "http", MCP_AUTH_DIAGNOSTICS_SCOPE_KEY: diagnostics}), + ) + ) try: server = MCPServer( - server_id="signed", name="signed", transport=transport, - url="https://up.example/mcp", auth_type="aws_sigv4", - aws_access_key_id="AKIDEXAMPLE", aws_secret_access_key="test-signing-secret", - aws_region_name="us-east-1", aws_service_name="execute-api", - command="python", args=["-c", "pass"], + server_id="signed", + name="signed", + transport=transport, + url="https://up.example/mcp", + auth_type="aws_sigv4", + aws_access_key_id="AKIDEXAMPLE", + aws_secret_access_key="test-signing-secret", + aws_region_name="us-east-1", + aws_service_name="execute-api", + command="python", + args=["-c", "pass"], ) client = await MCPServerManager()._create_mcp_client(server) if transport == "stdio": @@ -12784,12 +12780,16 @@ async def test_debug_reports_legacy_signing_and_non_http_transport(transport: Li async def test_temporary_server_discovery_reuses_resolved_metadata_without_publishing() -> None: manager: Final = MCPServerManager() server: Final = MCPServer( - server_id="temporary-oauth-discovery", name="temporary", url="https://idp.example.com/mcp", - transport=MCPTransport.http, auth_type=MCPAuth.true_passthrough, + server_id="temporary-oauth-discovery", + name="temporary", + url="https://idp.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.true_passthrough, ) manager._set_oauth_discovery_deferred(server.server_id, True) metadata: Final = MCPOAuthMetadata( - authorization_url="https://idp.example.com/authorize", token_url="https://idp.example.com/token", + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", registration_url="https://idp.example.com/register", ) with patch.object(manager, "_discover_oauth_metadata_for_server", AsyncMock(return_value=metadata)) as discovery: @@ -12809,13 +12809,18 @@ async def test_temporary_server_discovery_reuses_resolved_metadata_without_publi async def test_repeated_stale_oauth_discovery_is_bounded(auth_type: MCPAuth) -> None: manager: Final = MCPServerManager() server: Final = MCPServer( - server_id="repeated-stale", name="stale", url="https://idp.example.com/mcp", - transport=MCPTransport.http, auth_type=auth_type, oauth2_flow="authorization_code", + server_id="repeated-stale", + name="stale", + url="https://idp.example.com/mcp", + transport=MCPTransport.http, + auth_type=auth_type, + oauth2_flow="authorization_code", ) manager.registry[server.server_id] = server manager._set_oauth_discovery_deferred(server.server_id, True) metadata: Final = MCPOAuthMetadata( - authorization_url="https://idp.example.com/authorize", token_url="https://idp.example.com/token", + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", ) with ( patch.object(manager, "_discover_oauth_metadata_for_server", AsyncMock(return_value=metadata)) as discovery, @@ -12835,13 +12840,20 @@ async def test_repeated_stale_oauth_discovery_is_bounded(auth_type: MCPAuth) -> async def test_stale_discovery_falls_back_to_resolved_registered_server() -> None: manager: Final = MCPServerManager() original: Final = MCPServer( - server_id="resolved-replacement", name="replacement", url="https://old.example.com/mcp", - transport=MCPTransport.http, auth_type=MCPAuth.oauth2, oauth2_flow="authorization_code", + server_id="resolved-replacement", + name="replacement", + url="https://old.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + ) + replacement: Final = original.model_copy( + update={ + "url": "https://new.example.com/mcp", + "authorization_url": "https://new.example.com/authorize", + "token_url": "https://new.example.com/token", + } ) - replacement: Final = original.model_copy(update={ - "url": "https://new.example.com/mcp", "authorization_url": "https://new.example.com/authorize", - "token_url": "https://new.example.com/token", - }) manager.registry[original.server_id] = replacement assert await manager._rejoin_oauth_metadata_discovery(original, retry_stale=False) is replacement @@ -12849,8 +12861,11 @@ async def test_stale_discovery_falls_back_to_resolved_registered_server() -> Non def test_stale_discovery_cannot_overwrite_new_registered_server() -> None: manager: Final = MCPServerManager() original: Final = MCPServer( - server_id="stale-publication", name="publication", url="https://old.example.com/mcp", - transport=MCPTransport.http, auth_type=MCPAuth.oauth2, + server_id="stale-publication", + name="publication", + url="https://old.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, ) manager._set_oauth_discovery_deferred(original.server_id, True) original_slot: Final = manager._oauth_discovery_slot(original.server_id) @@ -12866,9 +12881,13 @@ def test_stale_discovery_cannot_overwrite_new_registered_server() -> None: async def test_temporary_oauth_discovery_expires_without_more_requests() -> None: manager: Final = MCPServerManager() server: Final = MCPServer( - server_id="expiring-session", name="temporary", url="https://idp.example.com/mcp", - transport=MCPTransport.http, auth_type=MCPAuth.true_passthrough, - authorization_url="https://idp.example.com/authorize", token_url="https://idp.example.com/token", + server_id="expiring-session", + name="temporary", + url="https://idp.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.true_passthrough, + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", ) manager._set_oauth_discovery_deferred(server.server_id, True) resolved: Final = await manager.ensure_oauth_metadata_discovered(server) @@ -12969,7 +12988,9 @@ async def test_openapi_health_reports_size_limit_as_unknown_and_caches_failure(r result = await manager.health_check_server(server.server_id) cached = await manager.health_check_server(server.server_id) assert result.status == "unknown" - assert result.health_check_error == "OpenAPI specification probe refused: Response exceeds the configured size limit" + assert ( + result.health_check_error == "OpenAPI specification probe refused: Response exceeds the configured size limit" + ) assert cached.health_check_error == result.health_check_error assert cached.last_health_check == result.last_health_check assert route.call_count == 1 @@ -12981,8 +13002,11 @@ async def test_openapi_health_cancellation_does_not_poison_cache(respx_mock, mon monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") manager = MCPServerManager() server = MCPServer( - server_id="cancelled-cache", name="cancelled-cache", transport=MCPTransport.http, - spec_path="https://93.184.216.34/cancelled-cache.json", auth_type=MCPAuth.none, + server_id="cancelled-cache", + name="cancelled-cache", + transport=MCPTransport.http, + spec_path="https://93.184.216.34/cancelled-cache.json", + auth_type=MCPAuth.none, ) manager.registry = {server.server_id: server} started = asyncio.Event() @@ -13040,11 +13064,18 @@ class _DiscoveryUpstream: return httpx.Response(202) self.requests = (*self.requests, (payload.method, request.headers.get("authorization", ""))) if payload.method == "initialize": - return httpx.Response(200, json={ - "jsonrpc": "2.0", "id": payload.id, - "result": {"protocolVersion": "2025-03-26", "serverInfo": {"name": "discovery", "version": "1"}, - "capabilities": {} if self.outcome == "unsupported" else {"prompts": {}, "resources": {}}}, - }) + return httpx.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload.id, + "result": { + "protocolVersion": "2025-03-26", + "serverInfo": {"name": "discovery", "version": "1"}, + "capabilities": {} if self.outcome == "unsupported" else {"prompts": {}, "resources": {}}, + }, + }, + ) self.entered.set() await self.release.wait() if self.outcome == "failure": @@ -13052,12 +13083,15 @@ class _DiscoveryUpstream: if self.outcome == "cancelled": raise asyncio.CancelledError() if self.outcome == "rejected": - return httpx.Response(200, json={"jsonrpc": "2.0", "id": payload.id, - "error": {"code": -32601, "message": "Unsupported"}}) + return httpx.Response( + 200, json={"jsonrpc": "2.0", "id": payload.id, "error": {"code": -32601, "message": "Unsupported"}} + ) result: Final = { "prompts/list": {"prompts": [{"name": "example", "description": "original"}]}, "resources/list": {"resources": [{"name": "example", "uri": "test://example", "description": "original"}]}, - "resources/templates/list": {"resourceTemplates": [{"name": "example", "uriTemplate": "test://{name}", "description": "original"}]}, + "resources/templates/list": { + "resourceTemplates": [{"name": "example", "uriTemplate": "test://{name}", "description": "original"}] + }, "tools/list": {"tools": []}, }[payload.method] return httpx.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": result}) @@ -13068,7 +13102,9 @@ class _DiscoveryUpstream: def _discovery_server() -> MCPServer: - return MCPServer(server_id="discovery", name="discovery", url="https://discovery.example/mcp", transport=MCPTransport.http) + return MCPServer( + server_id="discovery", name="discovery", url="https://discovery.example/mcp", transport=MCPTransport.http + ) @pytest.mark.asyncio @@ -13079,8 +13115,11 @@ async def test_discovery_cache_reuses_raw_results_and_expires(kind: str) -> None clock: Final = _DiscoveryClock() manager: Final = MCPServerManager(discovery_clock=clock) upstream: Final = _DiscoveryUpstream() - operation: Final = {"prompts": manager.get_prompts_from_server, "resources": manager.get_resources_from_server, - "templates": manager.get_resource_templates_from_server}[kind] + operation: Final = { + "prompts": manager.get_prompts_from_server, + "resources": manager.get_resources_from_server, + "templates": manager.get_resource_templates_from_server, + }[kind] server: Final = _discovery_server() with respx.mock(base_url="https://discovery.example") as router: router.route().mock(side_effect=upstream.respond) @@ -13109,8 +13148,11 @@ async def test_discovery_cache_empty_results_and_failures(kind: str, outcome: st manager: Final = MCPServerManager() upstream: Final = _DiscoveryUpstream() upstream.outcome = outcome - operation: Final = {"prompts": manager.get_prompts_from_server, "resources": manager.get_resources_from_server, - "templates": manager.get_resource_templates_from_server}[kind] + operation: Final = { + "prompts": manager.get_prompts_from_server, + "resources": manager.get_resources_from_server, + "templates": manager.get_resource_templates_from_server, + }[kind] with respx.mock(base_url="https://discovery.example") as router: router.route().mock(side_effect=upstream.respond) assert await operation(_discovery_server(), None) == [] @@ -13137,9 +13179,20 @@ async def test_discovery_cache_isolates_forwarded_credentials_and_shares_static_ assert len(await manager.get_prompts_from_server(server, user)) == 1 assert upstream.initializes == 1 for credential in ("first-secret", "second-secret", "first-secret"): - assert len(await manager.get_prompts_from_server(server, first_user, extra_headers={"Authorization": credential})) == 1 + assert ( + len( + await manager.get_prompts_from_server( + server, first_user, extra_headers={"Authorization": credential} + ) + ) + == 1 + ) assert upstream.initializes == 3 - assert {auth for method, auth in upstream.requests if method == "prompts/list"} == {"", "first-secret", "second-secret"} + assert {auth for method, auth in upstream.requests if method == "prompts/list"} == { + "", + "first-secret", + "second-secret", + } @pytest.mark.asyncio @@ -13151,7 +13204,9 @@ async def test_discovery_cache_coalesces_and_survives_waiter_cancellation() -> N upstream.release.clear() with respx.mock(base_url="https://discovery.example") as router: router.route().mock(side_effect=upstream.respond) - tasks: Final = tuple(asyncio.create_task(manager.get_prompts_from_server(_discovery_server(), None)) for _ in range(10)) + tasks: Final = tuple( + asyncio.create_task(manager.get_prompts_from_server(_discovery_server(), None)) for _ in range(10) + ) await asyncio.wait_for(upstream.entered.wait(), timeout=5) tasks[0].cancel() with pytest.raises(asyncio.CancelledError): @@ -13200,7 +13255,9 @@ async def test_discovery_cache_can_be_disabled(monkeypatch: pytest.MonkeyPatch) assert upstream.initializes == 2 -@pytest.mark.parametrize("value,expected", (("invalid", 60.0), ("nan", 60.0), ("inf", 60.0), ("-1", 60.0), ("12.5", 12.5))) +@pytest.mark.parametrize( + "value,expected", (("invalid", 60.0), ("nan", 60.0), ("inf", 60.0), ("-1", 60.0), ("12.5", 12.5)) +) def test_discovery_cache_ttl_validation(value: str, expected: float, monkeypatch: pytest.MonkeyPatch) -> None: from litellm.proxy._experimental.mcp_server.mcp_server_manager import _mcp_discovery_cache_ttl @@ -13318,9 +13375,15 @@ async def test_discovery_cache_tracks_resolved_credentials_across_workers() -> N source: Final = CredentialSource() managers: Final = (MCPServerManager(cred_provider=source), MCPServerManager(cred_provider=source)) server: Final = MCPServer( - server_id="discovery", name="discovery", url="https://discovery.example/mcp", transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, oauth2_flow="authorization_code", client_id="discovery-client", - authorization_url="https://discovery.example/authorize", token_url="https://discovery.example/token", + server_id="discovery", + name="discovery", + url="https://discovery.example/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + client_id="discovery-client", + authorization_url="https://discovery.example/authorize", + token_url="https://discovery.example/token", ) user: Final = UserAPIKeyAuth(user_id="same-user", api_key="same-key") upstream: Final = _DiscoveryUpstream() @@ -13339,11 +13402,15 @@ async def test_discovery_cache_tracks_resolved_credentials_across_workers() -> N with respx.mock(base_url="https://discovery.example") as router: router.route().mock(side_effect=respond) for manager in managers: - assert [item.name for item in await manager.get_prompts_from_server(server, user)] == ["discovery-account-a"] + assert [item.name for item in await manager.get_prompts_from_server(server, user)] == [ + "discovery-account-a" + ] assert upstream.initializes == 2 source.token = "token-b" for manager in managers: - assert [item.name for item in await manager.get_prompts_from_server(server, user)] == ["discovery-account-b"] + assert [item.name for item in await manager.get_prompts_from_server(server, user)] == [ + "discovery-account-b" + ] assert upstream.initializes == 4 source.token = None for manager in managers: @@ -13370,9 +13437,15 @@ async def test_discovery_resolves_stored_oauth_for_the_requesting_user() -> None store: Final = TokenStore() manager: Final = MCPServerManager(per_user_oauth_token_store=store) server: Final = MCPServer( - server_id="discovery", name="discovery", url="https://discovery.example/mcp", transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, oauth2_flow="authorization_code", client_id="discovery-client", - authorization_url="https://discovery.example/authorize", token_url="https://discovery.example/token", + server_id="discovery", + name="discovery", + url="https://discovery.example/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + client_id="discovery-client", + authorization_url="https://discovery.example/authorize", + token_url="https://discovery.example/token", ) user: Final = UserAPIKeyAuth(user_id="requesting-user") upstream: Final = _DiscoveryUpstream() diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index fad137af66b..61e11e72859 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -19,6 +19,7 @@ from litellm.proxy.hooks.proxy_track_cost_callback import ( run_spend_event, ) from litellm.proxy.route_llm_request import ProxyModelNotFoundError +from litellm.proxy.utils import ProxyUpdateSpend from litellm.proxy.spend_tracking.spend_event import SpendEventDecodeError, build_spend_event, decode_spend_event from litellm.proxy.spend_tracking.spend_event_producer import SpendEventProducer, UnixAddress from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload @@ -1885,9 +1886,15 @@ async def test_track_cost_callback_keeps_guardrail_cost_on_cache_hit(): } with ( - patch("litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock) as mock_increment, # test-quality-ok: the callback imports this from proxy_server inside its body, so there is no injection seam - patch("litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock), # test-quality-ok: same function-body import, no injection seam - patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging, # test-quality-ok: same function-body import, no injection seam + patch( # test-quality-ok: the callback imports this from proxy_server inside its body, so there is no injection seam + "litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock + ) as mock_increment, + patch( # test-quality-ok: same function-body import, no injection seam + "litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock + ), + patch( # test-quality-ok: same function-body import, no injection seam + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as mock_proxy_logging, ): mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() @@ -1912,14 +1919,15 @@ async def test_track_cost_callback_keeps_guardrail_cost_on_cache_hit(): ("allm_passthrough_route", True), ("aretrieve_batch", True), ("acompletion", False), - ("call_mcp_tool", False), + ("call_mcp_tool", True), (None, False), ], ) def test_should_track_cost_callback_pass_through_without_owner(call_type, expected): """Regression for LIT-3782: unauthenticated pass-through requests (auth=false) carry no key/user/team/end-user, yet must still be tracked so they land in - LiteLLM_SpendLogs. Other call types with no owner stay untracked. + LiteLLM_SpendLogs. Explicit MCP passthrough calls require the same handling. + Other call types with no owner stay untracked. aretrieve_batch is included for the same reason: CheckBatchCost's synthetic logging_obj for a completed managed batch only ever carries @@ -1939,10 +1947,26 @@ def test_should_track_cost_callback_pass_through_without_owner(call_type, expect ) +def test_should_track_cost_callback_respects_disabled_spend_updates(monkeypatch): + monkeypatch.setattr(ProxyUpdateSpend, "disable_spend_updates", staticmethod(lambda: True)) + + assert ( + _should_track_cost_callback( + user_api_key="key", + user_id="user", + team_id="team", + end_user_id="end-user", + call_type="call_mcp_tool", + ) + is False + ) + + @pytest.mark.parametrize( "call_type, expect_spend_log", [ ("pass_through_endpoint", True), + ("call_mcp_tool", True), ("aretrieve_batch", True), ("acompletion", False), (None, False), @@ -1953,8 +1977,8 @@ async def test_track_cost_callback_logs_unauthenticated_pass_through_request(cal """Regression for LIT-3782: a pass-through request with auth=false reaches the cost callback with no key/user/team/end-user. Before the fix the spend-log write was skipped and the request never appeared in request/usage logs. It - must now be written for pass-through call types while other unauthenticated - calls remain skipped. + must now be written for pass-through and MCP tool call types while other + unauthenticated calls remain skipped. aretrieve_batch is included because CheckBatchCost's completed-batch cost event reaches this same callback with no attributable key/user/team when