scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-02 02:11:58 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

fix(mcp): key the per-caller listed-tool slot by the hashed token

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yucheng 2026-09-27 02:46:18 +00:00
parent efc4f85fe6
commit 25797cbf2a
2 changed files with 39 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -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)

45
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
View file

@ -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()