From 97b7eaad5cd6d723c4849ab43e22ecdbfb9ce2d5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 17 Aug 2026 14:44:01 -0700 Subject: [PATCH] fix(mcp): keep issuer-anchored slots on reload, skip discovery for stamped M2M challenge Registry-swap reconciliation used bool(server.url) while registration uses _requires_oauth_discovery, dropping slots for issuer-anchored servers without a url. The preemptive 401 loop awaited discovery before the stamped client_credentials continue, so a deferred discovery failure could 503 requests whose challenge decision never reads metadata --- .../mcp_server/mcp_server_manager.py | 2 +- .../proxy/_experimental/mcp_server/server.py | 6 ++++ .../mcp_server/test_mcp_server.py | 16 ++++++++++ .../mcp_server/test_mcp_server_manager.py | 30 +++++++++++++++++++ 4 files changed, 53 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index cdaf3f2f206..91e5339dda2 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1691,7 +1691,7 @@ class MCPServerManager: def _reconcile_oauth_discovery_slots_for_servers(self, servers: Iterable[MCPServer]) -> None: """Align retry slots after an atomic registry replacement.""" for server in servers: - should_defer = bool(server.url) and _oauth_endpoints_unresolved(server) + should_defer = _requires_oauth_discovery(server.url, server.issuer_is_anchored, server) has_slot = self._oauth_discovery_slot(server.server_id) is not None if should_defer != has_slot: self._set_oauth_discovery_deferred(server.server_id, should_defer) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index e4ac40734dc..2b7f3e96371 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -3737,6 +3737,12 @@ if MCP_AVAILABLE: # preemptive challenge and let downstream authorization # return 403. continue + if server is not None and server.auth_type == MCPAuth.oauth2 and server.oauth2_flow == "client_credentials": + # Stamped M2M: the challenge decision below never reads discovered + # metadata, so deferred-discovery failures must not 503 this loop. + # Unstamped rows stay on the discover-first path because filling + # authorization_url/token_url can change their inferred flow. + continue if server is not None: server = await global_mcp_server_manager.ensure_oauth_metadata_discovered(server) if server and server.auth_type == MCPAuth.oauth2: 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 0193f2c9152..7b3b8c7a28c 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 @@ -7969,6 +7969,22 @@ class TestPreemptive401ModeAware: assert manager._oauth_discovery_slot(server.server_id) is None assert exc.value.status_code == 401 + @pytest.mark.asyncio + async def test_stamped_m2m_challenge_skips_deferred_discovery(self): + from litellm.proxy._experimental.mcp_server import server as server_module + + manager = server_module.global_mcp_server_manager + server = _make_oauth2_server("stamped_m2m", oauth2_flow="client_credentials") + + with patch.object( + manager, + "ensure_oauth_metadata_discovered", + new=AsyncMock(side_effect=HTTPException(status_code=503, detail="discovery down")), + ) as discovery: + await self._run(server, None, has_stored_token=False) + + discovery.assert_not_awaited() + @pytest.mark.asyncio async def test_gateway_managed_interactive_no_token_challenges_with_x_litellm_api_key(self): """No stored token, key in x-litellm-api-key (oauth2_headers empty): 401.""" 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 fe09303b298..76d6a178b9e 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 @@ -695,6 +695,36 @@ class TestMCPServerManager: assert resolved.token_url == "https://idp.example.com/token" assert manager._oauth_discovery_slot(replacement.server_id) is None + def test_registry_swap_reconcile_keeps_slot_for_issuer_anchored_server_without_url(self): + manager = MCPServerManager() + server = MCPServer( + server_id="anchored-no-url-1", + name="anchored_no_url", + url=None, + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + issuer="https://idp.example.com", + issuer_is_anchored=True, + ) + manager.registry[server.server_id] = server + manager._set_oauth_discovery_deferred(server.server_id, True) + + manager._reconcile_oauth_discovery_slots_for_servers([server]) + + assert manager._oauth_discovery_slot(server.server_id) is not None + + resolved = server.model_copy( + update={ + "authorization_url": "https://idp.example.com/authorize", + "token_url": "https://idp.example.com/token", + } + ) + manager.registry[resolved.server_id] = resolved + manager._reconcile_oauth_discovery_slots_for_servers([resolved]) + + assert manager._oauth_discovery_slot(server.server_id) is None + def _assert_oauth_discovery_state_removed(self, manager, server_id): assert manager._oauth_discovery_slot(server_id) is None