From fdb7da7c6117d5cecce4a1167b5f75f48714a5df Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 12 Sep 2026 19:41:20 +0000 Subject: [PATCH] fix(mcp): bound per-caller listed-tool catalogs to 256 identities per server Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/mcp_server_manager.py | 18 ++++++++--- .../proxy/_experimental/mcp_server/server.py | 2 +- .../guardrail_hooks/agent_365/agent_365.py | 2 +- .../mcp_server/test_mcp_server_manager.py | 30 +++++++++++++++++++ 4 files changed, 46 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index a1200f1784e..d439579a3ab 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -242,6 +242,9 @@ _user_env_vars_cache: Final[dict[tuple[str, str], tuple[dict[str, str], float]]] _USER_ENV_VARS_CACHE_TTL: Final = 60 # seconds _USER_ENV_VARS_CACHE_MAX_SIZE: Final = 4096 # cap to prevent unbounded growth +_NO_LISTED_TOOLS: Final[Mapping[str | None, Mapping[str, MCPTool]]] = MappingProxyType({}) +_LISTED_TOOLS_CALLERS_PER_SERVER: Final = 256 + # Auth types whose upstream OAuth endpoints (protected-resource + authorization-server metadata) the # gateway discovers from the upstream itself: interactive oauth2 and the two client-forwarded modes. # OBO/M2M endpoint discovery is decided separately via _obo_needs_endpoint_discovery. Shared by the @@ -4446,7 +4449,7 @@ class MCPServerManager: # through _create_prefixed_tools — that would add the prefix a second # time producing "test_petstore-test_petstore-getinventory". unprefixed_tools: Final = [ # mutable-ok: returned through the list[MCPTool] listing contract - t.model_copy(update={"name": t.name[len(registry_prefix) :]}) for t in tools + t.model_copy(update=MappingProxyType({"name": t.name[len(registry_prefix) :]})) for t in tools ] self._record_listed_tools(server, unprefixed_tools, user_api_key_auth) return tools if add_prefix else unprefixed_tools @@ -4523,9 +4526,16 @@ class MCPServerManager: ) -> None: identity: Final = self._listed_tools_identity(server, user_api_key_auth) listing: Final = MappingProxyType({tool.name: tool for tool in tools}) - self._listed_tools_by_server_id[server.server_id] = MappingProxyType( - {**self._listed_tools_by_server_id.get(server.server_id, {}), identity: listing} + existing: Final = self._listed_tools_by_server_id.get(server.server_id, _NO_LISTED_TOOLS) + shared: Final = existing.get(None) + callers: Final = tuple((key, value) for key, value in existing.items() if key not in (None, identity)) + evicted: Final = 0 if identity is None else max(len(callers) + 1 - _LISTED_TOOLS_CALLERS_PER_SERVER, 0) + entries: Final = ( + *(() if shared is None else ((None, shared),)), + *callers[evicted:], + (identity, listing), ) + self._listed_tools_by_server_id[server.server_id] = MappingProxyType(dict(entries)) def _discovery_key( self, @@ -5396,7 +5406,7 @@ class MCPServerManager: self, server: MCPServer, name: str, user_api_key_auth: UserAPIKeyAuth | None = None ) -> MCPTool | None: identity: Final = self._listed_tools_identity(server, user_api_key_auth) - listed: Final = self._listed_tools_by_server_id.get(server.server_id, {}).get(identity) + listed: Final = self._listed_tools_by_server_id.get(server.server_id, _NO_LISTED_TOOLS).get(identity) if not listed: return None return listed.get(name) or listed.get(strip_known_server_prefix(name, server)) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index dc191dbf8b8..62e13675a8b 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -4149,7 +4149,7 @@ if MCP_AVAILABLE: if server and ( (server.auth_type == MCPAuth.oauth2_token_exchange and not oauth2_headers) or ( - _get_mcp_servers_in_path(get_route_relative_request_path(scope)) == [server_name] + tuple(_get_mcp_servers_in_path(get_route_relative_request_path(scope)) or ()) == (server_name,) and not agent_365_subject_token_present(oauth2_headers) and agent_365_authorization_servers(server, user_api_key_auth) ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py b/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py index 24f1643d89a..58d8ba5f865 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py +++ b/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py @@ -666,7 +666,7 @@ def entra_assertion(value: object) -> str | None: def agent_365_subject_token_present(oauth2_headers: Mapping[str, str] | None) -> bool: """Whether the request's ``Authorization`` carries an Entra assertion the guardrail can exchange.""" - authorization: Final = (oauth2_headers or {}).get("Authorization", "") + authorization: Final = oauth2_headers.get("Authorization", "") if oauth2_headers else "" if not authorization.lower().startswith("bearer "): return False return entra_assertion(authorization[len("bearer ") :].strip()) is not None 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 4d1490b76a5..a504946a0d8 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 @@ -6665,6 +6665,36 @@ class TestMCPServerManager: for_bob = manager.get_listed_tool(shared, "echo", bob) assert for_bob is not None and for_bob.description == "everyone" + def test_per_caller_listed_tools_evict_oldest_caller_and_keep_shared(self): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import _LISTED_TOOLS_CALLERS_PER_SERVER + + manager = MCPServerManager() + server = MCPServer( + server_id="srv", + name="srv", + transport=MCPTransport.http, + url="http://srv", + auth_type=MCPAuth.oauth2_token_exchange, + ) + manager._create_prefixed_tools([MCPTool(name="read", description="shared", inputSchema={})], server) + callers = [UserAPIKeyAuth(user_id=f"u{i}", api_key=f"k{i}") for i in range(_LISTED_TOOLS_CALLERS_PER_SERVER + 1)] + for caller in callers: + manager._create_prefixed_tools( + [MCPTool(name="read", description=caller.user_id, inputSchema={})], server, user_api_key_auth=caller + ) + manager._create_prefixed_tools( + [MCPTool(name="read", description="u1 again", inputSchema={})], server, user_api_key_auth=callers[1] + ) + + assert manager.get_listed_tool(server, "srv-read", callers[0]) is None + second = manager.get_listed_tool(server, "srv-read", callers[1]) + assert second is not None and second.description == "u1 again" + newest = manager.get_listed_tool(server, "srv-read", callers[-1]) + assert newest is not None and newest.description == callers[-1].user_id + assert len(manager._listed_tools_by_server_id[server.server_id]) == _LISTED_TOOLS_CALLERS_PER_SERVER + 1 + shared = manager.get_listed_tool(server, "srv-read") + assert shared is not None and shared.description == "shared" + @pytest.mark.asyncio @pytest.mark.parametrize("add_prefix", [True, False]) async def test_openapi_listing_records_listed_tools(self, add_prefix):