From 4bfaf74edcdef5a24e37c583578070ced8ad9d27 Mon Sep 17 00:00:00 2001 From: yucheng Date: Tue, 15 Sep 2026 00:29:58 +0000 Subject: [PATCH] refactor(guardrails): keep Agent 365 PR to the guardrail, move MCP changes to stacked PRs 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> --- .../mcp_server/discoverable_endpoints.py | 24 +- .../mcp_server/mcp_server_manager.py | 237 ++------- .../_experimental/mcp_server/oauth_utils.py | 2 +- .../outbound_credentials/adapter.py | 8 +- .../proxy/_experimental/mcp_server/server.py | 69 +-- .../guardrail_hooks/agent_365/agent_365.py | 129 +---- litellm/proxy/utils.py | 4 - litellm/types/mcp.py | 4 - .../types/mcp_server/mcp_server_manager.py | 20 +- .../outbound_credentials/test_adapter.py | 16 - .../mcp_server/test_discoverable_endpoints.py | 102 ---- .../test_mcp_guardrail_usage_monitor.py | 1 - .../mcp_server/test_mcp_server.py | 447 +--------------- .../mcp_server/test_mcp_server_manager.py | 490 ------------------ .../mcp_server/test_openapi_tool_auth.py | 14 - .../guardrail_hooks/test_agent_365.py | 158 ------ .../utils/proxy_logging/test_mcp_bridging.py | 20 - 17 files changed, 89 insertions(+), 1656 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 5b0bcc254b2..bafe33d0a6b 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -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. diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 00f19bf964e..fb0c623473a 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -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--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 diff --git a/litellm/proxy/_experimental/mcp_server/oauth_utils.py b/litellm/proxy/_experimental/mcp_server/oauth_utils.py index 50eca9de2b5..39865a35ec6 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_utils.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_utils.py @@ -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 diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py index 2401ad4a776..77979a15199 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py @@ -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 = ( diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index ca12d9e9937..7feb1fd468d 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -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 diff --git a/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py b/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py index c332ac3c964..975d321104d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py +++ b/litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py @@ -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://``).""" - 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) - ) - ) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index b6908311219..d201ac4bc88 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -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(), ) diff --git a/litellm/types/mcp.py b/litellm/types/mcp.py index 87ffaa7cd9e..a59fcb1bcb5 100644 --- a/litellm/types/mcp.py +++ b/litellm/types/mcp.py @@ -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() diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 70ad101dcb0..985d31af997 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -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 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py index d650c5ea6c1..1b003e11993 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py @@ -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, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 975c4b03545..9ea870d3210 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -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 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_guardrail_usage_monitor.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_guardrail_usage_monitor.py index 956f87f0da1..24e6d2de10d 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_guardrail_usage_monitor.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_guardrail_usage_monitor.py @@ -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 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 59973b51054..f5e4a420496 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -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 ) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 2a1064ec510..d56f08c4e79 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -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): """ diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py index 67894d55e80..64614c094ba 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py @@ -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", diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_agent_365.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_agent_365.py index f73c0e87547..f9b7561b9d3 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_agent_365.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_agent_365.py @@ -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, - ) diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_mcp_bridging.py b/tests/test_litellm/proxy/utils/proxy_logging/test_mcp_bridging.py index 75dd11ea16b..56057dce7e0 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_mcp_bridging.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_mcp_bridging.py @@ -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 = {