diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 81f1d677d87..1a90435203a 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -251,7 +251,6 @@ _USER_ENV_VARS_CACHE_TTL: Final = 60 # seconds _USER_ENV_VARS_CACHE_MAX_SIZE: Final = 4096 # cap to prevent unbounded growth _ListedToolsByCaller: TypeAlias = Mapping[str | None, Mapping[str, MCPTool]] -_NO_LISTED_TOOLS: Final[_ListedToolsByCaller] = MappingProxyType({}) _LISTED_TOOLS_CALLERS_PER_SERVER: Final = 256 @@ -4617,7 +4616,7 @@ class MCPServerManager: ) -> None: identity: Final = self._listed_tools_identity(server, caller) listing: Final = MappingProxyType({tool.name: tool for tool in tools}) - existing: Final = self._listed_tools_by_server_id.get(server.server_id, _NO_LISTED_TOOLS) + existing: Final[_ListedToolsByCaller] = self._listed_tools_by_server_id.get(server.server_id, {}) 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) @@ -5505,7 +5504,7 @@ class MCPServerManager: def get_listed_tool(self, server: MCPServer, name: str, caller: ListedToolsCaller | None = None) -> MCPTool | None: identity: Final = self._listed_tools_identity(server, caller) - listed: Final = self._listed_tools_by_server_id.get(server.server_id, _NO_LISTED_TOOLS).get(identity) + listed: Final = self._listed_tools_by_server_id.get(server.server_id, {}).get(identity) if not listed: return None tool: Final = listed.get(name) or listed.get(strip_known_server_prefix(name, server)) 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 246db7ae7f1..e40cc943b45 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 @@ -6913,8 +6913,6 @@ class TestMCPServerManager: @pytest.mark.asyncio async def test_user_oauth_refresh_keeps_listed_tools(self): - """Tool definitions are server-wide, so one user's re-auth must not blank the metadata other - callers' tool calls hand to pre-call guardrails.""" manager = MCPServerManager() server = MCPServer(server_id="srv", name="srv", transport=MCPTransport.http, url="http://srv") manager._create_prefixed_tools([MCPTool(name="echo", description="shared", inputSchema={})], server) @@ -6997,7 +6995,6 @@ class TestMCPServerManager: ], ) def test_upstream_identity_inputs_keep_listed_tools_apart(self, server_kwargs, caller_a, caller_b): - """Whatever reaches upstream and can change its catalog must also split the listed-tool cache.""" manager = MCPServerManager() server = MCPServer( **{"server_id": "srv", "name": "srv", "transport": MCPTransport.http, "url": "http://srv", **server_kwargs} @@ -7037,8 +7034,6 @@ class TestMCPServerManager: ], ) def test_jwt_signer_makes_a_shared_server_list_per_caller(self, signer, static_headers, shared): - """MCPJWTSigner hands upstream a JWT naming the caller on an otherwise shared ``auth_type: none`` - server, so the upstream may tailor the catalog and the cache must not hand one caller another's.""" manager = MCPServerManager() server = MCPServer( server_id="srv", name="srv", transport=MCPTransport.http, url="http://srv", static_headers=static_headers @@ -7062,7 +7057,6 @@ class TestMCPServerManager: @pytest.mark.asyncio async def test_call_tool_hands_hooks_the_catalog_the_same_forwarded_headers_listed(self): - """Interleaved callers on a forwarded-header server: the hook must see the caller's own catalog.""" manager = MCPServerManager() server = MCPServer( server_id="catalog",