diff --git a/litellm/proxy/mcp_server/mcp_server_manager.py b/litellm/proxy/mcp_server/mcp_server_manager.py index eacba25257e..c5b3c100623 100644 --- a/litellm/proxy/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/mcp_server/mcp_server_manager.py @@ -1906,6 +1906,17 @@ def _listing_task_was_cancelled(error: asyncio.CancelledError) -> bool: return error.args != (SESSION_CLEANUP_TIMEOUT_MESSAGE,) +def is_anonymous_mcp_caller(user_api_key_auth: UserAPIKeyAuth | None) -> bool: + return not ( + user_api_key_auth + and ( + getattr(user_api_key_auth, "user_id", None) + or getattr(user_api_key_auth, "user_role", None) + or getattr(user_api_key_auth, "api_key", None) + ) + ) + + class MCPServerManager: _STDIO_ENV_TEMPLATE_PATTERN = re.compile(r"^\$\{(X-[^}]+)\}$") @@ -3775,14 +3786,7 @@ class MCPServerManager: ) ) - is_anonymous: Final = not ( - user_api_key_auth - and ( - getattr(user_api_key_auth, "user_id", None) - or getattr(user_api_key_auth, "user_role", None) - or getattr(user_api_key_auth, "api_key", None) - ) - ) + is_anonymous: Final = is_anonymous_mcp_caller(user_api_key_auth) if is_anonymous: anonymous_server_ids: Final = [ server.server_id diff --git a/litellm/proxy/mcp_server/server.py b/litellm/proxy/mcp_server/server.py index bbbbd09e527..302f48b11f1 100644 --- a/litellm/proxy/mcp_server/server.py +++ b/litellm/proxy/mcp_server/server.py @@ -593,6 +593,7 @@ if MCP_AVAILABLE: from litellm.proxy.mcp_server.mcp_server_manager import ( MCPServerManager, global_mcp_server_manager, + is_anonymous_mcp_caller, is_legacy_upstream_delegate, ) @@ -1885,6 +1886,24 @@ if MCP_AVAILABLE: ) return acting.model_copy(update={"object_permission": updated_op, "mcp_toolset_id": toolset_id}) + def _narrow_allowed_server_ids( + allowed_servers: Sequence[MCPServer], + toolset_allowed_server_ids: Set[str] | None, + ) -> frozenset[str]: + allowed_ids: Final = frozenset(srv.server_id for srv in allowed_servers) + if toolset_allowed_server_ids is None: + return allowed_ids + return allowed_ids & frozenset(toolset_allowed_server_ids) + + def _challenge_server_ids( + user_api_key_auth: UserAPIKeyAuth | None, + allowed_servers: Sequence[MCPServer], + toolset_allowed_server_ids: Set[str] | None, + ) -> frozenset[str] | None: + if is_anonymous_mcp_caller(user_api_key_auth): + return None if toolset_allowed_server_ids is None else frozenset(toolset_allowed_server_ids) + return _narrow_allowed_server_ids(allowed_servers, toolset_allowed_server_ids) + _apply_toolset_scope: Final = apply_toolset_scope async def _toolset_server_ids(toolset_id: str) -> set[str]: @@ -1902,7 +1921,7 @@ if MCP_AVAILABLE: mcp_server_auth_headers: dict[str, dict[str, str]] | None, user_api_key_auth: UserAPIKeyAuth | None, client_ip: str | None, - allowed_server_ids: set[str] | None = None, + allowed_server_ids: Set[str] | None = None, raw_headers: Mapping[str, str] | None = None, ) -> None: """Fail fast with HTTP 401 for MCP servers that need user auth but @@ -1910,11 +1929,11 @@ if MCP_AVAILABLE: (points clients at the gateway AS metadata) and pass-through OAuth (points clients at the upstream resource-metadata via our well-known). - ``allowed_server_ids`` may be passed by callers that have already - narrowed the authorized server set (e.g. toolset scoping); servers - not in that set are skipped so a client targeting a toolset that - excludes a passthrough server is not pushed into an OAuth flow for - a server it will be 403'd on immediately after authentication. + ``allowed_server_ids`` is the set of servers the caller's key may + reach, narrowed further by any active toolset scope; anonymous + cold-start callers are not narrowed by a key allowlist. Servers not in + that set are skipped so a client is not pushed into an OAuth flow + for a server it will be 403'd on immediately after authentication. """ for server_name in mcp_servers or []: server = operations.global_mcp_server_manager.get_mcp_server_by_name(server_name, client_ip=client_ip) @@ -2212,19 +2231,18 @@ if MCP_AVAILABLE: async def _check_passthrough_upstream_auth( scope: Scope, user_api_key_auth: UserAPIKeyAuth | None, - mcp_servers: list[str] | None, - client_ip: str | None, + allowed_servers: Sequence[MCPServer], *, oauth2_headers: Mapping[str, str] | None, mcp_server_auth_headers: dict[str, dict[str, str]] | None, raw_headers: Mapping[str, str] | None, requested_servers: Sequence[str] | None = None, + client_ip: str | None = None, ) -> None: """Probe pass-through and legacy delegate 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 - trigger an upstream probe against a server their key is not permitted for. + The caller passes the already-resolved allowed-server set so a user cannot + trigger a probe against a server their key is not permitted for. The MCP SDK commits HTTP 200 headers before invoking handlers, so a 401 can only be returned before that point. This function raises HTTPException(401) @@ -2236,19 +2254,11 @@ if MCP_AVAILABLE: if not oauth2_headers and not mcp_server_auth_headers: return - # Use the authorized server set, not the raw user-supplied names, so that - # a caller cannot force a probe to a server their key is not allowed to use. - allowed_servers: Final = await operations._get_allowed_mcp_servers( - user_api_key_auth=user_api_key_auth, - mcp_servers=mcp_servers, - client_ip=client_ip, - ) # A legacy delegate is probed only when the single requested name resolves to it the way admission # resolves it, so neither a multi-server connect nor an access group can be failed by one rejected # token; its challenge echoes the requested name, matching the tokenless challenge. - requested_scope: Final = requested_servers if requested_servers is not None else mcp_servers single_requested_name: Final = ( - requested_scope[0] if requested_scope is not None and len(requested_scope) == 1 else None + requested_servers[0] if requested_servers is not None and len(requested_servers) == 1 else None ) requested_server: Final = ( operations.global_mcp_server_manager.get_mcp_server_by_name(single_requested_name, client_ip=client_ip) @@ -2378,14 +2388,14 @@ if MCP_AVAILABLE: await _check_passthrough_upstream_auth( scope, context.user_api_key_auth, - authorized_names, - context.client_ip, + allowed, oauth2_headers=context.oauth2_headers, mcp_server_auth_headers={key: dict(value) for key, value in context.mcp_server_auth_headers.items()} if context.mcp_server_auth_headers is not None else None, raw_headers=context.raw_headers, requested_servers=context.mcp_servers, + client_ip=context.client_ip, ) return None @@ -2398,10 +2408,15 @@ if MCP_AVAILABLE: mcp_server_auth_headers: Mapping[str, Mapping[str, str]] | None, user_api_key_auth: UserAPIKeyAuth | None, client_ip: str | None, - allowed_server_ids: Set[str] | None, + toolset_allowed_server_ids: Set[str] | None, raw_headers: Mapping[str, str] | None, ) -> None: if not defer_upstream_probes and not targets_disabled_capability: + allowed_servers: Final = await operations._get_allowed_mcp_servers( + user_api_key_auth=user_api_key_auth, + mcp_servers=list(mcp_servers) if mcp_servers is not None else None, + client_ip=client_ip, + ) await _raise_preemptive_401_for_unauthenticated_servers( scope=scope, mcp_servers=list(mcp_servers) if mcp_servers is not None else None, @@ -2413,7 +2428,9 @@ if MCP_AVAILABLE: ), user_api_key_auth=user_api_key_auth, client_ip=client_ip, - allowed_server_ids=set(allowed_server_ids) if allowed_server_ids is not None else None, + allowed_server_ids=_challenge_server_ids( + user_api_key_auth, allowed_servers, toolset_allowed_server_ids + ), raw_headers=raw_headers, ) # Pre-flight auth check for pass-through servers. Must run after @@ -2422,8 +2439,7 @@ if MCP_AVAILABLE: await _check_passthrough_upstream_auth( scope, user_api_key_auth, - list(mcp_servers) if mcp_servers is not None else None, - client_ip, + allowed_servers, oauth2_headers=oauth2_headers, mcp_server_auth_headers=( {name: dict(headers) for name, headers in mcp_server_auth_headers.items()} @@ -2431,6 +2447,8 @@ if MCP_AVAILABLE: else None ), raw_headers=raw_headers, + requested_servers=mcp_servers, + client_ip=client_ip, ) async def _targets_disabled_capability(peek: _LazyPeekedBody | None) -> bool: @@ -2543,9 +2561,10 @@ if MCP_AVAILABLE: # https://datatracker.ietf.org/doc/html/rfc9728#name-www-authenticate-response # Must run after toolset scoping so the challenge set is derived - # from the fully-authorized server set: a passthrough server that - # the active toolset excludes should not trigger an OAuth flow - # for a server the caller will be 403'd on after authentication. + # from the fully-authorized server set (the key's allowed servers, + # narrowed by any active toolset scope): a server the key cannot + # reach should not trigger an OAuth flow for a server the caller + # will be 403'd on after authentication. await _preflight_streamable_http_upstream_auth( defer_upstream_probes=defer_upstream_probes, targets_disabled_capability=targets_disabled_capability, @@ -2555,7 +2574,7 @@ if MCP_AVAILABLE: mcp_server_auth_headers=mcp_server_auth_headers, user_api_key_auth=user_api_key_auth, client_ip=_client_ip, - allowed_server_ids=toolset_allowed_server_ids, + toolset_allowed_server_ids=toolset_allowed_server_ids, raw_headers=raw_headers, ) @@ -2904,9 +2923,15 @@ if MCP_AVAILABLE: # https://datatracker.ietf.org/doc/html/rfc9728#name-www-authenticate-response # Must run after toolset scoping so the challenge set is derived - # from the fully-authorized server set: a passthrough server that - # the active toolset excludes should not trigger an OAuth flow - # for a server the caller will be 403'd on after authentication. + # from the fully-authorized server set (the key's allowed servers, + # narrowed by any active toolset scope): a server the key cannot + # reach should not trigger an OAuth flow for a server the caller + # will be 403'd on after authentication. + allowed_servers: Final = await operations._get_allowed_mcp_servers( + user_api_key_auth=user_api_key_auth, + mcp_servers=mcp_servers, + client_ip=_sse_client_ip, + ) await _raise_preemptive_401_for_unauthenticated_servers( scope=scope, mcp_servers=mcp_servers, @@ -2914,7 +2939,9 @@ if MCP_AVAILABLE: mcp_server_auth_headers=mcp_server_auth_headers, user_api_key_auth=user_api_key_auth, client_ip=_sse_client_ip, - allowed_server_ids=toolset_allowed_server_ids, + allowed_server_ids=_challenge_server_ids( + user_api_key_auth, allowed_servers, toolset_allowed_server_ids + ), raw_headers=raw_headers, ) @@ -2927,11 +2954,12 @@ if MCP_AVAILABLE: await _check_passthrough_upstream_auth( scope, user_api_key_auth, - mcp_servers, - _sse_client_ip, + allowed_servers, oauth2_headers=oauth2_headers, mcp_server_auth_headers=mcp_server_auth_headers, raw_headers=raw_headers, + requested_servers=mcp_servers, + client_ip=_sse_client_ip, ) set_auth_context( user_api_key_auth=user_api_key_auth, diff --git a/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py b/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py index 621a7b65005..457f829a2d1 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py +++ b/tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py @@ -15,6 +15,7 @@ from fastapi import HTTPException from mcp import ReadResourceResult, Resource from mcp.server.context import ServerRequestContext from mcp.server.models import InitializationOptions +from mcp.server.streamable_http_manager import StreamableHTTPSessionManager from mcp.types import ( INVALID_REQUEST, METHOD_NOT_FOUND, @@ -7312,10 +7313,6 @@ async def test_legacy_delegate_bare_token_is_not_probed_upstream(): # test-qual with ( _patch_delegate_resolver(server, "delegate_test"), - patch( - "litellm.proxy.mcp_server.operations._get_allowed_mcp_servers", - new=AsyncMock(return_value=[server]), - ), patch( "litellm.proxy.mcp_server.server._probe_upstream_auth", new=AsyncMock(return_value=(401, 'Bearer realm="upstream", error="invalid_token"')), @@ -7324,8 +7321,7 @@ async def test_legacy_delegate_bare_token_is_not_probed_upstream(): # test-qual await _check_passthrough_upstream_auth( scope=scope, user_api_key_auth=UserAPIKeyAuth(), - mcp_servers=["delegate_test"], - client_ip=None, + allowed_servers=[server], oauth2_headers=None, mcp_server_auth_headers=None, raw_headers=None, @@ -7356,18 +7352,14 @@ async def test_legacy_delegate_dual_credentials_probe_only_the_upstream_token() with ( _patch_delegate_resolver(server, "delegate_test"), - patch( # test-quality-ok: isolate authorized-server resolution so this test targets the preflight boundary - "litellm.proxy.mcp_server.operations._get_allowed_mcp_servers", - new=AsyncMock(return_value=[server]), - ), respx.mock as upstream, ): probe: Final = upstream.post(server.url).respond(200) await _check_passthrough_upstream_auth( scope=scope, user_api_key_auth=UserAPIKeyAuth(user_id="admitted-user"), - mcp_servers=["delegate_test"], - client_ip=None, + allowed_servers=[server], + requested_servers=["delegate_test"], oauth2_headers={"Authorization": "Bearer upstream-token"}, mcp_server_auth_headers=None, raw_headers={"x-litellm-api-key": "sk-litellm-proxy-key", "Authorization": "Bearer upstream-token"}, @@ -7379,7 +7371,9 @@ async def test_legacy_delegate_dual_credentials_probe_only_the_upstream_token() assert all("sk-litellm-proxy-key" not in value for value in sent_headers.values()) -def _expired_delegate_token_preflight(mcp_servers: Sequence[str]) -> Coroutine[object, object, None]: +def _expired_delegate_token_preflight( + mcp_servers: Sequence[str], allowed_servers: Sequence[MCPServer] +) -> Coroutine[object, object, None]: from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.mcp_server.server import ( _check_passthrough_upstream_auth, @@ -7388,7 +7382,8 @@ def _expired_delegate_token_preflight(mcp_servers: Sequence[str]) -> Coroutine[o return _check_passthrough_upstream_auth( scope=_delegate_scope([(b"authorization", b"Bearer expired-upstream-token")]), user_api_key_auth=UserAPIKeyAuth().model_copy(update={"mcp_forwards_caller_authorization": True}), - mcp_servers=list(mcp_servers), + allowed_servers=allowed_servers, + requested_servers=mcp_servers, client_ip=None, oauth2_headers={"Authorization": "Bearer expired-upstream-token"}, mcp_server_auth_headers=None, @@ -7397,28 +7392,24 @@ def _expired_delegate_token_preflight(mcp_servers: Sequence[str]) -> Coroutine[o @contextlib.contextmanager -def _rejecting_delegate_upstream() -> Iterator[None]: +def _rejecting_delegate_upstream() -> Iterator[MCPServer]: server: Final = _delegate_auth_mcp_server() with ( _patch_delegate_resolver(server, "delegate_test"), - patch( # test-quality-ok: isolate authorized-server resolution so this test targets the preflight boundary - "litellm.proxy.mcp_server.operations._get_allowed_mcp_servers", - new=AsyncMock(return_value=[server]), - ), patch( # test-quality-ok: stands in for the upstream rejecting an expired token "litellm.proxy.mcp_server.server._probe_upstream_auth", new=AsyncMock(return_value=(401, 'Bearer error="invalid_token"')), ), ): - yield + yield server @pytest.mark.asyncio async def test_legacy_delegate_expired_upstream_token_gets_reauthorization_challenge() -> None: """An expired upstream token on a single-server route gets the RFC 9728 re-authorization challenge for the requested name, so the client signs in again instead of seeing an empty tool list.""" - with _rejecting_delegate_upstream(), pytest.raises(HTTPException) as exc_info: - await _expired_delegate_token_preflight(["delegate_test"]) + with _rejecting_delegate_upstream() as server, pytest.raises(HTTPException) as exc_info: + await _expired_delegate_token_preflight(["delegate_test"], [server]) assert exc_info.value.status_code == 401 challenge: Final = exc_info.value.headers["www-authenticate"] @@ -7431,8 +7422,8 @@ async def test_legacy_delegate_expired_upstream_token_gets_reauthorization_chall @pytest.mark.asyncio async def test_legacy_delegate_multi_server_scope_is_not_failed_by_one_upstream_token() -> None: """A multi-server scope is not probed, so one rejected upstream token cannot fail the whole connect.""" - with _rejecting_delegate_upstream(): - assert await _expired_delegate_token_preflight(["delegate_test", "other_server"]) is None + with _rejecting_delegate_upstream() as server: + assert await _expired_delegate_token_preflight(["delegate_test", "other_server"], [server]) is None @pytest.mark.asyncio @@ -7447,10 +7438,6 @@ async def test_legacy_delegate_probe_checks_the_callers_token_not_a_static_autho ) with ( _patch_delegate_resolver(server, "delegate_test"), - patch( # test-quality-ok: isolate authorized-server resolution so this test targets the preflight boundary - "litellm.proxy.mcp_server.operations._get_allowed_mcp_servers", - new=AsyncMock(return_value=[server]), - ), respx.mock as upstream, ): upstream.post(server.url).mock( @@ -7459,7 +7446,7 @@ async def test_legacy_delegate_probe_checks_the_callers_token_not_a_static_autho ) ) outcome: Final = await asyncio.gather( - _expired_delegate_token_preflight(["delegate_test"]), return_exceptions=True + _expired_delegate_token_preflight(["delegate_test"], [server]), return_exceptions=True ) assert getattr(outcome[0], "status_code", None) == 401 @@ -7557,10 +7544,6 @@ async def test_oauth_passthrough_preflight_preserves_status_contract(probe_statu ) with ( - patch( # test-quality-ok: isolate authorized-server resolution so this test exercises the preflight contract - "litellm.proxy.mcp_server.operations._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.mcp_server.server._probe_upstream_auth", new=AsyncMock(return_value=(probe_status, None)), @@ -7570,11 +7553,13 @@ async def test_oauth_passthrough_preflight_preserves_status_contract(probe_statu result = await _check_passthrough_upstream_auth( scope=scope, user_api_key_auth=UserAPIKeyAuth(user_id="admitted-user"), - mcp_servers=["passthrough_server"], - client_ip=None, + allowed_servers=[server], oauth2_headers={"Authorization": "Bearer upstream-token"}, mcp_server_auth_headers=None, - raw_headers={"x-litellm-api-key": "sk-litellm-proxy-key", "Authorization": "Bearer upstream-token"}, + raw_headers={ + "x-litellm-api-key": "sk-litellm-proxy-key", + "Authorization": "Bearer upstream-token", + }, ) assert result is None else: @@ -7582,11 +7567,13 @@ async def test_oauth_passthrough_preflight_preserves_status_contract(probe_statu await _check_passthrough_upstream_auth( scope=scope, user_api_key_auth=UserAPIKeyAuth(user_id="admitted-user"), - mcp_servers=["passthrough_server"], - client_ip=None, + allowed_servers=[server], oauth2_headers={"Authorization": "Bearer upstream-token"}, mcp_server_auth_headers=None, - raw_headers={"x-litellm-api-key": "sk-litellm-proxy-key", "Authorization": "Bearer upstream-token"}, + raw_headers={ + "x-litellm-api-key": "sk-litellm-proxy-key", + "Authorization": "Bearer upstream-token", + }, ) assert exc_info.value.status_code == expected_status if expected_status == 401: @@ -7609,10 +7596,6 @@ async def test_delegate_tokenless_request_not_probed(): with ( _patch_delegate_resolver(server, "delegate_test"), - patch( - "litellm.proxy.mcp_server.operations._get_allowed_mcp_servers", - new=AsyncMock(return_value=[server]), - ), patch( "litellm.proxy.mcp_server.server._probe_upstream_auth", new=AsyncMock(return_value=(401, None)), @@ -7621,8 +7604,7 @@ async def test_delegate_tokenless_request_not_probed(): await _check_passthrough_upstream_auth( scope=scope, user_api_key_auth=UserAPIKeyAuth(), - mcp_servers=["delegate_test"], - client_ip=None, + allowed_servers=[server], oauth2_headers=None, mcp_server_auth_headers=None, raw_headers=None, @@ -7645,10 +7627,6 @@ async def test_delegate_preflight_skipped_on_multi_server_routes(): with ( _patch_delegate_resolver(servers[0], "delegate_test", "other_server"), - patch( - "litellm.proxy.mcp_server.operations._get_allowed_mcp_servers", - new=AsyncMock(return_value=servers), - ), patch( "litellm.proxy.mcp_server.server._probe_upstream_auth", new=AsyncMock(return_value=(401, None)), @@ -7657,8 +7635,7 @@ async def test_delegate_preflight_skipped_on_multi_server_routes(): await _check_passthrough_upstream_auth( scope=scope, user_api_key_auth=UserAPIKeyAuth(), - mcp_servers=["delegate_test", "other_server"], - client_ip=None, + allowed_servers=servers, oauth2_headers=None, mcp_server_auth_headers=None, raw_headers=None, @@ -7691,10 +7668,6 @@ async def test_bare_authorization_never_probes_passthrough_servers(): with ( _patch_delegate_resolver(passthrough_server, "pt_server"), - patch( - "litellm.proxy.mcp_server.operations._get_allowed_mcp_servers", - new=AsyncMock(return_value=[passthrough_server]), - ), patch( "litellm.proxy.mcp_server.server._probe_upstream_auth", new=AsyncMock(return_value=(401, None)), @@ -7703,8 +7676,7 @@ async def test_bare_authorization_never_probes_passthrough_servers(): await _check_passthrough_upstream_auth( scope=scope, user_api_key_auth=UserAPIKeyAuth(), - mcp_servers=["pt_server"], - client_ip=None, + allowed_servers=[passthrough_server], oauth2_headers=None, mcp_server_auth_headers=None, raw_headers=None, @@ -7740,10 +7712,6 @@ async def test_delegate_not_probed_when_named_only_via_server_id(): with ( _patch_delegate_resolver(server, "delegate_test"), - patch( - "litellm.proxy.mcp_server.operations._get_allowed_mcp_servers", - new=AsyncMock(return_value=[server]), - ), patch( "litellm.proxy.mcp_server.server._probe_upstream_auth", new=AsyncMock(return_value=(401, None)), @@ -7752,8 +7720,7 @@ async def test_delegate_not_probed_when_named_only_via_server_id(): await _check_passthrough_upstream_auth( scope=scope, user_api_key_auth=UserAPIKeyAuth(user_id="u1", api_key="hashed-sk"), - mcp_servers=["delegate-secret-id"], - client_ip=None, + allowed_servers=[server], oauth2_headers=None, mcp_server_auth_headers=None, raw_headers=None, @@ -7780,10 +7747,6 @@ async def test_delegate_probe_not_fanned_out_to_access_group_members() -> None: ) with ( _patch_delegate_resolver(rejecting, "delegate_test"), - patch( # test-quality-ok: isolate authorized-server resolution so this test targets the preflight boundary - "litellm.proxy.mcp_server.operations._get_allowed_mcp_servers", - new=AsyncMock(return_value=[rejecting, accepting]), - ), respx.mock as upstream, ): upstream.post(rejecting.url).respond(401, headers={"www-authenticate": 'Bearer error="invalid_token"'}) @@ -7794,7 +7757,8 @@ async def test_delegate_probe_not_fanned_out_to_access_group_members() -> None: [(b"x-litellm-api-key", b"sk-litellm-proxy-key"), (b"authorization", b"Bearer upstream-token")] ), user_api_key_auth=UserAPIKeyAuth(user_id="admitted-user"), - mcp_servers=["prod_tools_group"], + allowed_servers=[rejecting, accepting], + requested_servers=["prod_tools_group"], client_ip=None, oauth2_headers={"Authorization": "Bearer upstream-token"}, mcp_server_auth_headers=None, @@ -10891,6 +10855,395 @@ class TestPreemptive401ModeAware: await self._run(delegate, self.LITELLM_KEY_HEADERS, has_stored_token=False) +class TestPreemptive401GatedByKeyAllowlist: + """The preemptive OAuth 401 challenge must be gated on the servers the + caller's key is allowed to reach, not only on toolset scoping. A key whose + object_permission.mcp_servers excludes an oauth2 server must fall through + to the existing 403 denial instead of being pushed into an OAuth flow it + can never complete.""" + + @staticmethod + def _auth() -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key="sk-litellm-virtual-key") + + @staticmethod + def _initialize_call(alias: str) -> tuple[dict[str, object], AsyncMock]: + body: Final = ( + b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":' + b'{"protocolVersion":"2024-11-05","capabilities":{},' + b'"clientInfo":{"name":"test","version":"0"}}}' + ) + receive: Final = AsyncMock(return_value={"type": "http.request", "body": body, "more_body": False}) + return {"type": "http", "method": "POST", "path": f"/mcp/{alias}", "headers": []}, receive + + @contextlib.contextmanager + def _patches( + self, + server: MCPServer, + *, + allowed_servers: list[MCPServer], + handle_target: StreamableHTTPSessionManager | None, + toolset_auth: UserAPIKeyAuth | None = None, + request_auth: UserAPIKeyAuth | None = None, + mcp_servers: list[str] | None = None, + oauth2_headers: dict[str, str] | None = None, + raw_headers: dict[str, str] | None = None, + ) -> Iterator[tuple[AsyncMock, AsyncMock]]: + def _resolve_allowed( + user_api_key_auth: UserAPIKeyAuth | None = None, + mcp_servers: list[str] | None = None, + client_ip: str | None = None, + ) -> list[MCPServer]: + if user_api_key_auth is None or user_api_key_auth.object_permission is None: + return allowed_servers + scoped_ids: Final = user_api_key_auth.object_permission.mcp_servers + if not scoped_ids: + return allowed_servers + return [srv for srv in allowed_servers if srv.server_id in scoped_ids] + + registry: Final = mcp_operations.global_mcp_server_manager.registry + registry[server.server_id] = server + try: + with contextlib.ExitStack() as stack: + stack.enter_context( + patch( # test-quality-ok: admission parse is unrelated to the challenge gate; the handler reads it as a module global + "litellm.proxy.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=( + request_auth if request_auth is not None else self._auth(), + None, + mcp_servers if mcp_servers is not None else [server.alias], + None, + oauth2_headers, + raw_headers, + ), + ) + ) + allowed_servers_lookup: Final = stack.enter_context( + patch( # test-quality-ok: allowed-server resolution hits the key/toolset DB; the test must control the allowlist + "litellm.proxy.mcp_server.operations._get_allowed_mcp_servers", + new_callable=AsyncMock, + side_effect=_resolve_allowed, + ) + ) + stack.enter_context( + patch( # test-quality-ok: module-level startup flag gates the transport; there is no injectable seam + "litellm.proxy.mcp_server.server._SESSION_MANAGERS_INITIALIZED", + True, + ) + ) + stack.enter_context( + patch.object( # test-quality-ok: no stored per-user token is the scenario; the real lookup hits the token cache/DB + mcp_operations.global_mcp_server_manager, + "has_user_oauth_token", + new_callable=AsyncMock, + return_value=False, + ) + ) + stack.enter_context( + patch.object( # test-quality-ok: the real cache warm-up calls the upstream server over the network + mcp_operations.global_mcp_server_manager, + "_ensure_upstream_initialize_instructions_cached", + new_callable=AsyncMock, + ) + ) + if toolset_auth is not None: + stack.enter_context( + patch( # test-quality-ok: toolset resolution hits the DB; the test controls the scoped auth directly + "litellm.proxy.mcp_server.server.apply_toolset_scope", + new_callable=AsyncMock, + return_value=toolset_auth, + ) + ) + stack.enter_context( + patch( # test-quality-ok: toolset resolution hits the DB; the test controls the scoped server ids + "litellm.proxy.mcp_server.server._toolset_server_ids", + new_callable=AsyncMock, + return_value={"id-other"}, + ) + ) + handle_request: Final = AsyncMock() + if handle_target is not None: + stack.enter_context( + patch.object( # test-quality-ok: session dispatch is a module-level global; rejected requests must not reach it + handle_target, + "handle_request", + new=handle_request, + ) + ) + yield allowed_servers_lookup, handle_request + finally: + registry.pop(server.server_id, None) + + @pytest.mark.asyncio + async def test_key_not_allowed_gets_403_not_401_streamable(self): + """Denied key on an oauth2 server must reach the scoped-access denial + (403) instead of an OAuth challenge (401) it can never satisfy.""" + from litellm.proxy.mcp_server import server as server_module + + server: Final = _make_oauth2_server("restricted") + scope, receive = self._initialize_call("restricted") + + with ( + self._patches(server, allowed_servers=[], handle_target=server_module.session_manager_stateful) as ( + _, + handle_request, + ), + pytest.raises(HTTPException) as exc, + ): + await server_module.handle_streamable_http_mcp(scope, receive, AsyncMock()) + + assert exc.value.status_code == 403 + handle_request.assert_not_awaited() + + @pytest.mark.asyncio + async def test_key_allowed_still_challenges_streamable(self): + """Key allowlist includes the oauth2 server with no stored token: + the 401 challenge still fires, so the gate is not over-broad.""" + from litellm.proxy.mcp_server import server as server_module + + server: Final = _make_oauth2_server("restricted") + scope, receive = self._initialize_call("restricted") + + with ( + self._patches(server, allowed_servers=[server], handle_target=server_module.session_manager_stateful) as ( + _, + handle_request, + ), + pytest.raises(HTTPException) as exc, + ): + await server_module.handle_streamable_http_mcp(scope, receive, AsyncMock()) + + assert exc.value.status_code == 401 + assert "www-authenticate" in {k.lower() for k in exc.value.headers} + handle_request.assert_not_awaited() + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("server_kind", "transport"), + [("pt", "streamable"), ("pt", "sse"), ("dl", "streamable"), ("dl", "sse")], + ) + async def test_keyless_cold_start_still_gets_oauth_challenge(self, server_kind: str, transport: str) -> None: + from litellm.proxy.mcp_server import server as server_module + + server: Final = MCPServer( + server_id=f"id-{server_kind}", + name=server_kind, + alias=server_kind, + server_name=server_kind, + url=None, + transport=MCPTransport.http, + auth_type=MCPAuth.none if server_kind == "pt" else MCPAuth.oauth2, + extra_headers=["Authorization"] if server_kind == "pt" else [], + oauth_passthrough=server_kind == "pt", + delegate_auth_to_upstream=server_kind == "dl", + oauth2_flow="authorization_code" if server_kind == "dl" else None, + authorization_url="https://idp.test/authorize" if server_kind == "dl" else None, + token_url="https://idp.test/token" if server_kind == "dl" else None, + mcp_info={"server_name": server_kind}, + ) + body: Final = ( + b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":' + b'{"protocolVersion":"2024-11-05","capabilities":{},' + b'"clientInfo":{"name":"test","version":"0"}}}' + ) + scope: Final = { + "type": "http", + "method": "POST" if transport == "streamable" else "GET", + "path": f"/mcp/{server_kind}", + "scheme": "http", + "server": ("127.0.0.1", 4000), + "root_path": "", + "query_string": b"", + "headers": [ + (b"host", b"127.0.0.1:4000"), + (b"accept", b"application/json, text/event-stream"), + (b"content-type", b"application/json"), + ], + } + receive: Final = AsyncMock(return_value={"type": "http.request", "body": body, "more_body": False}) + registry: Final = mcp_operations.global_mcp_server_manager.registry + expected_challenge: Final = ( + f'Bearer resource_metadata="http://127.0.0.1:4000/.well-known/oauth-protected-resource/mcp/{server_kind}"' + ) + + with ( + patch.dict(registry, {server.server_id: server}), + patch( # test-quality-ok: force the admitted anonymous cold-start auth context without DB or network + "litellm.proxy.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=(UserAPIKeyAuth(), None, [server.alias], None, None, None), + ), + patch( # test-quality-ok: module-level startup flag gates the transport, with no injectable seam + "litellm.proxy.mcp_server.server._SESSION_MANAGERS_INITIALIZED", True + ), + ): + if transport == "streamable": + handle_request: Final = AsyncMock() + with ( + patch.object( + server_module.session_manager_stateful, + "handle_request", + new=handle_request, + ), # test-quality-ok: rejected requests must not dispatch to the session manager + pytest.raises(HTTPException) as exc, + ): + await server_module.handle_streamable_http_mcp(scope, receive, AsyncMock()) + handle_request.assert_not_awaited() + else: + with ( + patch.object(server_module.sse, "connect_sse") as connect_sse, + pytest.raises(HTTPException) as exc, + ): + await server_module.handle_sse_mcp(scope, receive, AsyncMock()) + connect_sse.assert_not_called() + + assert exc.value.status_code == 401 + assert exc.value.headers is not None + assert exc.value.headers["www-authenticate"] == expected_challenge + + @pytest.mark.asyncio + async def test_streamable_probe_reuses_allowed_servers( # test-quality-ok: one lookup per request is the contract + self, monkeypatch: pytest.MonkeyPatch + ): + from litellm.proxy.mcp_server import server as server_module + + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"mcp_advertised_versions": ["2026-07-28"]}) + server: Final = MCPServer( + server_id="id-restricted", + name="restricted", + alias="restricted", + server_name="restricted", + url="https://restricted.test/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + extra_headers=["Authorization"], + oauth_passthrough=True, + ) + body: Final = json.dumps( + { + "jsonrpc": "2.0", + "id": 7, + "method": "tools/call", + "params": { + "name": "confirm", + "arguments": {}, + "_meta": { + "io.modelcontextprotocol/protocolVersion": "2026-07-28", + "io.modelcontextprotocol/clientCapabilities": {}, + }, + }, + } + ).encode() + scope: Final = { + "type": "http", + "method": "POST", + "path": "/mcp", + "scheme": "http", + "server": ("localhost", 8000), + "query_string": b"", + "root_path": "", + "headers": [ + (b"content-type", b"application/json"), + (b"x-litellm-api-key", b"sk-litellm-virtual-key"), + (b"authorization", b"Bearer upstream-token"), + (b"mcp-protocol-version", b"2026-07-28"), + (b"mcp-method", b"tools/call"), + (b"mcp-name", b"confirm"), + (b"accept", b"application/json, text/event-stream"), + ], + } + receive: Final = AsyncMock(return_value={"type": "http.request", "body": body, "more_body": False}) + + with ( + self._patches( + server, + allowed_servers=[server], + handle_target=server_module.session_manager_stateless, + mcp_servers=["restricted", "another"], + oauth2_headers={"Authorization": "Bearer upstream-token"}, + raw_headers={ + "x-litellm-api-key": "sk-litellm-virtual-key", + "Authorization": "Bearer upstream-token", + }, + ) as (allowed_servers_lookup, handle_request), + patch( + "litellm.proxy.mcp_server.server._probe_upstream_auth", + new_callable=AsyncMock, + return_value=(200, None), + ) as probe, + ): + await server_module.handle_streamable_http_mcp(scope, receive, AsyncMock()) + + allowed_servers_lookup.assert_awaited_once() + probe.assert_awaited_once_with(server.url, "Bearer upstream-token") + handle_request.assert_awaited_once() + + @pytest.mark.asyncio + async def test_toolset_excluding_server_gets_403_not_401_streamable(self): + """Key allowlist includes the server but the active toolset scope + excludes it: the caller-visible denial is the scoped-access 403, not + a 401 challenge.""" + from litellm.proxy._types import LiteLLM_ObjectPermissionTable + from litellm.proxy.mcp_server import server as server_module + + server: Final = _make_oauth2_server("restricted") + key_auth: Final = self._auth().model_copy( + update={ + "object_permission": LiteLLM_ObjectPermissionTable( + object_permission_id="key-scope", + mcp_servers=[server.server_id], + ) + } + ) + scoped_auth: Final = key_auth.model_copy( + update={ + "object_permission": LiteLLM_ObjectPermissionTable( + object_permission_id="toolset-scope", + mcp_servers=["id-other"], + ) + } + ) + scope, receive = self._initialize_call("restricted") + token: Final = server_module._mcp_active_toolset_id.set("toolset-1") + try: + with ( + self._patches( + server, + allowed_servers=[server], + handle_target=server_module.session_manager_stateful, + toolset_auth=scoped_auth, + request_auth=key_auth, + ) as (_, handle_request), + pytest.raises(HTTPException) as exc, + ): + await server_module.handle_streamable_http_mcp(scope, receive, AsyncMock()) + finally: + server_module._mcp_active_toolset_id.reset(token) + + assert exc.value.status_code == 403 + assert "www-authenticate" not in {k.lower() for k in exc.value.headers or {}} + handle_request.assert_not_awaited() + + @pytest.mark.asyncio + async def test_key_not_allowed_gets_403_not_401_sse(self): + """Same gate on the SSE transport: a denied key gets 403, not a 401.""" + from litellm.proxy.mcp_server import server as server_module + + server: Final = _make_oauth2_server("restricted") + scope: Final = {"type": "http", "method": "GET", "path": "/mcp/restricted", "headers": []} + + with ( + self._patches(server, allowed_servers=[], handle_target=None), + patch.object(server_module.sse, "connect_sse") as connect_sse, + pytest.raises(HTTPException) as exc, + ): + await server_module.handle_sse_mcp(scope, AsyncMock(), AsyncMock()) + + assert exc.value.status_code == 403 + connect_sse.assert_not_called() + + class TestSingleServerPreflightReachesIdJag: """The connect-time preflight is what turns a credential failure into an HTTP status the client can read. An oauth2_id_jag server has to reach it: its subject comes from the assertion stored at @@ -12265,6 +12618,7 @@ async def test_passthrough_probes_bind_each_token_to_its_authorized_server(allow ) for name in ("first", "second") ] + allowed_servers: Final = targets if allowed else [] mcp_operations.global_mcp_server_manager.registry.update({target.server_id: target for target in targets}) with respx.mock(assert_all_called=False) as upstream: first: Final = upstream.post(targets[0].url).respond(200) @@ -12279,8 +12633,7 @@ async def test_passthrough_probes_bind_each_token_to_its_authorized_server(allow mcp_servers=[target.server_id for target in targets] if allowed else [], ), ), - mcp_servers=["first", "second"], - client_ip=None, + allowed_servers=allowed_servers, oauth2_headers=None, mcp_server_auth_headers={ "first": {"x-other" if header_shape == "non_authorization" else "aUtHoRiZaTiOn": "Bearer first-token"}, diff --git a/tests/unit/proxy/mcp_server/test_mcp_stale_session.py b/tests/unit/proxy/mcp_server/test_mcp_stale_session.py index 88f18b609a4..7e3e701c896 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_stale_session.py +++ b/tests/unit/proxy/mcp_server/test_mcp_stale_session.py @@ -652,6 +652,11 @@ async def test_per_user_oauth_missing_stored_token_returns_preemptive_401(): "litellm.proxy.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name", return_value=oauth_server, ), + patch( # test-quality-ok: transport handler collaborators are module-level globals, no injectable seam + "litellm.proxy.mcp_server.operations._get_allowed_mcp_servers", + new_callable=AsyncMock, + return_value=[oauth_server], + ), patch.object( session_manager_stateless, "handle_request", @@ -740,6 +745,11 @@ async def test_admitted_subject_missing_stored_token_challenged_with_resource_me "litellm.proxy.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name", return_value=oauth_server, ), + patch( # test-quality-ok: transport handler collaborators are module-level globals, no injectable seam + "litellm.proxy.mcp_server.operations._get_allowed_mcp_servers", + new_callable=AsyncMock, + return_value=[oauth_server], + ), patch.object( session_manager_stateless, "handle_request", @@ -1035,6 +1045,11 @@ async def test_per_user_oauth_with_stored_token_skips_preemptive_401(): "litellm.proxy.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name", return_value=oauth_server, ), + patch( # test-quality-ok: transport handler collaborators are module-level globals, no injectable seam + "litellm.proxy.mcp_server.operations._get_allowed_mcp_servers", + new_callable=AsyncMock, + return_value=[oauth_server], + ), patch.object( session_manager_stateless, "handle_request", @@ -1138,6 +1153,11 @@ async def test_handle_streamable_http_mcp_delegated_server_without_token_returns "litellm.proxy.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name", return_value=delegated_server, ), + patch( # test-quality-ok: transport handler collaborators are module-level globals, no injectable seam + "litellm.proxy.mcp_server.operations._get_allowed_mcp_servers", + new_callable=AsyncMock, + return_value=[delegated_server], + ), patch.object( session_manager_stateful, "handle_request", @@ -1226,6 +1246,11 @@ async def test_handle_streamable_http_mcp_token_exchange_without_subject_returns "litellm.proxy.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name", return_value=obo_server, ), + patch( # test-quality-ok: transport handler collaborators are module-level globals, no injectable seam + "litellm.proxy.mcp_server.operations._get_allowed_mcp_servers", + new_callable=AsyncMock, + return_value=[obo_server], + ), patch.object( session_manager_stateful, "handle_request", @@ -1325,6 +1350,11 @@ async def test_handle_streamable_http_mcp_oauth_delegate_without_token_returns_g "litellm.proxy.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name", return_value=od_server, ), + patch( # test-quality-ok: transport handler collaborators are module-level globals, no injectable seam + "litellm.proxy.mcp_server.operations._get_allowed_mcp_servers", + new_callable=AsyncMock, + return_value=[od_server], + ), patch.object( session_manager_stateful, "handle_request", @@ -1582,6 +1612,11 @@ async def test_handle_streamable_http_mcp_true_passthrough_without_token_surface "litellm.proxy.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name", return_value=tp_server, ), + patch( # test-quality-ok: transport handler collaborators are module-level globals, no injectable seam + "litellm.proxy.mcp_server.operations._get_allowed_mcp_servers", + new_callable=AsyncMock, + return_value=[tp_server], + ), patch.object( session_manager_stateful, "handle_request", @@ -1650,6 +1685,11 @@ async def test_handle_streamable_http_mcp_true_passthrough_dcr_bridge_challenges "litellm.proxy.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name", return_value=bridge_server, ), + patch( # test-quality-ok: transport handler collaborators are module-level globals, no injectable seam + "litellm.proxy.mcp_server.operations._get_allowed_mcp_servers", + new_callable=AsyncMock, + return_value=[bridge_server], + ), patch.object( session_manager_stateful, "handle_request", @@ -1827,9 +1867,7 @@ async def test_streamable_http_skips_preflight_for_disabled_methods( # test-qua "tp_server", extra_headers=request_headers, ) - body: Final[bytes] = json.dumps( - {"jsonrpc": "2.0", "id": 1, "method": method, "params": params} - ).encode() + body: Final[bytes] = json.dumps({"jsonrpc": "2.0", "id": 1, "method": method, "params": params}).encode() receive: Final = AsyncMock(return_value={"type": "http.request", "body": body, "more_body": False}) send: Final = AsyncMock() user_auth: Final = UserAPIKeyAuth()