scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-17 23:51:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 139

refactor(guardrails): keep Agent 365 PR to the guardrail, move MCP changes to stacked PRs

Browse source 
The listed-tool metadata and per-caller catalog cache move to a follow-up PR
stacked on this one, and the RFC 9728 sign-in challenge and discovery metadata
move to a second PR stacked on that. On its own the guardrail evaluates the
tool name and arguments and the caller presents its Entra token

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yucheng 2026-09-15 00:29:58 +00:00
parent b030164f73
commit 4bfaf74edc
17 changed files with 89 additions and 1656 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

24
litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py
View file

@ -5,14 +5,13 @@ import secrets
import time
from collections.abc import Callable, Mapping
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict
from typing import TYPE_CHECKING, Any, Final, Literal, Optional
from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse
import httpx
from fastapi import APIRouter, Depends, Form, HTTPException, Request
from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Response
from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError
from typing_extensions import ReadOnly
from litellm._logging import verbose_logger
from litellm.caching.in_memory_cache import InMemoryCache
@ -82,10 +81,6 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import (
encrypt_value_helper,
)
from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
from litellm.proxy.guardrails.guardrail_hooks.agent_365.agent_365 import (
agent_365_authorization_servers,
agent_365_scopes_supported,
)
from litellm.types.mcp import MCPAuth, MCPCredentials
from litellm.types.mcp_server.mcp_server_manager import MCPServer, MCPTokenEndpointAuthMethod
@ -2403,7 +2398,7 @@ async def _build_oauth_protected_resource_response(
request: Request,
mcp_server_name: str | None,
use_standard_pattern: bool,
) -> Mapping[str, object]:
) -> dict:
"""
Build OAuth protected resource response with the appropriate URL pattern.
@ -2502,15 +2497,6 @@ async def _build_oauth_protected_resource_response(
if obo_response is not None:
return obo_response
agent_365_issuers: Final = agent_365_authorization_servers(mcp_server, None) if mcp_server else ()
if mcp_server is not None and agent_365_issuers:
agent_365_metadata: Final[_ProtectedResourceMetadata] = {
"authorization_servers": agent_365_issuers,
"resource": resource_url,
"scopes_supported": agent_365_scopes_supported(mcp_server, None),
}
return agent_365_metadata
if mcp_server is not None and mcp_server.advertises_gateway_authorization_server:
return {
"authorization_servers": [f"{request_base_url}/mcp"],
@ -2530,12 +2516,6 @@ async def _build_oauth_protected_resource_response(
}
class _ProtectedResourceMetadata(TypedDict):
authorization_servers: ReadOnly[tuple[str, ...]]
resource: ReadOnly[str]
scopes_supported: ReadOnly[tuple[str, ...]]
def _obo_protected_resource_response(mcp_server: MCPServer | None, resource_url: str) -> dict | None:
"""The OBO (token_exchange) PRM, or None when this server is not OBO / no issuer is configured.

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

@ -242,21 +242,6 @@ _user_env_vars_cache: Final[dict[tuple[str, str], tuple[dict[str, str], float]]]
_USER_ENV_VARS_CACHE_TTL: Final = 60 # seconds
_USER_ENV_VARS_CACHE_MAX_SIZE: Final = 4096 # cap to prevent unbounded growth
_ListedToolsByCaller: TypeAlias = Mapping[str | None, Mapping[str, MCPTool]]
_NO_LISTED_TOOLS: Final[_ListedToolsByCaller] = MappingProxyType({})
_LISTED_TOOLS_CALLERS_PER_SERVER: Final = 256
@dataclass(frozen=True, slots=True)
class ListedToolsCaller:
"""Request inputs that select which upstream catalog a caller was shown by tools/list."""
user_api_key_auth: UserAPIKeyAuth | None = None
mcp_auth_header: str | dict[str, str] | None = None
raw_headers: Mapping[str, str] | None = None
oauth2_headers: Mapping[str, str] | None = None
# Auth types whose upstream OAuth endpoints (protected-resource + authorization-server metadata) the
# gateway discovers from the upstream itself: interactive oauth2 and the two client-forwarded modes.
# OBO/M2M endpoint discovery is decided separately via _obo_needs_endpoint_discovery. Shared by the
@ -1134,25 +1119,6 @@ def _authorization_is_litellm_admission_credential(
return bool(user_api_key_auth and user_api_key_auth.api_key and not admission_header)
def _server_auth_header_for(
server: MCPServer,
mcp_server_auth_headers: Mapping[str, str | dict[str, str]] | None,
mcp_auth_header: str | dict[str, str] | None,
) -> str | dict[str, str] | None:
"""Server-specific ``x-mcp-<alias>-authorization`` header, else the deprecated global one."""
server_specific: Final = (
lookup_mcp_server_auth_in_headers(
mcp_server_auth_headers,
alias=server.alias,
server_name=server.server_name,
access_groups=server.access_groups,
)
if mcp_server_auth_headers
else None
)
return mcp_auth_header if server_specific is None else server_specific
def _format_byok_openapi_auth_header(mcp_server: MCPServer, mcp_auth_header: str) -> str:
"""Format a raw BYOK credential for OpenAPI tool ``Authorization`` injection.
@ -1979,7 +1945,6 @@ class MCPServerManager:
"gmail_send_email": "zapier_mcp_server",
}
"""
self._listed_tools_by_server_id: dict[str, _ListedToolsByCaller] = {} # mutable-ok: refreshed per tools/list
self._upstream_initialize_instructions_by_server_id: dict[str, str] = {}
# Per-server monotonic timestamp of last upstream prefetch attempt (success,
# empty result, or failure). Used to throttle re-probes for servers that do
@ -2675,7 +2640,7 @@ class MCPServerManager:
self._assign_unique_short_prefix(new_server)
_warn_legacy_delegate_auth_if_applicable(new_server, source="config")
_warn_config_id_jag_server_outruns_sso(new_server)
self._invalidate_server_definition_caches(server_id)
self._invalidate_discovery_lists(server_id)
self.config_mcp_servers[server_id] = new_server
self._set_oauth_discovery_deferred(
server_id,
@ -2877,7 +2842,7 @@ class MCPServerManager:
global_mcp_tool_registry,
)
self._invalidate_server_definition_caches(server.server_id)
self._invalidate_discovery_lists(server.server_id)
prefix_root: Final = normalize_server_name(get_server_prefix(server))
if server.spec_path and prefix_root:
openapi_key_prefix: Final = prefix_root + MCP_TOOL_PREFIX_SEPARATOR
@ -3254,7 +3219,7 @@ class MCPServerManager:
# env_vars_are_encrypted=False.
new_server: Final = await self.build_mcp_server_from_table(mcp_server, env_vars_are_encrypted=False)
self._assign_unique_short_prefix(new_server)
self._invalidate_server_definition_caches(mcp_server.server_id)
self._invalidate_discovery_lists(mcp_server.server_id)
self.registry[mcp_server.server_id] = new_server
await self._maybe_register_openapi_tools(new_server)
self.prime_oauth_metadata_discovery(new_server)
@ -3291,7 +3256,7 @@ class MCPServerManager:
previous_server=self.registry[mcp_server.server_id],
)
self._assign_unique_short_prefix(new_server)
self._invalidate_server_definition_caches(mcp_server.server_id)
self._invalidate_discovery_lists(mcp_server.server_id)
self.registry[mcp_server.server_id] = new_server
await self._maybe_register_openapi_tools(new_server)
self.prime_oauth_metadata_discovery(new_server)
@ -3750,7 +3715,19 @@ class MCPServerManager:
verbose_logger.warning("MCP Server %s not found", server_id)
return []
server_auth_header: Final = _server_auth_header_for(server, mcp_server_auth_headers, mcp_auth_header)
# Get server-specific auth header if available
server_auth_header: str | dict[str, str] | None = None
if mcp_server_auth_headers:
server_auth_header = lookup_mcp_server_auth_in_headers(
mcp_server_auth_headers,
alias=server.alias,
server_name=server.server_name,
access_groups=server.access_groups,
)
# Fall back to deprecated mcp_auth_header if no server-specific header found
if server_auth_header is None:
server_auth_header = mcp_auth_header
try:
tools: Final = await self._get_tools_from_server(
@ -3836,7 +3813,7 @@ class MCPServerManager:
def _build_stdio_env(
self,
server: MCPServer,
raw_headers: Mapping[str, str] | None = None,
raw_headers: dict[str, str] | None = None,
) -> dict[str, str] | None:
"""Resolve stdio env values, supporting header-driven placeholders."""
@ -4098,7 +4075,6 @@ class MCPServerManager:
oauth2_headers: dict[str, str] | None,
user_api_key_auth: UserAPIKeyAuth | None,
raw_headers: Mapping[str, str] | None = None,
resource_metadata_url: str | None = None,
) -> None:
"""Mint an exchange-backed server's upstream credential at the transport edge.
@ -4136,9 +4112,7 @@ class MCPServerManager:
if spec is None or not isinstance(spec.config, (TokenExchangeConfig, IdJagConfig)):
return
if subject_token is None and isinstance(spec.config, TokenExchangeConfig):
raise_token_exchange_challenge(
resolved_server, root_path=get_request_root_path(), resource_metadata_url=resource_metadata_url
)
raise_token_exchange_challenge(resolved_server, root_path=get_request_root_path())
match await self._cred_provider.resolve_credentials(to_subject(user_api_key_auth, subject_token), spec):
case Ok(_):
return
@ -4148,7 +4122,6 @@ class MCPServerManager:
resolved_server,
root_path=get_request_root_path(),
claims=err.unauthorized.claims,
resource_metadata_url=resource_metadata_url,
)
raise_public(err)
@ -4365,12 +4338,6 @@ class MCPServerManager:
verbose_logger.info("_get_tools_from_server for %s...", server.name)
client = None
listed_caller: Final = ListedToolsCaller(
user_api_key_auth=user_api_key_auth,
mcp_auth_header=mcp_auth_header,
raw_headers=raw_headers,
oauth2_headers=oauth2_headers,
)
try:
# Tool *listing* must not be blocked by missing per-user env vars —
@ -4447,25 +4414,29 @@ class MCPServerManager:
if server.spec_path:
# OpenAPI tools were stored in the registry under the prefix
# active at registration time — fetch by that same prefix.
registry_prefix: Final = normalize_server_name(get_server_prefix(server)) + MCP_TOOL_PREFIX_SEPARATOR
_tools: Final = global_mcp_tool_registry.list_tools(tool_prefix=registry_prefix)
_tools: Final = global_mcp_tool_registry.list_tools(tool_prefix=get_server_prefix(server))
tools = global_mcp_tool_registry.convert_tools_to_mcp_sdk_tool_type(_tools)
# OpenAPI tools are stored in the registry with their prefix already
# applied (e.g. "test_petstore-getinventory"). Do NOT pass them
# through _create_prefixed_tools — that would add the prefix a second
# time producing "test_petstore-test_petstore-getinventory".
unprefixed_tools: Final = [ # mutable-ok: returned through the list[MCPTool] listing contract
t.model_copy(update=MappingProxyType({"name": t.name[len(registry_prefix) :]})) for t in tools
]
self._record_listed_tools(server, unprefixed_tools, listed_caller)
return tools if add_prefix else unprefixed_tools
if not add_prefix:
prefix: Final = get_server_prefix(server)
sep: Final = MCP_TOOL_PREFIX_SEPARATOR
tools = [
(
t.model_copy(update={"name": t.name[len(prefix) + len(sep) :]})
if t.name.startswith(f"{prefix}{sep}")
else t
)
for t in tools
]
return tools
else:
tools = await self._fetch_tools_with_timeout(client, server.name)
self._remember_upstream_initialize_instructions(server, client)
prefixed_or_original_tools: Final = self._create_prefixed_tools(
tools, server, add_prefix=add_prefix, caller=listed_caller
)
prefixed_or_original_tools: Final = self._create_prefixed_tools(tools, server, add_prefix=add_prefix)
return prefixed_or_original_tools
@ -4509,86 +4480,6 @@ class MCPServerManager:
self._resource_discovery_cache.invalidate(server_id)
self._template_discovery_cache.invalidate(server_id)
def _invalidate_server_definition_caches(self, server_id: str) -> None:
self._invalidate_discovery_lists(server_id)
self._listed_tools_by_server_id.pop(server_id, None)
def _discovers_per_caller(self, server: MCPServer) -> bool:
return (
server.requires_per_user_auth
or self._references_per_user_env_var(server)
or server.delegate_auth_to_upstream
or server.auth_type in (MCPAuth.oauth2_token_exchange, MCPAuth.oauth2_id_jag)
or self._signs_caller_identity_upstream(server)
)
@staticmethod
def _signs_caller_identity_upstream(server: MCPServer) -> bool:
"""Whether MCPJWTSigner mints a per-caller ``Authorization`` for ``server``, so the upstream may
tailor its catalog to the caller even though the server itself is configured as shared."""
from litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer import ( # noqa: PLC0415 # lazy: guardrail package imports the proxy server
get_mcp_jwt_signer,
)
if get_mcp_jwt_signer() is None:
return False
return not any(k.lower() == "authorization" for k in (server.static_headers or {}))
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.
"""
if server.spec_path or caller is None:
return None
auth: Final = caller.user_api_key_auth
identity: Final = (
(auth.user_id, auth.api_key) if auth is not None and self._discovers_per_caller(server) else None
)
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 = (
self._extract_subject_token(caller.oauth2_headers, caller.raw_headers, auth)
if server.is_client_forwarded_token
else None
)
inputs: Final = (identity, caller.mcp_auth_header, forwarded, stdio_env, relayed_bearer)
if not any(inputs):
return None
material: Final = json.dumps(inputs, sort_keys=True, separators=(",", ":"))
return hashlib.sha256(material.encode()).hexdigest()
@staticmethod
def _forwarded_header_values(
server: MCPServer, raw_headers: Mapping[str, str] | None
) -> tuple[tuple[str, str], ...]:
if not raw_headers or not server.extra_headers:
return ()
forwarded_names: Final = frozenset(name.lower() for name in server.extra_headers)
return tuple(
sorted((name.lower(), value) for name, value in raw_headers.items() if name.lower() in forwarded_names)
)
def _record_listed_tools(
self, server: MCPServer, tools: Sequence[MCPTool], caller: ListedToolsCaller | None
) -> None:
identity: Final = self._listed_tools_identity(server, caller)
listing: Final = MappingProxyType({tool.name: tool for tool in tools})
existing: Final = self._listed_tools_by_server_id.get(server.server_id, _NO_LISTED_TOOLS)
shared: Final = existing.get(None)
callers: Final = tuple((key, value) for key, value in existing.items() if key not in (None, identity))
evicted: Final = 0 if identity is None else max(len(callers) + 1 - _LISTED_TOOLS_CALLERS_PER_SERVER, 0)
entries: Final = (
*(() if shared is None else ((None, shared),)),
*callers[evicted:],
(identity, listing),
)
self._listed_tools_by_server_id[server.server_id] = MappingProxyType(dict(entries))
def _discovery_key(
self,
server: MCPServer,
@ -4599,7 +4490,12 @@ class MCPServerManager:
subject_token: str | None,
credential_fingerprint: str | None = None,
) -> _DiscoveryKey:
per_user: Final = self._discovers_per_caller(server)
per_user: Final = (
server.requires_per_user_auth
or self._references_per_user_env_var(server)
or server.delegate_auth_to_upstream
or server.auth_type in (MCPAuth.oauth2_token_exchange, MCPAuth.oauth2_id_jag)
)
if not (per_user or mcp_auth_header or extra_headers or stdio_env or subject_token):
return server.server_id, None
identity: Final = (
@ -5412,13 +5308,7 @@ class MCPServerManager:
"attempts; the 3-character prefix space is too crowded."
)
def _create_prefixed_tools(
self,
tools: list[MCPTool],
server: MCPServer,
add_prefix: bool = True,
caller: ListedToolsCaller | None = None,
) -> list[MCPTool]:
def _create_prefixed_tools(self, tools: list[MCPTool], server: MCPServer, add_prefix: bool = True) -> list[MCPTool]:
"""
Create prefixed tools and update tool mapping.
@ -5450,21 +5340,9 @@ class MCPServerManager:
for spelling in iter_known_tool_name_spellings(original_name, server):
self.tool_name_to_mcp_server_name_mapping[spelling] = prefix
self._record_listed_tools(server, tools, caller)
verbose_logger.info("Successfully fetched %s tools from server %s", len(prefixed_tools), server.name)
return prefixed_tools
def get_listed_tool(self, server: MCPServer, name: str, caller: ListedToolsCaller | None = None) -> MCPTool | None:
identity: Final = self._listed_tools_identity(server, caller)
listed: Final = self._listed_tools_by_server_id.get(server.server_id, _NO_LISTED_TOOLS).get(identity)
if not listed:
return None
tool: Final = listed.get(name) or listed.get(strip_known_server_prefix(name, server))
if tool is None:
return None
description: Final = (server.tool_name_to_description or {}).get(tool.name)
return tool if description is None else tool.model_copy(update={"description": description})
def _create_prefixed_prompts(
self, prompts: Sequence[Prompt], server: MCPServer, add_prefix: bool = True
) -> list[Prompt]:
@ -5703,7 +5581,6 @@ class MCPServerManager:
server: MCPServer,
raw_headers: dict[str, str] | None = None,
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
tool: MCPTool | None = None,
) -> dict[str, Any]:
"""
Run pre-call checks and guardrail hooks for an MCP tool call.
@ -5717,9 +5594,6 @@ class MCPServerManager:
``pre_mcp_call`` evaluation (or a block) on the spend-log row the Guardrails
Monitor counts. It stays optional so callers that do no logging are unchanged.
``tool`` is the upstream tool definition when one was listed, so guardrails
can see its description and input schema, not just the name and arguments.
Returns a dict that may contain:
- "arguments": hook-modified tool arguments (only if changed)
- "extra_headers": headers injected by pre_mcp_call guardrail hooks
@ -5773,8 +5647,6 @@ class MCPServerManager:
"user_api_key_hash": (getattr(user_api_key_auth, "api_key_hash", None) if user_api_key_auth else None),
"incoming_bearer_token": incoming_bearer_token,
"headers": logging_safe_mcp_headers(raw_headers),
"tool_description": tool.description if tool is not None else None,
"tool_input_schema": tool.inputSchema if tool is not None else None,
}
# Create MCP request object for processing
@ -5829,7 +5701,6 @@ class MCPServerManager:
proxy_logging_obj: ProxyLogging,
start_time: datetime.datetime,
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
tool: MCPTool | None = None,
):
"""Create and return a during hook task for MCP tool calls.
@ -5844,8 +5715,6 @@ class MCPServerManager:
tool_name=name,
arguments=arguments,
server_name=server_name_from_prefix,
tool_description=tool.description if tool is not None else None,
tool_input_schema=tool.inputSchema if tool is not None else None,
start_time=start_time.timestamp() if start_time else None,
hidden_params=HiddenParams(),
)
@ -5977,7 +5846,21 @@ class MCPServerManager:
GuardrailRaisedException: If guardrails block the call
HTTPException: If an HTTP error occurs
"""
server_auth_header: Final = _server_auth_header_for(mcp_server, mcp_server_auth_headers, mcp_auth_header)
# Get server-specific auth header if available (case-insensitive)
# FIX: Added case-insensitive matching to handle auth header keys that may not match
# the exact case of server alias/name (e.g., '1litellmagcgateway' vs '1LiteLLMAGCGateway')
server_auth_header: dict[str, str] | str | None = None
if mcp_server_auth_headers:
server_auth_header = lookup_mcp_server_auth_in_headers(
mcp_server_auth_headers,
alias=mcp_server.alias,
server_name=mcp_server.server_name,
access_groups=mcp_server.access_groups,
)
# Fall back to deprecated mcp_auth_header if no server-specific header found
if server_auth_header is None:
server_auth_header = mcp_auth_header
# Extract subject token for OAuth2 Token Exchange (OBO) and ID-JAG flows
subject_token: str | None = None
@ -6413,12 +6296,6 @@ class MCPServerManager:
user_api_key_auth,
mcp_auth_header,
)
listed_caller: Final = ListedToolsCaller(
user_api_key_auth=user_api_key_auth,
mcp_auth_header=_server_auth_header_for(mcp_server, mcp_server_auth_headers, mcp_auth_header),
raw_headers=raw_headers,
oauth2_headers=oauth2_headers,
)
#########################################################
# Pre MCP Tool Call Hook
@ -6434,7 +6311,6 @@ class MCPServerManager:
server=mcp_server,
raw_headers=raw_headers,
litellm_logging_obj=litellm_logging_obj,
tool=self.get_listed_tool(mcp_server, name, listed_caller),
)
if "arguments" in hook_result:
arguments = hook_result["arguments"]
@ -6450,7 +6326,6 @@ class MCPServerManager:
proxy_logging_obj=proxy_logging_obj,
start_time=start_time,
litellm_logging_obj=litellm_logging_obj,
tool=self.get_listed_tool(mcp_server, name, listed_caller),
)
tasks.append(during_hook_task)
@ -6718,7 +6593,7 @@ class MCPServerManager:
for server_id in previous_registry.keys() | registered_registry.keys():
if previous_registry.get(server_id) != registered_registry.get(server_id):
self._invalidate_server_definition_caches(server_id)
self._invalidate_discovery_lists(server_id)
self.registry = registered_registry
# A discovery task may have published into ``previous_registry`` while
# this replacement was being staged. Reconcile every published entry

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

@ -199,7 +199,7 @@ def get_route_relative_request_path(scope: Scope) -> str:
:func:`litellm.proxy.auth.auth_utils.get_request_route`, which the rest of the MCP auth path
already routes through, so ``/litellmfoo`` is not truncated under ``root_path=/litellm``."""
raw_path = str(scope.get("_original_path") or scope.get("path", "") or "")
root_path = str(scope.get("app_root_path", scope.get("root_path")) or "").rstrip("/")
root_path = str(scope.get("app_root_path") or scope.get("root_path") or "").rstrip("/")
if root_path and (raw_path == root_path or raw_path.startswith(f"{root_path}/")):
return raw_path[len(root_path) :]
return raw_path

8
litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py
View file

@ -358,7 +358,6 @@ def raise_token_exchange_challenge(
*,
root_path: str,
claims: str | None = None,
resource_metadata_url: str | None = None,
) -> NoReturn:
"""Raise the RFC 9728 / RFC 6750 challenge an OBO (``token_exchange``) server returns when the
caller's subject token is missing or the IdP rejected it.
@ -376,13 +375,8 @@ def raise_token_exchange_challenge(
``error="invalid_token"`` and is byte-identical to the static one. Both the error value (one of
two literals) and the base64 claims draw from a fixed alphabet, so nothing from the IdP body
reaches the header unescaped.
``resource_metadata_url`` overrides the derived path with the absolute metadata URL matching the
route spelling the request arrived on: RFC 9728 §3.3 clients reject a ``resource`` that differs
from the URL they connected to, and a ``/{server}/mcp`` connect must not be sent to the
``/mcp/{server}`` document.
"""
resource_metadata: Final = resource_metadata_url or oauth_protected_resource_path(root_path, server)
resource_metadata: Final = oauth_protected_resource_path(root_path, server)
encoded_claims: Final = base64.b64encode(claims.encode()).decode() if claims else None
error: Final = "insufficient_claims" if encoded_claims else "invalid_token"
error_description: Final = (

69
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -57,7 +57,6 @@ from litellm.proxy._experimental.mcp_server.mcp_debug import (
from litellm.proxy._experimental.mcp_server.oauth_utils import (
_redact_mcp_resource_url,
get_byok_www_authenticate,
get_passthrough_resource_metadata_url,
get_passthrough_www_authenticate,
get_route_relative_request_path,
well_known_root_suffix,
@ -81,16 +80,12 @@ from litellm.proxy._types import (
UserAPIKeyAuth,
)
from litellm.proxy.auth.ip_address_utils import IPAddressUtils
from litellm.proxy.guardrails.guardrail_hooks.agent_365.agent_365 import (
agent_365_sign_in_required,
)
from litellm.proxy.litellm_pre_call_utils import (
LiteLLMProxyRequestSetup,
get_chain_id_from_headers,
)
from litellm.types.mcp import MCPAuth, MCPSpecVersion
from litellm.types.mcp_server.mcp_server_manager import MCPInfo, MCPServer
from litellm.types.mcp_server.tool_registry import MCPTool as RegisteredTool
from litellm.types.utils import CallTypes, StandardLoggingMCPToolCall
from litellm.utils import Rules, client, function_setup
@ -2782,11 +2777,6 @@ if MCP_AVAILABLE:
return managed_resource_templates
def _registered_tool_metadata(name: str, registered: RegisteredTool, server: MCPServer) -> MCPTool:
overrides: Final = server.tool_name_to_description
description: Final = overrides.get(name, registered.description) if overrides else registered.description
return MCPTool(name=name, description=description, inputSchema=registered.input_schema)
def _resolve_display_name_to_original(
name: str,
allowed_mcp_servers: list[MCPServer],
@ -3125,7 +3115,6 @@ if MCP_AVAILABLE:
server=mcp_server,
raw_headers=raw_headers,
litellm_logging_obj=litellm_logging_obj,
tool=_registered_tool_metadata(original_tool_name, local_tool, mcp_server),
)
# `pre_call_tool_check` may return guardrail-modified
# arguments; honor them on the local path too.
@ -3191,8 +3180,7 @@ if MCP_AVAILABLE:
# not in the registry either, `_handle_local_mcp_tool` below reports
# 404 and nothing runs, so demanding a server here would turn every
# unknown tool name into a misleading 503.
registered_local_tool: Final = global_mcp_tool_registry.get_tool(original_tool_name)
if registered_local_tool is not None:
if global_mcp_tool_registry.get_tool(original_tool_name) is not None:
# `mcp_server` is None here because the tool name is not in the
# tool -> server mapping, but the name still carries a prefix
# that the server-level check above compared against the
@ -3233,7 +3221,6 @@ if MCP_AVAILABLE:
server=prefix_server,
raw_headers=raw_headers,
litellm_logging_obj=litellm_logging_obj,
tool=_registered_tool_metadata(original_tool_name, registered_local_tool, prefix_server),
)
if "arguments" in hook_result:
arguments = hook_result["arguments"] # pyright: ignore[reportAny] # hook returns untyped args
@ -4020,21 +4007,6 @@ if MCP_AVAILABLE:
)
return user_api_key_auth.model_copy(update={"object_permission": updated_op})
async def _key_granted_single_server(
server: MCPServer,
mcp_servers: Sequence[str] | None,
user_api_key_auth: UserAPIKeyAuth | None,
client_ip: str | None,
) -> bool:
"""Sign-in challenges are issued only on a single-server connect the key's grant admits, so a key
without access gets the grant's 403 instead of a sign-in it could not use."""
if len(mcp_servers or []) != 1:
return False
allowed: Final = await _get_allowed_mcp_servers(
user_api_key_auth=user_api_key_auth, mcp_servers=mcp_servers, client_ip=client_ip
)
return any(granted.server_id == server.server_id for granted in allowed)
async def _raise_preemptive_401_for_unauthenticated_servers(
scope: Scope,
mcp_servers: list[str] | None,
@ -4157,24 +4129,8 @@ if MCP_AVAILABLE:
# (transport level, where WWW-Authenticate survives) with the RFC 9728 resource_metadata
# so the client discovers the IdP, SSOs, and retries with a subject token, which LiteLLM
# then exchanges. A tool-call-time 401 would be wrapped into a JSON-RPC error and the
# header lost, so the discovery flow needs this pre-emptive challenge. Servers gated by an
# Agent 365 guardrail (OBO to the evaluate API) get the same challenge, also when the only
# bearer is the LiteLLM key itself, which admits the caller but is not an exchangeable subject,
# and when Entra refuses the presented assertion (expired, wrong audience), so the client
# signs in again instead of failing every tool call. Only on the server's own route: the
# per-server metadata ``resource`` must equal the URL the client connected to (RFC 9728 3.3),
# which aggregate ``/mcp`` and multi-server connects never do.
granted_single_server = server is not None and await _key_granted_single_server(
server, mcp_servers, user_api_key_auth, client_ip
)
if server and (
(server.auth_type == MCPAuth.oauth2_token_exchange and not oauth2_headers)
or (
granted_single_server
and tuple(_get_mcp_servers_in_path(get_route_relative_request_path(scope)) or ()) == (server_name,)
and await agent_365_sign_in_required(server, user_api_key_auth, oauth2_headers)
)
):
# header lost, so the discovery flow needs this pre-emptive challenge.
if server and server.auth_type == MCPAuth.oauth2_token_exchange and not oauth2_headers:
from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( # noqa: PLC0415 # lazy: adapter pulls MCP subgraph
raise_token_exchange_challenge,
)
@ -4182,11 +4138,7 @@ if MCP_AVAILABLE:
get_request_root_path,
)
raise_token_exchange_challenge(
server,
root_path=get_request_root_path(),
resource_metadata_url=get_passthrough_resource_metadata_url(scope=scope, server_name=server_name),
)
raise_token_exchange_challenge(server, root_path=get_request_root_path())
# Exchange-backed modes (token_exchange's OBO mint, id_jag's stored-assertion mint): run
# the exchange here at the transport edge, so a rejected subject raises the RFC 9728
@ -4195,13 +4147,22 @@ if MCP_AVAILABLE:
# and what each mints from. Gated to single-server routes the key may reach; the
# multi-server aggregate keeps absorbing per-server auth failures so one bad server
# cannot 401 the whole connect.
if server and granted_single_server:
if (
server
and len(mcp_servers or []) == 1
and server.server_id
in frozenset(
allowed.server_id
for allowed in await _get_allowed_mcp_servers(
user_api_key_auth=user_api_key_auth, mcp_servers=mcp_servers, client_ip=client_ip
)
)
):
await global_mcp_server_manager.preflight_token_exchange(
server=server,
oauth2_headers=oauth2_headers,
user_api_key_auth=user_api_key_auth,
raw_headers=raw_headers,
resource_metadata_url=get_passthrough_resource_metadata_url(scope=scope, server_name=server_name),
)
# Pass-through OAuth: when the admin has opted a server into

129
litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py
View file

@ -19,10 +19,9 @@ from typing import TYPE_CHECKING, ClassVar, Final, Literal, NoReturn
import httpx
from fastapi import HTTPException
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError
from pydantic import TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.exceptions import Timeout as LitellmTimeout
from litellm.integrations.custom_guardrail import (
@ -36,7 +35,6 @@ from litellm.llms.custom_httpx.http_handler import (
httpxSpecialProvider,
)
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.mcp_server.mcp_server_manager import MCPServer
from litellm.types.proxy.guardrails.guardrail_hooks.agent_365 import (
AGENT_365_PROD_API_BASE,
AGENT_365_PROD_RESOURCE_APP_ID,
@ -51,11 +49,9 @@ if TYPE_CHECKING:
from litellm.types.utils import GuardrailStatus
TOKEN_ENDPOINT_TEMPLATE: Final = "https://login.microsoftonline.com/{tenant_id}/oauth2/v2.0/token"
ENTRA_ISSUER_TEMPLATE: Final = "https://login.microsoftonline.com/{tenant_id}/v2.0"
EVALUATE_PATH: Final = "/agents/tool-evaluation/evaluate"
MCP_SESSION_ID_HEADER: Final = "mcp-session-id"
DEFENDER_STATUS_EVALUATED: Final = "Evaluated"
GATEWAY_SCOPE_TEMPLATE: Final = "api://{client_id}/access_as_user"
_GATEWAY_OWNED_TOKEN_ERRORS: Final = frozenset(
{"invalid_client", "unauthorized_client", "invalid_scope", "invalid_resource"}
)
@ -64,7 +60,6 @@ _GATEWAY_OWNED_TOKEN_ERRORS: Final = frozenset(
_INVALID_ASSERTION_AADSTS_PREFIX: Final = "50027"
_AADSTS_CODES_ADAPTER: Final = TypeAdapter(tuple[int, ...])
_MCP_CALL_TYPES: Final[tuple[str, ...]] = ("mcp_call", "call_mcp_tool")
_TOOL_INPUT_SCHEMA_ADAPTER: Final = TypeAdapter(dict[str, object])
_OBO_CACHE_MAX_ENTRIES: Final = 1000
_DEFAULT_TOKEN_TTL_SECONDS: Final = 3599.0
_TOKEN_EXPIRY_SLACK_SECONDS: Final = 60.0
@ -86,11 +81,10 @@ def _parse_aadsts_codes(raw: object) -> tuple[int, ...]:
return ()
def _parse_tool_input_schema(raw: object) -> Mapping[str, object] | None:
try:
return _TOOL_INPUT_SCHEMA_ADAPTER.validate_python(raw)
except ValidationError:
return None
def entra_assertion(value: object) -> str | None:
"""``value`` when it is a compact JWS, the only bearer shape the OBO exchange accepts as its assertion.
A LiteLLM virtual key, session bearer, or opaque upstream token in ``Authorization`` yields ``None``."""
return value if isinstance(value, str) and value.count(".") == 2 else None
class _DefenderResult(TypedDict, total=False):
@ -105,14 +99,6 @@ class _EvaluateResponse(TypedDict, total=False):
correlationId: ReadOnly[str]
class _ToolReference(BaseModel):
model_config = ConfigDict(frozen=True)
name: str
description: str | None = None
input_schema: Mapping[str, object] | None = Field(default=None, serialization_alias="inputSchema")
class _UnavailableDetail(TypedDict):
error: ReadOnly[str]
message: ReadOnly[str]
@ -411,14 +397,8 @@ class Agent365Guardrail(CustomGuardrail):
arguments: Final = data.get("mcp_arguments")
server_name: Final = str(data.get("mcp_server_name") or "litellm")
agent_id: Final = self.agent_id or user_api_key_dict.key_alias
description: Final = data.get("mcp_tool_description")
tool_reference: Final = _ToolReference(
name=tool_name,
description=description if isinstance(description, str) and description else None,
input_schema=_parse_tool_input_schema(data.get("mcp_tool_input_schema")),
)
payload: Final[dict[str, object]] = { # mutable-ok: JSON body with optional fields added below
"tool": tool_reference.model_dump(by_alias=True, exclude_none=True),
"tool": {"name": tool_name},
"serverName": server_name,
"conversationId": self._resolve_conversation_id(data),
}
@ -458,18 +438,6 @@ class Agent365Guardrail(CustomGuardrail):
return call_id
return str(uuid.uuid4())
async def exchange_rejects_subject(self, assertion: str) -> bool:
"""Whether Entra refuses ``assertion`` as the On-Behalf-Of subject (expired, wrong audience, bad
signature). Gateway credential rejections and endpoint failures answer ``False``: the caller cannot fix
those by signing in again, so the tool call reports them."""
try:
await self._get_obo_token(assertion)
except Agent365TokenExchangeError as exc:
return not exc.gateway_owned
except (Agent365ThrottledError, Agent365MalformedResponseError, httpx.HTTPError, LitellmTimeout, TimeoutError):
return False
return False
async def _get_obo_token(self, assertion: str) -> str:
cache_key: Final = hashlib.sha256(assertion.encode("utf-8")).hexdigest()
now: Final = time.time()
@ -667,88 +635,3 @@ class Agent365Guardrail(CustomGuardrail):
guardrail_provider=self.guardrail_provider,
event_type=GuardrailEventHooks.pre_mcp_call,
)
def _applies_to_caller(guardrail: Agent365Guardrail, user_api_key_auth: "UserAPIKeyAuth") -> bool:
probe: Final[Mapping[str, object]] = {
"metadata": {
"user_api_key_metadata": user_api_key_auth.metadata,
"user_api_key_team_metadata": user_api_key_auth.team_metadata,
}
}
return guardrail.should_run_guardrail(data=dict(probe), event_type=GuardrailEventHooks.pre_mcp_call)
def _applicable_guardrails(
server: MCPServer, user_api_key_auth: "UserAPIKeyAuth | None"
) -> tuple[Agent365Guardrail, ...]:
"""Agent 365 guardrails whose sign-in the gateway advertises for ``server``: the ``default_on`` ones, minus
those the caller's key or team opted out of once the caller is known. A guardrail only a key or policy
selects still enforces at the tool call but never challenges, since the anonymous metadata fetch that
follows a challenge cannot see which key selected it and would advertise the wrong issuer. Only servers
that leave the caller's top-level ``Authorization`` with the gateway qualify: a forwarded API-key header
travels upstream in its own slot and does not displace the Entra assertion."""
if not server.keeps_caller_authorization:
return ()
advertised: Final = tuple(
callback
for callback in litellm.logging_callback_manager.get_custom_loggers_for_type(Agent365Guardrail)
if isinstance(callback, Agent365Guardrail) and callback.default_on
)
if user_api_key_auth is None:
return advertised
return tuple(g for g in advertised if _applies_to_caller(g, user_api_key_auth))
def entra_assertion(value: object) -> str | None:
"""``value`` when it is a compact JWS, the only bearer shape the OBO exchange accepts as its assertion.
A LiteLLM virtual key, session bearer, or opaque upstream token in ``Authorization`` yields ``None``."""
return value if isinstance(value, str) and value.count(".") == 2 else None
def _presented_assertion(oauth2_headers: Mapping[str, str] | None) -> str | None:
authorization: Final = oauth2_headers.get("Authorization", "") if oauth2_headers else ""
if not authorization.lower().startswith("bearer "):
return None
return entra_assertion(authorization[len("bearer ") :].strip())
async def agent_365_sign_in_required(
server: MCPServer, user_api_key_auth: "UserAPIKeyAuth | None", oauth2_headers: Mapping[str, str] | None
) -> bool:
"""Whether the connect must answer with the RFC 9728 sign-in challenge: an Agent 365 guardrail gates
``server`` for this caller and the request carries no Entra assertion, or one Entra will not exchange.
Decided at connect because a tool call's JSON-RPC error cannot carry ``WWW-Authenticate``."""
guardrails: Final = _applicable_guardrails(server, user_api_key_auth)
if not guardrails:
return False
assertion: Final = _presented_assertion(oauth2_headers)
if assertion is None:
return True
for guardrail in guardrails:
if await guardrail.exchange_rejects_subject(assertion):
return True
return False
def agent_365_authorization_servers(server: MCPServer, user_api_key_auth: "UserAPIKeyAuth | None") -> tuple[str, ...]:
"""Entra issuers an MCP client signs in with before calling ``server`` through an Agent 365 guardrail."""
return tuple(
dict.fromkeys(
ENTRA_ISSUER_TEMPLATE.format(tenant_id=g.tenant_id)
for g in _applicable_guardrails(server, user_api_key_auth)
)
)
def agent_365_scopes_supported(server: MCPServer, user_api_key_auth: "UserAPIKeyAuth | None") -> tuple[str, ...]:
"""Scopes the client requests from Entra for ``server``: the admin's ``scopes`` when set, otherwise the
``access_as_user`` scope of each gating guardrail's gateway app registration (``api://<client_id>``)."""
if server.scopes:
return tuple(server.scopes)
return tuple(
dict.fromkeys(
GATEWAY_SCOPE_TEMPLATE.format(client_id=g.client_id)
for g in _applicable_guardrails(server, user_api_key_auth)
)
)

4
litellm/proxy/utils.py
View file

@ -1245,8 +1245,6 @@ class ProxyLogging:
"user_api_key_request_route": kwargs.get("user_api_key_request_route"),
"mcp_tool_name": request_obj.tool_name, # Keep original for reference
"mcp_arguments": request_obj.arguments, # Keep original for reference
"mcp_tool_description": request_obj.tool_description,
"mcp_tool_input_schema": request_obj.tool_input_schema,
# Surface the per-MCP-server rate-limit identity so the
# ParallelRequestLimiterV3 hook can apply mcp_rpm_limit on the
# synthetic call_mcp_tool payload (otherwise a key with
@ -1468,8 +1466,6 @@ class ProxyLogging:
tool_name=kwargs.get("name", ""),
arguments=kwargs.get("arguments", {}),
server_name=kwargs.get("server_name"),
tool_description=kwargs.get("tool_description"),
tool_input_schema=kwargs.get("tool_input_schema"),
user_api_key_auth=user_api_key_auth_dict,
hidden_params=HiddenParams(),
)

4
litellm/types/mcp.py
View file

@ -391,8 +391,6 @@ class MCPPreCallRequestObject(BaseModel):
tool_name: str
arguments: dict[str, Any]
server_name: str | None = None
tool_description: str | None = None
tool_input_schema: Mapping[str, object] | None = None
user_api_key_auth: dict[str, Any] | None = None
hidden_params: HiddenParams = HiddenParams()
@ -416,8 +414,6 @@ class MCPDuringCallRequestObject(BaseModel):
tool_name: str
arguments: dict[str, Any]
server_name: str | None = None
tool_description: str | None = None
tool_input_schema: Mapping[str, object] | None = None
start_time: float | None = None
hidden_params: HiddenParams = HiddenParams()

20
litellm/types/mcp_server/mcp_server_manager.py
View file

@ -276,10 +276,10 @@ class MCPServer(BaseModel):
return self.per_server_oauth_discovery and self.auth_type == MCPAuth.oauth2 and not self.has_client_credentials
@property
def keeps_caller_authorization(self) -> bool:
"""Whether the caller's top-level ``Authorization`` stays with the gateway: the server neither relays
it upstream nor runs an OAuth mode that fills that slot itself, so a gateway guardrail may consume it
as the caller's own assertion. Forwarding a separate API-key header leaves the slot untouched."""
def advertises_gateway_authorization_server(self) -> bool:
"""Whether named discovery should advertise the aggregate gateway authorization server."""
if self.auth_type == MCPAuth.oauth2:
return self.is_gateway_managed_oauth2 and not self.uses_per_server_oauth_relay
if self.auth_type not in (
None,
MCPAuth.none,
@ -291,15 +291,9 @@ class MCPServer(BaseModel):
MCPAuth.aws_sigv4,
):
return False
return not any(header.lower() == "authorization" for header in (self.extra_headers or ()))
@property
def advertises_gateway_authorization_server(self) -> bool:
"""Whether named discovery should advertise the aggregate gateway authorization server."""
if self.auth_type == MCPAuth.oauth2:
return self.is_gateway_managed_oauth2 and not self.uses_per_server_oauth_relay
return self.keeps_caller_authorization and not any(
header.lower() in ("x-api-key", "api-key", "apikey") for header in (self.extra_headers or ())
return not any(
header.lower() in ("authorization", "x-api-key", "api-key", "apikey")
for header in (self.extra_headers or ())
)
@property

16
tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py
View file

@ -584,22 +584,6 @@ def test_raise_token_exchange_challenge_is_rfc9728_invalid_token():
assert "error_description=" in www
def test_raise_token_exchange_challenge_explicit_resource_metadata_url_wins():
from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import (
raise_token_exchange_challenge,
)
with pytest.raises(HTTPException) as exc_info:
raise_token_exchange_challenge(
_server(alias="obo-srv"),
root_path="/",
resource_metadata_url="https://gw.example.com/.well-known/oauth-protected-resource/obo-srv/mcp",
)
www = exc_info.value.headers["WWW-Authenticate"]
assert 'resource_metadata="https://gw.example.com/.well-known/oauth-protected-resource/obo-srv/mcp"' in www
assert "/mcp/obo-srv" not in www
def test_raise_token_exchange_challenge_includes_server_root_path(monkeypatch):
from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import (
raise_token_exchange_challenge,

102
tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
View file

@ -10,7 +10,6 @@ from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import HTTPException
from fastapi.encoders import jsonable_encoder
from litellm.types.mcp import MCPAuth
@ -7125,107 +7124,6 @@ async def test_build_oauth_protected_resource_response_obo_end_to_end():
global_mcp_server_manager.registry.clear()
@pytest.fixture
def agent_365_guardrail():
import litellm
from litellm.proxy.guardrails.guardrail_hooks.agent_365 import Agent365Guardrail
guardrail = Agent365Guardrail(
guardrail_name="agent-365-guard",
tenant_id="tenant-abc",
client_id="client-xyz",
client_secret="secret-123",
async_handler=AsyncMock(),
event_hook="pre_mcp_call",
default_on=True,
)
litellm.logging_callback_manager.add_litellm_callback(guardrail)
try:
yield guardrail
finally:
litellm.logging_callback_manager.remove_callback_from_list_by_object(
litellm.callbacks, guardrail, require_self=False
)
async def _agent_365_gated_prm(scopes, extra_headers=None):
from fastapi import Request
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
_build_oauth_protected_resource_response,
)
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
from litellm.proxy._types import MCPTransport
from litellm.types.mcp import MCPAuth
from litellm.types.mcp_server.mcp_server_manager import MCPServer
global_mcp_server_manager.registry.clear()
global_mcp_server_manager.registry["tools"] = MCPServer(
server_id="tools",
name="tools",
server_name="tools",
alias="tools",
transport=MCPTransport.http,
auth_type=MCPAuth.none,
scopes=scopes,
extra_headers=extra_headers,
)
mock_request = MagicMock(spec=Request)
mock_request.base_url = "https://litellm.example.com/"
mock_request.headers = {}
try:
return await _build_oauth_protected_resource_response(
request=mock_request, mcp_server_name="tools", use_standard_pattern=True
)
finally:
global_mcp_server_manager.registry.clear()
@pytest.mark.asyncio
async def test_agent_365_gated_server_prm_names_the_entra_tenant(agent_365_guardrail):
response = await _agent_365_gated_prm(scopes=["api://gateway-app/access_as_user"])
assert jsonable_encoder(response) == {
"authorization_servers": ["https://login.microsoftonline.com/tenant-abc/v2.0"],
"resource": "https://litellm.example.com/mcp/tools",
"scopes_supported": ["api://gateway-app/access_as_user"],
}
@pytest.mark.asyncio
async def test_agent_365_prm_defaults_scopeless_server_to_the_gateway_app_scope(agent_365_guardrail):
response = await _agent_365_gated_prm(scopes=None)
assert jsonable_encoder(response) == {
"authorization_servers": ["https://login.microsoftonline.com/tenant-abc/v2.0"],
"resource": "https://litellm.example.com/mcp/tools",
"scopes_supported": ["api://client-xyz/access_as_user"],
}
@pytest.mark.asyncio
async def test_agent_365_prm_survives_a_forwarded_upstream_api_key_header(agent_365_guardrail):
"""The forwarded ``x-api-key`` is the upstream's credential and rides in its own header, so the caller's
``Authorization`` still carries the Entra assertion and discovery must keep naming the Entra tenant."""
response = await _agent_365_gated_prm(scopes=None, extra_headers=["x-api-key"])
assert jsonable_encoder(response) == {
"authorization_servers": ["https://login.microsoftonline.com/tenant-abc/v2.0"],
"resource": "https://litellm.example.com/mcp/tools",
"scopes_supported": ["api://client-xyz/access_as_user"],
}
@pytest.mark.asyncio
async def test_key_selected_agent_365_guardrail_leaves_anonymous_prm_on_the_gateway_issuer(agent_365_guardrail):
"""A default-off guardrail gates only the keys that select it, so the anonymous discovery fetch must keep
pointing every other client at the gateway's own authorization server and scopes."""
agent_365_guardrail.default_on = False
response = await _agent_365_gated_prm(scopes=["mcp:read"])
assert jsonable_encoder(response) == {
"authorization_servers": ["https://litellm.example.com/mcp"],
"resource": "https://litellm.example.com/mcp/tools",
"scopes_supported": ["mcp:read"],
}
def _token_request(headers):
"""A real Starlette request with case-insensitive headers (matches production)."""
from starlette.requests import Request

1
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_guardrail_usage_monitor.py
View file

@ -48,7 +48,6 @@ def _bare_manager() -> MOD.MCPServerManager:
reaches the guardrail hooks; they have their own coverage elsewhere.
"""
mgr = MOD.MCPServerManager.__new__(MOD.MCPServerManager)
mgr._listed_tools_by_server_id = {}
mgr.check_allowed_or_banned_tools = lambda name, server: True
mgr.validate_allowed_params = lambda tool_name, arguments, server: None

447
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
View file

@ -6,7 +6,6 @@ from types import SimpleNamespace
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
from fastapi import HTTPException
from mcp import ReadResourceResult, Resource
@ -7053,12 +7052,9 @@ async def test_execute_mcp_tool_sets_model_in_model_call_details():
fake_server.server_name = "openapi-petstore"
fake_server.alias = None
fake_server.short_prefix = None
fake_server.tool_name_to_description = None
fake_tool = MagicMock()
fake_tool.name = "list_pets"
fake_tool.description = "test tool"
fake_tool.input_schema = {"type": "object"}
start_time = datetime.now(timezone.utc)
litellm_logging_obj, _ = function_setup(
@ -7109,141 +7105,6 @@ async def test_execute_mcp_tool_sets_model_in_model_call_details():
assert litellm_logging_obj.model == "MCP: list_pets"
@pytest.mark.asyncio
async def test_execute_mcp_tool_hands_openapi_registered_tool_metadata_to_pre_call_hooks():
"""OpenAPI-generated tools dispatch through the local registry, so the pre-call hooks must get the
registered description and input schema on that path too, even when no tools/list ran first."""
from litellm.proxy._experimental.mcp_server import server as mcp_module
petstore = MCPServer(
server_id="petstore-id",
name="petstore",
server_name="petstore",
transport=MCPTransport.http,
url=None,
spec_path="https://example.com/petstore.yaml",
)
schema = {"type": "object", "properties": {"limit": {"type": "integer"}}}
mcp_module.global_mcp_tool_registry.register_tool(
name="petstore-list_pets", description="List the pets", input_schema=schema, handler=lambda limit: "ok"
)
manager = mcp_module.global_mcp_server_manager
manager._listed_tools_by_server_id.pop(petstore.server_id, None)
pre_call_tool_check = AsyncMock(return_value={})
try:
with (
patch.object(manager, "_get_mcp_server_from_tool_name", return_value=petstore),
patch.object(manager, "pre_call_tool_check", new=pre_call_tool_check),
):
await mcp_module.execute_mcp_tool(
name="petstore-list_pets",
arguments={"limit": 10},
allowed_mcp_servers=[petstore],
start_time=datetime.now(),
user_api_key_auth=UserAPIKeyAuth(api_key="sk-user", user_id="alice"),
)
finally:
mcp_module.global_mcp_tool_registry.unregister_tools_with_prefix("petstore-")
handed_tool = pre_call_tool_check.call_args.kwargs["tool"]
assert (handed_tool.name, handed_tool.description, handed_tool.inputSchema) == (
"list_pets",
"List the pets",
schema,
)
@pytest.mark.asyncio
async def test_execute_mcp_tool_hands_openapi_hooks_the_admin_description_clients_saw():
"""tools/list shows the admin's tool_name_to_description wording, so the local-registry call path
must hand the pre-call hooks that same wording rather than the generated one."""
from litellm.proxy._experimental.mcp_server import server as mcp_module
petstore = MCPServer(
server_id="petstore-id",
name="petstore",
server_name="petstore",
transport=MCPTransport.http,
url=None,
spec_path="https://example.com/petstore.yaml",
tool_name_to_description={"getpetbyid": "ADMIN DESC"},
)
schema = {"type": "object", "properties": {"petId": {"type": "integer"}}}
mcp_module.global_mcp_tool_registry.register_tool(
name="petstore-getpetbyid", description="Find pet by ID", input_schema=schema, handler=lambda petId: "ok"
)
manager = mcp_module.global_mcp_server_manager
manager._listed_tools_by_server_id.pop(petstore.server_id, None)
pre_call_tool_check = AsyncMock(return_value={})
try:
with (
patch.object(manager, "_get_mcp_server_from_tool_name", return_value=petstore),
patch.object(manager, "pre_call_tool_check", new=pre_call_tool_check),
):
await mcp_module.execute_mcp_tool(
name="petstore-getpetbyid",
arguments={"petId": 1},
allowed_mcp_servers=[petstore],
start_time=datetime.now(),
user_api_key_auth=UserAPIKeyAuth(api_key="sk-user", user_id="alice"),
)
finally:
mcp_module.global_mcp_tool_registry.unregister_tools_with_prefix("petstore-")
handed_tool = pre_call_tool_check.call_args.kwargs["tool"]
assert (handed_tool.description, handed_tool.inputSchema) == ("ADMIN DESC", schema)
@pytest.mark.asyncio
async def test_execute_mcp_tool_hands_hooks_the_metadata_of_the_operation_it_runs_when_names_collide():
"""An OpenAPI operation whose name starts with its own server prefix must not be reported to the
pre-call hooks with the metadata of the shorter operation, since that is not the one that runs."""
from litellm.proxy._experimental.mcp_server import server as mcp_module
petstore = MCPServer(
server_id="petstore-id",
name="petstore",
server_name="petstore",
transport=MCPTransport.http,
url=None,
spec_path="https://example.com/petstore.yaml",
)
registry = mcp_module.global_mcp_tool_registry
registry.register_tool(name="petstore-get_pet", description="short", input_schema={}, handler=lambda: "short")
registry.register_tool(
name="petstore-petstore-get_pet",
description="long",
input_schema={"type": "object", "properties": {"petId": {"type": "integer"}}},
handler=lambda: "long",
)
manager = mcp_module.global_mcp_server_manager
pre_call_tool_check = AsyncMock(return_value={})
try:
with (
patch.object(manager, "_get_mcp_server_from_tool_name", return_value=petstore),
patch.object(manager, "pre_call_tool_check", new=pre_call_tool_check),
):
result = await mcp_module.execute_mcp_tool(
name="petstore-petstore-get_pet",
arguments={},
allowed_mcp_servers=[petstore],
start_time=datetime.now(),
user_api_key_auth=UserAPIKeyAuth(api_key="sk-user", user_id="alice"),
)
finally:
registry.unregister_tools_with_prefix("petstore-")
handed_tool = pre_call_tool_check.call_args.kwargs["tool"]
assert (handed_tool.description, handed_tool.inputSchema) == (
"long",
{"type": "object", "properties": {"petId": {"type": "integer"}}},
)
assert result.content[0].text == "long"
@pytest.mark.asyncio
async def test_execute_mcp_tool_rest_unresolved_prefixed_name_routes_to_requested_server():
"""A prefixed REST name that resolves to no tool must still dispatch to the server_id.
@ -8911,308 +8772,6 @@ class TestSingleServerPreflightReachesIdJag:
preflight.assert_not_awaited()
class TestAgent365ChallengeAtConnect:
"""A missing Entra bearer on an Agent 365 gated server is challenged at connect (RFC 9728), where the
WWW-Authenticate header survives, instead of only inside the tools/call JSON-RPC error."""
GATEWAY_SCOPE = "api://gateway-app/access_as_user"
ENTRA_BEARER = {"Authorization": "Bearer eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1LTEifQ.c2ln"}
@staticmethod
def _entra_response(status_code: int, body: dict[str, object]) -> httpx.Response:
return httpx.Response(
status_code=status_code,
json=body,
request=httpx.Request("POST", "https://login.microsoftonline.com/tenant-abc/oauth2/v2.0/token"),
)
def _server(self, scopes: list[str] | None, extra_headers: list[str] | None = None) -> MCPServer:
return MCPServer(
server_id="id-tools",
name="tools",
alias="tools",
server_name="tools",
url="https://tools.test/mcp",
transport=MCPTransport.http,
auth_type=MCPAuth.none,
scopes=scopes,
extra_headers=extra_headers,
mcp_info={"server_name": "tools"},
)
@pytest.fixture
def agent_365_guardrail(self):
import litellm
from litellm.proxy.guardrails.guardrail_hooks.agent_365 import Agent365Guardrail
handler = AsyncMock()
handler.post.return_value = self._entra_response(200, {"access_token": "obo-token", "expires_in": 3599})
guardrail = Agent365Guardrail(
guardrail_name="agent-365-guard",
tenant_id="tenant-abc",
client_id="client-xyz",
client_secret="secret-123",
async_handler=handler,
event_hook="pre_mcp_call",
default_on=True,
)
litellm.logging_callback_manager.add_litellm_callback(guardrail)
try:
yield guardrail
finally:
litellm.logging_callback_manager.remove_callback_from_list_by_object(
litellm.callbacks, guardrail, require_self=False
)
async def _connect(
self,
server: MCPServer,
oauth2_headers: dict[str, str] | None,
path: str = "/mcp/tools",
mount_scope: dict[str, str] | None = None,
granted: bool = True,
key_metadata: dict[str, object] | None = None,
) -> HTTPException | None:
from litellm.proxy._experimental.mcp_server import server as server_module
with (
patch.object( # test-quality-ok: route wiring must use the manager's configured server
server_module.global_mcp_server_manager, "get_mcp_server_by_name", return_value=server
),
patch.object( # test-quality-ok: allowed-set resolution needs the DB; the test controls its answer
server_module, "_get_allowed_mcp_servers", AsyncMock(return_value=[server] if granted else [])
),
):
try:
await server_module._raise_preemptive_401_for_unauthenticated_servers(
scope={
"type": "http",
"method": "POST",
"path": path,
"scheme": "https",
"server": ("gw.example.com", 443),
"headers": [],
**(mount_scope or {}),
},
mcp_servers=["tools"],
oauth2_headers=oauth2_headers,
mcp_server_auth_headers=None,
user_api_key_auth=UserAPIKeyAuth(
api_key="sk-litellm-virtual-key", user_id="u-1", metadata=key_metadata or {}
),
client_ip=None,
)
except HTTPException as challenge:
return challenge
return None
@pytest.mark.asyncio
async def test_no_bearer_gets_the_discovery_challenge(self, agent_365_guardrail):
challenge = await self._connect(self._server([self.GATEWAY_SCOPE]), None)
assert challenge is not None and challenge.status_code == 401
www_authenticate = (challenge.headers or {}).get("WWW-Authenticate", "")
assert 'error="invalid_token"' in www_authenticate
assert (
'resource_metadata="https://gw.example.com/.well-known/oauth-protected-resource/mcp/tools"'
in www_authenticate
)
@pytest.mark.asyncio
async def test_legacy_route_challenge_points_at_its_own_metadata(self, agent_365_guardrail):
"""RFC 9728 3.3: the metadata's ``resource`` must equal the URL the client connected to, so a
``/{server}/mcp`` connect is sent to the ``/{server}/mcp`` document, not the ``/mcp/{server}`` one."""
challenge = await self._connect(self._server([self.GATEWAY_SCOPE]), None, path="/tools/mcp")
assert challenge is not None and challenge.status_code == 401
www_authenticate = (challenge.headers or {}).get("WWW-Authenticate", "")
assert (
'resource_metadata="https://gw.example.com/.well-known/oauth-protected-resource/tools/mcp"'
in www_authenticate
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"server_root, mount_scope",
[
("", {"root_path": "/mcp", "app_root_path": ""}),
("/litellm", {"root_path": "/litellm/mcp", "app_root_path": "/litellm"}),
],
)
async def test_mounted_standard_route_is_challenged(self, agent_365_guardrail, server_root, mount_scope):
"""``/mcp/{server}`` is served by the ``/mcp`` Mount, which moves the mount prefix into
``root_path`` and leaves the app root (empty or SERVER_ROOT_PATH) in ``app_root_path``."""
with patch.dict(os.environ, {"SERVER_ROOT_PATH": server_root}):
challenge = await self._connect(
self._server([self.GATEWAY_SCOPE]), None, path=f"{server_root}/mcp/tools", mount_scope=mount_scope
)
assert challenge is not None and challenge.status_code == 401
www_authenticate = (challenge.headers or {}).get("WWW-Authenticate", "")
assert (
f'resource_metadata="https://gw.example.com{server_root}'
f'/.well-known/oauth-protected-resource{server_root}/mcp/tools"' in www_authenticate
)
@pytest.mark.asyncio
@pytest.mark.parametrize("path", ["/mcp", "/mcp/tools,other"])
async def test_aggregate_route_is_not_challenged_at_connect(self, agent_365_guardrail, path):
"""The per-server metadata's ``resource`` can never equal the aggregate ``/mcp`` URL the client
connected to (RFC 9728 3.3), and one guarded server must not 401 a multi-server connect, so the
Agent 365 challenge is left to tools/call there."""
assert await self._connect(self._server([self.GATEWAY_SCOPE]), None, path=path) is None
@pytest.mark.asyncio
async def test_entra_assertion_present_connects(self, agent_365_guardrail):
assert await self._connect(self._server([self.GATEWAY_SCOPE]), self.ENTRA_BEARER) is None
@pytest.mark.asyncio
@pytest.mark.parametrize(
"entra_body",
[
{"error": "invalid_grant", "error_description": "AADSTS700084: The refresh token was issued..."},
{
"error": "invalid_client",
"error_description": "AADSTS5002723: Invalid JWT token.",
"error_codes": [5002723],
},
],
ids=["expired-or-wrong-audience", "forged-reported-as-invalid_client"],
)
async def test_assertion_entra_refuses_to_exchange_is_challenged(self, agent_365_guardrail, entra_body):
"""An expired, wrong-audience, or forged Entra token looks like a valid one. Only the OBO exchange
can tell, and its verdict must arrive at connect, where WWW-Authenticate reaches the client,
rather than inside every tools/call JSON-RPC error. Entra files a forged assertion under
``invalid_client`` with an AADSTS50027xx sub-code, which must not read as a gateway secret problem."""
agent_365_guardrail.async_handler.post.return_value = self._entra_response(400, entra_body)
challenge = await self._connect(self._server([self.GATEWAY_SCOPE]), self.ENTRA_BEARER)
assert challenge is not None and challenge.status_code == 401
www_authenticate = (challenge.headers or {}).get("WWW-Authenticate", "")
assert 'error="invalid_token"' in www_authenticate
assert (
'resource_metadata="https://gw.example.com/.well-known/oauth-protected-resource/mcp/tools"'
in www_authenticate
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"entra_outcome",
[
{
"return_value": _entra_response(
401,
{
"error": "invalid_client",
"error_description": "AADSTS7000215: Invalid client secret",
"error_codes": [7000215],
},
)
},
{"return_value": _entra_response(503, {"error": "temporarily_unavailable"})},
{"side_effect": httpx.ConnectError("dns")},
],
ids=["gateway-credentials-rejected", "entra-5xx", "entra-unreachable"],
)
async def test_gateway_side_exchange_failures_do_not_send_the_caller_to_sign_in(
self, agent_365_guardrail, entra_outcome
):
"""Signing in again cannot fix the gateway's own client secret or an Entra outage, so those stay
with the tool call, which reports them as the guardrail's unavailable path."""
agent_365_guardrail.async_handler.post = AsyncMock(**entra_outcome)
assert await self._connect(self._server([self.GATEWAY_SCOPE]), self.ENTRA_BEARER) is None
@pytest.mark.asyncio
async def test_key_selected_default_off_guardrail_never_challenges(self):
"""The anonymous metadata fetch that follows a challenge cannot see which key selected the guardrail
and would advertise the gateway issuer, so a challenge here would send the client to the wrong IdP.
The guardrail still enforces at tools/call."""
import litellm
from litellm.proxy.guardrails.guardrail_hooks.agent_365 import Agent365Guardrail
guardrail = Agent365Guardrail(
guardrail_name="agent-365-guard",
tenant_id="tenant-abc",
client_id="client-xyz",
client_secret="secret-123",
async_handler=AsyncMock(),
event_hook="pre_mcp_call",
default_on=False,
)
litellm.logging_callback_manager.add_litellm_callback(guardrail)
try:
with patch( # test-quality-ok: key-selected guardrails read the proxy server premium global, no injection seam
"litellm.proxy.proxy_server.premium_user", True
):
assert (
await self._connect(
self._server([self.GATEWAY_SCOPE]),
None,
key_metadata={"guardrails": ["agent-365-guard"]},
)
is None
)
finally:
litellm.logging_callback_manager.remove_callback_from_list_by_object(
litellm.callbacks, guardrail, require_self=False
)
@pytest.mark.asyncio
async def test_litellm_key_in_authorization_is_still_challenged(self, agent_365_guardrail):
"""A LiteLLM virtual key admits the caller but is no Entra assertion, so the tools/call would fail
401 inside JSON-RPC with the WWW-Authenticate header lost. The connect must challenge instead."""
challenge = await self._connect(
self._server([self.GATEWAY_SCOPE]), {"Authorization": "Bearer sk-litellm-virtual-key"}
)
assert challenge is not None and challenge.status_code == 401
www_authenticate = (challenge.headers or {}).get("WWW-Authenticate", "")
assert 'error="invalid_token"' in www_authenticate
assert (
'resource_metadata="https://gw.example.com/.well-known/oauth-protected-resource/mcp/tools"'
in www_authenticate
)
@pytest.mark.asyncio
async def test_scopeless_server_is_still_challenged(self, agent_365_guardrail):
challenge = await self._connect(self._server(None), None)
assert challenge is not None and challenge.status_code == 401
assert 'error="invalid_token"' in (challenge.headers or {}).get("WWW-Authenticate", "")
@pytest.mark.asyncio
async def test_server_forwarding_an_upstream_api_key_header_is_still_challenged(self, agent_365_guardrail):
"""``x-api-key`` travels upstream in its own header and leaves the caller's ``Authorization`` free for
the Entra assertion, so a key-only connect must still be sent to sign in."""
challenge = await self._connect(self._server(None, extra_headers=["x-api-key"]), None)
assert challenge is not None and challenge.status_code == 401
www_authenticate = (challenge.headers or {}).get("WWW-Authenticate", "")
assert 'error="invalid_token"' in www_authenticate
assert (
'resource_metadata="https://gw.example.com/.well-known/oauth-protected-resource/mcp/tools"'
in www_authenticate
)
@pytest.mark.asyncio
async def test_server_relaying_the_caller_authorization_is_not_challenged(self, agent_365_guardrail):
"""Forwarding ``Authorization`` hands the caller's bearer to the upstream, so the gateway holds no Entra
assertion of its own to exchange and must not advertise a sign-in it cannot consume."""
assert await self._connect(self._server(None, extra_headers=["Authorization"]), None) is None
@pytest.mark.asyncio
async def test_no_registered_guardrail_means_no_challenge(self):
assert await self._connect(self._server([self.GATEWAY_SCOPE]), None) is None
@pytest.mark.asyncio
async def test_key_without_the_server_grant_is_not_sent_to_sign_in(self, agent_365_guardrail):
"""Signing in cannot earn a key a server it was never granted, so the connect must fall through to
the ordinary 403 grant denial instead of leading with an Entra challenge the caller cannot use."""
assert await self._connect(self._server([self.GATEWAY_SCOPE]), None, granted=False) is None
def _make_obo_server(alias: str) -> MCPServer:
return MCPServer(
server_id=f"id-{alias}",
@ -9284,11 +8843,7 @@ class TestOboPreflightScopedToAllowedServers:
_, preflight = await self._run(requested, allowed=[requested], user_api_key_auth=key)
preflight.assert_awaited_once_with(
server=requested,
oauth2_headers=self.SUBJECT_HEADERS,
user_api_key_auth=key,
raw_headers=None,
resource_metadata_url="/.well-known/oauth-protected-resource/mcp/obo_tools",
server=requested, oauth2_headers=self.SUBJECT_HEADERS, user_api_key_auth=key, raw_headers=None
)

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

@ -34,7 +34,6 @@ from pydantic import AnyUrl, TypeAdapter
from litellm.constants import MCP_METADATA_TIMEOUT
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
ListedToolsCaller,
MCPServerManager,
_deserialize_json_dict,
_flow_endpoints_missing,
@ -5092,36 +5091,6 @@ class TestMCPServerManager:
)
assert server2.requires_per_user_auth is False
@pytest.mark.parametrize(
"auth_type, extra_headers, keeps_authorization, advertises_gateway",
[
(MCPAuth.none, None, True, True),
(MCPAuth.api_key, None, True, True),
(MCPAuth.none, ["x-api-key"], True, False),
(MCPAuth.none, ["API-Key"], True, False),
(MCPAuth.none, ["Authorization"], False, False),
(MCPAuth.none, ["x-api-key", "authorization"], False, False),
(MCPAuth.oauth_delegate, None, False, False),
(MCPAuth.true_passthrough, None, False, False),
(MCPAuth.oauth2_token_exchange, None, False, False),
],
)
def test_forwarded_api_key_header_keeps_caller_authorization_but_not_gateway_discovery(
self, auth_type, extra_headers, keeps_authorization, advertises_gateway
):
"""A forwarded API-key header is the upstream's own credential and leaves the caller's top-level
``Authorization`` with the gateway, while still ruling out the gateway's aggregate OAuth discovery."""
server = MCPServer(
server_id="s",
name="s",
transport=MCPTransport.http,
auth_type=auth_type,
url="http://s.example",
extra_headers=extra_headers,
)
assert server.keeps_caller_authorization is keeps_authorization
assert server.advertises_gateway_authorization_server is advertises_gateway
@pytest.mark.asyncio
async def test_register_openapi_tools_includes_static_headers(self, tmp_path):
"""Ensure OpenAPI-to-MCP tool calls include server.static_headers (Issue #19341)."""
@ -6602,465 +6571,6 @@ class TestMCPServerManager:
# Verify the MCP client call was awaited exactly once
assert mock_client.call_tool.await_count == 1
@staticmethod
def _manager_ready_for_call_tool(listed_tools: list[MCPTool]) -> tuple[MCPServerManager, MagicMock]:
from mcp.types import CallToolResult
manager = MCPServerManager()
server = MCPServer(
server_id="test-server",
name="test-server",
transport=MCPTransport.http,
url="http://test-server.com",
)
manager.registry = {"test-server": server}
manager.tool_name_to_mcp_server_name_mapping["test_tool"] = "test-server"
manager.tool_name_to_mcp_server_name_mapping["test-server-test_tool"] = "test-server"
manager._create_prefixed_tools(listed_tools, server)
mock_client = AsyncMock()
mock_client.call_tool.return_value = MagicMock(spec=CallToolResult, content=[], isError=False)
manager._create_mcp_client = AsyncMock(return_value=mock_client)
proxy_logging_obj = MagicMock()
proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock(return_value={})
proxy_logging_obj._convert_mcp_to_llm_format = MagicMock(return_value={})
proxy_logging_obj.pre_call_hook = AsyncMock(return_value={})
proxy_logging_obj.during_call_hook = AsyncMock(return_value=None)
return manager, proxy_logging_obj
@staticmethod
def _unrestricted_auth() -> MagicMock:
user_api_key_auth = MagicMock()
user_api_key_auth.object_permission = None
user_api_key_auth.object_permission_id = None
return user_api_key_auth
@pytest.mark.asyncio
async def test_call_tool_hands_listed_tool_description_and_schema_to_pre_call_hooks(self):
schema = {"type": "object", "properties": {"param": {"type": "string"}}, "required": ["param"]}
listed = [MCPTool(name="test_tool", description="Runs the test tool", inputSchema=schema)]
manager, proxy_logging_obj = self._manager_ready_for_call_tool(listed)
await manager.call_tool(
server_name="test-server",
name="test_tool",
arguments={"param": "value"},
user_api_key_auth=self._unrestricted_auth(),
proxy_logging_obj=proxy_logging_obj,
)
hook_kwargs = proxy_logging_obj._create_mcp_request_object_from_kwargs.call_args.args[0]
assert (hook_kwargs["tool_description"], hook_kwargs["tool_input_schema"]) == ("Runs the test tool", schema)
@pytest.mark.asyncio
async def test_call_tool_hands_listed_tool_metadata_to_during_call_hooks_through_real_conversion(self):
schema = {"type": "object", "properties": {"param": {"type": "string"}}}
listed = [MCPTool(name="test_tool", description="Runs the test tool", inputSchema=schema)]
manager, _ = self._manager_ready_for_call_tool(listed)
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
proxy_logging_obj.pre_call_hook = AsyncMock(return_value={})
proxy_logging_obj.during_call_hook = AsyncMock(return_value=None)
await manager.call_tool(
server_name="test-server",
name="test_tool",
arguments={"param": "value"},
user_api_key_auth=UserAPIKeyAuth(api_key="sk-test"),
proxy_logging_obj=proxy_logging_obj,
)
during_data = proxy_logging_obj.during_call_hook.call_args.kwargs["data"]
assert (during_data["mcp_tool_description"], during_data["mcp_tool_input_schema"]) == (
"Runs the test tool",
schema,
)
@pytest.mark.asyncio
async def test_call_tool_passes_no_tool_metadata_when_tool_was_never_listed(self):
manager, proxy_logging_obj = self._manager_ready_for_call_tool(
[MCPTool(name="other_tool", description="Unrelated", inputSchema={"type": "object"})]
)
await manager.call_tool(
server_name="test-server",
name="test_tool",
arguments={"param": "value"},
user_api_key_auth=self._unrestricted_auth(),
proxy_logging_obj=proxy_logging_obj,
)
hook_kwargs = proxy_logging_obj._create_mcp_request_object_from_kwargs.call_args.args[0]
assert (hook_kwargs["tool_description"], hook_kwargs["tool_input_schema"]) == (None, None)
def test_get_listed_tool_resolves_prefixed_name_and_latest_listing(self):
manager = MCPServerManager()
server = MCPServer(server_id="srv", name="srv", transport=MCPTransport.http, url="http://srv")
manager._create_prefixed_tools([MCPTool(name="echo", description="v1", inputSchema={})], server)
manager._create_prefixed_tools([MCPTool(name="echo", description="v2", inputSchema={})], server)
by_prefixed_name = manager.get_listed_tool(server, "srv-echo")
assert by_prefixed_name is not None and by_prefixed_name.description == "v2"
assert manager.get_listed_tool(server, "missing") is None
def test_get_listed_tool_uses_admin_description_override_clients_saw(self):
manager = MCPServerManager()
server = MCPServer(
server_id="srv",
name="srv",
transport=MCPTransport.http,
url="http://srv",
tool_name_to_description={"echo": "Admin wording"},
)
schema = {"type": "object", "properties": {"text": {"type": "string"}}}
manager._create_prefixed_tools(
[
MCPTool(name="echo", description="Upstream wording", inputSchema=schema),
MCPTool(name="ping", description="Untouched", inputSchema={}),
],
server,
)
overridden = manager.get_listed_tool(server, "srv-echo")
assert overridden is not None
assert (overridden.name, overridden.description, overridden.inputSchema) == ("echo", "Admin wording", schema)
untouched = manager.get_listed_tool(server, "ping")
assert untouched is not None and untouched.description == "Untouched"
def test_server_definition_change_drops_listed_tools(self):
manager = MCPServerManager()
server = MCPServer(server_id="srv", name="srv", transport=MCPTransport.http, url="http://srv")
other = MCPServer(server_id="other", name="other", transport=MCPTransport.http, url="http://other")
manager._create_prefixed_tools([MCPTool(name="echo", description="old", inputSchema={})], server)
manager._create_prefixed_tools([MCPTool(name="ping", description="kept", inputSchema={})], other)
manager._invalidate_server_definition_caches(server.server_id)
assert manager.get_listed_tool(server, "echo") is None
kept = manager.get_listed_tool(other, "ping")
assert kept is not None and kept.description == "kept"
@pytest.mark.asyncio
async def test_user_oauth_refresh_keeps_listed_tools(self):
"""Tool definitions are server-wide, so one user's re-auth must not blank the metadata other
callers' tool calls hand to pre-call guardrails."""
manager = MCPServerManager()
server = MCPServer(server_id="srv", name="srv", transport=MCPTransport.http, url="http://srv")
manager._create_prefixed_tools([MCPTool(name="echo", description="shared", inputSchema={})], server)
await manager.invalidate_user_oauth_token_cache("alice", server.server_id)
listed = manager.get_listed_tool(server, "echo")
assert listed is not None and listed.description == "shared"
def test_per_caller_server_keeps_listed_tools_per_identity(self):
manager = MCPServerManager()
server = MCPServer(
server_id="srv",
name="srv",
transport=MCPTransport.http,
url="http://srv",
auth_type=MCPAuth.oauth2_token_exchange,
)
alice = UserAPIKeyAuth(user_id="alice", api_key="hashed-alice")
bob = UserAPIKeyAuth(user_id="bob", api_key="hashed-bob")
alice_schema = {"type": "object", "properties": {"path": {"type": "string"}}}
bob_schema = {"type": "object", "properties": {"path": {"type": "string"}, "site": {"type": "string"}}}
manager._create_prefixed_tools(
[MCPTool(name="read", description="alice view", inputSchema=alice_schema)],
server,
caller=ListedToolsCaller(user_api_key_auth=alice),
)
manager._create_prefixed_tools(
[MCPTool(name="read", description="bob view", inputSchema=bob_schema)],
server,
caller=ListedToolsCaller(user_api_key_auth=bob),
)
alice_tool = manager.get_listed_tool(server, "srv-read", ListedToolsCaller(user_api_key_auth=alice))
bob_tool = manager.get_listed_tool(server, "srv-read", ListedToolsCaller(user_api_key_auth=bob))
assert alice_tool is not None and (alice_tool.description, alice_tool.inputSchema) == (
"alice view",
alice_schema,
)
assert bob_tool is not None and (bob_tool.description, bob_tool.inputSchema) == ("bob view", bob_schema)
carol = ListedToolsCaller(user_api_key_auth=UserAPIKeyAuth(user_id="carol", api_key="k"))
assert manager.get_listed_tool(server, "srv-read", carol) is None
shared = MCPServer(server_id="shared", name="shared", transport=MCPTransport.http, url="http://shared")
manager._create_prefixed_tools(
[MCPTool(name="echo", description="everyone", inputSchema={})],
shared,
caller=ListedToolsCaller(user_api_key_auth=alice),
)
for_bob = manager.get_listed_tool(shared, "echo", ListedToolsCaller(user_api_key_auth=bob))
assert for_bob is not None and for_bob.description == "everyone"
@pytest.mark.parametrize(
("server_kwargs", "caller_a", "caller_b"),
[
pytest.param(
{"extra_headers": ["X-Workspace"]},
ListedToolsCaller(raw_headers={"x-workspace": "A"}),
ListedToolsCaller(raw_headers={"X-Workspace": "B"}),
id="forwarded-header",
),
pytest.param(
{"auth_type": MCPAuth.true_passthrough},
ListedToolsCaller(raw_headers={"authorization": "Bearer upstream-a"}),
ListedToolsCaller(raw_headers={"authorization": "Bearer upstream-b"}),
id="anonymous-passthrough-bearer",
),
pytest.param(
{"auth_type": MCPAuth.bearer_token},
ListedToolsCaller(mcp_auth_header="byok-a"),
ListedToolsCaller(mcp_auth_header="byok-b"),
id="per-server-auth-header",
),
pytest.param(
{"transport": MCPTransport.stdio, "command": "srv", "env": {"WS": "${X-WS}"}},
ListedToolsCaller(raw_headers={"X-WS": "A"}),
ListedToolsCaller(raw_headers={"X-WS": "B"}),
id="header-driven-stdio-env",
),
],
)
def test_upstream_identity_inputs_keep_listed_tools_apart(self, server_kwargs, caller_a, caller_b):
"""Whatever reaches upstream and can change its catalog must also split the listed-tool cache."""
manager = MCPServerManager()
server = MCPServer(
**{"server_id": "srv", "name": "srv", "transport": MCPTransport.http, "url": "http://srv", **server_kwargs}
)
manager._create_prefixed_tools(
[MCPTool(name="turn", description="Catalog A", inputSchema={})], server, caller=caller_a
)
manager._create_prefixed_tools(
[MCPTool(name="turn", description="Catalog B", inputSchema={})], server, caller=caller_b
)
for_a = manager.get_listed_tool(server, "srv-turn", caller_a)
for_b = manager.get_listed_tool(server, "srv-turn", caller_b)
assert for_a is not None and for_a.description == "Catalog A"
assert for_b is not None and for_b.description == "Catalog B"
assert manager.get_listed_tool(server, "srv-turn", ListedToolsCaller()) is None
def test_shared_server_ignores_headers_it_never_forwards(self):
manager = MCPServerManager()
server = MCPServer(server_id="srv", name="srv", transport=MCPTransport.http, url="http://srv")
manager._create_prefixed_tools(
[MCPTool(name="turn", description="everyone", inputSchema={})],
server,
caller=ListedToolsCaller(raw_headers={"authorization": "Bearer sk-litellm", "x-workspace": "A"}),
)
other = ListedToolsCaller(raw_headers={"authorization": "Bearer sk-other", "x-workspace": "B"})
listed = manager.get_listed_tool(server, "turn", other)
assert listed is not None and listed.description == "everyone"
@pytest.mark.parametrize(
("signer", "static_headers", "shared"),
[
pytest.param(MagicMock(), None, False, id="signer-mints-per-caller-authorization"),
pytest.param(MagicMock(), {"Authorization": "Bearer admin-token"}, True, id="static-authorization-wins"),
pytest.param(None, None, True, id="no-signer-stays-shared"),
],
)
def test_jwt_signer_makes_a_shared_server_list_per_caller(self, signer, static_headers, shared):
"""MCPJWTSigner hands upstream a JWT naming the caller on an otherwise shared ``auth_type: none``
server, so the upstream may tailor the catalog and the cache must not hand one caller another's."""
manager = MCPServerManager()
server = MCPServer(
server_id="srv", name="srv", transport=MCPTransport.http, url="http://srv", static_headers=static_headers
)
alice = ListedToolsCaller(user_api_key_auth=UserAPIKeyAuth(user_id="alice", api_key="hashed-alice"))
bob = ListedToolsCaller(user_api_key_auth=UserAPIKeyAuth(user_id="bob", api_key="hashed-bob"))
with patch( # test-quality-ok: the signer is a process-wide singleton the manager reads, no injection seam
"litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer.get_mcp_jwt_signer",
return_value=signer,
):
manager._create_prefixed_tools(
[MCPTool(name="turn", description="alice view", inputSchema={})], server, caller=alice
)
for_bob = manager.get_listed_tool(server, "srv-turn", bob)
if shared:
assert for_bob is not None and for_bob.description == "alice view"
else:
assert for_bob is None
@pytest.mark.asyncio
async def test_call_tool_hands_hooks_the_catalog_the_same_forwarded_headers_listed(self):
"""Interleaved callers on a forwarded-header server: the hook must see the caller's own catalog."""
manager = MCPServerManager()
server = MCPServer(
server_id="catalog",
name="catalog",
transport=MCPTransport.http,
url="http://catalog",
extra_headers=["X-Workspace"],
)
manager.registry = {"catalog": server}
catalogs = {
"A": [
MCPTool(
name="turn", description="Catalog A", inputSchema={"properties": {"turn": {"description": "A"}}}
)
],
"B": [
MCPTool(
name="turn", description="Catalog B", inputSchema={"properties": {"turn": {"description": "B"}}}
)
],
}
mock_client = AsyncMock()
mock_client.call_tool.return_value = MagicMock(spec=CallToolResult, content=[], isError=False)
manager._create_mcp_client = AsyncMock(return_value=mock_client)
manager._fetch_tools_with_timeout = AsyncMock(side_effect=lambda client, name: catalogs[client.workspace])
for workspace in ("A", "B"):
manager._create_mcp_client.return_value.workspace = workspace
await manager._get_tools_from_server(
server=server,
extra_headers={"X-Workspace": workspace},
raw_headers={"x-workspace": workspace, "authorization": "Bearer sk-litellm"},
user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm", user_id="shared-key"),
)
proxy_logging_obj = MagicMock()
proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock(return_value={})
proxy_logging_obj._convert_mcp_to_llm_format = MagicMock(return_value={})
proxy_logging_obj.pre_call_hook = AsyncMock(return_value={})
proxy_logging_obj.during_call_hook = AsyncMock(return_value=None)
await manager.call_tool(
server_name="catalog",
name="catalog-turn",
arguments={"turn": "A-1"},
user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm", user_id="shared-key"),
proxy_logging_obj=proxy_logging_obj,
raw_headers={"x-workspace": "A", "authorization": "Bearer sk-litellm"},
)
hook_kwargs = proxy_logging_obj._create_mcp_request_object_from_kwargs.call_args.args[0]
assert (hook_kwargs["tool_description"], hook_kwargs["tool_input_schema"]) == (
"Catalog A",
{"properties": {"turn": {"description": "A"}}},
)
def test_per_caller_listed_tools_evict_oldest_caller_and_keep_shared(self):
from litellm.proxy._experimental.mcp_server.mcp_server_manager import _LISTED_TOOLS_CALLERS_PER_SERVER
manager = MCPServerManager()
server = MCPServer(
server_id="srv",
name="srv",
transport=MCPTransport.http,
url="http://srv",
auth_type=MCPAuth.oauth2_token_exchange,
)
manager._create_prefixed_tools([MCPTool(name="read", description="shared", inputSchema={})], server)
callers = [
ListedToolsCaller(user_api_key_auth=UserAPIKeyAuth(user_id=f"u{i}", api_key=f"k{i}"))
for i in range(_LISTED_TOOLS_CALLERS_PER_SERVER + 1)
]
for caller in callers:
manager._create_prefixed_tools(
[MCPTool(name="read", description=caller.user_api_key_auth.user_id, inputSchema={})],
server,
caller=caller,
)
manager._create_prefixed_tools(
[MCPTool(name="read", description="u1 again", inputSchema={})], server, caller=callers[1]
)
assert manager.get_listed_tool(server, "srv-read", callers[0]) is None
second = manager.get_listed_tool(server, "srv-read", callers[1])
assert second is not None and second.description == "u1 again"
newest = manager.get_listed_tool(server, "srv-read", callers[-1])
assert newest is not None and newest.description == callers[-1].user_api_key_auth.user_id
assert len(manager._listed_tools_by_server_id[server.server_id]) == _LISTED_TOOLS_CALLERS_PER_SERVER + 1
shared = manager.get_listed_tool(server, "srv-read")
assert shared is not None and shared.description == "shared"
@pytest.mark.asyncio
@pytest.mark.parametrize("add_prefix", [True, False])
async def test_openapi_listing_records_listed_tools(self, add_prefix):
from litellm.proxy._experimental.mcp_server.tool_registry import global_mcp_tool_registry
server = MCPServer(
server_id="petstore-id",
name="petstore",
alias="petstore",
transport=MCPTransport.http,
url=None,
spec_path="/spec.yaml",
)
manager = MCPServerManager()
manager._create_mcp_client = AsyncMock(return_value=AsyncMock())
async def _handler(**kwargs):
return None
global_mcp_tool_registry.unregister_tools_with_prefix("petstore-")
global_mcp_tool_registry.register_tool(
name="petstore-list_pets",
description="List pets",
input_schema={"type": "object", "properties": {"limit": {"type": "integer"}}},
handler=_handler,
)
try:
listed = await manager._get_tools_from_server(server=server, add_prefix=add_prefix)
finally:
global_mcp_tool_registry.unregister_tools_with_prefix("petstore-")
assert [t.name for t in listed] == ["petstore-list_pets" if add_prefix else "list_pets"]
for name in ("list_pets", "petstore-list_pets"):
tool = manager.get_listed_tool(server, name)
assert tool is not None and tool.description == "List pets"
assert tool.inputSchema["properties"] == {"limit": {"type": "integer"}}
@pytest.mark.asyncio
async def test_openapi_listing_ignores_overlapping_server_prefix(self):
from litellm.proxy._experimental.mcp_server.tool_registry import global_mcp_tool_registry
server = MCPServer(
server_id="pet-id",
name="pet",
alias="pet",
transport=MCPTransport.http,
url=None,
spec_path="/spec.yaml",
)
manager = MCPServerManager()
manager._create_mcp_client = AsyncMock(return_value=AsyncMock())
async def _handler(**kwargs):
return None
for prefix in ("pet-", "petstore-"):
global_mcp_tool_registry.unregister_tools_with_prefix(prefix)
global_mcp_tool_registry.register_tool(
name="pet-petstore-list",
description="Local pet tool",
input_schema={"type": "object", "properties": {"limit": {"type": "integer"}}},
handler=_handler,
)
global_mcp_tool_registry.register_tool(
name="petstore-list",
description="Foreign petstore tool",
input_schema={"type": "object", "properties": {"status": {"type": "string"}}},
handler=_handler,
)
try:
listed = await manager._get_tools_from_server(server=server, add_prefix=True)
finally:
for prefix in ("pet-", "petstore-"):
global_mcp_tool_registry.unregister_tools_with_prefix(prefix)
assert [t.name for t in listed] == ["pet-petstore-list"]
tool = manager.get_listed_tool(server, "petstore-list")
assert tool is not None and tool.description == "Local pet tool"
assert tool.inputSchema["properties"] == {"limit": {"type": "integer"}}
@pytest.mark.asyncio
async def test_get_allowed_mcp_servers_with_user_api_key_auth(self):
"""

14
tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py
View file

@ -40,12 +40,9 @@ async def test_openapi_local_tool_runs_pre_call_tool_check():
fake_server.server_name = "openapi-petstore"
fake_server.alias = None
fake_server.short_prefix = None
fake_server.tool_name_to_description = None
fake_tool = MagicMock()
fake_tool.name = "list_pets"
fake_tool.description = "test tool"
fake_tool.input_schema = {"type": "object"}
pre_call = AsyncMock(return_value={})
handle_local = AsyncMock(return_value=[])
@ -124,12 +121,9 @@ async def test_openapi_local_tool_blocked_when_pre_call_check_raises():
fake_server.server_name = "openapi-petstore"
fake_server.alias = None
fake_server.short_prefix = None
fake_server.tool_name_to_description = None
fake_tool = MagicMock()
fake_tool.name = "delete_pet"
fake_tool.description = "test tool"
fake_tool.input_schema = {"type": "object"}
pre_call = AsyncMock(
side_effect=HTTPException(status_code=403, detail="not allowed")
@ -192,8 +186,6 @@ async def test_openapi_local_tool_denied_when_server_not_resolvable():
fake_tool = MagicMock()
fake_tool.name = "list_pets"
fake_tool.description = "test tool"
fake_tool.input_schema = {"type": "object"}
pre_call = AsyncMock(return_value={})
handle_local = AsyncMock(return_value=[])
@ -278,8 +270,6 @@ async def test_openapi_local_tool_injects_resolved_oauth_token():
fake_tool = MagicMock()
fake_tool.name = "get_values"
fake_tool.description = "test tool"
fake_tool.input_schema = {"type": "object"}
captured: dict = {}
async def handle_local(_name, _arguments):
@ -626,8 +616,6 @@ async def test_per_server_auth_header_reaches_both_openapi_dispatch_arms(dispatc
if dispatch_arm == "local_registry":
fake_tool = MagicMock()
fake_tool.name = "list_reports"
fake_tool.description = "test tool"
fake_tool.input_schema = {"type": "object"}
with (
patch.object(manager, "_get_mcp_server_from_tool_name", return_value=server),
patch.object(mcp_module.global_mcp_tool_registry, "get_tool", return_value=fake_tool),
@ -699,8 +687,6 @@ async def test_local_dispatch_reports_the_outcome_instead_of_success(failure: st
fake_tool = MagicMock()
fake_tool.name = "list_reports"
fake_tool.description = "test tool"
fake_tool.input_schema = {"type": "object"}
fake_tool.handler = raising_handler
server = MCPServer(
server_id="srv-openapi",

158
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_agent_365.py
View file

@ -1,9 +1,7 @@
import time
import uuid
from collections.abc import Iterator
from types import SimpleNamespace
from typing import Any, Final
from unittest.mock import patch
import httpx
import pytest
@ -22,18 +20,12 @@ from litellm.proxy.guardrails.guardrail_hooks.agent_365 import (
guardrail_initializer_registry,
initialize_guardrail,
)
from litellm.proxy.guardrails.guardrail_hooks.agent_365.agent_365 import (
agent_365_authorization_servers,
agent_365_scopes_supported,
)
from litellm.proxy.utils import ProxyLogging
from litellm.types.guardrails import (
GuardrailEventHooks,
LitellmParams,
SupportedGuardrailIntegrations,
)
from litellm.types.mcp import MCPAuth, MCPTransport
from litellm.types.mcp_server.mcp_server_manager import MCPServer
from litellm.types.proxy.guardrails.guardrail_hooks.agent_365 import (
AGENT_365_PROD_API_BASE,
AGENT_365_PROD_RESOURCE_APP_ID,
@ -308,29 +300,6 @@ class TestAllowFlow:
assert evaluate_call.json["conversationId"] == "sess-123"
assert evaluate_call.json["agentId"] == "agent-007"
@pytest.mark.asyncio
async def test_evaluate_payload_includes_listed_tool_metadata(self):
handler: Final = FakeHandler([_token_response(), _allow_response()])
guardrail: Final = _make_guardrail(handler)
schema: Final = {"type": "object", "properties": {"to": {"type": "string"}}, "required": ["to"]}
await _run(guardrail, _mcp_data(mcp_tool_description="Send an email", mcp_tool_input_schema=schema))
assert handler.calls[1].json["tool"] == {
"name": "send_email",
"description": "Send an email",
"inputSchema": schema,
}
@pytest.mark.asyncio
@pytest.mark.parametrize(
("description", "schema"),
[(None, None), ("", None), (None, ["not", "a", "schema"]), (42, "type: object")],
)
async def test_evaluate_payload_omits_missing_or_malformed_tool_metadata(self, description, schema):
handler: Final = FakeHandler([_token_response(), _allow_response()])
guardrail: Final = _make_guardrail(handler)
await _run(guardrail, _mcp_data(mcp_tool_description=description, mcp_tool_input_schema=schema))
assert handler.calls[1].json["tool"] == {"name": "send_email"}
@pytest.mark.asyncio
async def test_agent_id_falls_back_to_key_alias(self):
handler: Final = FakeHandler([_token_response(), _allow_response()])
@ -1039,130 +1008,3 @@ class TestFinalArgumentsEvaluated:
)
assert result["modified_arguments"] == {"turn": "please [REWRITE_ME_REDACTED] now"}
assert handler.calls[1].json["arguments"] == {"turn": "please [REWRITE_ME_REDACTED] now"}
ENTRA_ISSUER: Final = "https://login.microsoftonline.com/tenant-abc/v2.0"
GATEWAY_SCOPE: Final = "api://gateway-app/access_as_user"
def _mcp_server(auth_type: MCPAuth = MCPAuth.none, scopes: list[str] | None = None, **fields: Any) -> MCPServer:
return MCPServer(
server_id="tools-id",
name="tools",
server_name="tools",
transport=MCPTransport.http,
url="https://tools.test/mcp",
auth_type=auth_type,
scopes=scopes,
**fields,
)
@pytest.fixture
def registered_guardrail() -> Iterator[Agent365Guardrail]:
guardrail: Final = _make_guardrail(FakeHandler([]))
litellm.logging_callback_manager.add_litellm_callback(guardrail)
try:
yield guardrail
finally:
litellm.logging_callback_manager.remove_callback_from_list_by_object(
litellm.callbacks, guardrail, require_self=False
)
class TestAgent365AuthorizationServers:
def test_names_the_guardrail_tenant_for_a_scoped_gateway_signed_in_server(self, registered_guardrail):
assert agent_365_authorization_servers(_mcp_server(scopes=[GATEWAY_SCOPE]), None) == (ENTRA_ISSUER,)
assert agent_365_authorization_servers(
_mcp_server(MCPAuth.api_key, scopes=[GATEWAY_SCOPE], auth_value="k"), None
) == (ENTRA_ISSUER,)
@pytest.mark.parametrize("scopes", [None, []], ids=["unset", "empty"])
def test_scopeless_server_signs_in_with_the_gateway_app_scope(self, registered_guardrail, scopes):
server: Final = _mcp_server(scopes=scopes)
assert agent_365_authorization_servers(server, None) == (ENTRA_ISSUER,)
assert agent_365_scopes_supported(server, None) == ("api://client-xyz/access_as_user",)
def test_admin_scopes_override_the_default_gateway_scope(self, registered_guardrail):
assert agent_365_scopes_supported(_mcp_server(scopes=[GATEWAY_SCOPE]), None) == (GATEWAY_SCOPE,)
def test_no_default_scope_when_no_guardrail_gates_the_server(self):
assert agent_365_scopes_supported(_mcp_server(scopes=None), None) == ()
def test_silent_when_no_guardrail_is_registered(self):
assert agent_365_authorization_servers(_mcp_server(scopes=[GATEWAY_SCOPE]), None) == ()
@pytest.mark.parametrize(
"server",
[
_mcp_server(MCPAuth.oauth2, scopes=[GATEWAY_SCOPE]),
_mcp_server(MCPAuth.oauth2_token_exchange, scopes=[GATEWAY_SCOPE], token_exchange_endpoint="https://i/t"),
_mcp_server(MCPAuth.oauth2_id_jag, scopes=[GATEWAY_SCOPE]),
_mcp_server(MCPAuth.true_passthrough, scopes=[GATEWAY_SCOPE]),
_mcp_server(MCPAuth.oauth_delegate, scopes=[GATEWAY_SCOPE]),
_mcp_server(MCPAuth.none, scopes=[GATEWAY_SCOPE], extra_headers=["Authorization"]),
],
ids=["oauth2", "token_exchange", "id_jag", "true_passthrough", "oauth_delegate", "forwards_authorization"],
)
def test_leaves_servers_whose_own_auth_mode_owns_sign_in_alone(self, registered_guardrail, server):
assert agent_365_authorization_servers(server, None) == ()
@pytest.mark.parametrize("header", ["x-api-key", "API-Key", "apikey"])
def test_forwarded_api_key_header_leaves_authorization_to_entra(self, registered_guardrail, header):
"""An upstream API key rides in its own header, so the caller's ``Authorization`` still carries the
Entra assertion and a key-only client must be told where to sign in."""
server: Final = _mcp_server(MCPAuth.none, scopes=None, extra_headers=[header])
assert agent_365_authorization_servers(server, None) == (ENTRA_ISSUER,)
assert agent_365_scopes_supported(server, None) == ("api://client-xyz/access_as_user",)
def test_dedupes_guardrails_sharing_a_tenant(self, registered_guardrail):
twin: Final = _make_guardrail(FakeHandler([]))
twin.guardrail_name = "agent-365-twin"
litellm.logging_callback_manager.add_litellm_callback(twin)
try:
assert agent_365_authorization_servers(_mcp_server(scopes=[GATEWAY_SCOPE]), None) == (ENTRA_ISSUER,)
finally:
litellm.logging_callback_manager.remove_callback_from_list_by_object(
litellm.callbacks, twin, require_self=False
)
def test_key_selected_guardrail_never_advertises_sign_in(self):
"""The challenge a guarded key would get and the anonymous metadata fetch that follows it must name
the same issuer. The anonymous fetch cannot see the key, so neither side advertises Entra."""
guardrail: Final = _make_guardrail(FakeHandler([]))
guardrail.default_on = False
litellm.logging_callback_manager.add_litellm_callback(guardrail)
server: Final = _mcp_server(scopes=[GATEWAY_SCOPE])
plain_key: Final = UserAPIKeyAuth(api_key="sk-plain", user_id="u-1")
guarded_key: Final = UserAPIKeyAuth(
api_key="sk-guarded", user_id="u-2", metadata={"guardrails": ["agent-365-guard"]}
)
try:
with patch( # test-quality-ok: key-selected guardrails read the proxy server premium global, no injection seam
"litellm.proxy.proxy_server.premium_user", True
):
assert agent_365_authorization_servers(server, plain_key) == ()
assert agent_365_authorization_servers(server, guarded_key) == ()
assert agent_365_authorization_servers(server, None) == ()
assert agent_365_scopes_supported(_mcp_server(scopes=None), None) == ()
finally:
litellm.logging_callback_manager.remove_callback_from_list_by_object(
litellm.callbacks, guardrail, require_self=False
)
@pytest.mark.parametrize(
"metadata",
[{"opted_out_global_guardrails": ["agent-365-guard"]}, {"disable_global_guardrails": True}],
ids=["opted-out", "globals-disabled"],
)
def test_key_opted_out_of_the_default_on_guardrail_is_not_challenged(self, registered_guardrail, metadata):
server: Final = _mcp_server(scopes=[GATEWAY_SCOPE])
opted_out: Final = UserAPIKeyAuth(api_key="sk-out", user_id="u-3", metadata=metadata)
team_opted_out: Final = UserAPIKeyAuth(api_key="sk-team", user_id="u-4", team_metadata=metadata)
assert agent_365_authorization_servers(server, opted_out) == ()
assert agent_365_authorization_servers(server, team_opted_out) == ()
assert agent_365_authorization_servers(server, UserAPIKeyAuth(api_key="sk-in", user_id="u-5")) == (
ENTRA_ISSUER,
)

20
tests/test_litellm/proxy/utils/proxy_logging/test_mcp_bridging.py
View file

@ -381,26 +381,6 @@ def test_create_mcp_request_object_from_kwargs_full(proxy_logging, make_user_api
assert snapshot == {"tool_name": "calc", "arguments": {"x": 1}, "server_name": "math", "auth_user_id": "u-1"}
def test_mcp_tool_metadata_flows_from_kwargs_to_synthetic_data(proxy_logging):
schema = {"type": "object", "properties": {"x": {"type": "integer"}}}
obj = proxy_logging._create_mcp_request_object_from_kwargs(
kwargs={
"name": "calc",
"arguments": {"x": 1},
"tool_description": "Adds numbers",
"tool_input_schema": schema,
}
)
out = proxy_logging._convert_mcp_to_llm_format(request_obj=obj, kwargs={})
assert (out["mcp_tool_description"], out["mcp_tool_input_schema"]) == ("Adds numbers", schema)
def test_mcp_tool_metadata_absent_when_tool_was_never_listed(proxy_logging):
obj = proxy_logging._create_mcp_request_object_from_kwargs(kwargs={"name": "calc", "arguments": {}})
out = proxy_logging._convert_mcp_to_llm_format(request_obj=obj, kwargs={})
assert (out["mcp_tool_description"], out["mcp_tool_input_schema"]) == (None, None)
def test_create_mcp_request_object_from_kwargs_empty(proxy_logging):
obj = proxy_logging._create_mcp_request_object_from_kwargs(kwargs={})
snapshot = {