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

fix(mcp): key the listed-tools cache by the OBO subject token

Browse source 
token_exchange servers list upstream with the caller's own Entra bearer, so two callers on one
LiteLLM key with different subjects were sharing a catalog slot

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yucheng 2026-09-15 02:09:05 +00:00
parent 8cfbf5dc2e
commit 0ab0e49e6a
2 changed files with 20 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -4575,10 +4575,11 @@ class MCPServerManager:
def _listed_tools_identity(self, server: MCPServer, caller: ListedToolsCaller | None) -> str | None:
"""Key the listed-tool cache by every request input that can change the upstream catalog.
Forwarded headers, header-driven stdio env, a relayed caller bearer, and the
server-specific auth header all reach upstream, so two callers differing in any of
them may be shown different tools. Shared servers with none of those stay on the
shared (``None``) slot. OpenAPI servers list from the process-wide registry.
Forwarded headers, header-driven stdio env, the caller bearer (forwarded as-is or
exchanged as the OBO subject), and the server-specific auth header all reach
upstream, so two callers differing in any of them may be shown different tools. Shared
servers with none of those stay on the shared (``None``) slot. OpenAPI servers list from
the process-wide registry.
"""
if server.spec_path or caller is None:
return None
@ -4589,12 +4590,12 @@ class MCPServerManager:
forwarded: Final = self._forwarded_header_values(server, caller.raw_headers)
header_env: Final = self._build_stdio_env(server, caller.raw_headers)
stdio_env: Final = None if header_env == self._build_stdio_env(server) else header_env
relayed_bearer: Final = (
caller_bearer: Final = (
self._extract_subject_token(caller.oauth2_headers, caller.raw_headers, auth)
if server.is_client_forwarded_token
if server.is_client_forwarded_token or server.auth_type == MCPAuth.oauth2_token_exchange
else None
)
inputs: Final = (identity, caller.mcp_auth_header, forwarded, stdio_env, relayed_bearer)
inputs: Final = (identity, caller.mcp_auth_header, forwarded, stdio_env, caller_bearer)
if not any(inputs):
return None
material: Final = json.dumps(inputs, sort_keys=True, separators=(",", ":"))

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

@ -6986,6 +6986,18 @@ class TestMCPServerManager:
ListedToolsCaller(mcp_auth_header="byok-b"),
id="per-server-auth-header",
),
pytest.param(
{"auth_type": MCPAuth.oauth2_token_exchange},
ListedToolsCaller(
user_api_key_auth=UserAPIKeyAuth(user_id="team-bot", api_key="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"),
raw_headers={"x-litellm-api-key": "sk-shared", "authorization": "Bearer entra-bob"},
),
id="shared-key-different-obo-subjects",
),
pytest.param(
{"transport": MCPTransport.stdio, "command": "srv", "env": {"WS": "${X-WS}"}},
ListedToolsCaller(raw_headers={"X-WS": "A"}),