diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index c273c715ba6..b12e68ab87d 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -541,6 +541,9 @@ class AnthropicChatCompletion(BaseLLM): async def acompletion_dispatch() -> "ModelResponse | CustomStreamWrapper": """Translate then send, so the provider config can inline remote media off the event loop.""" + from litellm.proxy._types import ProxyException + from litellm.responses.mcp.request_context import authorize_provider_mcp_request + request_headers, data = finish_request( await config.async_transform_request( model=model, @@ -550,6 +553,19 @@ class AnthropicChatCompletion(BaseLLM): headers=headers, ) ) + try: + await authorize_provider_mcp_request( + TypeAdapter(dict[str, object]).validate_python(data), + TypeAdapter(dict[str, object]).validate_python(litellm_params), + provider_bound=True, + ) + except ProxyException as error: + raise litellm.PermissionDeniedError( + message=error.message, + model=model, + llm_provider="anthropic", + response=httpx.Response(403, request=httpx.Request("POST", api_base)), + ) from error if ( stream is True ): # if function call - fake the streaming (need complete blocks for output parsing in openai format) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 4a10db31606..5193787b92e 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -606,6 +606,40 @@ _DECODED_BODY_STALE_HEADERS: Final[frozenset[str]] = frozenset({"content-encodin _OPTIONAL_PARAMS: Final = TypeAdapter(Mapping[str, object]) _JSON_OBJECT: Final = TypeAdapter(dict[str, object]) _JSON_OBJECTS: Final = TypeAdapter(list[dict[str, object]]) + + +async def _admit_hosted_mcp_chat_request( + provider_config: BaseConfig, + data: Mapping[str, object], + litellm_params: Mapping[str, object], + *, + model: str, + custom_llm_provider: str, + api_base: str, +) -> None: + """Check provider-run MCP attachments in the final chat body against the caller's tool grants. + + Only chat configs that forward MCP server URLs to the provider (Claude Platform on Bedrock) + need this; the Anthropic chat handler runs the same check on its own path.""" + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + from litellm.proxy._types import ProxyException + from litellm.responses.mcp.request_context import authorize_provider_mcp_request + + if not (isinstance(provider_config, AnthropicConfig) and provider_config.runs_hosted_mcp()): + return + try: + await authorize_provider_mcp_request( + _JSON_OBJECT.validate_python(data), _JSON_OBJECT.validate_python(litellm_params), provider_bound=True + ) + except ProxyException as error: + raise litellm.PermissionDeniedError( + message=error.message, + model=model, + llm_provider=custom_llm_provider, + response=httpx.Response(403, request=httpx.Request("POST", api_base)), + ) from error + + _DECODED_JSON: Final = TypeAdapter(object) @@ -1015,6 +1049,21 @@ class BaseLLMHTTPHandler: shared_session=shared_session, ) + async def admit_then_dispatch( + data: dict[str, object], # mutable-ok: async_completion takes dict + signed_headers: dict[str, object], # mutable-ok: async_completion takes dict + signed_json_body: bytes | None, + ) -> ModelResponse | CustomStreamWrapper: + await _admit_hosted_mcp_chat_request( + provider_config, + data, + cast("Mapping[str, object]", litellm_params), # cast-ok: completion's litellm_params is a bare dict + model=model, + custom_llm_provider=custom_llm_provider, + api_base=api_base, + ) + return await dispatch_async(data, signed_headers, signed_json_body) + if acompletion is True and provider_config.uses_async_transform_request: async def transform_then_dispatch() -> ModelResponse | CustomStreamWrapper: @@ -1033,7 +1082,7 @@ class BaseLLMHTTPHandler: if isinstance(provider_config, SignsRequestsWithAWS) else asyncio.to_thread(sign_and_log, transformed) ) - return await dispatch_async(*signed_request) + return await admit_then_dispatch(*signed_request) return transform_then_dispatch() @@ -1048,7 +1097,7 @@ class BaseLLMHTTPHandler: ) if acompletion is True: - return dispatch_async(data, signed_headers, signed_json_body) + return admit_then_dispatch(data, signed_headers, signed_json_body) if stream is True: data = self._add_stream_param_to_request_body( @@ -2693,6 +2742,9 @@ class BaseLLMHTTPHandler: litellm_params=litellm_params, headers=headers, ) + from litellm.responses.mcp.request_context import authorize_provider_mcp_request + + await authorize_provider_mcp_request(_JSON_OBJECT.validate_python(request_body), kwargs, provider_bound=True) logging_obj.stream = stream logging_obj.model_call_details.update(request_body) @@ -3297,6 +3349,13 @@ class BaseLLMHTTPHandler: extra_body=extra_body, request=data, model=model, litellm_params=litellm_params ) ) + from litellm.responses.mcp.request_context import authorize_provider_mcp_request + + await authorize_provider_mcp_request( + data, + {"litellm_metadata": litellm_metadata}, + provider_bound=True, + ) stream = bool(stream or data.get("stream")) # Preserve the OpenAI-style request context (not sent to the provider) for streaming diff --git a/litellm/proxy/mcp_server/auth/provider_mcp.py b/litellm/proxy/mcp_server/auth/provider_mcp.py new file mode 100644 index 00000000000..411849d263e --- /dev/null +++ b/litellm/proxy/mcp_server/auth/provider_mcp.py @@ -0,0 +1,177 @@ +from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final, NoReturn + +from pydantic import TypeAdapter, ValidationError + +from litellm.models.mcp_server import LiteLLM_MCPServerTable +from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth +from litellm.proxy.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler +from litellm.proxy.mcp_server.catalog import catalog_operation, global_manager +from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler +from litellm.types.mcp_server.provider_registration import canonical_provider_mcp_url + +_OBJECT: Final = TypeAdapter(dict[str, object]) +_OBJECTS: Final = TypeAdapter(tuple[dict[str, object], ...]) +_NAMES: Final = TypeAdapter(tuple[str, ...]) + + +@dataclass(frozen=True, slots=True) +class ProviderMCPAttachment: + url: str + tools: frozenset[str] + + +@dataclass(frozen=True, slots=True) +class MCPAdmissionError: + message: str + + +def _names(value: object) -> frozenset[str] | MCPAdmissionError: + try: + names: Final = _NAMES.validate_python(value) + except ValidationError: + return MCPAdmissionError("Provider MCP attachments require an explicit allowed_tools list") + if any(not name.strip() or any(char in name for char in "*?[]") for name in names): + return MCPAdmissionError("Provider MCP tool names must be exact names without wildcards") + return frozenset(names) + + +def _attachment(url: object, tools: frozenset[str] | MCPAdmissionError) -> ProviderMCPAttachment | MCPAdmissionError: + if isinstance(tools, MCPAdmissionError): + return tools + try: + canonical: Final = canonical_provider_mcp_url(url) + except ValueError: + return MCPAdmissionError("Provider MCP attachments require a registered HTTPS server_url") + return ProviderMCPAttachment(canonical, tools) + + +def _openai_attachment(tool: Mapping[str, object]) -> ProviderMCPAttachment | MCPAdmissionError: + if tool.get("connector_id") is not None: + return MCPAdmissionError("Provider MCP connector IDs are not supported by URL registrations") + allowed: Final = tool.get("allowed_tools") + names: Final = _OBJECT.validate_python(allowed).get("tool_names") if isinstance(allowed, Mapping) else allowed + if isinstance(allowed, Mapping) and not names: + # OpenAI reads an empty tool_names filter as no filter at all, i.e. every tool on the server. + return MCPAdmissionError("Provider MCP tool filters require a nonempty tool_names list") + return _attachment(tool.get("server_url"), _names(names)) + + +def _anthropic_attachment( + server: Mapping[str, object], toolsets: Mapping[str, Mapping[str, object]] +) -> ProviderMCPAttachment | MCPAdmissionError: + name: Final = server.get("name") + if server.get("type") != "url" or not isinstance(name, str): + return MCPAdmissionError("Provider MCP servers must have type=url and a name") + toolset: Final = toolsets.get(name) + if toolset is not None and server.get("tool_configuration") is not None: + return MCPAdmissionError("Provider MCP servers require one unambiguous tool configuration") + if toolset is not None: + default: Final = _OBJECT.validate_python(toolset.get("default_config") or {}) + if default.get("enabled") is not False: + return MCPAdmissionError("Provider MCP toolsets require default_config.enabled=false") + configs: Final = TypeAdapter(dict[str, dict[str, object]]).validate_python(toolset.get("configs") or {}) + if any("enabled" in config and type(config["enabled"]) is not bool for config in configs.values()): + return MCPAdmissionError("Provider MCP tool enabled settings must be booleans") + enabled: Final = tuple(name for name, config in configs.items() if config.get("enabled") is True) + return _attachment(server.get("url"), _names(enabled)) + config: Final = _OBJECT.validate_python(server.get("tool_configuration") or {}) + if config.get("enabled") is False: + return _attachment(server.get("url"), frozenset()) + names: Final = _names(config.get("allowed_tools")) + if names == frozenset(): + return MCPAdmissionError("Enabled legacy MCP connectors require a nonempty allowed_tools list") + return _attachment(server.get("url"), names) + + +async def _attachments( + body: Mapping[str, object], *, provider_bound: bool +) -> tuple[ProviderMCPAttachment | MCPAdmissionError, ...]: + tools: Final = _OBJECTS.validate_python(body.get("tools") or ()) + mcp_tools: Final = tuple(tool for tool in tools if tool.get("type") == "mcp") + gateway: Final = ( + (False,) * len(mcp_tools) + if provider_bound + else await LiteLLM_Proxy_MCP_Handler.routes_through_gateway(mcp_tools) + ) + openai: Final = tuple(_openai_attachment(tool) for tool, routed in zip(mcp_tools, gateway) if not routed) + servers: Final = ( + *_OBJECTS.validate_python(body.get("mcp_servers") or ()), + *(tool for tool in tools if tool.get("type") == "url"), + ) + toolsets: Final = tuple(tool for tool in tools if tool.get("type") == "mcp_toolset") + server_names: Final = frozenset(_NAMES.validate_python(tuple(server.get("name") for server in servers))) + toolset_names: Final = _NAMES.validate_python(tuple(tool.get("mcp_server_name") for tool in toolsets)) + if len(server_names) != len(servers) or any(name not in server_names for name in toolset_names): + return (MCPAdmissionError("Provider MCP server names must be unique and toolsets must reference them"),) + indexed_toolsets: Final = MappingProxyType({name: tool for name, tool in zip(toolset_names, toolsets)}) + if len(indexed_toolsets) != len(toolsets): + return (MCPAdmissionError("Provider MCP servers require one unambiguous tool configuration"),) + return (*openai, *(_anthropic_attachment(server, indexed_toolsets) for server in servers)) + + +def _raise_denied(error: MCPAdmissionError) -> NoReturn: + raise ProxyException( + message=error.message, + type=ProxyErrorTypes.tool_access_denied, + param="tools", + code=403, + ) + + +@catalog_operation(global_manager) +async def authorize_provider_mcp( + body: Mapping[str, object], auth: UserAPIKeyAuth, *, provider_bound: bool = False +) -> None: + try: + attachments: Final = await _attachments(body, provider_bound=provider_bound) + except (ValidationError, TypeError): + _raise_denied(MCPAdmissionError("Invalid provider MCP tool configuration")) + if not attachments: + return + current_auth: Final = await MCPRequestHandler.refresh_catalog_authority(auth) + if current_auth is None: + _raise_denied(MCPAdmissionError("Provider MCP requests require an authenticated caller")) + manager: Final = global_manager() + servers: Final = tuple(manager.catalog.provider_registrations().values()) + # Uniqueness of provider URLs is enforced on write, but a request must not depend on it: + # every registration of a URL is considered, so a duplicate can never shadow a grant. + registrations: Final = { + url: tuple(server for server in servers if server.url == url) for url in {s.url for s in servers} + } + allowed_servers: Final = frozenset(await manager.get_allowed_mcp_servers(current_auth)) + for attachment in attachments: + await _authorize_attachment(attachment, registrations, allowed_servers, current_auth) + + +async def _granted_tools(registration: LiteLLM_MCPServerTable, auth: UserAPIKeyAuth) -> frozenset[str]: + permitted: Final = await MCPRequestHandler.get_allowed_tools_for_server( + registration.server_id, auth, fail_closed=True + ) + ceiling: Final = frozenset(registration.allowed_tools or ()) + return ceiling if permitted is None else ceiling.intersection(permitted) + + +async def _authorize_attachment( + attachment: ProviderMCPAttachment | MCPAdmissionError, + registrations: Mapping[str | None, tuple[LiteLLM_MCPServerTable, ...]], + allowed_servers: frozenset[str], + auth: UserAPIKeyAuth, +) -> None: + """Admit the attachment when any registration of its URL that the key or team is granted + allows every requested tool.""" + if isinstance(attachment, MCPAdmissionError): + _raise_denied(attachment) + granted: Final = tuple( + registration + for registration in registrations.get(attachment.url, ()) + if registration.server_id in allowed_servers + ) + if not granted: + _raise_denied(MCPAdmissionError("Provider MCP server is not registered or not permitted for this key/team")) + for registration in granted: + if attachment.tools <= await _granted_tools(registration, auth): + return + _raise_denied(MCPAdmissionError("Provider MCP tools exceed the registered server or key/team permissions")) diff --git a/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py index 482c2ad2587..6d2115e1471 100644 --- a/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py @@ -3246,6 +3246,7 @@ class MCPRequestHandler: user_api_key_auth: UserAPIKeyAuth | None = None, *, keyless_source: bool = False, + fail_closed: bool = False, ) -> list[str] | None: """ Get list of allowed tool names for a specific server based on key/team permissions. @@ -3254,6 +3255,8 @@ class MCPRequestHandler: Args: server_id: Server ID to check permissions for user_api_key_auth: User auth + fail_closed: Deny every tool ([]) when the lookup fails, instead of the legacy allow-all. + Provider-native admission sets it, because None there means every registered tool. Returns: List[str] if restrictions exist, None if no restrictions (allow all) @@ -3295,7 +3298,8 @@ class MCPRequestHandler: # keyless_source AND the marker are needed: each source resolves through an UNMARKED auth, so # without keyless_source a fault under a source returns None and wins the union as allow-all. deny_all: Final = ( - unreadable_entitlement + fail_closed + or unreadable_entitlement or keyless_source or is_mcp_admitted_user_subject(user_api_key_auth) or personal_key_inherits_team_grants(user_api_key_auth) diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 1104aee353e..4ca9eb7e7da 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -117,6 +117,15 @@ async def _served_mcp_path_names( return await served_names(names) if names else frozenset[str]() +def _provider_url(tool: Mapping[str, object]) -> str | None: + from litellm.types.mcp_server.provider_registration import canonical_provider_mcp_url + + try: + return canonical_provider_mcp_url(tool.get("server_url")) + except ValueError: + return None + + class LiteLLM_Proxy_MCP_Handler: """ Helper class with static methods for MCP integration with Responses API. @@ -159,13 +168,15 @@ class LiteLLM_Proxy_MCP_Handler: served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] = _gateway_served_names, ) -> SplitTools: items: Final = tuple(tools or ()) - served: Final = await _served_mcp_path_names(items, served_names) + if not any(proxy_path_mcp_name(tool) is not None for tool in items): + return LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(items) + routed: Final = await LiteLLM_Proxy_MCP_Handler.routes_through_gateway(items, served_names) return LiteLLM_Proxy_MCP_Handler._parse_mcp_tools( [ - {**tool, "server_url": f"{LITELLM_PROXY_MCP_SERVER_URL_PREFIX}{name}"} - if (name := proxy_path_mcp_name(tool)) in served + {**tool, "server_url": f"{LITELLM_PROXY_MCP_SERVER_URL_PREFIX}{proxy_path_mcp_name(tool)}"} + if through_gateway and not names_gateway_explicitly(tool) else tool - for tool in items + for tool, through_gateway in zip(items, routed) ] ) @@ -179,7 +190,16 @@ class LiteLLM_Proxy_MCP_Handler: ) -> tuple[bool, ...]: items: Final = tuple(tools or ()) served: Final = await _served_mcp_path_names(items, served_names) - return tuple(names_gateway_explicitly(tool) or proxy_path_mcp_name(tool) in served for tool in items) + if not served: + return tuple(map(names_gateway_explicitly, items)) + provider_urls: Final = frozenset( + server.url for server in global_manager().catalog.provider_registrations().values() + ) + return tuple( + names_gateway_explicitly(tool) + or (proxy_path_mcp_name(tool) in served and _provider_url(tool) not in provider_urls) + for tool in items + ) @staticmethod async def _apply_toolset_permissions( diff --git a/litellm/responses/mcp/request_context.py b/litellm/responses/mcp/request_context.py index 3254557c1f0..4bc50021518 100644 --- a/litellm/responses/mcp/request_context.py +++ b/litellm/responses/mcp/request_context.py @@ -11,6 +11,7 @@ still executes the tool, just with no credentials. from collections.abc import Iterable, Mapping from copy import deepcopy from dataclasses import dataclass +from itertools import chain from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final @@ -27,6 +28,66 @@ class _AuthCarryingMetadata(TypedDict): user_api_key_auth: ReadOnly[NotRequired["UserAPIKeyAuth | None"]] +PROVIDER_MCP_ENFORCEMENT_SETTING: Final = "enforce_provider_mcp_permissions" + + +def provider_mcp_enforcement_enabled() -> bool: + """Provider-native MCP checks are opt-in. Off (the default) forwards provider-run MCP attachments + unchecked, exactly as before registrations existed. Read per request, so a config reload applies.""" + from litellm.proxy import proxy_server + + value: Final = proxy_server.general_settings.get(PROVIDER_MCP_ENFORCEMENT_SETTING) + return value is True or (isinstance(value, str) and value.strip().lower() == "true") + + +async def authorize_provider_mcp_request( + body: Mapping[str, object], context: Mapping[str, object], *, provider_bound: bool = False +) -> None: + buckets: Final = tuple( + TypeAdapter(dict[str, object]).validate_python(context[key]) + for key in ("litellm_metadata", "metadata") + if isinstance(context.get(key), Mapping) + ) + auth: Final = next((bucket["user_api_key_auth"] for bucket in buckets if "user_api_key_auth" in bucket), None) + if auth is None or not provider_mcp_enforcement_enabled(): + return + extra: Final = body.get("extra_body") + effective: Final = ( + {**body, **TypeAdapter(dict[str, object]).validate_python(extra)} if isinstance(extra, Mapping) else body + ) + tools: Final = (*_tool_list(effective.get("tools")), *_input_item_tools(effective.get("input"))) + has_native_tools: Final = any( + isinstance(tool, Mapping) + and TypeAdapter(dict[str, object]).validate_python(tool).get("type") in ("mcp", "mcp_toolset", "url") + for tool in tools + ) + if not has_native_tools and not effective.get("mcp_servers"): + return + from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth + + if not isinstance(auth, UserAPIKeyAuth): + return + from litellm.proxy.mcp_server.auth.provider_mcp import authorize_provider_mcp + + await authorize_provider_mcp({**effective, "tools": list(tools)}, auth, provider_bound=provider_bound) + + +def _tool_list(value: object) -> tuple[object, ...]: + return TypeAdapter(tuple[object, ...]).validate_python(value) if isinstance(value, (list, tuple)) else () + + +def _input_item_tools(value: object) -> tuple[object, ...]: + """Tool definitions carried inside Responses input items, such as ``tool_search_output``; the provider + loads them like top-level tools, so they need the same admission.""" + return tuple( + chain.from_iterable( + _tool_list(TypeAdapter(dict[str, object]).validate_python(item).get("tools")) + for item in _tool_list(value) + if isinstance(item, Mapping) + ) + ) + + @dataclass(frozen=True, slots=True) class MCPRequestContext: """Everything a gateway handler must forward to MCP tool listing and execution.""" diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 288e4a7119d..4976f9d994b 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -2214,13 +2214,27 @@ def with_token_limit_checks( async def _admit_frame_turn( - turn_admission: ResponsesWsTurnAdmission | None, raw_message: str + turn_admission: ResponsesWsTurnAdmission | None, + raw_message: str, + *, + user_api_key_auth: UserAPIKeyAuth | None = None, ) -> _AdmittedTurn | _RefusedTurn: frame: Final = _json_object_or_none(raw_message) - if turn_admission is None or frame is None or frame.get("type") != "response.create": + if frame is None or frame.get("type") != "response.create": return _AdmittedTurn() try: - return _AdmittedTurn(await turn_admission()) + if user_api_key_auth is not None: + from litellm.responses.mcp.request_context import authorize_provider_mcp_request + + nested: Final = frame.get("response") + bodies: Final = (frame, nested) if _is_json_object(nested) else (frame,) + for body in bodies: + await authorize_provider_mcp_request( + body, + {"litellm_metadata": {"user_api_key_auth": user_api_key_auth}}, + provider_bound=True, + ) + return _AdmittedTurn(await turn_admission()) if turn_admission is not None else _AdmittedTurn() except RateLimitError as refusal: return _RefusedTurn(refusal) except Exception as e: @@ -2235,7 +2249,14 @@ def turn_refusal_frame(refusal: RateLimitError) -> str: def _turn_admission_error_frame(failure: Exception, status: int) -> str: - return json.dumps({"type": "error", "status": status, "error": {"type": "server_error", "message": str(failure)}}) + _, error_type, error_code = _error_event_fields(failure) + return json.dumps( + { + "type": "error", + "status": status if status != 500 else _status_code_for_error_fields(error_type, error_code), + "error": {"type": error_type or "server_error", "message": str(failure)}, + } + ) def _turn_rejection_frame(rejection: Exception) -> str: @@ -2652,12 +2673,12 @@ class ResponsesWebSocketStreaming: guarded: Final = await self._guard_frame(message) if guarded is None: return - admitted: Final = await self._admit_or_reject_turn(guarded, message, turn_admission) - if admitted is None: - return if guarded.pii_tokens: self._remember_pii_tokens(guarded.pii_tokens) masked: Final = await self._mask_response_create(guarded.frame) + admitted: Final = await self._admit_or_reject_turn(guarded, masked, turn_admission) + if admitted is None: + return self._record_client_frame(masked, guarded.guardrail_information, admitted.reservation) await self.backend_ws.send(masked) @@ -3060,7 +3081,7 @@ class ResponsesWebSocketStreaming: async def _admit_or_reject_turn( self, guarded: GuardedResponsesWsFrame, message: str, turn_admission: ResponsesWsTurnAdmission | None ) -> _AdmittedTurn | None: - verdict: Final = await _admit_frame_turn(turn_admission, message) + verdict: Final = await _admit_frame_turn(turn_admission, message, user_api_key_auth=self.user_api_key_dict) match verdict: case _AdmittedTurn(): return verdict diff --git a/litellm/utils.py b/litellm/utils.py index 53b02ebddac..e39b17da5f5 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1788,6 +1788,21 @@ def check_coroutine(value) -> bool: return get_coroutine_checker().is_async_callable(value) +def _chat_deployment_runs_hosted_mcp(kwargs: Mapping[str, object]) -> bool: + """Whether this chat deployment forwards MCP server URLs for the provider to call (hosted MCP), + as Claude Platform on Bedrock does, so its tool grants must be checked before dispatch.""" + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + + model: Final = kwargs.get("model") + provider: Final = _deployment_provider(kwargs) + if not isinstance(model, str) or provider not in LlmProviders._value2member_map_: + return False + config: Final = ProviderConfigManager.get_provider_chat_config( + model=litellm.get_llm_provider(model=model, custom_llm_provider=provider)[0], provider=LlmProviders(provider) + ) + return isinstance(config, AnthropicConfig) and config.runs_hosted_mcp() + + def _deployment_provider(kwargs: Mapping[str, object]) -> str: model: Final = kwargs.get("model") custom_llm_provider: Final = kwargs.get("custom_llm_provider") @@ -1833,6 +1848,13 @@ async def async_pre_call_deployment_hook(kwargs: Mapping[str, object], call_type if result is not None: modified_kwargs = result + if call_type in (CallTypes.aresponses.value, CallTypes.anthropic_messages.value) or ( + call_type == CallTypes.acompletion.value + and (_deployment_provider(modified_kwargs) == "anthropic" or _chat_deployment_runs_hosted_mcp(modified_kwargs)) + ): + from litellm.responses.mcp.request_context import authorize_provider_mcp_request + + await authorize_provider_mcp_request(TypeAdapter(dict[str, object]).validate_python(modified_kwargs), kwargs) return _unstreamed_when_post_call_guardrails_apply(modified_kwargs, typed_call_type) diff --git a/tests/integration/mcp/test_mcp_llm_endpoints.py b/tests/integration/mcp/test_mcp_llm_endpoints.py index 0c97ae62d07..d685af24cc8 100644 --- a/tests/integration/mcp/test_mcp_llm_endpoints.py +++ b/tests/integration/mcp/test_mcp_llm_endpoints.py @@ -47,7 +47,9 @@ pytestmark: Final = [pytest.mark.xdist_group("mcp_llm_private_gateway"), pytest. @contextmanager def _private_gateway( - directory: Path, guardrails: Sequence[Mapping[str, JsonValue]] = () + directory: Path, + guardrails: Sequence[Mapping[str, JsonValue]] = (), + extra_general_settings: Mapping[str, JsonValue] | None = None, ) -> Iterator[tuple[Gateway, Mapping[str, str]]]: stock: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) config: Final = directory / "stored-responses.yaml" @@ -56,7 +58,11 @@ def _private_gateway( yaml.safe_dump( { **stock, - "general_settings": {**object_value(stock["general_settings"]), "store_prompts_in_spend_logs": True}, + "general_settings": { + **object_value(stock["general_settings"]), + "store_prompts_in_spend_logs": True, + **(extra_general_settings or {}), + }, "litellm_settings": { **object_value(stock["litellm_settings"]), "cost_margin_config": { @@ -97,6 +103,22 @@ def _llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gat yield served +@pytest.fixture(scope="module") +def _enforcing_llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, Mapping[str, str]]]: + with _private_gateway( + tmp_path_factory.mktemp("mcp-llm-enforcing-gateway"), + extra_general_settings={"enforce_provider_mcp_permissions": True}, + ) as served: + yield served + + +@pytest.fixture +def enforcing_gateway(_enforcing_llm_gateway: tuple[Gateway, Mapping[str, str]]) -> Iterator[Gateway]: + """A proxy that opted in to provider-native MCP checks.""" + with _serving(_enforcing_llm_gateway) as private: + yield private + + @pytest.fixture def gateway(_llm_gateway: tuple[Gateway, Mapping[str, str]]) -> Iterator[Gateway]: with _serving(_llm_gateway) as private: @@ -750,6 +772,7 @@ CALLER_FUNCTION_TOOLS: Final[Mapping[Surface, Mapping[str, object]]] = { def test_tools_that_do_not_name_the_gateway_reach_the_provider_unchanged_and_never_the_peer( gateway: Gateway, surface: Surface ) -> None: + """Provider-native MCP checks are opt-in, so by default these reach the provider exactly as before.""" with _rig(gateway, surface) as rig: remote_mcp_tools: Final = [ { @@ -778,6 +801,66 @@ def test_tools_that_do_not_name_the_gateway_reach_the_provider_unchanged_and_nev assert tool_calls(rig.peer.drain()) == (), "a tool that does not name the gateway reached the peer" +@pytest.mark.parametrize("surface", ("chat", "responses", "messages")) +def test_non_gateway_tools_follow_provider_native_admission_policy( + enforcing_gateway: Gateway, surface: Surface +) -> None: + with _rig(enforcing_gateway, surface) as rig: + remote_mcp_tools: Final = [ + { + "type": "mcp", + "server_label": "docs", + "server_url": "https://mcp.example.com/sse", + "require_approval": "never", + }, + { + "type": "mcp", + "server_label": "repos", + "server_url": f"https://mcp.example.com/mcp/unserved{uuid.uuid4().hex[:8]}", + "require_approval": "never", + }, + ] + response: Final = rig.send(_granted_key(rig), [CALLER_FUNCTION_TOOLS[surface], *remote_mcp_tools]) + assert response.status_code == (200 if surface == "chat" else 403), response.text + bodies: Final = tuple( + JSON_OBJECT.validate_json(request.body) for request in rig.wire.drain() if request.method == "POST" + ) + assert tool_calls(rig.peer.drain()) == (), "a tool that does not name the gateway reached the peer" + if surface != "chat": + assert "explicit allowed_tools list" in response.text, response.text + assert bodies == (), bodies + return + assert len(bodies) == 1, bodies + forwarded: Final = bodies[0].get("tools") + assert isinstance(forwarded, list) and len(forwarded) == 3, forwarded + assert [tool for tool in forwarded if isinstance(tool, dict) and tool.get("type") == "mcp"] == remote_mcp_tools + assert "lookup" in json.dumps([tool for tool in forwarded if tool not in remote_mcp_tools]), forwarded + + +@pytest.mark.parametrize("stream", (False, True)) +def test_chat_to_anthropic_denies_native_mcp_before_provider_dispatch( + enforcing_gateway: Gateway, stream: bool +) -> None: + with _rig(enforcing_gateway, "chat", upstream="messages") as rig: + response: Final = rig.send( + _granted_key(rig), + [CALLER_FUNCTION_TOOLS["chat"]], + stream=stream, + mcp_servers=[ + { + "type": "url", + "name": "unregistered", + "url": "https://native.example/mcp", + "tool_configuration": {"allowed_tools": ["search"]}, + } + ], + ) + assert response.status_code == 403, response.text + assert "not registered or not permitted" in response.text, response.text + assert tuple(request for request in rig.wire.drain() if request.method == "POST") == () + assert tool_calls(rig.peer.drain()) == () + + def test_streaming_chat_executes_the_tool_once_and_streams_the_follow_up(gateway: Gateway) -> None: with _rig(gateway, "chat") as rig: key: Final = _granted_key(rig) diff --git a/tests/unit/proxy/mcp_server/auth/test_provider_mcp.py b/tests/unit/proxy/mcp_server/auth/test_provider_mcp.py new file mode 100644 index 00000000000..1ee7addb812 --- /dev/null +++ b/tests/unit/proxy/mcp_server/auth/test_provider_mcp.py @@ -0,0 +1,822 @@ +import asyncio +from collections.abc import Mapping +from copy import deepcopy +from types import SimpleNamespace +from typing import Final +from unittest.mock import AsyncMock + +import pytest +import respx + +import litellm +from litellm.caching.caching import Cache +from litellm.caching.caching_handler import _PENDING_CACHE_WRITES +from litellm.caching.dual_cache import DualCache +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy import proxy_server +from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable, ProxyException, UserAPIKeyAuth +from litellm.proxy.mcp_server import mcp_server_manager +from litellm.proxy.mcp_server.auth.provider_mcp import authorize_provider_mcp +from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler +from litellm.responses.mcp.request_context import authorize_provider_mcp_request +from litellm.types.mcp import MCPTransport +from litellm.types.mcp_server.mcp_server_manager import MCPServer +from litellm.utils import async_pre_call_deployment_hook + + +def _openai(tools: object = ("search",), url: str = "https://native.example/mcp") -> dict[str, object]: + return {"tools": [{"type": "mcp", "server_label": "arbitrary", "server_url": url, "allowed_tools": tools}]} + + +def _anthropic(legacy: bool = False) -> dict[str, object]: + return { + "mcp_servers": [ + { + "type": "url", + "name": "docs", + "url": "https://native.example/mcp", + **({"tool_configuration": {"allowed_tools": ["search"]}} if legacy else {}), + } + ], + **( + {} + if legacy + else { + "tools": [ + { + "type": "mcp_toolset", + "mcp_server_name": "docs", + "default_config": {"enabled": False}, + "configs": {"search": {"enabled": True}, "delete": {"enabled": False}}, + } + ] + } + ), + } + + +@pytest.fixture +async def native_auth(monkeypatch: pytest.MonkeyPatch) -> UserAPIKeyAuth: + manager: Final = mcp_server_manager.MCPServerManager() + cache: Final = DualCache() + monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager) + monkeypatch.setattr(proxy_server, "prisma_client", None) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + await manager.load_servers_from_config( + { + "docs": { + "server_id": "native-id", + "alias": "native-alias", + "execution_mode": "provider", + "url": "https://native.example/mcp", + "allowed_tools": ["search", "fetch"], + } + } + ) + await cache.async_set_cache( + key="team_id:team", + value=LiteLLM_TeamTable( + team_id="team", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="team-grant", + mcp_servers=["docs"], + mcp_tool_permissions={"docs": ["search"]}, + ), + ), + ) + team: Final = await cache.async_get_cache(key="team_id:team") + database: Final = SimpleNamespace(litellm_teamtable=SimpleNamespace(find_unique=AsyncMock(return_value=team))) + monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=database, writer_db=database)) + monkeypatch.setattr( + proxy_server, "general_settings", {"supported_db_objects": [], "enforce_provider_mcp_permissions": True} + ) + return UserAPIKeyAuth( + api_key="test-key", + team_id="team", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="key-grant", + mcp_servers=["native-alias"], + mcp_tool_permissions={"native-alias": ["search", "fetch"]}, + ), + ) + + +@pytest.mark.parametrize( + "body", [_openai(), _openai({"tool_names": ["search"], "read_only": True}), _anthropic(), _anthropic(True)] +) +async def test_allowed_native_attachment_is_unchanged(body: Mapping[str, object], native_auth: UserAPIKeyAuth) -> None: + before: Final = deepcopy(body) + await authorize_provider_mcp(body, native_auth) + assert body == before + assert mcp_server_manager.global_mcp_server_manager.get_registry() == {} + + +@pytest.mark.parametrize( + "body", + [ + _openai(["fetch"]), + _openai(["delete"]), + _openai(None), + _openai({"read_only": True}), + _openai(["*"]), + _openai(["search"], "https://unregistered.example/mcp"), + _openai(["search"], "http://native.example/mcp"), + _openai(["search"], "https://native.example/mcp?endpoint=other"), + _openai(["search"], "https://native.example/mcp/"), + _openai(["search"], "https://native.example@other.example/mcp"), + {"mcp_servers": [{"type": "url", "name": "docs", "url": "https://native.example/mcp"}]}, + { + "mcp_servers": [ + { + "type": "url", + "name": "docs", + "url": "https://native.example/mcp", + "tool_configuration": {"allowed_tools": []}, + } + ] + }, + { + **_anthropic(), + "tools": [{"type": "mcp_toolset", "mcp_server_name": "docs", "configs": {"search": {"enabled": True}}}], + }, + { + **_anthropic(), + "tools": [{"type": "mcp_toolset", "mcp_server_name": "missing", "default_config": {"enabled": False}}], + }, + { + **_anthropic(), + "tools": [{"type": "mcp_toolset", "mcp_server_name": "docs", "default_config": {"enabled": False}}] * 2, + }, + { + **_anthropic(), + "mcp_servers": [{"type": "url", "name": "docs", "url": "https://native.example/mcp"}] * 2, + }, + {**_anthropic(True), "tools": _anthropic()["tools"]}, + { + **_anthropic(), + "tools": [ + { + "type": "mcp_toolset", + "mcp_server_name": "docs", + "default_config": {"enabled": False}, + "configs": {"search": {"enabled": "true"}}, + } + ], + }, + {"tools": [{"type": "mcp", "connector_id": "connector_example", "allowed_tools": ["search"]}]}, + ], +) +async def test_native_attachment_denies_unbounded_or_unauthorized_scope( + body: Mapping[str, object], native_auth: UserAPIKeyAuth +) -> None: + with pytest.raises(ProxyException) as denied: + await authorize_provider_mcp(body, native_auth) + assert denied.value.code == "403" + + +@pytest.mark.parametrize("servers", [["no-mcp-servers"], ["unrelated"]]) +async def test_server_denial_cannot_be_bypassed_by_tool_name(servers: list[str], native_auth: UserAPIKeyAuth) -> None: + auth: Final = native_auth.model_copy( + update={ + "object_permission": LiteLLM_ObjectPermissionTable( + object_permission_id="denied", + mcp_servers=servers, + ) + } + ) + with pytest.raises(ProxyException) as denied: + await authorize_provider_mcp(_openai(), auth) + assert denied.value.code == "403" + + +async def test_duplicate_url_registration_neither_shadows_nor_widens_the_granted_one( + native_auth: UserAPIKeyAuth, monkeypatch: pytest.MonkeyPatch +) -> None: + """Two database rows share one URL, as a database without the partial index can hold. The ungranted + duplicate is listed last, where a one-registration-per-URL map would have picked it.""" + from litellm.models.mcp_server import LiteLLM_MCPServerTable + + granted: Final = LiteLLM_MCPServerTable( + server_id="native-id", server_name="docs", alias="native-alias", execution_mode="provider", + transport=None, url="https://native.example/mcp", allowed_tools=["search", "fetch"], + ) + duplicate: Final = LiteLLM_MCPServerTable( + server_id="duplicate-id", server_name="duplicate", execution_mode="provider", + transport=None, url="https://native.example/mcp", allowed_tools=["delete"], + ) + team: Final = await proxy_server.user_api_key_cache.async_get_cache(key="team_id:team") + tables: Final = SimpleNamespace( + litellm_mcpservertable=SimpleNamespace(find_many=AsyncMock(return_value=[granted, duplicate])), + litellm_config=SimpleNamespace(find_unique=AsyncMock(return_value=SimpleNamespace(reload_revision=1))), + litellm_teamtable=SimpleNamespace(find_unique=AsyncMock(return_value=team)), + ) + monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=tables, writer_db=tables)) + monkeypatch.setattr(proxy_server, "general_settings", {"enforce_provider_mcp_permissions": True}) + monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", mcp_server_manager.MCPServerManager()) + + await authorize_provider_mcp(_openai(["search"]), native_auth) + with pytest.raises(ProxyException) as denied: + await authorize_provider_mcp(_openai(["delete"]), native_auth) + assert denied.value.code == "403" + + +async def test_empty_openai_tool_name_filter_is_denied_because_openai_reads_it_as_every_tool( + native_auth: UserAPIKeyAuth, +) -> None: + with pytest.raises(ProxyException) as denied: + await authorize_provider_mcp(_openai({"tool_names": []}), native_auth) + assert (denied.value.code, denied.value.message) == ( + "403", + "Provider MCP tool filters require a nonempty tool_names list", + ) + + +async def test_failed_tool_grant_lookup_denies_instead_of_allowing_every_registered_tool( + native_auth: UserAPIKeyAuth, monkeypatch: pytest.MonkeyPatch +) -> None: + """The key reaches the server directly, but its toolset grant cannot be read: the tool lookup fails + while server admission still succeeds, and that failure must not widen to every registered tool.""" + team: Final = await proxy_server.user_api_key_cache.async_get_cache(key="team_id:team") + toolset_store_down: Final = SimpleNamespace( + litellm_teamtable=SimpleNamespace(find_unique=AsyncMock(return_value=team)), + litellm_mcptoolsettable=SimpleNamespace(find_many=AsyncMock(side_effect=RuntimeError("toolset store down"))), + ) + monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=toolset_store_down, writer_db=toolset_store_down)) + auth: Final = native_auth.model_copy( + update={ + "requires_fresh_policy": True, + "object_permission": LiteLLM_ObjectPermissionTable( + object_permission_id="key-grant", mcp_servers=["native-alias"], mcp_toolsets=["research"] + ), + } + ) + with pytest.raises(ProxyException) as denied: + await authorize_provider_mcp(_openai(["search"]), auth) + assert denied.value.code == "403" + + +@pytest.mark.parametrize("tools,admitted", [(["search"], True), (["delete"], False)]) +async def test_mcp_tools_carried_in_responses_input_items_are_admitted_like_top_level_tools( + native_auth: UserAPIKeyAuth, tools: list[str], admitted: bool +) -> None: + body: Final = {"input": [{"type": "tool_search_output", "tools": _openai(tools)["tools"]}]} + context: Final = {"litellm_metadata": {"user_api_key_auth": native_auth}} + if admitted: + await authorize_provider_mcp_request(body, context) + return + with pytest.raises(ProxyException) as denied: + await authorize_provider_mcp_request(body, context) + assert denied.value.code == "403" + + +@pytest.mark.parametrize( + "body", + [ + _openai(["search"], "https://unregistered.example/mcp"), + _openai(["delete"]), + _openai({"tool_names": []}), + {"tools": [{"type": "mcp", "server_label": "docs", "server_url": "https://native.example/mcp"}]}, + ], + ids=["unregistered_url", "ungranted_tool", "empty_tool_name_filter", "unbounded_attachment"], +) +async def test_provider_mcp_checks_are_off_unless_the_proxy_opts_in( + native_auth: UserAPIKeyAuth, monkeypatch: pytest.MonkeyPatch, body: dict[str, object] +) -> None: + """Without enforce_provider_mcp_permissions, provider-run MCP passes through as it did before.""" + context: Final = {"litellm_metadata": {"user_api_key_auth": native_auth}} + monkeypatch.setattr(proxy_server, "general_settings", {"supported_db_objects": []}) + assert await authorize_provider_mcp_request(body, context) is None + monkeypatch.setattr( + proxy_server, "general_settings", {"supported_db_objects": [], "enforce_provider_mcp_permissions": "true"} + ) + with pytest.raises(ProxyException) as denied: + await authorize_provider_mcp_request(body, context) + assert denied.value.code == "403" + + +async def test_ungranted_provider_tool_reaches_provider_unchanged_when_checks_are_off( + native_auth: UserAPIKeyAuth, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + import json + + monkeypatch.setattr(proxy_server, "general_settings", {"supported_db_objects": []}) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + litellm.in_memory_llm_clients_cache.flush_cache() + upstream: Final = respx_mock.post(url__regex=r"https://native-provider.example/.*").respond( + 200, + json={ + "id": "resp_native", "object": "response", "created_at": 1, "model": "gpt-5.6-luna", + "status": "completed", "output": [], "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + }, + ) + attachment: Final = _openai(["delete"], "https://unregistered.example/mcp") + await litellm.aresponses( + model="openai/gpt-5.6-luna", + input="Use the tool", + api_key="test-provider-key", + api_base="https://native-provider.example/v1", + litellm_metadata={"user_api_key_auth": native_auth}, + **attachment, + ) + assert json.loads(upstream.calls.last.request.content)["tools"] == attachment["tools"] + + +async def test_canonical_host_and_default_port_match(native_auth: UserAPIKeyAuth) -> None: + body: Final = _openai(["search"], "https://NATIVE.example:443/mcp") + await authorize_provider_mcp(body, native_auth) + assert body == _openai(["search"], "https://NATIVE.example:443/mcp") + + +async def test_toolsets_match_only_their_named_server_in_a_mixed_configuration(native_auth: UserAPIKeyAuth) -> None: + body: Final = { + "mcp_servers": [ + { + "type": "url", + "name": "legacy", + "url": "https://native.example/mcp", + "tool_configuration": {"allowed_tools": ["search"]}, + }, + {"type": "url", "name": "modern", "url": "https://native.example/mcp"}, + ], + "tools": [ + { + "type": "mcp_toolset", + "mcp_server_name": "modern", + "default_config": {"enabled": False}, + "configs": {"search": {"enabled": True}}, + } + ], + } + original: Final = deepcopy(body) + await authorize_provider_mcp(body, native_auth) + assert body == original + + +@pytest.mark.parametrize( + "call_type,model", + [("aresponses", "openai/gpt-5.6-luna"), ("acompletion", "anthropic/claude-haiku-5-5")], +) +async def test_extra_body_cannot_replace_authorized_attachment( + native_auth: UserAPIKeyAuth, call_type: str, model: str +) -> None: + kwargs: Final = { + "model": model, + **_openai(), + "extra_body": _openai(["delete"]), + "litellm_metadata": {"user_api_key_auth": native_auth}, + } + with pytest.raises(ProxyException) as denied: + await async_pre_call_deployment_hook(kwargs, call_type) + assert denied.value.code == "403" + + +@pytest.mark.parametrize("provider_bound", [False, True]) +async def test_registration_revocation_is_seen_on_next_admission( + native_auth: UserAPIKeyAuth, provider_bound: bool +) -> None: + await authorize_provider_mcp(_openai(), native_auth, provider_bound=provider_bound) + await mcp_server_manager.global_mcp_server_manager.load_servers_from_config( + { + "docs": { + "server_id": "native-id", + "execution_mode": "provider", + "url": "https://native.example/mcp", + "allowed_tools": [], + } + } + ) + with pytest.raises(ProxyException) as denied: + await authorize_provider_mcp(_openai(), native_auth, provider_bound=provider_bound) + assert denied.value.code == "403" + + +async def test_sdk_without_proxy_identity_preserves_native_passthrough(native_auth: UserAPIKeyAuth) -> None: + body: Final = _openai(["anything"], "https://unregistered.example/mcp") + await authorize_provider_mcp_request(body, {}) + assert body == _openai(["anything"], "https://unregistered.example/mcp") + + +async def test_provider_url_wins_over_gateway_path_shortcut(native_auth: UserAPIKeyAuth) -> None: + manager: Final = mcp_server_manager.global_mcp_server_manager + await manager.load_servers_from_config( + { + "native": { + "server_id": "provider-path", + "execution_mode": "provider", + "url": "https://native.example/mcp/gateway", + "allowed_tools": ["search"], + } + } + ) + + async def gateway_names(names: object) -> frozenset[str]: + return frozenset({"gateway"}) + + tools: Final = [{"type": "mcp", "server_url": "https://native.example/mcp/gateway", "allowed_tools": ["search"]}] + assert await LiteLLM_Proxy_MCP_Handler.routes_through_gateway(tools, gateway_names) == (False,) + assert await LiteLLM_Proxy_MCP_Handler.split_mcp_tools(tools, gateway_names) == ([], tools) + + +async def test_explicit_gateway_references_retain_existing_path(native_auth: UserAPIKeyAuth) -> None: + tools: Final = [{"type": "mcp", "server_url": "litellm_proxy/mcp/gateway"}] + await authorize_provider_mcp({"tools": tools}, native_auth) + assert await LiteLLM_Proxy_MCP_Handler.split_mcp_tools(tools) == (tools, []) + + +@pytest.mark.parametrize("api", ["aresponses", "anthropic_messages", "acompletion"]) +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("configured", [False, True]) +async def test_denied_native_request_never_reaches_provider( + native_auth: UserAPIKeyAuth, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + api: str, + stream: bool, + configured: bool, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + litellm.in_memory_llm_clients_cache.flush_cache() + upstream: Final = respx_mock.post(url__regex=r"https://native-provider.example/.*").respond(500) + attachment: Final = ( + _openai(["delete"]) + if api == "aresponses" + else { + "mcp_servers": [ + { + "type": "url", + "name": "docs", + "url": "https://unregistered.example/mcp", + "tool_configuration": {"allowed_tools": ["search"]}, + } + ], + } + ) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "native-model", + "litellm_params": { + "model": "openai/gpt-5.6-luna" if api == "aresponses" else "anthropic/claude-haiku-5-5", + "api_key": "test-provider-key", + "api_base": "https://native-provider.example/v1", + **(attachment if configured else {}), + }, + } + ], + num_retries=0, + ) + request: Final = { + "model": "native-model", + "stream": stream, + "metadata" if api == "acompletion" else "litellm_metadata": {"user_api_key_auth": native_auth}, + **({} if configured else attachment), + } + call: Final = ( + router.aresponses(input="Use the tool", **request) + if api == "aresponses" + else router.anthropic_messages(messages=[{"role": "user", "content": "Use the tool"}], max_tokens=32, **request) + if api == "anthropic_messages" + else router.acompletion(messages=[{"role": "user", "content": "Use the tool"}], max_tokens=32, **request) + ) + with pytest.raises(ProxyException) as denied: + await call + assert str(getattr(denied.value, "status_code", getattr(denied.value, "code", ""))) == "403" + assert upstream.call_count == 0 + + +CLAUDE_PLATFORM_PARAMS: Final = { + "model": "bedrock/claude_platform/claude-sonnet-4-6", + "api_key": "test-provider-key", + "api_base": "https://native-provider.example/v1/messages", + "workspace_id": "wrkspc_test", + "aws_region_name": "us-west-2", +} +CLAUDE_PLATFORM_REPLY: Final = { + "id": "msg_native", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [{"type": "text", "text": "done"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, +} + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("routed", [False, True]) +async def test_claude_platform_chat_with_ungranted_provider_tools_never_reaches_bedrock( + native_auth: UserAPIKeyAuth, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, stream: bool, routed: bool +) -> None: + """Claude Platform on Bedrock runs MCP servers named in a chat request, so chat must be checked + like the Anthropic provider, both before dispatch (routed) and on the final body (direct).""" + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + litellm.in_memory_llm_clients_cache.flush_cache() + upstream: Final = respx_mock.post(url__regex=r"https://native-provider.example/.*").respond( + 200, json=CLAUDE_PLATFORM_REPLY + ) + request: Final = { + "messages": [{"role": "user", "content": "Use the tool"}], + "max_tokens": 32, + "stream": stream, + "metadata": {"user_api_key_auth": native_auth}, + **_openai(["delete"]), + } + call: Final = ( + litellm.Router( + model_list=[{"model_name": "platform", "litellm_params": dict(CLAUDE_PLATFORM_PARAMS)}], num_retries=0 + ).acompletion(model="platform", **request) + if routed + else litellm.acompletion(**CLAUDE_PLATFORM_PARAMS, **request) + ) + with pytest.raises((ProxyException, litellm.PermissionDeniedError)) as denied: + await call + assert str(getattr(denied.value, "status_code", getattr(denied.value, "code", ""))) == "403" + assert upstream.call_count == 0 + + +async def test_claude_platform_chat_with_granted_provider_tools_reaches_bedrock_with_its_mcp_server( + native_auth: UserAPIKeyAuth, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + import json + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + litellm.in_memory_llm_clients_cache.flush_cache() + upstream: Final = respx_mock.post(url__regex=r"https://native-provider.example/.*").respond( + 200, json=CLAUDE_PLATFORM_REPLY + ) + await litellm.acompletion( + **CLAUDE_PLATFORM_PARAMS, + messages=[{"role": "user", "content": "Use the tool"}], + max_tokens=32, + metadata={"user_api_key_auth": native_auth}, + **_openai(["search"]), + ) + sent: Final = json.loads(upstream.calls.last.request.content) + assert sent["mcp_servers"] == [ + { + "type": "url", + "url": "https://native.example/mcp", + "name": "arbitrary", + "tool_configuration": {"allowed_tools": ["search"]}, + } + ] + + +@pytest.mark.parametrize("api", ["aresponses", "anthropic_messages", "acompletion"]) +async def test_allowed_native_request_reaches_provider_unchanged( + native_auth: UserAPIKeyAuth, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + api: str, +) -> None: + import json + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + litellm.in_memory_llm_clients_cache.flush_cache() + attachment: Final = _openai(["search"]) if api == "aresponses" else _anthropic(legacy=api == "acompletion") + response: Final = ( + { + "id": "resp_native", + "object": "response", + "created_at": 1, + "model": "gpt-5.6-luna", + "status": "completed", + "output": [], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + } + if api == "aresponses" + else { + "id": "msg_native", + "type": "message", + "role": "assistant", + "model": "claude-haiku-5-5", + "content": [{"type": "text", "text": "done"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + } + ) + upstream: Final = respx_mock.post(url__regex=r"https://native-provider.example/.*").respond(200, json=response) + request: Final = { + "api_key": "test-provider-key", + "api_base": "https://native-provider.example/v1", + "litellm_metadata": {"user_api_key_auth": native_auth}, + **attachment, + } + if api == "aresponses": + await litellm.aresponses(model="openai/gpt-5.6-luna", input="Use the tool", **request) + elif api == "anthropic_messages": + await litellm.anthropic_messages( + model="anthropic/claude-haiku-5-5", + messages=[{"role": "user", "content": "Use the tool"}], + max_tokens=32, + **request, + ) + else: + await litellm.acompletion( + model="anthropic/claude-haiku-5-5", + messages=[{"role": "user", "content": "Use the tool"}], + max_tokens=32, + **request, + ) + sent: Final = json.loads(upstream.calls.last.request.content) + assert {key: sent[key] for key in attachment} == attachment + assert upstream.call_count == 1 + + +@pytest.mark.parametrize("api", ["aresponses", "acompletion"]) +async def test_cached_response_does_not_bypass_revoked_mcp_permission( + native_auth: UserAPIKeyAuth, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + api: str, +) -> None: + monkeypatch.setattr(litellm, "cache", Cache(type="local")) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + upstream: Final = respx_mock.post(url__regex=r"https://native-provider.example/.*").respond( + 200, + json={ + "id": "resp_cached", + "object": "response", + "created_at": 1, + "model": "gpt-5.6-luna", + "status": "completed", + "output": [ + { + "type": "message", + "id": "msg_cached", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "tool result", "annotations": []}], + } + ], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + } + if api == "aresponses" + else { + "id": "msg_cached", + "type": "message", + "role": "assistant", + "model": "claude-haiku-5-5", + "content": [{"type": "text", "text": "tool result"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + request: Final = { + "caching": True, + "api_key": "test-provider-key", + "api_base": "https://native-provider.example/v1", + **( + {"model": "openai/gpt-5.6-luna", "input": "Use the tool", **_openai(["search"])} + if api == "aresponses" + else { + "model": "anthropic/claude-haiku-5-5", + "messages": [{"role": "user", "content": "Use the tool"}], + "tools": [ + { + "type": "url", + "name": "docs", + "url": "https://native.example/mcp", + "tool_configuration": {"allowed_tools": ["search"]}, + } + ], + } + ), + } + call: Final = litellm.aresponses if api == "aresponses" else litellm.acompletion + await call(**request, litellm_metadata={"user_api_key_auth": native_auth}) + await asyncio.gather(*tuple(_PENDING_CACHE_WRITES)) + await call(**request, litellm_metadata={"user_api_key_auth": native_auth}) + assert upstream.call_count == 1 + denied_auth: Final = native_auth.model_copy( + update={ + "object_permission": LiteLLM_ObjectPermissionTable( + object_permission_id="revoked", + mcp_servers=["no-mcp-servers"], + ) + } + ) + with pytest.raises(ProxyException) as denied: + await call(**request, litellm_metadata={"user_api_key_auth": denied_auth}) + assert denied.value.code == "403" + assert upstream.call_count == 1 + + +class _LateMCPAttachment(CustomLogger): + async def async_pre_request_hook(self, model: str, messages: list[object], kwargs: dict[str, object]) -> dict: + return { + **kwargs, + "mcp_servers": [ + { + "type": "url", + "name": "docs", + "url": "https://unregistered.example/mcp", + "tool_configuration": {"allowed_tools": ["search"]}, + } + ], + "tools": [], + } + + +class _LateChatMCPAttachment(CustomLogger): + def log_pre_api_call(self, model: str, messages: list, kwargs: dict) -> None: + body: Final = kwargs["additional_args"]["complete_input_dict"] # mutable-ok: exercise payload mutation by a callback + body.update( + { + "mcp_servers": [ + { + "type": "url", + "name": "docs", + "url": "https://unregistered.example/mcp", + "tool_configuration": {"allowed_tools": ["search"]}, + } + ] + } + ) + + +@pytest.mark.parametrize("api", ["anthropic_messages", "acompletion", "claude_platform_chat"]) +@pytest.mark.parametrize("stream", [False, True]) +async def test_final_body_check_blocks_attachment_added_after_deployment_hook( + native_auth: UserAPIKeyAuth, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + api: str, + stream: bool, +) -> None: + callback: Final = _LateChatMCPAttachment() if api in ("acompletion", "claude_platform_chat") else _LateMCPAttachment() + monkeypatch.setattr(litellm, "callbacks", [callback]) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + upstream: Final = respx_mock.post(url__regex=r"https://native-provider.example/.*").respond(500) + call: Final = litellm.anthropic_messages if api == "anthropic_messages" else litellm.acompletion + deployment: Final = ( + CLAUDE_PLATFORM_PARAMS + if api == "claude_platform_chat" + else { + "model": "anthropic/claude-haiku-5-5", + "api_key": "test-provider-key", + "api_base": "https://native-provider.example/v1", + } + ) + with pytest.raises((ProxyException, litellm.PermissionDeniedError)) as denied: + await call( + **deployment, + messages=[{"role": "user", "content": "hello"}], + max_tokens=32, + stream=stream, + litellm_metadata={"user_api_key_auth": native_auth}, + ) + assert str(getattr(denied.value, "status_code", getattr(denied.value, "code", ""))) == "403" + assert upstream.call_count == 0 + + +async def test_gateway_registration_does_not_authorize_provider_execution(native_auth: UserAPIKeyAuth) -> None: + manager: Final = mcp_server_manager.global_mcp_server_manager + manager.catalog.publish_config_registrations({}) + manager.config_mcp_servers = { + "gateway": MCPServer( + server_id="gateway", + name="gateway", + server_name="gateway", + url="https://native.example/mcp", + transport=MCPTransport.http, + allowed_tools=["search"], + ) + } + auth: Final = native_auth.model_copy( + update={ + "team_id": None, + "object_permission": LiteLLM_ObjectPermissionTable( + object_permission_id="gateway-grant", + mcp_servers=["gateway"], + ), + } + ) + with pytest.raises(ProxyException) as denied: + await authorize_provider_mcp(_openai(), auth) + assert denied.value.code == "403" + + +async def test_connector_id_cannot_override_a_registered_url(native_auth: UserAPIKeyAuth) -> None: + with pytest.raises(ProxyException) as denied: + await authorize_provider_mcp( + { + "tools": [ + { + "type": "mcp", + "server_url": "https://native.example/mcp", + "connector_id": "connector_other", + "allowed_tools": ["search"], + } + ] + }, + native_auth, + ) + assert denied.value.code == "403" diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index 7e044b55f5a..c9447ea2aa3 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -55,6 +55,7 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging from litellm.responses.streaming_iterator import ResponsesWsTurnAdmission from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck + from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth class TestResponsesAPIWebSocketSupport: @@ -3212,6 +3213,127 @@ def _make_streaming(**kwargs): return ResponsesWebSocketStreaming(**kwargs) +@pytest.fixture +async def provider_mcp_auth(monkeypatch: pytest.MonkeyPatch) -> "UserAPIKeyAuth": + from litellm.proxy import proxy_server + from litellm.proxy._types import LiteLLM_ObjectPermissionTable, UserAPIKeyAuth + from litellm.proxy.mcp_server import mcp_server_manager + + manager: Final = mcp_server_manager.MCPServerManager() + monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager) + monkeypatch.setattr(proxy_server, "prisma_client", None) + monkeypatch.setattr(proxy_server, "general_settings", {"enforce_provider_mcp_permissions": True}) + await manager.load_servers_from_config( + { + "docs": { + "server_id": "native-id", + "execution_mode": "provider", + "url": "https://native.example/mcp", + "allowed_tools": ["search", "fetch"], + } + } + ) + return UserAPIKeyAuth( + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="key-grant", + mcp_servers=["native-id"], + mcp_tool_permissions={"native-id": ["search"]}, + ) + ) + + +def _native_mcp_tools(tool: str = "search", url: str = "https://native.example/mcp") -> list[dict[str, object]]: + return [{"type": "mcp", "server_label": "client_label", "server_url": url, "allowed_tools": [tool]}] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("shape", ["flat", "nested", "mixed"]) +async def test_native_websocket_checks_mcp_permissions_on_first_and_later_frames( + provider_mcp_auth: "UserAPIKeyAuth", shape: str +) -> None: + allowed: Final = {"input": "search docs", "tools": _native_mcp_tools()} + denied: Final = {**allowed, "tools": _native_mcp_tools("fetch")} + frames: Final = tuple( + json.dumps( + { + "type": "response.create", + **({"response": body} if shape == "nested" else body), + **({"response": allowed} if shape == "mixed" else {}), + } + ) + for body in (denied, allowed, denied, allowed) + ) + client: Final = _ClientSocket(*frames[1:]) + backend: Final = AsyncMock() + handler: Final = _make_streaming( + websocket=client, backend_ws=backend, user_api_key_dict=provider_mcp_auth, first_message=frames[0] + ) + + await handler.client_to_backend() + + assert [call.args[0] for call in backend.send.await_args_list] == [frames[1], frames[3]] + assert [(frame["status"], frame["error"]["type"]) for frame in client.sent] == [ + (403, "tool_access_denied"), (403, "tool_access_denied") + ] + assert client.close_code is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("source", ["client", "deployment_default", "deployment_override"]) +@pytest.mark.parametrize("denial", ["url", "tool", "unbounded"]) +async def test_native_websocket_authorizes_effective_mcp_tools_before_turn_reservation( + provider_mcp_auth: "UserAPIKeyAuth", source: str, denial: str +) -> None: + from litellm.types.responses.streaming_websocket import ResponsesWebSocketRequestDefaults + + selected: Final = { + "url": _native_mcp_tools(url="https://unregistered.example/mcp"), + "tool": _native_mcp_tools("fetch"), + "unbounded": [{"type": "mcp", "server_label": "client_label", "server_url": "https://native.example/mcp"}], + }[denial] + body: Final = { + "type": "response.create", + "input": "search docs", + **({"tools": selected} if source == "client" else {}), + **({"tools": _native_mcp_tools()} if source == "deployment_override" else {}), + } + client: Final = _ClientSocket() + backend: Final = AsyncMock() + reserve: Final = AsyncMock(return_value={}) + handler: Final = _make_streaming( + websocket=client, + backend_ws=backend, + user_api_key_dict=provider_mcp_auth, + request_defaults=ResponsesWebSocketRequestDefaults( + fill_missing={"tools": selected} if source == "deployment_default" else {}, + overrides={"tools": selected} if source == "deployment_override" else {}, + ), + ) + + await handler._forward_guarded_frame(json.dumps(body), reserve) + + backend.send.assert_not_awaited() + reserve.assert_not_awaited() + assert [(frame["status"], frame["error"]["type"]) for frame in client.sent] == [(403, "tool_access_denied")] + + +@pytest.mark.asyncio +async def test_native_websocket_rechecks_registration_on_an_open_connection(provider_mcp_auth: "UserAPIKeyAuth") -> None: + from litellm.proxy.mcp_server import mcp_server_manager + + frame: Final = json.dumps({"type": "response.create", "input": "search docs", "tools": _native_mcp_tools()}) + client: Final = _ClientSocket() + backend: Final = AsyncMock() + handler: Final = _make_streaming(websocket=client, backend_ws=backend, user_api_key_dict=provider_mcp_auth) + + await handler._forward_guarded_frame(frame) + await mcp_server_manager.global_mcp_server_manager.load_servers_from_config({}) + await handler._forward_guarded_frame(frame) + + backend.send.assert_awaited_once_with(frame) + assert [(event["status"], event["error"]["type"]) for event in client.sent] == [(403, "tool_access_denied")] + + class TestNativeWebSocketGuardrailMasking: """Exercises the input/output PII masking hooks on ResponsesWebSocketStreaming."""