fix(mcp): bind discovery caching to credentials and active callers

litellm/experimental_mcp_client/Readme.md

@@ -2,10 +2,5 @@
 
 LiteLLM MCP Client is a client that allows you to use MCP tools with LiteLLM.
 
-## Gateway discovery caching
 
-The MCP gateway caches each upstream server's prompt, resource, and resource-template lists for 60 seconds per worker. Set `LITELLM_MCP_DISCOVERY_CACHE_TTL` to a nonnegative number of seconds to change the lifetime, or `0` to disable caching. Invalid values use the 60-second default
 
-Discovery results may remain unchanged until that lifetime expires. Server configuration updates invalidate the affected server's entries. Concurrent requests for the same list share one upstream fetch. Each list cache holds at most 1,024 entries per worker
-
-User-dependent upstream authentication uses separate cache entries. Gateway access checks still run for every request. Successful empty lists and unsupported capabilities are cached; failed requests retain the existing empty-list response and are retried on the next request

litellm/experimental_mcp_client/client.py

@@ -4,6 +4,8 @@ LiteLLM Proxy uses this MCP Client to connnect to other MCP servers.
 
 import asyncio
 import base64
+import hashlib
+import json
 import os
 from collections.abc import Awaitable, Callable, Generator
 from contextlib import AbstractAsyncContextManager
@@ -343,6 +345,22 @@ class MCPClient:
         if auth_value:
             self.update_auth_value(auth_value)
 
+    async def discovery_auth_fingerprint(self) -> str:
+        request: Final = httpx.Request("POST", self.server_url or "http://localhost/", headers=self._get_auth_headers())
+        if self._resolved_auth is None:
+            return self._hash_discovery_auth(request)
+        flow: Final = self._resolved_auth.async_auth_flow(request)
+        try:
+            authenticated: Final = await flow.__anext__()
+            return self._hash_discovery_auth(authenticated)
+        finally:
+            await flow.aclose()
+
+    @staticmethod
+    def _hash_discovery_auth(request: httpx.Request) -> str:
+        material: Final = json.dumps((str(request.url), tuple(sorted(request.headers.multi_items()))))
+        return hashlib.sha256(material.encode()).hexdigest()
+
     def _create_transport_context(
         self,
     ) -> tuple[_TransportContext, httpx.AsyncClient | None]:

litellm/proxy/_experimental/mcp_server/mcp_server_manager.py

@@ -1681,6 +1681,7 @@ def _record_mcp_guardrail_evaluations(
 
 _DiscoveryItem = TypeVar("_DiscoveryItem", bound=BaseModel)
 _DiscoveryKey: TypeAlias = tuple[str, str | None]
+_DISCOVERY_CACHE_LIMIT: Final = 1024
 
 
 @dataclass(frozen=True, slots=True)
@@ -1695,6 +1696,7 @@ class _DiscoveryCache(Generic[_DiscoveryItem]):
         self._clock = clock
         self._entries: Mapping[_DiscoveryKey, _DiscoveryEntry[_DiscoveryItem]] = MappingProxyType({})
         self._pending: Mapping[_DiscoveryKey, asyncio.Task[list[_DiscoveryItem]]] = MappingProxyType({})
+        self._waiters: Mapping[asyncio.Task[list[_DiscoveryItem]], int] = MappingProxyType({})
 
     def invalidate(self, server_id: str) -> None:
         self._entries = MappingProxyType({key: entry for key, entry in self._entries.items() if key[0] != server_id})
@@ -1715,11 +1717,34 @@ class _DiscoveryCache(Generic[_DiscoveryItem]):
             return tuple(item.model_copy(deep=True) for item in entry.items)
         pending: Final = self._pending.get(key)
         if pending is not None:
-            return tuple(item.model_copy(deep=True) for item in await asyncio.shield(pending))
+            return await self._await_fetch(key, pending)
+        if len(self._pending) >= _DISCOVERY_CACHE_LIMIT:
+            return tuple(await fetch())
         task: Final = asyncio.create_task(self._fetch(key, fetch))
         self._pending = MappingProxyType({**self._pending, key: task})
         task.add_done_callback(self._observe_completion)
-        return tuple(item.model_copy(deep=True) for item in await asyncio.shield(task))
+        return await self._await_fetch(key, task)
+
+    async def _await_fetch(
+        self, key: _DiscoveryKey, task: asyncio.Task[list[_DiscoveryItem]]
+    ) -> tuple[_DiscoveryItem, ...]:
+        self._waiters = MappingProxyType({**self._waiters, task: self._waiters.get(task, 0) + 1})
+        try:
+            return tuple(item.model_copy(deep=True) for item in await asyncio.shield(task))
+        finally:
+            remaining: Final = self._waiters[task] - 1
+            if remaining:
+                self._waiters = MappingProxyType({**self._waiters, task: remaining})
+            else:
+                self._waiters = MappingProxyType(
+                    {pending: count for pending, count in self._waiters.items() if pending is not task}
+                )
+                if self._pending.get(key) is task:
+                    self._pending = MappingProxyType(
+                        {entry_key: pending for entry_key, pending in self._pending.items() if entry_key != key}
+                    )
+                if not task.done():
+                    task.cancel()
 
     async def _fetch(
         self, key: _DiscoveryKey, fetch: Callable[[], Awaitable[list[_DiscoveryItem]]]
@@ -1735,7 +1760,7 @@ class _DiscoveryCache(Generic[_DiscoveryItem]):
                     {
                         entry_key: entry
                         for entry_key, entry in (
-                            *live_entries[-1023:],
+                            *live_entries[-(_DISCOVERY_CACHE_LIMIT - 1) :],
                             (
                                 key,
                                 _DiscoveryEntry(now + self._ttl, tuple(item.model_copy(deep=True) for item in items)),
@@ -4484,6 +4509,7 @@ class MCPServerManager:
         extra_headers: dict[str, str] | None,
         stdio_env: dict[str, str] | None,
         subject_token: str | None,
+        credential_fingerprint: str | None = None,
     ) -> _DiscoveryKey:
         per_user: Final = (
             server.requires_per_user_auth
@@ -4499,7 +4525,9 @@ class MCPServerManager:
             else None
         )
         material: Final = json.dumps(
-            (identity, mcp_auth_header, extra_headers, stdio_env, subject_token), sort_keys=True, separators=(",", ":")
+            (identity, mcp_auth_header, extra_headers, stdio_env, subject_token, credential_fingerprint),
+            sort_keys=True,
+            separators=(",", ":"),
         )
         return server.server_id, hashlib.sha256(material.encode()).hexdigest()
 
@@ -4524,18 +4552,20 @@ class MCPServerManager:
             )
             stdio_env: Final = self._build_stdio_env(server, raw_headers)
             subject_token: Final = self._obo_subject_token(server, raw_headers, user_api_key_auth)
+            client: Final = await self._create_mcp_client(
+                server=server,
+                mcp_auth_header=mcp_auth_header,
+                extra_headers=headers,
+                stdio_env=stdio_env,
+                subject_token=subject_token,
+                user_api_key_auth=user_api_key_auth,
+            )
+            credential_fingerprint: Final = await client.discovery_auth_fingerprint()
             key: Final = self._discovery_key(
-                server, user_api_key_auth, mcp_auth_header, headers, stdio_env, subject_token
+                server, user_api_key_auth, mcp_auth_header, headers, stdio_env, subject_token, credential_fingerprint
             )
 
             async def fetch() -> list[Prompt]:
-                client: Final = await self._create_mcp_client(
-                    server=server,
-                    mcp_auth_header=mcp_auth_header,
-                    extra_headers=headers,
-                    stdio_env=stdio_env,
-                    subject_token=subject_token,
-                )
                 return await client.list_prompts(raise_on_error=True)
 
             items: Final = await self._prompt_discovery_cache.get(key, fetch)
@@ -4565,18 +4595,20 @@ class MCPServerManager:
             )
             stdio_env: Final = self._build_stdio_env(server, raw_headers)
             subject_token: Final = self._obo_subject_token(server, raw_headers, user_api_key_auth)
+            client: Final = await self._create_mcp_client(
+                server=server,
+                mcp_auth_header=mcp_auth_header,
+                extra_headers=headers,
+                stdio_env=stdio_env,
+                subject_token=subject_token,
+                user_api_key_auth=user_api_key_auth,
+            )
+            credential_fingerprint: Final = await client.discovery_auth_fingerprint()
             key: Final = self._discovery_key(
-                server, user_api_key_auth, mcp_auth_header, headers, stdio_env, subject_token
+                server, user_api_key_auth, mcp_auth_header, headers, stdio_env, subject_token, credential_fingerprint
             )
 
             async def fetch() -> list[Resource]:
-                client: Final = await self._create_mcp_client(
-                    server=server,
-                    mcp_auth_header=mcp_auth_header,
-                    extra_headers=headers,
-                    stdio_env=stdio_env,
-                    subject_token=subject_token,
-                )
                 return await client.list_resources(raise_on_error=True)
 
             items: Final = await self._resource_discovery_cache.get(key, fetch)
@@ -4606,18 +4638,20 @@ class MCPServerManager:
             )
             stdio_env: Final = self._build_stdio_env(server, raw_headers)
             subject_token: Final = self._obo_subject_token(server, raw_headers, user_api_key_auth)
+            client: Final = await self._create_mcp_client(
+                server=server,
+                mcp_auth_header=mcp_auth_header,
+                extra_headers=headers,
+                stdio_env=stdio_env,
+                subject_token=subject_token,
+                user_api_key_auth=user_api_key_auth,
+            )
+            credential_fingerprint: Final = await client.discovery_auth_fingerprint()
             key: Final = self._discovery_key(
-                server, user_api_key_auth, mcp_auth_header, headers, stdio_env, subject_token
+                server, user_api_key_auth, mcp_auth_header, headers, stdio_env, subject_token, credential_fingerprint
             )
 
             async def fetch() -> list[ResourceTemplate]:
-                client: Final = await self._create_mcp_client(
-                    server=server,
-                    mcp_auth_header=mcp_auth_header,
-                    extra_headers=headers,
-                    stdio_env=stdio_env,
-                    subject_token=subject_token,
-                )
                 return await client.list_resource_templates(raise_on_error=True)
 
             items: Final = await self._template_discovery_cache.get(key, fetch)
@@ -6112,6 +6146,7 @@ class MCPServerManager:
         failure is logged, never raised, because the DB write already succeeded and the TTL remains
         the backstop.
         """
+        self._invalidate_discovery_lists(server_id)
         try:
             await self._per_user_oauth_token_store.invalidate(user_id, server_id)
         except Exception as exc:  # noqa: BLE001 - cache drop is best-effort; TTL is the backstop

tests/test_litellm/experimental_mcp_client/test_mcp_client.py

@@ -1912,3 +1912,25 @@ def test_client_import_before_proxy_credentials_succeeds_in_fresh_process():
     )
     assert result.returncode == 0, result.stderr
     assert result.stdout.strip() == "MCPServerManager"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("resolved", (False, True))
+async def test_discovery_auth_fingerprint_tracks_effective_credentials(resolved: bool) -> None:
+    from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import StaticHeaderAuth
+
+    def client(token: str) -> MCPClient:
+        return MCPClient(
+            server_url="https://example.com/mcp",
+            auth_type=MCPAuth.api_key,
+            auth_value=None if resolved else token,
+            resolved_auth=StaticHeaderAuth(token) if resolved else None,
+        )
+
+    original: Final = await client("private-original-credential").discovery_auth_fingerprint()
+    repeated: Final = await client("private-original-credential").discovery_auth_fingerprint()
+    replaced: Final = await client("private-replaced-credential").discovery_auth_fingerprint()
+    assert original == repeated
+    assert original != replaced
+    assert len(original) == 64
+    assert "private-original-credential" not in original

tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py

@@ -3708,6 +3708,7 @@ class TestMCPServerManager:
         mock_prompt = Prompt(name="hello", description="Say hi")
         mock_client = AsyncMock()
         mock_client.list_prompts = AsyncMock(return_value=[mock_prompt])
+        mock_client.discovery_auth_fingerprint = AsyncMock(return_value="test-credential-hash")
 
         with patch.object(
             manager,
@@ -3779,6 +3780,7 @@ class TestMCPServerManager:
         mock_client = AsyncMock()
         mock_resources = [Resource(name="file", uri="https://example.com/file")]
         mock_client.list_resources = AsyncMock(return_value=mock_resources)
+        mock_client.discovery_auth_fingerprint = AsyncMock(return_value="test-credential-hash")
         prefixed_resources = [Resource(name="alias-server-file", uri="https://example.com/file")]
 
         with (
@@ -3826,6 +3828,7 @@ class TestMCPServerManager:
             )
         ]
         mock_client.list_resource_templates = AsyncMock(return_value=mock_templates)
+        mock_client.discovery_auth_fingerprint = AsyncMock(return_value="test-credential-hash")
         expected_templates = [
             ResourceTemplate(
                 name="template",
@@ -3855,6 +3858,7 @@ class TestMCPServerManager:
             extra_headers=None,
             stdio_env=None,
             subject_token=None,
+            user_api_key_auth=None,
         )
         mock_client.list_resource_templates.assert_awaited_once()
         assert result == expected_templates
@@ -13231,3 +13235,151 @@ async def test_discovery_cache_retries_cancelled_fetches() -> None:
     with pytest.raises(asyncio.CancelledError):
         await cache.get(("server", None), cancelled)
     assert [item.name for item in await cache.get(("server", None), supported)] == ["recovered"]
+
+
+@pytest.mark.asyncio
+async def test_discovery_cache_cancels_fetch_when_last_waiter_leaves() -> None:
+    from litellm.proxy._experimental.mcp_server.mcp_server_manager import _DiscoveryCache
+
+    cache: Final = _DiscoveryCache[Prompt](60, _DiscoveryClock())
+    entered: Final = asyncio.Event()
+    stopped: Final = asyncio.Event()
+    release: Final = asyncio.Event()
+
+    async def fetch() -> list[Prompt]:
+        entered.set()
+        try:
+            await release.wait()
+            return [Prompt(name="result")]
+        finally:
+            stopped.set()
+
+    tasks: Final = tuple(asyncio.create_task(cache.get(("server", None), fetch)) for _ in range(3))
+    await asyncio.wait_for(entered.wait(), timeout=5)
+    for task in tasks:
+        task.cancel()
+    outcomes: Final = await asyncio.gather(*tasks, return_exceptions=True)
+    assert all(isinstance(outcome, asyncio.CancelledError) for outcome in outcomes)
+    try:
+        await asyncio.wait_for(stopped.wait(), timeout=1)
+    finally:
+        release.set()
+
+
+@pytest.mark.asyncio
+async def test_discovery_cache_bounds_detached_fetches_without_dropping_results() -> None:
+    from litellm.proxy._experimental.mcp_server.mcp_server_manager import _DiscoveryCache
+
+    cache: Final = _DiscoveryCache[Prompt](60, _DiscoveryClock())
+    entered: Final[asyncio.Queue[None]] = asyncio.Queue()
+    release: Final = asyncio.Event()
+
+    async def blocked() -> list[Prompt]:
+        await entered.put(None)
+        await release.wait()
+        return [Prompt(name="blocked")]
+
+    tasks: Final = tuple(asyncio.create_task(cache.get((str(index), None), blocked)) for index in range(1024))
+    try:
+        for _ in tasks:
+            await asyncio.wait_for(entered.get(), timeout=5)
+        active_tasks: Final = frozenset(asyncio.all_tasks())
+
+        async def overflow() -> list[Prompt]:
+            assert frozenset(asyncio.all_tasks()) <= active_tasks
+            return [Prompt(name="overflow")]
+
+        result: Final = await cache.get(("overflow", None), overflow)
+        assert [item.name for item in result] == ["overflow"]
+    finally:
+        release.set()
+        outcomes: Final = await asyncio.gather(*tasks)
+        assert all(result[0].name == "blocked" for result in outcomes)
+
+
+@pytest.mark.asyncio
+async def test_discovery_cache_tracks_resolved_credentials_across_workers() -> None:
+    import respx
+    from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import StaticHeaderAuth
+    from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import UpstreamCredentialProvider
+    from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok, Result
+    from litellm.proxy._experimental.mcp_server.outbound_credentials.types import CredError, ServerSpec, Subject
+
+    class CredentialSource(UpstreamCredentialProvider):
+        def __init__(self) -> None:
+            super().__init__()
+            self.token: str | None = "token-a"
+
+        async def resolve_credentials(self, subject: Subject, server: ServerSpec) -> Result[httpx.Auth, CredError]:
+            if self.token is None:
+                return Error(CredError.of_unauthorized("Credential revoked"))
+            return Ok(StaticHeaderAuth("Bearer " + self.token))
+
+    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",
+    )
+    user: Final = UserAPIKeyAuth(user_id="same-user", api_key="same-key")
+    upstream: Final = _DiscoveryUpstream()
+
+    async def respond(request: httpx.Request) -> httpx.Response:
+        response: Final = await upstream.respond(request)
+        if '"prompts/list"' not in request.content.decode():
+            return response
+        from mcp.types import JSONRPCMessage, JSONRPCRequest
+
+        payload: Final = JSONRPCMessage.model_validate_json(request.content).root
+        assert isinstance(payload, JSONRPCRequest)
+        name: Final = {"Bearer token-a": "account-a", "Bearer token-b": "account-b"}[request.headers["authorization"]]
+        return httpx.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": {"prompts": [{"name": name}]}})
+
+    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 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 upstream.initializes == 4
+        source.token = None
+        for manager in managers:
+            assert await manager.get_prompts_from_server(server, user) == []
+        assert upstream.initializes == 4
+
+
+@pytest.mark.asyncio
+async def test_discovery_resolves_stored_oauth_for_the_requesting_user() -> None:
+    import respx
+    from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import OAuthToken
+
+    class TokenStore:
+        def __init__(self) -> None:
+            self.calls: tuple[tuple[str, str], ...] = ()
+
+        async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None:
+            self.calls = (*self.calls, (user_id, server_id))
+            return OAuthToken(access_token="stored-token")
+
+        async def invalidate(self, user_id: str, server_id: str) -> None:
+            return 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",
+    )
+    user: Final = UserAPIKeyAuth(user_id="requesting-user")
+    upstream: Final = _DiscoveryUpstream()
+    with respx.mock(base_url="https://discovery.example") as router:
+        router.route().mock(side_effect=upstream.respond)
+        assert len(await manager.get_prompts_from_server(server, user)) == 1
+        assert len(await manager.get_prompts_from_server(server, user)) == 1
+    assert store.calls == (("requesting-user", "discovery"), ("requesting-user", "discovery"))
+    assert upstream.initializes == 1
+    assert ("prompts/list", "Bearer stored-token") in upstream.requests