From 9437ce11feb8c7106ea6dc32670558801894d1ab Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Thu, 25 Jun 2026 16:24:50 -0700 Subject: [PATCH] feat(mcp): route the preemptive 401 existence check through the v2 resolver The discovery-phase 401 no longer calls v1's _get_user_oauth_extra_headers_from_db to decide whether a migrated server has a token; it asks the v2 resolver via a new has_user_oauth_token manager method (to_server_spec + to_subject + resolve_credentials, Ok means a token exists). With this, every authorization_code resolution runs through the v2 resolver: the call_tool egress, the listing connection, and the discovery challenge. Delegate servers short-circuit before the check (the client completes PKCE with the upstream). The challenge itself still emits the RFC 8414 authorization_uri form; the format unification stays a follow-up. --- .../mcp_server/mcp_server_manager.py | 16 ++++++ .../outbound_credentials/resolver.py | 10 ++++ .../proxy/_experimental/mcp_server/server.py | 12 ++--- .../outbound_credentials/test_resolver.py | 25 ++++++++++ .../mcp_server/test_mcp_server_manager.py | 49 ++++++++++++++++++- 5 files changed, 104 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 397fd63757e..d1bd6e2ea84 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -3680,6 +3680,22 @@ class MCPServerManager: return mcp_server + async def has_user_oauth_token( + self, server: MCPServer, user_api_key_auth: Optional[UserAPIKeyAuth] + ) -> bool: + """Whether the v2 resolver can produce a per-user token for this server right now. + + This is the preemptive 401's existence check, routed through the same resolver that drives + the egress so every authorization_code resolution (egress and the discovery challenge) runs + through v2. Returns False for a server the resolver does not own (a None spec). + """ + spec = to_server_spec(server) + if spec is None: + return False + return await self._cred_provider.has_user_token( + to_subject(user_api_key_auth, None), spec + ) + async def _resolve_oauth2_headers_for_tool_call( self, mcp_server: MCPServer, diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py index 9163bfdedaa..e6a9a90fdd2 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py @@ -99,6 +99,16 @@ class UpstreamCredentialProvider: return _not_implemented(AuthSpecKind.aws_sigv4) assert_never(server.config) + async def has_user_token(self, subject: Subject, server: ServerSpec) -> bool: + """Whether a usable per-user token exists for this server (the preemptive 401's check). + + Reads from the same per-user store as the ``authorization_code`` arm, so the discovery + challenge and the egress agree on whether the user is authorized. Returns a typed ``bool`` + (no ``httpx.Auth``), unlike ``resolve_credentials``. A non-per-user mode has no token in the + store, so it reads as False without a per-mode branch here. + """ + return await self._authz_token(subject, server) is not None + def _api_key(self, config: ApiKeyConfig) -> Result[httpx.Auth, CredError]: match config.key_source: case SharedKey() as source: diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 61376067f23..e5e5b709121 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -3420,12 +3420,6 @@ if MCP_AVAILABLE: # If no stored token exists, fail fast with 401 so clients can # kick off PKCE/interactive OAuth flow immediately. if server.needs_user_oauth_token: - stored_oauth_headers = await _get_user_oauth_extra_headers_from_db( - server=server, - user_api_key_auth=user_api_key_auth, - ) - if stored_oauth_headers: - continue if getattr(server, "delegate_auth_to_upstream", False) is True: # Delegate-auth servers run upstream PKCE: challenge with # the proxied resource_metadata (RFC 9728), not the @@ -3440,6 +3434,12 @@ if MCP_AVAILABLE: detail="Unauthorized", headers={"www-authenticate": www_authenticate}, ) + # The v2 resolver owns the existence check, so every authorization_code + # resolution (egress and this discovery challenge) runs through it. + if await global_mcp_server_manager.has_user_oauth_token( + server, user_api_key_auth + ): + continue request = StarletteRequest(scope) base_url = get_request_base_url(request) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py index 89e4397c3bc..73e9a52b937 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py @@ -166,6 +166,31 @@ async def test_authorization_code_isolates_by_subject(): assert isinstance(bob, Error) and bob.error.tag == "unauthorized" +@pytest.mark.asyncio +async def test_has_user_token_reflects_the_stored_token(): + present = UpstreamCredentialProvider( + oauth_token_store=_FakeTokenStore( + {("alice", "s"): OAuthToken(access_token="at")} + ) + ) + absent = UpstreamCredentialProvider(oauth_token_store=_FakeTokenStore({})) + spec = _spec(AuthorizationCodeConfig()) + subject = Subject(tenant_id="", subject_id="alice") + assert await present.has_user_token(subject, spec) is True + assert await absent.has_user_token(subject, spec) is False + + +@pytest.mark.asyncio +async def test_has_user_token_false_for_a_non_per_user_mode(): + # A none-mode server has no per-user token to check. + provider = UpstreamCredentialProvider() + spec = _spec(NoneConfig()) + assert ( + await provider.has_user_token(Subject(tenant_id="", subject_id="a"), spec) + is False + ) + + _STUBBED = [ ("api_key_byok", ApiKeyConfig(key_source=Byok())), ("passthrough", PassthroughConfig()), 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 97d379cb8b1..e44a345609c 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 @@ -2275,7 +2275,7 @@ class TestMCPServerManager: @pytest.mark.asyncio async def test_resolve_oauth2_headers_swallows_lookup_exception(self): - """Returns supplied headers (None) when the stored-token lookup raises.""" + """Returns supplied headers (None) when the v1 stored-token lookup raises (delegate path).""" from litellm.proxy._types import UserAPIKeyAuth manager = MCPServerManager() @@ -2284,6 +2284,7 @@ class TestMCPServerManager: name="oauth-srv", transport=MCPTransport.http, auth_type=MCPAuth.oauth2, + delegate_auth_to_upstream=True, # non-migrated, so it reaches the v1 lookup ) user_auth = UserAPIKeyAuth(api_key="sk-test", user_id="alice") @@ -2296,6 +2297,51 @@ class TestMCPServerManager: ) assert result is None + @pytest.mark.asyncio + async def test_has_user_oauth_token_delegates_to_provider(self): + """has_user_oauth_token maps the server and delegates the verdict to the v2 resolver.""" + from litellm.proxy._types import UserAPIKeyAuth + + for verdict in (True, False): + + class _Provider: + async def has_user_token(self, subject, spec): + return verdict + + manager = MCPServerManager(cred_provider=_Provider()) + server = MCPServer( + server_id="s", + name="n", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + ) + user_auth = UserAPIKeyAuth(api_key="sk", user_id="alice") + assert await manager.has_user_oauth_token(server, user_auth) is verdict + + @pytest.mark.asyncio + async def test_has_user_oauth_token_short_circuits_for_unmigrated_server(self): + """A server the resolver does not own (None spec, e.g. delegate) is False without a call.""" + from litellm.proxy._types import UserAPIKeyAuth + + calls: list = [] + + class _Provider: + async def has_user_token(self, subject, spec): + calls.append(spec) + return True + + manager = MCPServerManager(cred_provider=_Provider()) + server = MCPServer( + server_id="s", + name="n", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + delegate_auth_to_upstream=True, + ) + user_auth = UserAPIKeyAuth(api_key="sk", user_id="alice") + assert await manager.has_user_oauth_token(server, user_auth) is False + assert calls == [] # short-circuited on the None spec, never hit the resolver + @pytest.mark.asyncio async def test_resolve_oauth2_headers_no_user_id(self): """Skip lookup entirely when user_api_key_auth has no user_id.""" @@ -2384,7 +2430,6 @@ class TestMCPServerManager: # Unprefixed resolution resolved_server_unpref = manager._get_mcp_server_from_tool_name("create_zap") - print(resolved_server_unpref) assert resolved_server_unpref is not None assert resolved_server_unpref.server_id == server.server_id