From 25797cbf2a99583f92c7c8bdf0ef12797beec317 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sun, 27 Sep 2026 02:46:18 +0000 Subject: [PATCH] fix(mcp): key the per-caller listed-tool slot by the hashed token Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/mcp_server_manager.py | 2 +- .../mcp_server/test_mcp_server_manager.py | 45 ++++++++++++++++--- 2 files changed, 39 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 074a30eb0db..4c14b3ac1b8 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -4606,7 +4606,7 @@ class MCPServerManager: return None auth: Final = caller.user_api_key_auth identity: Final = ( - (auth.user_id, auth.api_key) if auth is not None and self._discovers_per_caller(server) else None + (auth.user_id, auth.token) if auth is not None and self._discovers_per_caller(server) else None ) forwarded: Final = self._forwarded_header_values(server, caller.raw_headers) header_env: Final = self._build_stdio_env(server, caller.raw_headers) 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 e44c66f6e35..019b93925ea 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 @@ -6931,8 +6931,8 @@ class TestMCPServerManager: url="http://srv", auth_type=MCPAuth.oauth2_token_exchange, ) - alice = UserAPIKeyAuth(user_id="alice", api_key="hashed-alice") - bob = UserAPIKeyAuth(user_id="bob", api_key="hashed-bob") + alice = UserAPIKeyAuth(user_id="alice", token="hashed-alice") + bob = UserAPIKeyAuth(user_id="bob", token="hashed-bob") alice_schema = {"type": "object", "properties": {"path": {"type": "string"}}} bob_schema = {"type": "object", "properties": {"path": {"type": "string"}, "site": {"type": "string"}}} manager._create_prefixed_tools( @@ -6953,7 +6953,7 @@ class TestMCPServerManager: alice_schema, ) assert bob_tool is not None and (bob_tool.description, bob_tool.input_schema) == ("bob view", bob_schema) - carol = ListedToolsCaller(user_api_key_auth=UserAPIKeyAuth(user_id="carol", api_key="k")) + carol = ListedToolsCaller(user_api_key_auth=UserAPIKeyAuth(user_id="carol", token="k")) assert manager.get_listed_tool(server, "srv-read", carol) is None shared = MCPServer(server_id="shared", name="shared", transport=MCPTransport.http, url="http://shared") @@ -6989,11 +6989,11 @@ class TestMCPServerManager: pytest.param( {"auth_type": MCPAuth.oauth2_token_exchange}, ListedToolsCaller( - user_api_key_auth=UserAPIKeyAuth(user_id="team-bot", api_key="hashed-shared"), + user_api_key_auth=UserAPIKeyAuth(user_id="team-bot", token="hashed-shared"), raw_headers={"x-litellm-api-key": "sk-shared", "authorization": "Bearer entra-alice"}, ), ListedToolsCaller( - user_api_key_auth=UserAPIKeyAuth(user_id="team-bot", api_key="hashed-shared"), + user_api_key_auth=UserAPIKeyAuth(user_id="team-bot", token="hashed-shared"), raw_headers={"x-litellm-api-key": "sk-shared", "authorization": "Bearer entra-bob"}, ), id="shared-key-different-obo-subjects", @@ -7106,8 +7106,8 @@ class TestMCPServerManager: server = MCPServer( server_id="srv", name="srv", transport=MCPTransport.http, url="http://srv", static_headers=static_headers ) - alice = ListedToolsCaller(user_api_key_auth=UserAPIKeyAuth(user_id="alice", api_key="hashed-alice")) - bob = ListedToolsCaller(user_api_key_auth=UserAPIKeyAuth(user_id="bob", api_key="hashed-bob")) + alice = ListedToolsCaller(user_api_key_auth=UserAPIKeyAuth(user_id="alice", token="hashed-alice")) + bob = ListedToolsCaller(user_api_key_auth=UserAPIKeyAuth(user_id="bob", token="hashed-bob")) with patch( # test-quality-ok: the signer is a process-wide singleton the manager reads, no injection seam "litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer.get_mcp_jwt_signer", @@ -7123,6 +7123,37 @@ class TestMCPServerManager: else: assert for_bob is None + def test_per_caller_slot_identity_is_the_token_not_the_api_key(self): + """Two callers sharing user_id split on the hashed token the admission validator stamps; + the raw api_key never enters the identity, so a caller carrying only that token lands on + the same slot.""" + from litellm.proxy._types import hash_token + + manager = MCPServerManager() + server = MCPServer(server_id="srv", name="srv", transport=MCPTransport.http, url="http://srv") + alice = ListedToolsCaller( + user_api_key_auth=UserAPIKeyAuth(user_id="same-user", api_key="sk-alpha") + ) + bob = ListedToolsCaller( + user_api_key_auth=UserAPIKeyAuth(user_id="same-user", api_key="sk-beta") + ) + + with patch( # test-quality-ok: the signer is a process-wide singleton the manager reads, no injection seam + "litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer.get_mcp_jwt_signer", + return_value=MagicMock(), + ): + manager._create_prefixed_tools( + [MCPTool(name="turn", description="slot a", inputSchema={})], server, caller=alice + ) + assert manager.get_listed_tool(server, "srv-turn", bob) is None + + same_token = ListedToolsCaller( + user_api_key_auth=UserAPIKeyAuth(user_id="same-user", token=hash_token("sk-alpha")) + ) + listed = manager.get_listed_tool(server, "srv-turn", same_token) + + assert listed is not None and listed.description == "slot a" + @pytest.mark.asyncio async def test_call_tool_hands_hooks_the_catalog_the_same_forwarded_headers_listed(self): manager = MCPServerManager()