fix(mcp): address delegate review findings on probes, master-key JWTs and introspection

The upstream-token probe now sends the same headers a tool call sends, including the server's
resolved static headers, and steps aside when they can't be resolved yet. A legacy delegate is
probed only when the caller asked for a single server, judged by the requested scope rather than
the reachable one

A JWT that verifies against the master key is LiteLLM's own, even once expired, so it takes normal
admission and is never forwarded upstream

OAuth2 introspection runs once per request inside the normal auth check. A provider that does not
vouch for the token raises OAuth2TokenValidationError, which admission reads as anonymous
This commit is contained in:
shrey kharbanda 2026-10-28 17:42:18 -07:00
parent 65988f3eb8
commit 0676e437d7
6 changed files with 429 additions and 138 deletions

View file

@ -12,6 +12,11 @@ from litellm.llms.custom_httpx.http_handler import (
from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth
class OAuth2TokenValidationError(ValueError):
"""The token-info or introspection endpoint did not vouch for the token: it rejected it, reported it
inactive, or could not be reached."""
class Oauth2Handler:
"""
Handles OAuth2 token validation.
@ -207,7 +212,7 @@ class Oauth2Handler:
)
except httpx.HTTPStatusError as e:
# This will catch any 4xx or 5xx errors
raise ValueError(f"Oauth 2.0 Token validation failed: {e}")
raise OAuth2TokenValidationError(f"Oauth 2.0 Token validation failed: {e}")
except Exception as e:
# This will catch any other errors (like network issues)
raise ValueError(f"An error occurred during token validation: {e}")
raise OAuth2TokenValidationError(f"An error occurred during token validation: {e}")

View file

@ -10,6 +10,7 @@ from itertools import chain
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal, TypeAlias
import jwt
from fastapi import HTTPException
from pydantic import SecretStr, TypeAdapter
from starlette.datastructures import Headers
@ -45,6 +46,7 @@ from litellm.proxy.agent_endpoints.auth.agent_access_groups import (
from litellm.proxy.agent_endpoints.auth.agent_caller import agent_caller_auth
from litellm.proxy.agent_endpoints.auth.managed_authorization import managed_agent_policy
from litellm.proxy.auth.ip_address_utils import IPAddressUtils
from litellm.proxy.auth.oauth2_check import OAuth2TokenValidationError
from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy module exports
_get_bearer_token_or_received_api_key, # pyright: ignore[reportPrivateUsage] # shared x-litellm-api-key parser lives with user_api_key_auth
_run_centralized_common_checks, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
@ -250,10 +252,25 @@ def _caused_by(exc: BaseException, error_type: type[BaseException]) -> bool:
return any(isinstance(link, error_type) for link in _cause_chain(exc))
def _signed_with(token: str, secret: str) -> bool:
"""True when ``token`` is a JWT whose HS256 signature verifies with ``secret``, whatever its claims say."""
try:
jwt.decode(
token,
secret,
algorithms=["HS256"],
options={"verify_exp": False, "verify_nbf": False, "verify_iat": False, "verify_aud": False},
)
except jwt.PyJWTError:
return False
return True
async def _bearer_is_litellm_credential(token: str) -> bool:
"""True when ``token`` is, or may be, a credential LiteLLM issued: the master key, a sealed gateway
bearer, or a virtual key the key table holds (a lookup fault counts as held). Only ``sk-`` tokens are
looked up, because no other token can be a virtual key."""
"""True when ``token`` is, or may be, a credential LiteLLM issued: the master key or a JWT it signed
(dashboard sessions, BYOK sessions), a sealed gateway bearer, or a virtual key the key table holds (a
lookup fault counts as held). Only ``sk-`` tokens are looked up, because no other token can be a
virtual key."""
from litellm.proxy.mcp_server.bridge_token_flow import ( # noqa: PLC0415 # bridge flow imports this module
is_sealed_gateway_bearer,
virtual_key_may_exist,
@ -263,7 +280,7 @@ async def _bearer_is_litellm_credential(token: str) -> bool:
)
from litellm.proxy.proxy_server import master_key # noqa: PLC0415 # startup owns the master key
if master_key and secrets.compare_digest(token.encode(), master_key.encode()):
if master_key and (secrets.compare_digest(token.encode(), master_key.encode()) or _signed_with(token, master_key)):
return True
if is_sealed_gateway_bearer(token):
return True
@ -298,16 +315,13 @@ def _identity_check(token: str, route: str, settings: Mapping[str, object]) -> _
return None
async def _verify_identity(check: _IdentityCheck, token: str, request: Request) -> None:
from litellm.proxy.auth.oauth2_check import Oauth2Handler # noqa: PLC0415 # proxy import cycle
async def _verify_identity(check: Literal["oauth2_proxy", "jwt"], token: str, request: Request) -> None:
from litellm.proxy.auth.oauth2_proxy_hook import ( # noqa: PLC0415 # proxy import cycle
handle_oauth2_proxy_request,
)
from litellm.proxy.proxy_server import jwt_handler # noqa: PLC0415 # startup owns the active auth configuration
match check:
case "oauth2":
await Oauth2Handler.check_oauth2_token(token=token)
case "oauth2_proxy":
await handle_oauth2_proxy_request(request=request)
case "jwt":
@ -320,8 +334,10 @@ async def _upstream_bearer_identity(token: str, litellm_api_key: str, request: R
"""The LiteLLM user an upstream bearer also proves, or ``None`` when it proves none.
Custom auth decides alone, so any rejection is ``None``. Otherwise the identity check LiteLLM would run
(OAuth2 introspection, reverse-proxy headers, or a JWT) must pass first. Past that, only a user LiteLLM
does not know is ``None``; every other refusal (budget, blocked team, disallowed route) propagates."""
must vouch for the caller. Reverse-proxy headers and a JWT are verified first, and OAuth2 introspection
runs once inside the full check, where a provider that does not vouch is ``None``. Past that, only a
user LiteLLM does not know is ``None``; every other refusal (budget, blocked team, disallowed route)
propagates."""
from litellm.proxy.auth.auth_utils import get_request_route # noqa: PLC0415 # proxy import cycle
from litellm.proxy.mcp_server.bridge_token_flow import ( # noqa: PLC0415 # bridge flow imports this module
custom_auth_configured,
@ -340,14 +356,15 @@ async def _upstream_bearer_identity(token: str, litellm_api_key: str, request: R
)
if check is None:
return None
try:
await _verify_identity(check, token, request)
except Exception: # noqa: BLE001 # any verification failure means LiteLLM does not vouch for this bearer
return None
if check != "oauth2":
try:
await _verify_identity(check, token, request)
except Exception: # noqa: BLE001 # any verification failure means LiteLLM does not vouch for this bearer
return None
try:
return await user_api_key_auth(api_key=litellm_api_key, request=request)
except (HTTPException, ProxyException) as exc:
if _caused_by(exc, UserNotFoundError):
if _caused_by(exc, UserNotFoundError) or _caused_by(exc, OAuth2TokenValidationError):
return None
raise

View file

@ -4176,6 +4176,12 @@ class MCPServerManager:
referenced: Final = collect_env_var_references(strings=static_headers.values())
return bool(referenced & user_var_names)
async def resolved_static_headers(
self, server: MCPServer, user_api_key_auth: UserAPIKeyAuth | None
) -> dict[str, str] | None:
"""``server.static_headers`` interpolated exactly as a tool call sends them."""
return await self._resolve_static_headers_with_env_vars(server, user_api_key_auth)
async def _resolve_static_headers_with_env_vars(
self,
server: MCPServer,

View file

@ -2104,6 +2104,7 @@ if MCP_AVAILABLE:
url: str,
auth_header: str,
timeout: float = 5.0,
headers: Mapping[str, str] | None = None,
) -> tuple[int, str | None]:
"""JSON-RPC initialize-probe the upstream URL to check whether the token is accepted.
@ -2136,6 +2137,7 @@ if MCP_AVAILABLE:
probe_headers: Final = {
"Accept": "application/json, text/event-stream",
**({"Authorization": auth_header} if auth_header else {}),
**(headers or {}),
}
try:
resp: Final = await client.post(
@ -2154,6 +2156,30 @@ if MCP_AVAILABLE:
verbose_logger.debug("_probe_upstream_auth: probe to %s failed (%s), allowing request through", url, exc)
return 200, None
async def _probe_with_outbound_headers(
server: MCPServer,
auth_header: str,
caller_headers: Mapping[str, str],
user_api_key_auth: UserAPIKeyAuth | None,
) -> tuple[int, str | None]:
"""Probe with the headers a real tool call sends: the caller's prepared headers with the server's
resolved static headers on top, so a static ``Authorization`` wins here too. When those cannot be
resolved yet, such as a per-user variable the caller has not set, the probe fails open like a network
error and leaves the error to the call."""
try:
static_headers: Final = await global_mcp_server_manager.resolved_static_headers(server, user_api_key_auth)
except Exception as exc: # noqa: BLE001 # the probe fails open on anything it cannot reproduce
verbose_logger.debug("Skipping upstream auth probe for %s (%s)", server.name, type(exc).__name__)
return 200, None
outbound: Final = merge_mcp_headers(extra_headers=caller_headers, static_headers=static_headers) or {}
authorization: Final = next(
(value for name, value in reversed(outbound.items()) if name.lower() == "authorization"), auth_header
)
companion_headers: Final = {name: value for name, value in outbound.items() if name.lower() != "authorization"}
if not companion_headers:
return await _probe_upstream_auth(server.url or "", authorization)
return await _probe_upstream_auth(server.url or "", authorization, headers=companion_headers)
async def _check_passthrough_upstream_auth(
scope: Scope,
user_api_key_auth: UserAPIKeyAuth | None,
@ -2163,6 +2189,7 @@ if MCP_AVAILABLE:
oauth2_headers: Mapping[str, str] | None,
mcp_server_auth_headers: dict[str, dict[str, str]] | None,
raw_headers: Mapping[str, str] | None,
requested_servers: Sequence[str] | None = None,
) -> None:
"""Probe pass-through and legacy delegate upstream servers in parallel before the MCP session starts.
@ -2189,7 +2216,10 @@ if MCP_AVAILABLE:
)
# A legacy delegate is probed only on a single-server route, so one rejected token cannot 401 a
# multi-server connect; its challenge echoes the requested name, matching the tokenless challenge.
single_requested_name: Final = mcp_servers[0] if mcp_servers is not None and len(mcp_servers) == 1 else None
requested_scope: Final = requested_servers if requested_servers is not None else mcp_servers
single_requested_name: Final = (
requested_scope[0] if requested_scope is not None and len(requested_scope) == 1 else None
)
prepared_headers: Final = tuple(
(
srv,
@ -2216,8 +2246,13 @@ if MCP_AVAILABLE:
)
for srv, (auth, extra) in prepared_headers
)
probe_targets: Final[tuple[tuple[MCPServer, str, str], ...]] = tuple(
(srv, auth_header, srv.name if srv.is_oauth_passthrough else single_requested_name or srv.name)
probe_targets: Final[tuple[tuple[MCPServer, str, str, Mapping[str, str]], ...]] = tuple(
(
srv,
auth_header,
srv.name if srv.is_oauth_passthrough else single_requested_name or srv.name,
headers or {},
)
for srv, headers in probe_headers
if (
auth_header := next(
@ -2229,9 +2264,12 @@ if MCP_AVAILABLE:
return
probe_results: Final = await asyncio.gather(
*[_probe_upstream_auth(srv.url or "", auth_header) for srv, auth_header, _ in probe_targets]
*[
_probe_with_outbound_headers(srv, auth_header, headers, user_api_key_auth)
for srv, auth_header, _, headers in probe_targets
]
)
for (_srv, _, challenge_server_name), (probe_status, _) in zip(probe_targets, probe_results):
for (_srv, _, challenge_server_name, _), (probe_status, _) in zip(probe_targets, probe_results):
if probe_status == 401:
# Token is missing or expired: keep pass-through clients on the
# protected-resource discovery flow so they re-authorize against
@ -2315,6 +2353,7 @@ if MCP_AVAILABLE:
if context.mcp_server_auth_headers is not None
else None,
raw_headers=context.raw_headers,
requested_servers=context.mcp_servers,
)
return None

View file

@ -1599,9 +1599,7 @@ class TestMCPRequestHandler:
team_id="test-team-id",
)
with patch(
"litellm.proxy.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth"
) as mock_user_api_key_auth:
with patch("litellm.proxy.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth") as mock_user_api_key_auth:
mock_user_api_key_auth.return_value = mock_auth_result
# Call the method
@ -2775,6 +2773,40 @@ class TestMCPDelegateAuthToUpstream:
assert auth_result.mcp_forwards_caller_authorization is False
assert oauth2_headers is None
@pytest.mark.parametrize(
"signing_secret, expires_in, forwarded",
[
("sk-master-secret", timedelta(hours=1), False),
("sk-master-secret", timedelta(hours=-1), False),
("upstream-idp-signing-secret", timedelta(hours=1), True),
],
ids=["master-key-signed", "master-key-signed-expired", "signed-elsewhere"],
)
async def test_delegate_master_key_signed_jwt_is_never_forwarded(
self, signing_secret: str, expires_in: timedelta, forwarded: bool
) -> None:
"""LiteLLM signs dashboard and OAuth tokens with the master key, so such a JWT is LiteLLM's own even
once expired: it takes normal admission and never reaches the upstream. A JWT signed by anyone
else is the upstream's token."""
import jwt
from litellm.proxy import proxy_server
token: Final = jwt.encode(
{"sub": "dashboard-user", "exp": datetime.now(timezone.utc) + expires_in}, signing_secret, algorithm="HS256"
)
admission: Final = AsyncMock(return_value=UserAPIKeyAuth(user_id="dashboard-user"))
with (
patch.object(proxy_server, "master_key", "sk-master-secret"),
TestMCPDelegateAuthToUpstream._delegate_admission(admission),
):
auth_result, *_ = await MCPRequestHandler.process_mcp_request(
TestMCPDelegateAuthToUpstream._bearer_scope(token)
)
assert auth_result.mcp_forwards_caller_authorization is forwarded
assert admission.await_count == (0 if forwarded else 1)
@pytest.mark.parametrize(
"token, ui_token_lookup",
[
@ -2912,28 +2944,26 @@ class TestMCPDelegateAuthToUpstream:
):
yield
@pytest.mark.parametrize("mechanism", ["oauth2", "oauth2_proxy"])
@pytest.mark.parametrize(
"verify_effect, allowed_server_ids, expected_user_id",
[
(None, ["test-server-id"], "idp-user"),
(None, [], None),
(ValueError("Token is not active"), ["test-server-id"], None),
(ValueError("Missing trusted proxy header"), ["test-server-id"], None),
],
ids=["granted-user", "user-not-granted", "identity-rejected"],
)
async def test_delegate_idp_identity_follows_the_jwt_rules(
async def test_delegate_proxy_header_identity_follows_the_jwt_rules(
self,
mechanism: Literal["oauth2", "oauth2_proxy"],
verify_effect: Exception | None,
allowed_server_ids: Sequence[str],
expected_user_id: str | None,
) -> None:
"""OAuth2 introspection and trusted reverse-proxy headers admit a granted user by name, and a user
who is not granted or not vouched for anonymously. The bearer is forwarded either way."""
"""Trusted reverse-proxy headers admit a granted user by name, and a user who is not granted or not
vouched for anonymously. The bearer is forwarded either way."""
admission: Final = AsyncMock(return_value=UserAPIKeyAuth(user_id="idp-user"))
with (
TestMCPDelegateAuthToUpstream._idp_auth_enabled(mechanism, AsyncMock(side_effect=verify_effect)),
TestMCPDelegateAuthToUpstream._idp_auth_enabled("oauth2_proxy", AsyncMock(side_effect=verify_effect)),
TestMCPDelegateAuthToUpstream._delegate_admission(admission, allowed_server_ids),
):
auth_result, _, _, _, oauth2_headers, _ = await MCPRequestHandler.process_mcp_request(
@ -2944,6 +2974,55 @@ class TestMCPDelegateAuthToUpstream:
assert auth_result.mcp_forwards_caller_authorization is True
assert oauth2_headers == {"Authorization": "Bearer opaque-idp-token"}
@pytest.mark.parametrize(
"active, allowed_server_ids, expected_user_id",
[(True, ["test-server-id"], "idp-user"), (True, [], None), (False, ["test-server-id"], None)],
ids=["granted-user", "user-not-granted", "token-inactive"],
)
async def test_delegate_oauth2_introspects_once_and_follows_the_jwt_rules(
self, active: bool, allowed_server_ids: Sequence[str], expected_user_id: str | None
) -> None:
"""The provider is asked about the token once per request. A token it vouches for admits a granted
user by name, and one it reports inactive is forwarded anonymously instead of failing the request."""
import httpx
from litellm.proxy import proxy_server
from litellm.types.mcp import MCPAuth
async def introspect(outbound: httpx.Request, **kwargs: object) -> httpx.Response:
return httpx.Response(200, request=outbound, json={"active": active, "sub": "idp-user"})
with (
patch.dict(proxy_server.general_settings, {"enable_oauth2_auth": True}),
patch.object(proxy_server, "premium_user", True),
patch.object(proxy_server, "user_custom_auth", None),
patch.object(proxy_server, "master_key", "sk-fixture-master"),
patch.object(proxy_server, "prisma_client", None),
patch.object(proxy_server, "user_api_key_cache", DualCache()),
patch.dict(
os.environ,
{
"OAUTH_TOKEN_INFO_ENDPOINT": "https://idp.fixture.test/introspect",
"OAUTH_CLIENT_ID": "fixture-client",
"OAUTH_CLIENT_SECRET": "fixture-client-secret",
},
),
patch("httpx.AsyncClient.send", side_effect=introspect) as network,
patch("litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr,
):
mock_mgr.get_mcp_server_by_name.return_value = TestMCPDelegateAuthToUpstream._make_server(
auth_type=MCPAuth.oauth2, delegate_auth_to_upstream=True
)
mock_mgr.get_allowed_mcp_servers = AsyncMock(return_value=list(allowed_server_ids))
auth_result, _, _, _, oauth2_headers, _ = await MCPRequestHandler.process_mcp_request(
TestMCPDelegateAuthToUpstream._bearer_scope("opaque-idp-token")
)
assert network.call_count == 1
assert auth_result.user_id == expected_user_id
assert auth_result.mcp_forwards_caller_authorization is True
assert oauth2_headers == {"Authorization": "Bearer opaque-idp-token"}
@pytest.mark.parametrize("mechanism", ["oauth2", "oauth2_proxy"])
async def test_delegate_idp_user_refusal_propagates(self, mechanism: Literal["oauth2", "oauth2_proxy"]) -> None:
"""Once the provider vouches for the caller, LiteLLM's budget and permission refusals apply."""
@ -2966,7 +3045,9 @@ class TestMCPDelegateAuthToUpstream:
introspect: Final = AsyncMock()
with (
TestMCPDelegateAuthToUpstream._idp_auth_enabled("oauth2", introspect, premium=False),
TestMCPDelegateAuthToUpstream._delegate_admission(AsyncMock(return_value=UserAPIKeyAuth(user_id="idp-user"))),
TestMCPDelegateAuthToUpstream._delegate_admission(
AsyncMock(return_value=UserAPIKeyAuth(user_id="idp-user"))
),
):
auth_result, *_ = await MCPRequestHandler.process_mcp_request(
TestMCPDelegateAuthToUpstream._bearer_scope("opaque-idp-token")
@ -8005,9 +8086,7 @@ class TestMCPDcrBridgeDualCredential:
new_callable=AsyncMock,
return_value=UserAPIKeyAuth(user_id=presented_user, api_key="any-hash"),
),
patch(
"litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager"
) as mock_mgr,
patch("litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr,
patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY),
):
mock_mgr.get_mcp_server_by_name.return_value = self._DELEGATE._bridge_delegate_server()
@ -8058,7 +8137,9 @@ class TestMCPDcrBridgeDualCredential:
@pytest.mark.parametrize("root_path", ["", "/litellm"])
@pytest.mark.parametrize("invalid_token", [False, True])
def test_gateway_challenge_metadata_contains_the_public_resource(
public_path: str, root_path: str, invalid_token: bool,
public_path: str,
root_path: str,
invalid_token: bool,
) -> None:
from mcp.shared.auth_utils import check_resource_allowed
from starlette.requests import Request
@ -8071,8 +8152,12 @@ def test_gateway_challenge_metadata_contains_the_public_resource(
)
scope: Final[Scope] = {
"type": "http", "method": "POST", "scheme": "https", "path": "/mcp",
"root_path": root_path, "_original_path": f"{root_path}{public_path}",
"type": "http",
"method": "POST",
"scheme": "https",
"path": "/mcp",
"root_path": root_path,
"_original_path": f"{root_path}{public_path}",
"headers": [(b"host", b"gateway.example.com")],
}
request: Final = Request(scope)
@ -8081,8 +8166,7 @@ def test_gateway_challenge_metadata_contains_the_public_resource(
metadata_url: Final = header.split('resource_metadata="', 1)[1].split('"', 1)[0]
toolset: Final = public_path in ("/toolset/analysis/mcp", "/analysis/mcp")
assert metadata_url == (
f"https://gateway.example.com{root_path}/.well-known/oauth-protected-resource"
+ ("" if toolset else "/mcp")
f"https://gateway.example.com{root_path}/.well-known/oauth-protected-resource" + ("" if toolset else "/mcp")
)
metadata: Final = (
oauth_protected_resource_root(request) if toolset else _build_aggregate_protected_resource_response(request)
@ -8285,9 +8369,7 @@ class TestAggregateGatewayDcrChallenge:
):
with (
patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()),
patch(
"litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager"
) as mock_mgr,
patch("litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr,
):
mock_mgr.get_mcp_server_by_name.return_value = server
with pytest.raises(HTTPException) as exc_info:
@ -8333,9 +8415,7 @@ class TestAggregateGatewayDcrChallenge:
with (
patch.dict(os.environ, {"SERVER_ROOT_PATH": "/litellm"}),
patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()),
patch(
"litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager"
) as mock_mgr,
patch("litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr,
):
mock_mgr.get_mcp_server_by_name.return_value = server
with pytest.raises(HTTPException) as exc_info:
@ -8366,9 +8446,7 @@ class TestAggregateGatewayDcrChallenge:
):
with (
patch(self._AUTH_PATCH_TARGET, side_effect=self._auth_401()),
patch(
"litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager"
) as mock_mgr,
patch("litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr,
):
mock_mgr.get_mcp_server_by_name.return_value = resolved
mock_mgr.get_mcp_server_by_identifier.return_value = resolved
@ -8413,9 +8491,7 @@ class TestAggregateGatewayDcrChallenge:
(None, None),
]
for resolved, expected in cases:
with patch(
"litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager"
) as mock_mgr:
with patch("litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr:
mock_mgr.get_mcp_server_by_identifier.return_value = resolved
assert _gateway_dcr_challenge_target("/mcp/srv", None, None) == expected, resolved
assert _gateway_dcr_challenge_target("/mcp/a,b", None, None) is None
@ -11311,14 +11387,15 @@ async def test_catalog_refresh_uses_current_virtual_key_policy_and_keeps_session
from litellm.proxy import proxy_server
from litellm.types.mcp_server.mcp_toolset import MCPToolset
permission = LiteLLM_ObjectPermissionTable(
object_permission_id="current-policy", mcp_toolsets=["session-toolset"]
)
permission = LiteLLM_ObjectPermissionTable(object_permission_id="current-policy", mcp_toolsets=["session-toolset"])
toolset: Final = MCPToolset(
toolset_id="session-toolset", toolset_name="session-toolset",
toolset_id="session-toolset",
toolset_name="session-toolset",
tools=[{"server_id": "current-server", "tool_name": "read"}],
)
tables: Final = SimpleNamespace(litellm_mcptoolsettable=SimpleNamespace(find_many=AsyncMock(return_value=[toolset])))
tables: Final = SimpleNamespace(
litellm_mcptoolsettable=SimpleNamespace(find_many=AsyncMock(return_value=[toolset]))
)
monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=tables, writer_db=tables))
monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache())
current = UserAPIKeyAuth(
@ -11345,10 +11422,13 @@ async def test_catalog_refresh_uses_current_virtual_key_policy_and_keeps_session
caller.mcp_toolset_id = "session-toolset"
refreshed = await MCPRequestHandler.refresh_catalog_authority(caller)
reload_key.assert_awaited_once_with("owned-key-hash", check_db_only=True)
assert refreshed.object_permission == permission.model_copy(update={
"mcp_servers": ["current-server"], "mcp_toolsets": [],
"mcp_tool_permissions": {"current-server": ["read"]},
})
assert refreshed.object_permission == permission.model_copy(
update={
"mcp_servers": ["current-server"],
"mcp_toolsets": [],
"mcp_tool_permissions": {"current-server": ["read"]},
}
)
assert refreshed.object_permission_id == "current-policy"
assert (refreshed.team_id, refreshed.org_id, refreshed.project_id, refreshed.user_id) == (
"new-team",
@ -12452,9 +12532,7 @@ async def test_personal_key_guardrail_auth_preserves_native_tool_name_matching_s
from litellm.proxy.mcp_server.mcp_server_manager import global_mcp_server_manager
async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable:
return _personal_key_team_granting_tools(
team_id, ("math-multiply",) if team_id == "team-b" else ("other",)
)
return _personal_key_team_granting_tools(team_id, ("math-multiply",) if team_id == "team-b" else ("other",))
monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object)
monkeypatch.setattr(auth_checks, "get_org_object", AsyncMock(return_value=None))
@ -12493,9 +12571,7 @@ async def test_personal_key_mcp_billing_inherits_adopted_organization_limits(
)
async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable:
return _personal_key_team_granting_tools(team_id, ("divide",)).model_copy(
update={"organization_id": "org-1"}
)
return _personal_key_team_granting_tools(team_id, ("divide",)).model_copy(update={"organization_id": "org-1"})
org_loader: Final = AsyncMock(return_value=organization)
monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object)
@ -12541,16 +12617,12 @@ async def test_personal_key_mcp_billing_keeps_key_organization_limits(
)
async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable:
return _personal_key_team_granting_tools(team_id, ("divide",)).model_copy(
update={"organization_id": "org-1"}
)
return _personal_key_team_granting_tools(team_id, ("divide",)).model_copy(update={"organization_id": "org-1"})
org_loader: Final = AsyncMock(return_value=organization)
monkeypatch.setattr(auth_checks, "get_team_object", _get_team_object)
monkeypatch.setattr(auth_checks, "get_org_object_for_request", org_loader)
auth: Final = personal_key_owner_teams.model_copy(
update={"org_id": "org-k", "organization_rpm_limit": 9}
)
auth: Final = personal_key_owner_teams.model_copy(update={"org_id": "org-k", "organization_rpm_limit": 9})
server: Final = global_mcp_server_manager.registry["srv-b"]
with patch.object(global_mcp_server_manager, "get_mcp_server_from_tool_name", return_value=server):
@ -12682,9 +12754,7 @@ async def test_personal_key_mcp_billing_uses_fresh_team_and_membership_for_stric
}
)
async def _get_team_membership(
user_id: str, team_id: str, **kwargs: object
) -> LiteLLM_TeamMembership:
async def _get_team_membership(user_id: str, team_id: str, **kwargs: object) -> LiteLLM_TeamMembership:
database_value: Final = bool(kwargs.get("check_db_only"))
return LiteLLM_TeamMembership(
user_id=user_id,

View file

@ -7335,16 +7335,19 @@ async def test_legacy_delegate_bare_token_is_not_probed_upstream(): # test-qual
@pytest.mark.asyncio
async def test_legacy_delegate_dual_credentials_probe_only_the_upstream_token() -> None: # test-quality-ok: the credential sent to the upstream probe is the security contract and has no return value
@pytest.mark.usefixtures("httpx_transport")
async def test_legacy_delegate_dual_credentials_probe_only_the_upstream_token() -> None:
"""With a separate x-litellm-api-key, the preflight probes the upstream with its own token and never
with the LiteLLM key."""
import respx
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.mcp_server.server import (
_check_passthrough_upstream_auth,
)
server = _delegate_auth_mcp_server()
scope = _delegate_scope(
server: Final = _delegate_auth_mcp_server()
scope: Final = _delegate_scope(
[
(b"x-litellm-api-key", b"sk-litellm-proxy-key"),
(b"authorization", b"Bearer upstream-token"),
@ -7356,11 +7359,9 @@ async def test_legacy_delegate_dual_credentials_probe_only_the_upstream_token()
"litellm.proxy.mcp_server.operations._get_allowed_mcp_servers",
new=AsyncMock(return_value=[server]),
),
patch( # test-quality-ok: the probe's outbound credential is the security contract under test
"litellm.proxy.mcp_server.server._probe_upstream_auth",
new=AsyncMock(return_value=(200, None)),
) as probe,
respx.mock as upstream,
):
probe: Final = upstream.post(server.url).respond(200)
await _check_passthrough_upstream_auth(
scope=scope,
user_api_key_auth=UserAPIKeyAuth(user_id="admitted-user"),
@ -7371,7 +7372,10 @@ async def test_legacy_delegate_dual_credentials_probe_only_the_upstream_token()
raw_headers={"x-litellm-api-key": "sk-litellm-proxy-key", "Authorization": "Bearer upstream-token"},
)
assert probe.await_args.args == ("http://upstream:9401/mcp", "Bearer upstream-token")
assert probe.call_count == 1
sent_headers: Final = probe.calls[0].request.headers
assert sent_headers["Authorization"] == "Bearer upstream-token"
assert all("sk-litellm-proxy-key" not in value for value in sent_headers.values())
def _expired_delegate_token_preflight(mcp_servers: Sequence[str]) -> Coroutine[object, object, None]:
@ -7416,7 +7420,9 @@ async def test_legacy_delegate_expired_upstream_token_gets_reauthorization_chall
assert exc_info.value.status_code == 401
challenge: Final = exc_info.value.headers["www-authenticate"]
assert 'error="invalid_token"' in challenge
assert 'resource_metadata="http://localhost:4000/.well-known/oauth-protected-resource/mcp/delegate_test"' in challenge
assert (
'resource_metadata="http://localhost:4000/.well-known/oauth-protected-resource/mcp/delegate_test"' in challenge
)
@pytest.mark.asyncio
@ -7426,6 +7432,131 @@ async def test_legacy_delegate_multi_server_scope_is_not_failed_by_one_upstream_
assert await _expired_delegate_token_preflight(["delegate_test", "other_server"]) is None
@pytest.mark.asyncio
@pytest.mark.usefixtures("httpx_transport")
@pytest.mark.parametrize(
"static_headers, expected_authorization",
[
({"X-Tenant": "tenant-1"}, "Bearer expired-upstream-token"),
({"X-Tenant": "tenant-1", "Authorization": "Bearer server-wide-token"}, "Bearer server-wide-token"),
],
ids=["companion-header", "static-authorization-wins"],
)
async def test_legacy_delegate_probe_sends_the_headers_a_tool_call_sends(
static_headers: dict[str, str], expected_authorization: str
) -> None:
"""An upstream that needs a static header next to the token sees it in the probe too, so a token the
upstream accepts on real calls is not rejected before the session opens."""
import respx
server: Final = _delegate_auth_mcp_server().model_copy(update={"static_headers": static_headers})
with (
patch( # test-quality-ok: isolate authorized-server resolution so this test targets the preflight boundary
"litellm.proxy.mcp_server.operations._get_allowed_mcp_servers",
new=AsyncMock(return_value=[server]),
),
respx.mock as upstream,
):
probe: Final = upstream.post(server.url).mock(
side_effect=lambda request: httpx.Response(200 if request.headers.get("X-Tenant") == "tenant-1" else 401)
)
assert await _expired_delegate_token_preflight(["delegate_test"]) is None
assert probe.call_count == 1
assert probe.calls[0].request.headers["Authorization"] == expected_authorization
@pytest.mark.asyncio
@pytest.mark.usefixtures("httpx_transport")
async def test_legacy_delegate_probe_steps_aside_when_static_headers_need_unset_user_values() -> None:
"""A static header built from a per-user value the caller has not set cannot be reproduced, so the
probe lets the request through and the tool call reports the missing value."""
import respx
server: Final = _delegate_auth_mcp_server().model_copy(
update={
"static_headers": {"X-Tenant-Token": "${TENANT_TOKEN}"},
"env_vars": [{"name": "TENANT_TOKEN", "value": "", "scope": "user"}],
}
)
with (
patch( # test-quality-ok: isolate authorized-server resolution so this test targets the preflight boundary
"litellm.proxy.mcp_server.operations._get_allowed_mcp_servers",
new=AsyncMock(return_value=[server]),
),
respx.mock as upstream,
):
probe: Final = upstream.post(server.url).respond(
401, headers={"www-authenticate": 'Bearer error="invalid_token"'}
)
assert await _expired_delegate_token_preflight(["delegate_test"]) is None
assert probe.call_count == 0
@pytest.mark.asyncio
@pytest.mark.usefixtures("httpx_transport")
@pytest.mark.parametrize(
"requested, expected_status",
[(("delegate_test",), 401), (("delegate_test", "other_server"), None)],
ids=["single-server", "two-requested-one-reachable"],
)
async def test_modern_interaction_probes_legacy_delegate_only_for_a_single_requested_server(
requested: tuple[str, ...], expected_status: int | None
) -> None:
"""The single-server rule follows what the caller asked for, not what it can reach, so a request
naming two servers is not failed by one expired delegate token."""
import respx
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.mcp_server import server
from litellm.proxy.mcp_server.contracts import OperationContext
target: Final = _delegate_auth_mcp_server()
scope: Final = _delegate_scope(
[
(b"authorization", b"Bearer expired-upstream-token"),
(b"mcp-protocol-version", b"2026-07-28"),
(b"mcp-method", b"tools/call"),
(b"mcp-name", b"delegate-add"),
]
)
context: Final = OperationContext(
_caller=UserAPIKeyAuth().model_copy(update={"mcp_forwards_caller_authorization": True}),
mcp_servers=requested,
oauth2_headers={"Authorization": "Bearer expired-upstream-token"},
raw_headers={"authorization": "Bearer expired-upstream-token"},
)
body: Final = json.dumps(
{
"jsonrpc": "2.0",
"id": 1,
"method": "tools/call",
"params": {
"name": "delegate-add",
"arguments": {"a": 2, "b": 3},
"_meta": {
"io.modelcontextprotocol/protocolVersion": "2026-07-28",
"io.modelcontextprotocol/clientCapabilities": {},
},
},
}
).encode()
with (
patch( # test-quality-ok: isolate authorized-server resolution so this test targets the preflight boundary
"litellm.proxy.mcp_server.operations._get_allowed_mcp_servers",
new=AsyncMock(return_value=[target]),
),
respx.mock as upstream,
):
upstream.post(target.url).respond(401, headers={"www-authenticate": 'Bearer error="invalid_token"'})
outcome: Final = await asyncio.gather(
server._preflight_modern_interaction(scope, body, context), return_exceptions=True
)
assert getattr(outcome[0], "status_code", None) == expected_status
@pytest.mark.parametrize(
"probe_status, expected_status",
[(200, None), (401, 401), (403, 403)],
@ -8969,8 +9100,7 @@ async def test_get_allowed_mcp_servers_from_mcp_server_names_unknown_name_fails_
]
with patch(
"litellm.proxy.mcp_server.auth.user_api_key_auth_mcp."
"MCPRequestHandler.get_mcp_servers_from_access_groups",
"litellm.proxy.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler.get_mcp_servers_from_access_groups",
new_callable=AsyncMock,
return_value=[],
):
@ -9028,8 +9158,7 @@ async def test_get_allowed_mcp_servers_from_mcp_server_names_known_alias_returns
]
with patch(
"litellm.proxy.mcp_server.auth.user_api_key_auth_mcp."
"MCPRequestHandler.get_mcp_servers_from_access_groups",
"litellm.proxy.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler.get_mcp_servers_from_access_groups",
new_callable=AsyncMock,
return_value=[],
):
@ -9061,8 +9190,7 @@ async def test_get_allowed_mcp_servers_from_mcp_server_names_mixed_known_and_unk
]
with patch(
"litellm.proxy.mcp_server.auth.user_api_key_auth_mcp."
"MCPRequestHandler.get_mcp_servers_from_access_groups",
"litellm.proxy.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler.get_mcp_servers_from_access_groups",
new_callable=AsyncMock,
return_value=[],
):
@ -9094,8 +9222,7 @@ async def test_get_allowed_mcp_servers_from_mcp_server_names_access_group_resolv
]
with patch(
"litellm.proxy.mcp_server.auth.user_api_key_auth_mcp."
"MCPRequestHandler.get_mcp_servers_from_access_groups",
"litellm.proxy.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler.get_mcp_servers_from_access_groups",
new_callable=AsyncMock,
return_value=["id-b"],
):
@ -9363,15 +9490,11 @@ async def test_stateful_mcp_tracing_follows_each_message_exclusion(
is_otel_v2_enabled.cache_clear()
notified: Final = asyncio.Event()
async def call_tool(
ctx: ServerRequestContext[object, object], params: CallToolRequestParams
) -> CallToolResult:
async def call_tool(ctx: ServerRequestContext[object, object], params: CallToolRequestParams) -> CallToolResult:
with tracer.start_as_current_span("tool-handler"):
return CallToolResult(content=[TextContent(type="text", text="done")])
async def notification(
ctx: ServerRequestContext[object, object], params: NotificationParams
) -> HandlerResult:
async def notification(ctx: ServerRequestContext[object, object], params: NotificationParams) -> HandlerResult:
with tracer.start_as_current_span("notification-handler"):
notified.set()
return None
@ -9389,51 +9512,85 @@ async def test_stateful_mcp_tracing_follows_each_message_exclusion(
try:
async with manager.run():
async with httpx2.AsyncClient(transport=httpx2.ASGITransport(app=app)) as client:
initialized: Final = await client.post(initial_url, headers=headers, json={
"jsonrpc": "2.0", "id": 0, "method": "initialize", "params": {
"protocolVersion": "2025-11-25", "capabilities": {},
"clientInfo": {"name": "trace-exclusion", "version": "1"},
},
})
assert initialized.status_code == 200, initialized.text
assert any(
span.name == "initialize" for span in exporter.get_finished_spans()
) is not initialized_excluded
exporter.clear()
session_headers: Final = {**headers, "MCP-Session-Id": initialized.headers["mcp-session-id"]}
ready: Final = await client.post(initial_url, headers=session_headers,
json={"jsonrpc": "2.0", "method": "notifications/initialized"})
assert ready.status_code == 202, ready.text
response: Final = await client.post(current_url, headers=session_headers, json={
"jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {
"name": "probe", "arguments": {}, "_meta": {
"traceparent": "00-11111111111111111111111111111111-2222222222222222-01",
initialized: Final = await client.post(
initial_url,
headers=headers,
json={
"jsonrpc": "2.0",
"id": 0,
"method": "initialize",
"params": {
"protocolVersion": "2025-11-25",
"capabilities": {},
"clientInfo": {"name": "trace-exclusion", "version": "1"},
},
},
})
)
assert initialized.status_code == 200, initialized.text
assert (
any(span.name == "initialize" for span in exporter.get_finished_spans()) is not initialized_excluded
)
exporter.clear()
session_headers: Final = {**headers, "MCP-Session-Id": initialized.headers["mcp-session-id"]}
ready: Final = await client.post(
initial_url, headers=session_headers, json={"jsonrpc": "2.0", "method": "notifications/initialized"}
)
assert ready.status_code == 202, ready.text
response: Final = await client.post(
current_url,
headers=session_headers,
json={
"jsonrpc": "2.0",
"id": 1,
"method": "tools/call",
"params": {
"name": "probe",
"arguments": {},
"_meta": {
"traceparent": "00-11111111111111111111111111111111-2222222222222222-01",
},
},
},
)
assert response.status_code == 200 and '"text":"done"' in response.text, response.text
names: Final = {span.name for span in exporter.get_finished_spans()}
assert ("tools/call probe" in names) is initialized_excluded
assert ("tool-handler" in names) is initialized_excluded
exporter.clear()
notification_response: Final = await client.post(current_url, headers=session_headers,
json={"jsonrpc": "2.0", "method": "notifications/exclusion-probe"})
notification_response: Final = await client.post(
current_url,
headers=session_headers,
json={"jsonrpc": "2.0", "method": "notifications/exclusion-probe"},
)
assert notification_response.status_code == 202, notification_response.text
await notified.wait()
assert any(
span.name == "notification-handler" for span in exporter.get_finished_spans()
) is initialized_excluded
assert (
any(span.name == "notification-handler" for span in exporter.get_finished_spans())
is initialized_excluded
)
exporter.clear()
modern_response: Final = await client.post(current_url, headers={
"Accept": "application/json, text/event-stream", "MCP-Protocol-Version": "2026-07-28",
"MCP-Method": "tools/call", "MCP-Name": "probe",
}, json={
"jsonrpc": "2.0", "id": 2, "method": "tools/call", "params": {
"name": "probe", "arguments": {}, "_meta": {
PROTOCOL_VERSION_META_KEY: "2026-07-28", CLIENT_CAPABILITIES_META_KEY: {},
modern_response: Final = await client.post(
current_url,
headers={
"Accept": "application/json, text/event-stream",
"MCP-Protocol-Version": "2026-07-28",
"MCP-Method": "tools/call",
"MCP-Name": "probe",
},
json={
"jsonrpc": "2.0",
"id": 2,
"method": "tools/call",
"params": {
"name": "probe",
"arguments": {},
"_meta": {
PROTOCOL_VERSION_META_KEY: "2026-07-28",
CLIENT_CAPABILITIES_META_KEY: {},
},
},
},
})
)
assert modern_response.status_code == 200, modern_response.text
assert '"text":"done"' in modern_response.text, modern_response.text
modern_names: Final = {span.name for span in exporter.get_finished_spans()}
@ -9536,8 +9693,7 @@ async def test_get_allowed_mcp_servers_includes_active_servers_submitted_by_user
return_value={"submitted-1": submitted_server},
),
patch(
"litellm.proxy.mcp_server.auth.user_api_key_auth_mcp."
"MCPRequestHandler.get_allowed_mcp_servers",
"litellm.proxy.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler.get_allowed_mcp_servers",
AsyncMock(return_value=[]),
),
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
@ -10497,9 +10653,7 @@ async def test_handle_list_tools_all_faults_on_continuation_stays_degraded(
outcomes={"a": ServerListFault(tag="upstream_error", status_code=429)},
)
result: Final = await _list_tools_with_listing(
_mcp_request_ctx, listing, PaginatedRequestParams(cursor="cursor-1")
)
result: Final = await _list_tools_with_listing(_mcp_request_ctx, listing, PaginatedRequestParams(cursor="cursor-1"))
assert isinstance(result, ListToolsResult)
assert result.tools == []