mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
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:
parent
65988f3eb8
commit
0676e437d7
6 changed files with 429 additions and 138 deletions
|
|
@ -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}")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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 == []
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue