diff --git a/litellm/proxy/_experimental/mcp_server/capabilities.py b/litellm/proxy/_experimental/mcp_server/capabilities.py index 1ee263633be..e12c2346d1a 100644 --- a/litellm/proxy/_experimental/mcp_server/capabilities.py +++ b/litellm/proxy/_experimental/mcp_server/capabilities.py @@ -15,6 +15,8 @@ from litellm.types.mcp import ( MCP_LEGACY_VERSIONS, MCPAdvertisedVersion, MCPAdvertisedVersions, + MCPDisableableCapability, + MCPDisabledCapabilities, MCPSpecVersion, MCPTransport, ) @@ -30,6 +32,12 @@ GATEWAY_OPERATIONS: Final = frozenset( "resources/templates/list", } ) +_CAPABILITY_OPERATIONS: Final[Mapping[MCPDisableableCapability, frozenset[str]]] = MappingProxyType( + { + "prompts": frozenset({"prompts/list", "prompts/get"}), + "resources": frozenset({"resources/list", "resources/read", "resources/templates/list"}), + } +) @dataclass(frozen=True, slots=True) @@ -60,6 +68,7 @@ REVISION_SUPPORT: Final[Mapping[str, RevisionSupport]] = MappingProxyType( _COMPLETED_REVISIONS: Final = tuple(version for version, support in REVISION_SUPPORT.items() if support.completed) TRANSLATION_PAIRS: Final = frozenset(product(_COMPLETED_REVISIONS, repeat=2)) _ADVERTISED_VERSIONS: Final[TypeAdapter[tuple[MCPAdvertisedVersion, ...]]] = TypeAdapter(MCPAdvertisedVersions) +_DISABLED_CAPABILITIES: Final[TypeAdapter[MCPDisabledCapabilities]] = TypeAdapter(MCPDisabledCapabilities) def configured_versions() -> tuple[str, ...]: @@ -69,6 +78,22 @@ def configured_versions() -> tuple[str, ...]: return _ADVERTISED_VERSIONS.validate_python(MCP_LEGACY_VERSIONS if configured is None else configured) +def disabled_capabilities() -> frozenset[MCPDisableableCapability]: + from litellm.proxy.proxy_server import general_settings_view + + return frozenset( + _DISABLED_CAPABILITIES.validate_python(general_settings_view().get("mcp_disabled_capabilities") or ()) + ) + + +def disabled_operations(disabled: frozenset[MCPDisableableCapability]) -> frozenset[str]: + return frozenset[str]().union(*(_CAPABILITY_OPERATIONS[name] for name in disabled)) + + +def enabled_gateway_operations(disabled: frozenset[MCPDisableableCapability]) -> frozenset[str]: + return GATEWAY_OPERATIONS - disabled_operations(disabled) + + def build_discovery( *, configured: tuple[str, ...], @@ -118,8 +143,13 @@ def build_discovery( class GatewayVersionPolicy: - def __init__(self, versions: Callable[[], tuple[str, ...]] = configured_versions) -> None: + def __init__( + self, + versions: Callable[[], tuple[str, ...]] = configured_versions, + disabled: Callable[[], frozenset[MCPDisableableCapability]] = disabled_capabilities, + ) -> None: self._versions = versions + self._disabled = disabled async def __call__(self, ctx: ServerRequestContext[object, object], call_next: CallNext) -> HandlerResult: versions: Final = self._versions() @@ -143,7 +173,7 @@ class GatewayVersionPolicy: configured=versions, revision=initialized.protocol_version, transport=MCPTransport.http, - authorized_operations=GATEWAY_OPERATIONS, + authorized_operations=enabled_gateway_operations(self._disabled()), upstream_versions=frozenset(HANDSHAKE_PROTOCOL_VERSIONS), capabilities=initialized.capabilities, instructions=initialized.instructions, diff --git a/litellm/proxy/_experimental/mcp_server/operations.py b/litellm/proxy/_experimental/mcp_server/operations.py index 578f54fdd68..1cd14e032a3 100644 --- a/litellm/proxy/_experimental/mcp_server/operations.py +++ b/litellm/proxy/_experimental/mcp_server/operations.py @@ -64,9 +64,10 @@ from litellm.proxy._experimental.mcp_server.byok_credential_cache import ( get_cached_byok_credential, ) from litellm.proxy._experimental.mcp_server.capabilities import ( - GATEWAY_OPERATIONS, build_discovery, configured_versions, + disabled_capabilities, + enabled_gateway_operations, ) from litellm.proxy._experimental.mcp_server.catalog import catalog_operation, global_manager from litellm.proxy._experimental.mcp_server.contracts import ( @@ -3092,6 +3093,8 @@ async def _execute_list_prompts( if context.mcp_proxy_mode: _reject_mcp_proxy_operation() + if "prompts" in disabled_capabilities(): + return ListPromptsResult(prompts=[]) try: catalog: Final = MCPServerTargetCatalog(global_mcp_server_manager, listing=list_gateway_catalog) result: Final = await catalog.list(context, ListPromptsRequest(params=params)) @@ -3114,6 +3117,8 @@ async def _execute_get_prompt( ) -> GetPromptResult | InputRequiredResult: if context.mcp_proxy_mode: _reject_mcp_proxy_operation() + if "prompts" in disabled_capabilities(): + _reject_disabled_capability("prompts") ( user_api_key_auth, mcp_auth_header, @@ -3152,6 +3157,8 @@ async def _execute_list_resources( if context.mcp_proxy_mode: _reject_mcp_proxy_operation() + if "resources" in disabled_capabilities(): + return ListResourcesResult(resources=[]) try: catalog: Final = MCPServerTargetCatalog(global_mcp_server_manager, listing=list_gateway_catalog) result: Final = await catalog.list(context, ListResourcesRequest(params=params)) @@ -3180,6 +3187,8 @@ async def _execute_list_resource_templates( if context.mcp_proxy_mode: _reject_mcp_proxy_operation() + if "resources" in disabled_capabilities(): + return ListResourceTemplatesResult(resource_templates=[]) try: catalog: Final = MCPServerTargetCatalog(global_mcp_server_manager, listing=list_gateway_catalog) result: Final = await catalog.list(context, ListResourceTemplatesRequest(params=params)) @@ -3202,6 +3211,8 @@ async def _execute_read_resource( ) -> ReadResourceResult | InputRequiredResult: if context.mcp_proxy_mode: _reject_mcp_proxy_operation() + if "resources" in disabled_capabilities(): + _reject_disabled_capability("resources") ( user_api_key_auth, mcp_auth_header, @@ -3236,6 +3247,13 @@ def _reject_mcp_proxy_operation() -> NoReturn: raise MCPError(code=METHOD_NOT_FOUND, message="Operation unavailable on /mcp/proxy") +def _reject_disabled_capability(capability: str) -> NoReturn: + from mcp.shared.exceptions import MCPError + from mcp.types import METHOD_NOT_FOUND + + raise MCPError(code=METHOD_NOT_FOUND, message=f"MCP {capability} are disabled on this gateway") + + def prepare_context( user_api_key_auth: UserAPIKeyAuth | None = None, mcp_auth_header: str | None = None, @@ -3375,10 +3393,14 @@ class GatewayOperations: async def _execute(self, operation: GatewayOperation, context: OperationContext) -> GatewayResult: match operation: case DiscoverRequest(): + disabled: Final = disabled_capabilities() listings: Final = ( () if context.mcp_proxy_mode - else (ListPromptsRequest(), ListResourcesRequest(), ListResourceTemplatesRequest()) + else ( + *(() if "prompts" in disabled else (ListPromptsRequest(),)), + *(() if "resources" in disabled else (ListResourcesRequest(), ListResourceTemplatesRequest())), + ) ) memo_token: Final = _mcp_server_admission_memo.set({}) tasks: Final = ( @@ -3405,7 +3427,7 @@ class GatewayOperations: configured=configured_versions(), revision=context.protocol_version or "2025-11-25", transport=MCPTransport.http, - authorized_operations=GATEWAY_OPERATIONS, + authorized_operations=enabled_gateway_operations(disabled), upstream_versions=frozenset(MCP_LEGACY_VERSIONS), capabilities=ServerCapabilities( tools=ToolsCapability() diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index e5f9a211945..76156e9878e 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -13,7 +13,7 @@ import os import time import types from collections import Counter -from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterable, Mapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterable, Mapping, Sequence, Set from typing import TYPE_CHECKING, Final, NoReturn, Protocol import httpx @@ -569,7 +569,12 @@ if MCP_AVAILABLE: from litellm.proxy._experimental.mcp_server.auth.litellm_auth_handler import ( MCPAuthenticatedUser, ) - from litellm.proxy._experimental.mcp_server.capabilities import GatewayVersionPolicy, configured_versions + from litellm.proxy._experimental.mcp_server.capabilities import ( + GatewayVersionPolicy, + configured_versions, + disabled_capabilities, + disabled_operations, + ) from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( MCPServerManager, global_mcp_server_manager, @@ -1393,6 +1398,12 @@ if MCP_AVAILABLE: except (json.JSONDecodeError, UnicodeDecodeError, TypeError): return False + def _jsonrpc_request_method(body: bytes) -> str | None: + try: + return JSONRPCRequest.model_validate_json(body).method + except ValidationError: + return None + def _extract_initialize_client_info(body: bytes) -> Implementation | None: try: return InitializeRequest.model_validate_json(body, by_name=False).params.client_info @@ -2285,6 +2296,87 @@ if MCP_AVAILABLE: ) return None + async def _preflight_streamable_http_upstream_auth( + defer_upstream_probes: bool, + targets_disabled_capability: bool, + scope: Scope, + mcp_servers: Sequence[str] | None, + oauth2_headers: Mapping[str, str] | None, + mcp_server_auth_headers: Mapping[str, Mapping[str, str]] | None, + user_api_key_auth: UserAPIKeyAuth | None, + client_ip: str | None, + allowed_server_ids: Set[str] | None, + raw_headers: Mapping[str, str] | None, + ) -> None: + if not defer_upstream_probes and not targets_disabled_capability: + await _raise_preemptive_401_for_unauthenticated_servers( + scope=scope, + mcp_servers=list(mcp_servers) if mcp_servers is not None else None, + oauth2_headers=dict(oauth2_headers) if oauth2_headers is not None else None, + mcp_server_auth_headers=( + {name: dict(headers) for name, headers in mcp_server_auth_headers.items()} + if mcp_server_auth_headers is not None + else None + ), + user_api_key_auth=user_api_key_auth, + client_ip=client_ip, + allowed_server_ids=set(allowed_server_ids) if allowed_server_ids is not None else None, + raw_headers=raw_headers, + ) + # Pre-flight auth check for pass-through servers. Must run after + # toolset scoping so the probe list is derived from the fully-authorized + # server set, not the raw user-supplied names. + await _check_passthrough_upstream_auth( + scope, + user_api_key_auth, + list(mcp_servers) if mcp_servers is not None else None, + client_ip, + oauth2_headers=oauth2_headers, + mcp_server_auth_headers=( + {name: dict(headers) for name, headers in mcp_server_auth_headers.items()} + if mcp_server_auth_headers is not None + else None + ), + raw_headers=raw_headers, + ) + + async def _targets_disabled_capability(peek: _LazyPeekedBody | None) -> bool: + if peek is None: + return False + disabled_ops: Final = disabled_operations(disabled_capabilities()) + if not disabled_ops: + return False + return _jsonrpc_request_method(await peek.full_body()) in disabled_ops + + async def _preflight_streamable_http_modern_interaction( + peek: _LazyPeekedBody | None, + defer_upstream_probes: bool, + targets_disabled_capability: bool, + scope: Scope, + user_api_key_auth: UserAPIKeyAuth | None, + mcp_auth_header: str | None, + mcp_servers: Sequence[str] | None, + mcp_server_auth_headers: Mapping[str, Mapping[str, str]] | None, + oauth2_headers: Mapping[str, str] | None, + raw_headers: Mapping[str, str] | None, + client_ip: str | None, + ) -> JSONResponse | None: + if peek is None or not defer_upstream_probes or targets_disabled_capability: + return None + return await _preflight_modern_interaction( + scope, + await peek.full_body(), + OperationContext( + _caller=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_servers=tuple(mcp_servers) if mcp_servers is not None else None, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + client_ip=client_ip, + ), + ) + async def handle_streamable_http_mcp(scope: Scope, receive: Receive, send: Send) -> None: """Handle MCP requests through StreamableHTTP.""" try: @@ -2329,6 +2421,7 @@ if MCP_AVAILABLE: defer_upstream_probes: Final = request_headers.get( "mcp-protocol-version" ) == "2026-07-28" and request_headers.get("mcp-method") in {"tools/call", "prompts/get", "resources/read"} + targets_disabled_capability: Final = await _targets_disabled_capability(peek) # Extract client IP for MCP access control _client_ip: Final = IPAddressUtils.get_mcp_client_ip(StarletteRequest(scope)) @@ -2357,30 +2450,18 @@ if MCP_AVAILABLE: # from the fully-authorized server set: a passthrough server that # the active toolset excludes should not trigger an OAuth flow # for a server the caller will be 403'd on after authentication. - if not defer_upstream_probes: - await _raise_preemptive_401_for_unauthenticated_servers( - scope=scope, - mcp_servers=mcp_servers, - oauth2_headers=oauth2_headers, - mcp_server_auth_headers=mcp_server_auth_headers, - user_api_key_auth=user_api_key_auth, - client_ip=_client_ip, - allowed_server_ids=toolset_allowed_server_ids, - raw_headers=raw_headers, - ) - - # Pre-flight auth check for pass-through servers. Must run after - # toolset scoping so the probe list is derived from the fully-authorized - # server set, not the raw user-supplied names. - await _check_passthrough_upstream_auth( - scope, - user_api_key_auth, - mcp_servers, - _client_ip, - oauth2_headers=oauth2_headers, - mcp_server_auth_headers=mcp_server_auth_headers, - raw_headers=raw_headers, - ) + await _preflight_streamable_http_upstream_auth( + defer_upstream_probes=defer_upstream_probes, + targets_disabled_capability=targets_disabled_capability, + scope=scope, + mcp_servers=mcp_servers, + oauth2_headers=oauth2_headers, + mcp_server_auth_headers=mcp_server_auth_headers, + user_api_key_auth=user_api_key_auth, + client_ip=_client_ip, + allowed_server_ids=toolset_allowed_server_ids, + raw_headers=raw_headers, + ) # Inject masked debug headers when client sends x-litellm-mcp-debug: true _debug_headers: Final = MCPDebug.maybe_build_debug_headers( @@ -2444,23 +2525,22 @@ if MCP_AVAILABLE: return session_id = _get_session_id_from_scope(scope) - if peek is not None and defer_upstream_probes: - rejection: Final = await _preflight_modern_interaction( - scope, - await peek.full_body(), - OperationContext( - _caller=user_api_key_auth, - mcp_auth_header=mcp_auth_header, - mcp_servers=tuple(mcp_servers) if mcp_servers is not None else None, - mcp_server_auth_headers=mcp_server_auth_headers, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, - client_ip=_client_ip, - ), - ) - if rejection is not None: - await rejection(scope, receive, send) - return + rejection: Final = await _preflight_streamable_http_modern_interaction( + peek=peek, + defer_upstream_probes=defer_upstream_probes, + targets_disabled_capability=targets_disabled_capability, + scope=scope, + user_api_key_auth=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_servers=mcp_servers, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + client_ip=_client_ip, + ) + if rejection is not None: + await rejection(scope, receive, send) + return scope_response: Final = ( OAuthScopeResponse( diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 6608e152466..682a5a56ee9 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -62,6 +62,7 @@ from litellm.types.mcp import ( MCPAuth, MCPAuthType, MCPCredentials, + MCPDisabledCapabilities, MCPTransport, MCPTransportType, MCPUpstreamProtocol, @@ -2885,6 +2886,13 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): description="MCP revisions enabled by the gateway. Defaults to all completed legacy revisions. " "Modern protocol serving requires explicit opt-in. Apps/Tasks remain disabled.", ) + mcp_disabled_capabilities: MCPDisabledCapabilities | None = Field( + None, + description="MCP capabilities the gateway stops serving on every route, any of 'prompts' and 'resources'. " + "A disabled capability is left out of the initialize and server/discover responses, its list operations " + "return an empty list without contacting any upstream MCP server, and its get/read operations answer " + "METHOD_NOT_FOUND. Unset means prompts and resources stay advertised.", + ) mcp_allowed_clients: list[MCPAllowedClient] | None = Field( None, description="MCP client applications admitted by the gateway, each an {alias, value} pair where alias is the name shown in the dashboard and logs and value is the identity that must match exactly. When set, every MCP request must carry a client identity equal to one of the values: a JWT caller is identified by the claim named in litellm_jwtauth.mcp_client_id_jwt_field, any other caller by the header named in mcp_client_id_header. A request with no resolvable identity, or an unlisted one, is rejected with 403. Unset means every client is admitted.", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4459172cf58..e5bceb25ef3 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6927,6 +6927,12 @@ class ProxyConfig: TypeAdapter(MCPAdvertisedVersions).validate_python(general_settings["mcp_advertised_versions"]) + disabled_mcp_capabilities: Final = typed_general_settings.get("mcp_disabled_capabilities") + if disabled_mcp_capabilities is not None: + from litellm.types.mcp import MCPDisabledCapabilities + + TypeAdapter(MCPDisabledCapabilities).validate_python(disabled_mcp_capabilities) + if os.getenv("NUM_WORKERS", "1") != "1" and redis_usage_cache is None: warn_login_counters_are_per_worker(os.getenv("NUM_WORKERS", "1")) if declared_proxy_ranges(general_settings) is None: @@ -18814,6 +18820,7 @@ _GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES: Final[Mapping[str, str]] = MappingPro "always_include_stream_usage": "Boolean", "forward_client_headers_to_llm_api": "Boolean", "mcp_required_fields": "List", + "mcp_disabled_capabilities": "List", "cancel_on_disconnect": "Boolean", "disable_auto_add_proxy_admin_to_teams": "Boolean", "apply_user_budget_to_team_keys": "Boolean", diff --git a/litellm/types/mcp.py b/litellm/types/mcp.py index 5dcbf77993a..72e564b6919 100644 --- a/litellm/types/mcp.py +++ b/litellm/types/mcp.py @@ -73,6 +73,8 @@ def validate_mcp_protocol_transport(protocol_version: MCPUpstreamProtocol, trans MCPAdvertisedVersion: TypeAlias = MCPLegacyVersion | Literal["2026-07-28"] MCPAdvertisedVersions: TypeAlias = Annotated[tuple[MCPAdvertisedVersion, ...], Field(min_length=1)] +MCPDisableableCapability: TypeAlias = Literal["prompts", "resources"] +MCPDisabledCapabilities: TypeAlias = tuple[MCPDisableableCapability, ...] MCPSpecVersionType = Literal[ MCPSpecVersion.nov_2024, MCPSpecVersion.mar_2025, diff --git a/tests/unit/proxy/_experimental/mcp_server/test_capabilities.py b/tests/unit/proxy/_experimental/mcp_server/test_capabilities.py index f8e3655b2f2..39a6b1fcf15 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_capabilities.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_capabilities.py @@ -4,7 +4,15 @@ import pytest from mcp import Client from mcp.server import Server from mcp.shared.exceptions import MCPError -from mcp.types import PromptsCapability, ResourcesCapability, ServerCapabilities, ToolsCapability +from mcp.types import ( + ListPromptsResult, + ListResourcesResult, + PaginatedRequestParams, + PromptsCapability, + ResourcesCapability, + ServerCapabilities, + ToolsCapability, +) from mcp_types.version import HANDSHAKE_PROTOCOL_VERSIONS from litellm.proxy._experimental.mcp_server.capabilities import ( @@ -13,6 +21,8 @@ from litellm.proxy._experimental.mcp_server.capabilities import ( TRANSLATION_PAIRS, GatewayVersionPolicy, build_discovery, + disabled_capabilities, + enabled_gateway_operations, ) from litellm.types.mcp import MCPTransport @@ -27,7 +37,9 @@ def test_discovery_only_exposes_authorized_completed_support(revision, transport authorized_operations=frozenset({"tools/list", "tools/call"}), upstream_versions=frozenset(HANDSHAKE_PROTOCOL_VERSIONS), capabilities=ServerCapabilities( - tools=ToolsCapability(), prompts=PromptsCapability(), resources=ResourcesCapability(), + tools=ToolsCapability(), + prompts=PromptsCapability(), + resources=ResourcesCapability(), extensions={"io.modelcontextprotocol/ui": {}}, ), client_extensions=frozenset({"io.modelcontextprotocol/ui"}), @@ -59,18 +71,26 @@ def test_unproven_translation_never_advertises_operations(upstream): @pytest.mark.parametrize("revision", ["2026-07-28", "unknown", "2024-11-05"]) def test_unadvertised_revision_never_gains_capabilities(revision): result = build_discovery( - configured=("2025-11-25",), revision=revision, transport=MCPTransport.http, - authorized_operations=GATEWAY_OPERATIONS, upstream_versions=frozenset(HANDSHAKE_PROTOCOL_VERSIONS), + configured=("2025-11-25",), + revision=revision, + transport=MCPTransport.http, + authorized_operations=GATEWAY_OPERATIONS, + upstream_versions=frozenset(HANDSHAKE_PROTOCOL_VERSIONS), capabilities=ServerCapabilities(tools=ToolsCapability()), ) assert result.capabilities.tools is None def test_discovery_results_do_not_share_mutable_capabilities(): - capabilities = ServerCapabilities(tools=ToolsCapability(), prompts=PromptsCapability(), resources=ResourcesCapability()) + capabilities = ServerCapabilities( + tools=ToolsCapability(), prompts=PromptsCapability(), resources=ResourcesCapability() + ) args = dict( - configured=HANDSHAKE_PROTOCOL_VERSIONS, revision="2025-11-25", transport=MCPTransport.http, - upstream_versions=frozenset(HANDSHAKE_PROTOCOL_VERSIONS), capabilities=capabilities, + configured=HANDSHAKE_PROTOCOL_VERSIONS, + revision="2025-11-25", + transport=MCPTransport.http, + upstream_versions=frozenset(HANDSHAKE_PROTOCOL_VERSIONS), + capabilities=capabilities, ) allowed = build_discovery(**args, authorized_operations=GATEWAY_OPERATIONS) denied = build_discovery(**args, authorized_operations=frozenset()) @@ -101,7 +121,9 @@ async def test_version_policy_gates_the_actual_sdk_handshake(versions, accepted) result = await client.session.send_ping() assert result is not None else: - with pytest.RaisesGroup(pytest.RaisesExc(MCPError, match="Unsupported MCP protocol version"), flatten_subgroups=True): + with pytest.RaisesGroup( + pytest.RaisesExc(MCPError, match="Unsupported MCP protocol version"), flatten_subgroups=True + ): async with Client(server, mode="legacy"): pytest.fail("The excluded revision must not initialize") @@ -140,3 +162,54 @@ def test_modern_protocol_is_rejected_on_legacy_sse_paths(path: str, monkeypatch: unsupported_protocol_version({"path": path, "headers": [(b"mcp-protocol-version", b"2026-07-28")]}) == "2026-07-28" ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "disabled,prompts_advertised,resources_advertised", + [ + (frozenset(), True, True), + (frozenset({"prompts"}), False, True), + (frozenset({"resources"}), True, False), + (frozenset({"prompts", "resources"}), False, False), + ], +) +async def test_version_policy_hides_disabled_capabilities_from_the_handshake( + disabled, prompts_advertised, resources_advertised +) -> None: + server: Final = Server("test-gateway", version="1") + + async def list_prompts(ctx: object, params: PaginatedRequestParams) -> ListPromptsResult: + return ListPromptsResult(prompts=[]) + + async def list_resources(ctx: object, params: PaginatedRequestParams) -> ListResourcesResult: + return ListResourcesResult(resources=[]) + + server.add_request_handler("prompts/list", PaginatedRequestParams, list_prompts) + server.add_request_handler("resources/list", PaginatedRequestParams, list_resources) + server.middleware.append(GatewayVersionPolicy(lambda: ("2025-11-25",), lambda: disabled)) + async with Client(server, mode="legacy") as client: + assert (client.server_capabilities.prompts is not None) is prompts_advertised + assert (client.server_capabilities.resources is not None) is resources_advertised + + +@pytest.mark.parametrize( + "configured,expected", + [ + (None, frozenset()), + (["prompts"], frozenset({"prompts"})), + (["prompts", "resources"], frozenset({"prompts", "resources"})), + ], +) +def test_disabled_capabilities_follow_general_settings(configured, expected, monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy import proxy_server + + monkeypatch.setitem(proxy_server.general_settings, "mcp_disabled_capabilities", configured) + assert disabled_capabilities() == expected + + +def test_enabled_gateway_operations_only_drop_the_disabled_capability() -> None: + assert enabled_gateway_operations(frozenset()) == GATEWAY_OPERATIONS + remaining: Final = enabled_gateway_operations(frozenset({"resources"})) + assert {"tools/list", "tools/call", "prompts/list", "prompts/get"} <= remaining + assert not any(operation.startswith("resources/") for operation in remaining) diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_stale_session.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_stale_session.py index bf1f0f97a18..66d853bdd44 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_stale_session.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_stale_session.py @@ -8,6 +8,8 @@ they may send a stale `mcp-session-id` header. This test verifies that: """ import asyncio +import json +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -1737,3 +1739,150 @@ async def test_handle_streamable_http_mcp_true_passthrough_with_token_skips_prob assert mock_handle_request.await_count == 1 probe_client.post.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "disabled,method,params,expected_preflight_calls,extra_request_headers,expected_modern_preflight_calls", + [ + (("prompts", "resources"), "prompts/list", {}, 0, (), 0), + (("prompts", "resources"), "prompts/get", {"name": "x"}, 0, (), 0), + (("prompts", "resources"), "resources/list", {}, 0, (), 0), + (("prompts", "resources"), "resources/read", {"uri": "file:///x"}, 0, (), 0), + (("prompts", "resources"), "resources/templates/list", {}, 0, (), 0), + (("prompts", "resources"), "tools/list", {}, 1, (), 0), + (None, "prompts/list", {}, 1, (), 0), + (("prompts",), "resources/list", {}, 1, (), 0), + ( + ("prompts", "resources"), + "prompts/get", + {"name": "x"}, + 0, + ((b"mcp-protocol-version", b"2026-07-28"), (b"mcp-method", b"prompts/get")), + 0, + ), + ( + ("prompts", "resources"), + "resources/read", + {"uri": "file:///x"}, + 0, + ((b"mcp-protocol-version", b"2026-07-28"), (b"mcp-method", b"resources/read")), + 0, + ), + ( + ("prompts", "resources"), + "tools/call", + {"name": "tp_server-add", "arguments": {}}, + 0, + ((b"mcp-protocol-version", b"2026-07-28"), (b"mcp-method", b"tools/call")), + 1, + ), + ( + None, + "prompts/get", + {"name": "x"}, + 0, + ((b"mcp-protocol-version", b"2026-07-28"), (b"mcp-method", b"prompts/get")), + 1, + ), + ( + ("prompts", "resources"), + "prompts/get", + {"name": "x", "arguments": {"pad": "a" * 5000}}, + 0, + ((b"mcp-protocol-version", b"2026-07-28"), (b"mcp-method", b"prompts/get")), + 0, + ), + (None, "prompts/get", {"name": "x", "arguments": {"pad": "a" * 5000}}, 1, (), 0), + ], +) +async def test_streamable_http_skips_preflight_for_disabled_methods( # test-quality-ok: no injectable preflight seam + disabled: tuple[str, ...] | None, + method: str, + params: dict[str, object], + expected_preflight_calls: int, + extra_request_headers: tuple[tuple[bytes, bytes], ...], + expected_modern_preflight_calls: int, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server.server import ( + handle_streamable_http_mcp, + session_manager_stateless, + ) + + monkeypatch.setitem(proxy_server.general_settings, "mcp_advertised_versions", ["2026-07-28"]) + if disabled is None: + monkeypatch.delitem(proxy_server.general_settings, "mcp_disabled_capabilities", raising=False) + else: + monkeypatch.setitem(proxy_server.general_settings, "mcp_disabled_capabilities", list(disabled)) + + probe_client: Final = MagicMock() + probe_client.post = AsyncMock() + request_headers: Final = [ + (b"authorization", b"Bearer upstream-token"), + *extra_request_headers, + ] + scope: Final = _passthrough_mode_scope( + "tp_server", + extra_headers=request_headers, + ) + body: Final[bytes] = json.dumps( + {"jsonrpc": "2.0", "id": 1, "method": method, "params": params} + ).encode() + receive: Final = AsyncMock(return_value={"type": "http.request", "body": body, "more_body": False}) + send: Final = AsyncMock() + user_auth: Final = UserAPIKeyAuth() + user_auth.user_id = None + tp_server: Final = _build_passthrough_mode_server("tp_server", MCPAuth.true_passthrough) + + with ( + patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=(user_auth, None, ["tp_server"], None, None, None), + ), + patch("litellm.proxy._experimental.mcp_server.server.set_auth_context"), + patch( + "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", + True, + ), + patch( + "litellm.proxy._experimental.mcp_server.server._handle_stale_mcp_session", + new_callable=AsyncMock, + return_value=False, + ), + patch( + "litellm.proxy._experimental.mcp_server.server._raise_preemptive_401_for_unauthenticated_servers", + new_callable=AsyncMock, + ) as mock_preemptive_401, + patch( + "litellm.proxy._experimental.mcp_server.server._check_passthrough_upstream_auth", + new_callable=AsyncMock, + ) as mock_passthrough_auth, + patch( + "litellm.proxy._experimental.mcp_server.server._preflight_modern_interaction", + new_callable=AsyncMock, + return_value=None, + ) as mock_modern_preflight, + patch( + "litellm.proxy._experimental.mcp_server.server.get_async_httpx_client", + return_value=probe_client, + ), + patch( + "litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager.get_mcp_server_by_name", + return_value=tp_server, + ), + patch.object( + session_manager_stateless, + "handle_request", + new_callable=AsyncMock, + ) as mock_handle_request, + patch.object(session_manager_stateless, "_server_instances", {}), + ): + await handle_streamable_http_mcp(scope, receive, send) + + assert mock_preemptive_401.await_count == expected_preflight_calls + assert mock_passthrough_auth.await_count == expected_preflight_calls + assert mock_modern_preflight.await_count == expected_modern_preflight_calls + assert mock_handle_request.await_count == 1 diff --git a/tests/unit/proxy/_experimental/mcp_server/test_operations.py b/tests/unit/proxy/_experimental/mcp_server/test_operations.py index e96664aae84..5d8db1e6e2b 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_operations.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_operations.py @@ -48,8 +48,11 @@ async def _served_catalog_tool() -> str: @pytest.mark.parametrize("mode", ["direct", "virtual", "proxy"]) @pytest.mark.parametrize("verdict", ["allow", "block", "obo_error", "evaluate_error", "missing"]) async def test_mapped_caller_bearer_reaches_agent_evaluation( - monkeypatch: pytest.MonkeyPatch, _mcp_request_ctx: Callable[..., ServerRequestContext], - mode: str, verdict: str, caplog: pytest.LogCaptureFixture, + monkeypatch: pytest.MonkeyPatch, + _mcp_request_ctx: Callable[..., ServerRequestContext], + mode: str, + verdict: str, + caplog: pytest.LogCaptureFixture, ) -> None: from urllib.parse import parse_qs @@ -66,18 +69,31 @@ async def test_mapped_caller_bearer_reaches_agent_evaluation( from litellm.proxy.guardrails.guardrail_hooks.agent_365.agent_365 import Agent365Guardrail bearer: Final = "signed.caller.assertion" - request: Final = Request({ - "type": "http", "method": "POST", "path": "/mcp", - "headers": [] if verdict == "missing" else [(b"authorization", f"Bearer {bearer}".encode())], - }) + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": "/mcp", + "headers": [] if verdict == "missing" else [(b"authorization", f"Bearer {bearer}".encode())], + } + ) auth: Final = UserAPIKeyAuth( - api_key="stored-key-hash", user_id="mapped-caller", jwt_claims={"sub": "mapped-caller"}, - object_permission={"object_permission_id": "mapped", "mcp_servers": ["guarded"], "mcp_tool_search_enabled": mode == "virtual"}, + api_key="stored-key-hash", + user_id="mapped-caller", + jwt_claims={"sub": "mapped-caller"}, + object_permission={ + "object_permission_id": "mapped", + "mcp_servers": ["guarded"], + "mcp_tool_search_enabled": mode == "virtual", + }, ) oauth, raw, _, _ = MCPRequestHandler.scrub_gateway_admission_credentials( - admitted=False, admitted_credential=bearer, + admitted=False, + admitted_credential=bearer, oauth2_headers={"Authorization": f"Bearer {bearer}"}, - raw_headers=dict(request.headers), mcp_auth_header=None, mcp_server_auth_headers=None, + raw_headers=dict(request.headers), + mcp_auth_header=None, + mcp_server_auth_headers=None, ) assert "authorization" not in raw provider_requests: Final[asyncio.Queue[str]] = asyncio.Queue() @@ -92,9 +108,13 @@ async def test_mapped_caller_bearer_reaches_agent_evaluation( assert exchange.headers["authorization"] == "Bearer delegated-token" if verdict == "evaluate_error": return httpx.Response(503, json={"error": "unavailable"}) - return httpx.Response(200, json={ - "allowed": verdict == "allow", "defender": {"status": "Evaluated", "verdict": "Allow" if verdict == "allow" else "Block"}, - }) + return httpx.Response( + 200, + json={ + "allowed": verdict == "allow", + "defender": {"status": "Evaluated", "verdict": "Allow" if verdict == "allow" else "Block"}, + }, + ) executions: Final[asyncio.Queue[str]] = asyncio.Queue() @@ -107,36 +127,55 @@ async def test_mapped_caller_bearer_reaches_agent_evaluation( handler: Final = AsyncHTTPHandler(transport=httpx.MockTransport(provider)) guardrail: Final = Agent365Guardrail( - guardrail_name="agent-evaluation", tenant_id="tenant", client_id="gateway", client_secret="secret", - async_handler=handler, event_hook="pre_mcp_call", default_on=True, + guardrail_name="agent-evaluation", + tenant_id="tenant", + client_id="gateway", + client_secret="secret", + async_handler=handler, + event_hook="pre_mcp_call", + default_on=True, ) monkeypatch.setattr(litellm, "callbacks", [guardrail]) managed: Final = MCPServer( - server_id="guarded", name="guarded", transport=MCPTransport.http, - spec_path="/catalog.yaml", allow_all_keys=True, + server_id="guarded", + name="guarded", + transport=MCPTransport.http, + spec_path="/catalog.yaml", + allow_all_keys=True, ) monkeypatch.setitem(operations.global_mcp_server_manager.registry, managed.server_id, managed) global_mcp_tool_registry.register_tool( - name="guarded-echo", description="Echo", input_schema={"type": "object"}, handler=tool, + name="guarded-echo", + description="Echo", + input_schema={"type": "object"}, + handler=tool, ) previous: Final = auth_context_var.set(None) proxy_mode: Final = _mcp_proxy_mode.set(mode == "proxy") - tool_id: Final = mcp_proxy_tool_id(with_mcp_proxy_identity( - MCPTool(name="guarded-echo", input_schema={"type": "object"}), managed.server_id, - )) + tool_id: Final = mcp_proxy_tool_id( + with_mcp_proxy_identity( + MCPTool(name="guarded-echo", input_schema={"type": "object"}), + managed.server_id, + ) + ) params: Final = { "direct": CallToolRequestParams(name="guarded-echo", arguments={}), - "virtual": CallToolRequestParams(name="mcp_tool_call", arguments={"tool_name": "guarded-echo", "arguments": {}}), + "virtual": CallToolRequestParams( + name="mcp_tool_call", arguments={"tool_name": "guarded-echo", "arguments": {}} + ), "proxy": CallToolRequestParams(name="call_tool", arguments={"tool_id": tool_id, "arguments": {}}), }[mode] try: gateway.set_auth_context( - auth, mcp_servers=[managed.server_id], oauth2_headers=oauth, + auth, + mcp_servers=[managed.server_id], + oauth2_headers=oauth, raw_headers={"authorization": "Bearer previous.caller.assertion"} if verdict == "missing" else raw, ) async with gateway._legacy_operation_context(_mcp_request_ctx(request=request), trace=False) as context: result: Final = await GatewayOperations().execute( - CallToolRequest(params=params), context, + CallToolRequest(params=params), + context, ) assert isinstance(result, CallToolResult) assert result.is_error is (verdict != "allow") @@ -162,12 +201,16 @@ async def test_listing_records_only_tools_the_caller_received( ) -> None: manager: Final = operations.global_mcp_server_manager server: Final = MCPServer( - server_id="served-catalog", name="served-catalog", transport=MCPTransport.http, - spec_path="/catalog.yaml", allow_all_keys=True, + server_id="served-catalog", + name="served-catalog", + transport=MCPTransport.http, + spec_path="/catalog.yaml", + allow_all_keys=True, allowed_tools=["echo"] if restriction == "server" else None, ) auth: Final = UserAPIKeyAuth( - api_key="sk-served-catalog", user_id="lister", + api_key="sk-served-catalog", + user_id="lister", object_permission={ "object_permission_id": "served-permission", "mcp_servers": [server.server_id], @@ -181,18 +224,24 @@ async def test_listing_records_only_tools_the_caller_received( monkeypatch.setattr(litellm, "callbacks", [capture]) for name in ("echo", "status"): global_mcp_tool_registry.register_tool( - name=f"served-catalog-{name}", description=f"{name} description", - input_schema={"type": "object"}, handler=_served_catalog_tool, + name=f"served-catalog-{name}", + description=f"{name} description", + input_schema={"type": "object"}, + handler=_served_catalog_tool, ) try: if surface == "mcp": listing: Final = await operations._list_mcp_tools( - user_api_key_auth=auth, mcp_servers=[server.server_id], record_listing=True, + user_api_key_auth=auth, + mcp_servers=[server.server_id], + record_listing=True, ) assert [tool.name for tool in listing.tools] == ["served-catalog-echo"] else: rest_listing: Final = await rest_endpoints._get_tools_for_single_server( - server, None, user_api_key_auth=auth, + server, + None, + user_api_key_auth=auth, ) assert [tool.name for tool in rest_listing] == ["echo"] granted: Final = auth.model_copy(update={"object_permission": None}) @@ -203,7 +252,10 @@ async def test_listing_records_only_tools_the_caller_received( assert (served.description, served.input_schema) == ("echo description", {"type": "object"}) server.allowed_tools = None result: Final = await manager.call_tool( - server_name=server.server_id, name="status", arguments={}, user_api_key_auth=granted, + server_name=server.server_id, + name="status", + arguments={}, + user_api_key_auth=granted, proxy_logging_obj=ProxyLogging(user_api_key_cache=UserApiKeyCache()), ) assert result.is_error is False @@ -403,6 +455,54 @@ def _catalog_case(method): return cases[method] +@pytest.mark.asyncio +@pytest.mark.parametrize( + "disabled,operation", + [ + ("prompts", "prompts/list"), + ("prompts", "prompts/get"), + ("resources", "resources/list"), + ("resources", "resources/templates/list"), + ("resources", "resources/read"), + ], +) +async def test_disabled_capability_never_reaches_upstream( + disabled: str, operation: str, monkeypatch: pytest.MonkeyPatch +) -> None: + from unittest.mock import MagicMock + + from mcp.shared.exceptions import MCPError + from mcp.types import METHOD_NOT_FOUND + + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server import catalog + + monkeypatch.setitem(proxy_server.general_settings, "mcp_disabled_capabilities", [disabled]) + request, _client_method, manager_method, upstream_value, result_field = _catalog_case(operation) + server: Final = MCPServer( + server_id="catalog-disabled", name="catalog", server_name="catalog", transport=MCPTransport.http + ) + upstream: Final = AsyncMock(return_value=upstream_value) + manager: Final = operations.global_mcp_server_manager + context: Final = prepare_context(UserAPIKeyAuth(api_key=hash_token("sk-disabled")), mcp_servers=[server.server_id]) + with ( + patch.object(operations, "_get_allowed_mcp_servers", AsyncMock(return_value=[server])), + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock(enforce_mcp_server_rate_limits=AsyncMock())), + patch("litellm.proxy.proxy_server.prisma_client", None), + patch.dict(manager.registry, {server.server_id: server}), + patch.object(manager, manager_method, upstream), + patch.object(catalog, "fetch_optional_catalog_page", upstream), + ): + if result_field is None: + with pytest.raises(MCPError) as rejected: + await GatewayOperations().execute(request, context) + assert rejected.value.error.code == METHOD_NOT_FOUND + else: + result = await GatewayOperations().execute(request, context) + assert getattr(result, result_field) == [] + upstream.assert_not_awaited() + + def _mcp_rate_limited_proxy_logging() -> ProxyLogging: proxy_logging: Final = ProxyLogging(user_api_key_cache=UserApiKeyCache()) proxy_logging.proxy_hook_mapping["parallel_request_limiter"] = PROXY_MaxParallelRequestsHandler_v3( @@ -452,9 +552,7 @@ async def test_mcp_server_rpm_limits_every_catalog_operation(operation: str) -> resources=[types.Resource(name="document", uri="https://example.com/document")] ), "resources/templates/list": types.ListResourceTemplatesResult( - resource_templates=[ - types.ResourceTemplate(name="document", uri_template="https://example.com/{name}") - ] + resource_templates=[types.ResourceTemplate(name="document", uri_template="https://example.com/{name}")] ), "prompts/get": GetPromptResult(messages=[]), "resources/read": types.ReadResourceResult(contents=[]), @@ -537,9 +635,7 @@ async def test_tools_call_warmup_does_not_consume_mcp_server_rpm() -> None: ) caller: Final = UserAPIKeyAuth(api_key=hash_token("sk-catalog-warmup")) proxy_logging: Final = _mcp_rate_limited_proxy_logging() - upstream: Final = AsyncMock( - return_value=[types.Tool(name="echo", inputSchema={"type": "object"})] - ) + upstream: Final = AsyncMock(return_value=[types.Tool(name="echo", inputSchema={"type": "object"})]) with ( patch.object(operations, "_get_allowed_mcp_servers", AsyncMock(return_value=[server])), patch.object(operations.global_mcp_server_manager, "server_exposes_tool", return_value=False), @@ -565,7 +661,6 @@ async def test_tools_call_warmup_does_not_consume_mcp_server_rpm() -> None: assert [tool.name for tool in listing.tools] == ["echo"] - def _scoped_rig( monkeypatch: pytest.MonkeyPatch, ) -> tuple[MCPServer, MCPServer, AsyncMock, AsyncMock]: @@ -585,9 +680,7 @@ def _scoped_rig( manager.config_mcp_servers = {server.server_id: server for server in (memory, other)} monkeypatch.setattr(proxy_server, "prisma_client", None) monkeypatch.setattr(operations, "global_mcp_server_manager", manager) - monkeypatch.setattr( - manager, "get_allowed_mcp_servers", AsyncMock(return_value=[memory.server_id, other.server_id]) - ) + monkeypatch.setattr(manager, "get_allowed_mcp_servers", AsyncMock(return_value=[memory.server_id, other.server_id])) monkeypatch.setattr(global_mcp_tool_registry, "published_tools", {}) listing: Final = AsyncMock() monkeypatch.setattr(operations, "_get_tools_from_mcp_servers", listing) @@ -647,6 +740,7 @@ async def test_bare_tool_name_without_a_single_server_scope_is_still_unresolved_ listing.assert_not_awaited() handled.assert_not_awaited() + @pytest.mark.asyncio async def test_tools_call_pre_call_check_enforces_mcp_server_rpm() -> None: from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager @@ -794,8 +888,13 @@ async def test_catalog_helpers_preserve_context_results_and_failure_policy(metho if collection: helper = getattr(operations, "_list_mcp_" + collection) result = await helper( - user_api_key_auth=caller, mcp_auth_header=None, mcp_servers=["catalog"], - mcp_server_auth_headers=None, oauth2_headers=None, raw_headers=headers, client_ip="192.0.2.41", + user_api_key_auth=caller, + mcp_auth_header=None, + mcp_servers=["catalog"], + mcp_server_auth_headers=None, + oauth2_headers=None, + raw_headers=headers, + client_ip="192.0.2.41", ) assert result == (payload if state == "success" else []) else: @@ -1107,7 +1206,9 @@ async def test_discovery_preserves_caller_scope_and_proxy_restrictions(): from mcp.types import DiscoverRequest, ListToolsResult, Tool listed = AsyncMock(return_value=ListToolsResult(tools=[Tool(name="allowed", input_schema={"type": "object"})])) - context = prepare_context(UserAPIKeyAuth(user_id="scoped"), mcp_servers=["only-this"], mcp_proxy_mode=True, protocol_version="2025-06-18") + context = prepare_context( + UserAPIKeyAuth(user_id="scoped"), mcp_servers=["only-this"], mcp_proxy_mode=True, protocol_version="2025-06-18" + ) with patch("litellm.proxy._experimental.mcp_server.operations._execute_handle_list_tools", listed): result = await GatewayOperations().execute(DiscoverRequest(), context) assert result.capabilities.tools is not None @@ -1134,16 +1235,34 @@ async def test_discovery_denial_cannot_advertise_tools(): @pytest.mark.parametrize("available", ["none", "resources", "templates", "prompts"]) async def test_discovery_lists_each_capability_with_the_same_caller(available): from mcp.types import ( - DiscoverRequest, ListToolsResult, ListPromptsResult, ListResourcesResult, - ListResourceTemplatesResult, Prompt, Resource, ResourceTemplate, + DiscoverRequest, + ListToolsResult, + ListPromptsResult, + ListResourcesResult, + ListResourceTemplatesResult, + Prompt, + Resource, + ResourceTemplate, ) from litellm.proxy._experimental.mcp_server import operations context = prepare_context(UserAPIKeyAuth(user_id="scoped"), mcp_servers=["authorized"]) tools = AsyncMock(return_value=ListToolsResult(tools=[])) - prompts = AsyncMock(return_value=ListPromptsResult(prompts=[Prompt(name="allowed")] if available == "prompts" else [])) - resources = AsyncMock(return_value=ListResourcesResult(resources=[Resource(name="allowed", uri="test://allowed")] if available == "resources" else [])) - templates = AsyncMock(return_value=ListResourceTemplatesResult(resource_templates=[ResourceTemplate(name="allowed", uri_template="test://{id}")] if available == "templates" else [])) + prompts = AsyncMock( + return_value=ListPromptsResult(prompts=[Prompt(name="allowed")] if available == "prompts" else []) + ) + resources = AsyncMock( + return_value=ListResourcesResult( + resources=[Resource(name="allowed", uri="test://allowed")] if available == "resources" else [] + ) + ) + templates = AsyncMock( + return_value=ListResourceTemplatesResult( + resource_templates=[ResourceTemplate(name="allowed", uri_template="test://{id}")] + if available == "templates" + else [] + ) + ) with ( patch.object(operations, "_execute_handle_list_tools", tools), patch.object(operations, "_execute_list_prompts", prompts), @@ -1158,6 +1277,50 @@ async def test_discovery_lists_each_capability_with_the_same_caller(available): assert listing.await_args.args[0] is context +@pytest.mark.asyncio +@pytest.mark.parametrize("disabled", ["prompts", "resources"]) +async def test_discovery_skips_disabled_capability_listings(disabled: str, monkeypatch: pytest.MonkeyPatch) -> None: + from mcp.types import ( + DiscoverRequest, + ListPromptsResult, + ListResourcesResult, + ListResourceTemplatesResult, + ListToolsResult, + Prompt, + Resource, + ResourceTemplate, + ) + + from litellm.proxy import proxy_server + + monkeypatch.setitem(proxy_server.general_settings, "mcp_disabled_capabilities", [disabled]) + context: Final = prepare_context(UserAPIKeyAuth(user_id="scoped"), mcp_servers=["authorized"]) + tools: Final = AsyncMock(return_value=ListToolsResult(tools=[])) + prompts: Final = AsyncMock(return_value=ListPromptsResult(prompts=[Prompt(name="allowed")])) + resources: Final = AsyncMock( + return_value=ListResourcesResult(resources=[Resource(name="allowed", uri="test://allowed")]) + ) + templates: Final = AsyncMock( + return_value=ListResourceTemplatesResult( + resource_templates=[ResourceTemplate(name="allowed", uri_template="test://{id}")] + ) + ) + with ( + patch.object(operations, "_execute_handle_list_tools", tools), + patch.object(operations, "_execute_list_prompts", prompts), + patch.object(operations, "_execute_list_resources", resources), + patch.object(operations, "_execute_list_resource_templates", templates), + ): + result = await GatewayOperations().execute(DiscoverRequest(), context) + assert (result.capabilities.prompts is None) == (disabled == "prompts") + assert (result.capabilities.resources is None) == (disabled == "resources") + if disabled == "prompts": + prompts.assert_not_awaited() + else: + resources.assert_not_awaited() + templates.assert_not_awaited() + + @pytest.mark.asyncio async def test_discovery_shares_one_server_admission_across_catalog_listings() -> None: from mcp import types @@ -1214,6 +1377,7 @@ async def test_discovery_shares_one_server_admission_across_catalog_listings() - AsyncMock(side_effect=fetch_resources), AsyncMock(side_effect=fetch_resource_templates), ) + async def fetch_optional_page( _context: OperationContext, request: types.ListPromptsRequest | types.ListResourcesRequest | types.ListResourceTemplatesRequest, @@ -1257,13 +1421,24 @@ async def test_discovery_shares_one_server_admission_across_catalog_listings() - @pytest.mark.asyncio @pytest.mark.parametrize("outcome", ["success", "failure", "cancel"]) async def test_discovery_concurrent_listings_drain_on_failure_and_cancellation(outcome): - from mcp.types import DiscoverRequest, ListToolsResult, ListPromptsResult, ListResourcesResult, ListResourceTemplatesResult + from mcp.types import ( + DiscoverRequest, + ListToolsResult, + ListPromptsResult, + ListResourcesResult, + ListResourceTemplatesResult, + ) from litellm.proxy._experimental.mcp_server import operations ready = [asyncio.Event() for _ in range(4)] closed = [asyncio.Event() for _ in range(4)] release = asyncio.Event() - responses = (ListToolsResult(tools=[]), ListPromptsResult(prompts=[]), ListResourcesResult(resources=[]), ListResourceTemplatesResult(resource_templates=[])) + responses = ( + ListToolsResult(tools=[]), + ListPromptsResult(prompts=[]), + ListResourcesResult(resources=[]), + ListResourceTemplatesResult(resource_templates=[]), + ) def listing(index): async def run(*args, **kwargs): @@ -1277,6 +1452,7 @@ async def test_discovery_concurrent_listings_drain_on_failure_and_cancellation(o return responses[index] finally: closed[index].set() + return run with ( @@ -1285,7 +1461,9 @@ async def test_discovery_concurrent_listings_drain_on_failure_and_cancellation(o patch.object(operations, "_execute_list_resources", side_effect=listing(2)), patch.object(operations, "_execute_list_resource_templates", side_effect=listing(3)), ): - task = asyncio.create_task(GatewayOperations().execute(DiscoverRequest(), prepare_context(UserAPIKeyAuth(user_id="scoped")))) + task = asyncio.create_task( + GatewayOperations().execute(DiscoverRequest(), prepare_context(UserAPIKeyAuth(user_id="scoped"))) + ) try: await asyncio.gather(*(event.wait() for event in ready)) if outcome == "cancel": @@ -1314,7 +1492,8 @@ async def test_tools_listing_preserves_explicit_spend_log_policy(log_enabled): listing = AsyncMock(return_value=operations.AggregateToolListing(tools=[], outcomes={})) with patch.object(operations, "_list_mcp_tools", listing): result = await operations._execute_handle_list_tools( - prepare_context(UserAPIKeyAuth(user_id="caller")), PaginatedRequestParams(), + prepare_context(UserAPIKeyAuth(user_id="caller")), + PaginatedRequestParams(), log_list_tools_to_spendlogs=log_enabled, ) assert result.tools == [] @@ -1351,7 +1530,10 @@ async def test_list_mcp_tools_records_the_catalog_only_when_asked( @pytest.mark.asyncio async def test_discovery_keeps_one_catalog_revision_across_concurrent_listings(monkeypatch): from mcp.types import ( - DiscoverRequest, ListToolsResult, ListPromptsResult, ListResourcesResult, + DiscoverRequest, + ListToolsResult, + ListPromptsResult, + ListResourcesResult, ListResourceTemplatesResult, ) from litellm.proxy import proxy_server @@ -1365,8 +1547,12 @@ async def test_discovery_keeps_one_catalog_revision_across_concurrent_listings(m monkeypatch.setattr(proxy_server, "prisma_client", None) monkeypatch.setattr(operations, "global_mcp_server_manager", manager) observed = [] - responses = (ListToolsResult(tools=[]), ListPromptsResult(prompts=[]), - ListResourcesResult(resources=[]), ListResourceTemplatesResult(resource_templates=[])) + responses = ( + ListToolsResult(tools=[]), + ListPromptsResult(prompts=[]), + ListResourcesResult(resources=[]), + ListResourceTemplatesResult(resource_templates=[]), + ) def listing(index): async def run(*args, **kwargs): @@ -1374,6 +1560,7 @@ async def test_discovery_keeps_one_catalog_revision_across_concurrent_listings(m observed.append(manager.get_mcp_server_by_id(original.server_id).name) manager.registry = {updated.server_id: updated} return responses[index] + return run with ( @@ -1402,8 +1589,12 @@ async def test_local_handler_freshness_tracks_registered_owner_with_overlapping_ monkeypatch.setattr(proxy_server, "prisma_client", None) manager = MCPServerManager() now = datetime.now(timezone.utc) - private = MCPServer(server_id="private", name="billing", alias="billing", transport=MCPTransport.http, updated_at=now) - allowed = MCPServer(server_id="allowed", name="billing_admin", alias="billing-admin", transport=MCPTransport.http, updated_at=now) + private = MCPServer( + server_id="private", name="billing", alias="billing", transport=MCPTransport.http, updated_at=now + ) + allowed = MCPServer( + server_id="allowed", name="billing_admin", alias="billing-admin", transport=MCPTransport.http, updated_at=now + ) manager.config_mcp_servers = {server.server_id: server for server in (private, allowed)} monkeypatch.setattr(operations, "global_mcp_server_manager", manager) monkeypatch.setattr(global_mcp_tool_registry, "published_tools", {}) @@ -1411,7 +1602,9 @@ async def test_local_handler_freshness_tracks_registered_owner_with_overlapping_ global_mcp_tool_registry.register_tool("billing-admin-export", "export", {}, handler, server_id=private.server_id) async with manager.catalog.operation(): - manager.config_mcp_servers[changed_server_id] = manager.config_mcp_servers[changed_server_id].model_copy(update={"updated_at": now + timedelta(seconds=1)}) + manager.config_mcp_servers[changed_server_id] = manager.config_mcp_servers[changed_server_id].model_copy( + update={"updated_at": now + timedelta(seconds=1)} + ) if changed_server_id == "private": with pytest.raises(HTTPException) as denied: await operations._handle_local_mcp_tool("billing-admin-export", {}) @@ -1540,9 +1733,15 @@ async def test_continuation_preserves_current_authority_unavailable_error(monkey monkeypatch.setattr(proxy_server, "prisma_client", None) caller = UserAPIKeyAuth(api_key="sk-owned-key-without-database") caller.via_virtual_key = True - request = {"prompts/list": ListPromptsRequest, "resources/list": ListResourcesRequest, "resources/templates/list": ListResourceTemplatesRequest}[method] + request = { + "prompts/list": ListPromptsRequest, + "resources/list": ListResourcesRequest, + "resources/templates/list": ListResourceTemplatesRequest, + }[method] with pytest.raises(MCPError, match="Server misconfigured: no database connection"): - await GatewayOperations().execute(request(params=PaginatedRequestParams(cursor="existing-state")), prepare_context(caller)) + await GatewayOperations().execute( + request(params=PaginatedRequestParams(cursor="existing-state")), prepare_context(caller) + ) @pytest.mark.asyncio @@ -1552,13 +1751,17 @@ async def test_virtual_tool_catalog_rejects_a_cursor_and_preserves_its_complete_ from litellm.proxy._types import LiteLLM_ObjectPermissionTable - caller = UserAPIKeyAuth(object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="search", mcp_tool_search_enabled=True)) + caller = UserAPIKeyAuth( + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="search", mcp_tool_search_enabled=True) + ) context = prepare_context(caller) result = await GatewayOperations().execute(ListToolsRequest(), context) assert result.tools assert result.next_cursor is None with pytest.raises(MCPError, match="fresh listing"): - await GatewayOperations().execute(ListToolsRequest(params=PaginatedRequestParams(cursor="existing-state")), context) + await GatewayOperations().execute( + ListToolsRequest(params=PaginatedRequestParams(cursor="existing-state")), context + ) @pytest.mark.asyncio @@ -1573,7 +1776,9 @@ async def test_optional_catalog_preserves_revoked_user_error(monkeypatch, reques monkeypatch.setattr(proxy_server, "general_settings", {"supported_db_objects": []}) table = SimpleNamespace(find_unique=AsyncMock(return_value=None)) - monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(writer_db=SimpleNamespace(litellm_usertable=table))) + monkeypatch.setattr( + proxy_server, "prisma_client", SimpleNamespace(writer_db=SimpleNamespace(litellm_usertable=table)) + ) monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache()) caller = UserAPIKeyAuth(user_id="revoked-catalog-user") caller.mcp_admitted_user_subject = True @@ -1603,7 +1808,9 @@ async def test_tool_continuation_failure_requires_a_fresh_listing(monkeypatch: p @pytest.mark.asyncio @pytest.mark.parametrize("gateway", [False, True]) -async def test_initial_tool_listing_preserves_legacy_error_fallback(monkeypatch: pytest.MonkeyPatch, gateway: bool) -> None: +async def test_initial_tool_listing_preserves_legacy_error_fallback( + monkeypatch: pytest.MonkeyPatch, gateway: bool +) -> None: from mcp.types import ListToolsRequest fetch: Final = AsyncMock(side_effect=RuntimeError("catalog temporarily unavailable")) @@ -1624,12 +1831,20 @@ def test_discovery_extra_headers_exclude_gateway_admission_key(header_name: str) caller_key: Final = "Bearer sk-admission-only" raw_headers: Final = {"x-litellm-api-key": caller_key, "x-tenant": "tenant-control"} server: Final = MCPServer( - server_id="header-boundary", name="header-boundary", transport=MCPTransport.http, - url="https://example.invalid/mcp", auth_type=MCPAuth.none, + server_id="header-boundary", + name="header-boundary", + transport=MCPTransport.http, + url="https://example.invalid/mcp", + auth_type=MCPAuth.none, extra_headers=[header_name, "X-Tenant"], ) auth_header, extra_headers = operations._prepare_mcp_server_headers( - server, None, None, None, raw_headers, UserAPIKeyAuth(api_key="sk-admission-only"), + server, + None, + None, + None, + raw_headers, + UserAPIKeyAuth(api_key="sk-admission-only"), ) assert auth_header is None assert extra_headers == {"X-Tenant": "tenant-control"} @@ -1639,17 +1854,27 @@ def test_discovery_extra_headers_exclude_gateway_admission_key(header_name: str) @pytest.mark.asyncio @pytest.mark.parametrize("request_kind", ["absent", "non_starlette", "without_bearer", "admission_key"]) async def test_guardrail_context_never_recovers_a_previous_request_bearer( - _mcp_request_ctx: Callable[..., ServerRequestContext], request_kind: str, monkeypatch: pytest.MonkeyPatch, + _mcp_request_ctx: Callable[..., ServerRequestContext], + request_kind: str, + monkeypatch: pytest.MonkeyPatch, ) -> None: from mcp.server.auth.middleware.auth_context import auth_context_var from starlette.requests import Request from litellm.proxy._experimental.mcp_server import server as gateway request: Final = ( - None if request_kind == "absent" else object() if request_kind == "non_starlette" else Request({ - "type": "http", "method": "POST", "path": "/mcp", - "headers": [(b"authorization", b"Bearer sk-current-key")] if request_kind == "admission_key" else [], - }) + None + if request_kind == "absent" + else object() + if request_kind == "non_starlette" + else Request( + { + "type": "http", + "method": "POST", + "path": "/mcp", + "headers": [(b"authorization", b"Bearer sk-current-key")] if request_kind == "admission_key" else [], + } + ) ) previous: Final = auth_context_var.set(None) try: @@ -1661,7 +1886,9 @@ async def test_guardrail_context_never_recovers_a_previous_request_bearer( monkeypatch.setattr(litellm, "callbacks", [capture]) async with gateway._legacy_operation_context(_mcp_request_ctx(request=request), trace=False) as context: await operations.global_mcp_server_manager.pre_call_tool_check( - name="echo", arguments={}, server_name="guarded", + name="echo", + arguments={}, + server_name="guarded", user_api_key_auth=context.user_api_key_auth, proxy_logging_obj=ProxyLogging(user_api_key_cache=UserApiKeyCache()), server=MCPServer(server_id="guarded", name="guarded", transport=MCPTransport.http, allow_all_keys=True), diff --git a/tests/unit/proxy/test_proxy_config_unit_test.py b/tests/unit/proxy/test_proxy_config_unit_test.py index bee3c3a9eaa..28e017bb050 100644 --- a/tests/unit/proxy/test_proxy_config_unit_test.py +++ b/tests/unit/proxy/test_proxy_config_unit_test.py @@ -1,4 +1,6 @@ import os +from pathlib import Path +from typing import Final import traceback from unittest import mock import pytest @@ -340,3 +342,16 @@ class TestYamlStorePromptsDbOverride: proxy_config = ProxyConfig() assert proxy_config.settings.source("store_prompts_in_spend_logs") == "unset" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("disabled", [["prompt"], "prompts", ["prompts", "tools"]]) +async def test_load_config_rejects_unknown_mcp_disabled_capability(tmp_path: Path, disabled: object) -> None: + import yaml + from pydantic import ValidationError + + config_path: Final = tmp_path / "config.yaml" + config_path.write_text(yaml.safe_dump({"general_settings": {"mcp_disabled_capabilities": disabled}})) + + with pytest.raises(ValidationError): + await ProxyConfig().load_config(router=None, config_file_path=str(config_path)) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx index 4a486bfc648..210e2173457 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx @@ -147,6 +147,45 @@ describe("MCPNetworkSettings", () => { expect(updateConfigFieldSetting).not.toHaveBeenCalled(); }); + it("loads both capability switches on and writes nothing on Save when the proxy reports no stored value", async () => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([ + { field_name: "mcp_disabled_capabilities", field_value: null }, + ]); + renderSettings(); + const switches = await screen.findAllByRole("switch"); + expect(switches.map((s) => s.getAttribute("aria-checked"))).toEqual(["true", "true"]); + await userEvent.click(screen.getByRole("button", { name: /Save/ })); + + await waitFor(() => expect(toast.success).toHaveBeenCalledWith("MCP network settings saved")); + expect(deleteConfigFieldSetting).not.toHaveBeenCalled(); + expect(updateConfigFieldSetting).not.toHaveBeenCalled(); + }); + + it("stores the capabilities an admin switches off", async () => { + renderSettings(); + await userEvent.click((await screen.findAllByRole("switch"))[0]); + await userEvent.click(screen.getByRole("button", { name: /Save/ })); + + await waitFor(() => + expect(updateConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_disabled_capabilities", ["prompts"]), + ); + expect(deleteConfigFieldSetting).not.toHaveBeenCalled(); + }); + + it("deletes the stored setting when the admin switches the last disabled capability back on", async () => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([ + { field_name: "mcp_disabled_capabilities", field_value: ["resources"] }, + ]); + renderSettings(); + const switches = await screen.findAllByRole("switch"); + expect(switches.map((s) => s.getAttribute("aria-checked"))).toEqual(["true", "false"]); + await userEvent.click(switches[1]); + await userEvent.click(screen.getByRole("button", { name: /Save/ })); + + await waitFor(() => expect(deleteConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_disabled_capabilities")); + expect(updateConfigFieldSetting).not.toHaveBeenCalled(); + }); + it("labels the section Allowed Clients and renders each stored client as a card showing alias and value", async () => { vi.mocked(getGeneralSettingsCall).mockResolvedValue([ { field_name: "mcp_allowed_clients", field_value: [ANTIGRAVITY, CODEX] }, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx index c421f1a932e..45180203a58 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx @@ -13,6 +13,7 @@ import { } from "@/components/ui/dialog"; import { Input } from "@/components/ui/input"; import { Label } from "@/components/ui/label"; +import { Switch } from "@/components/ui/switch"; import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; import { toast } from "@/lib/toast"; import { @@ -159,6 +160,22 @@ const AllowedClientDialog: React.FC = ({ draft, onChan ); }; +type DisableableCapability = "prompts" | "resources"; + +const CAPABILITY_TOGGLES: { capability: DisableableCapability; label: string; description: string }[] = [ + { capability: "prompts", label: "Prompts", description: "Advertise prompts and answer prompts/list and prompts/get" }, + { + capability: "resources", + label: "Resources", + description: "Advertise resources and answer resources/list, resources/templates/list and resources/read", + }, +]; + +const parseDisabledCapabilities = (value: unknown): DisableableCapability[] => + Array.isArray(value) + ? value.filter((name): name is DisableableCapability => name === "prompts" || name === "resources") + : []; + const MCPNetworkSettings: React.FC = ({ accessToken }) => { const [loading, setLoading] = useState(true); const [saving, setSaving] = useState(false); @@ -168,9 +185,12 @@ const MCPNetworkSettings: React.FC = ({ accessToken }) const [storedRanges, setStoredRanges] = useState(null); const [storedClients, setStoredClients] = useState(ABSENT); const [storedClientIdHeader, setStoredClientIdHeader] = useState(null); + const [disabledCapabilities, setDisabledCapabilities] = useState([]); + const [storedDisabledCapabilities, setStoredDisabledCapabilities] = useState(null); const [currentIp, setCurrentIp] = useState(null); const [rangeDraft, setRangeDraft] = useState(""); const [clientDraft, setClientDraft] = useState(null); + const capabilityToggleId = useId(); useEffect(() => { loadSettings(); @@ -196,6 +216,11 @@ const MCPNetworkSettings: React.FC = ({ accessToken }) setClientIdHeader(field.field_value); setStoredClientIdHeader(field.field_value); } + if (field.field_name === "mcp_disabled_capabilities") { + const stored = parseDisabledCapabilities(field.field_value); + setDisabledCapabilities(stored); + setStoredDisabledCapabilities(stored.length > 0 ? stored : null); + } } } catch (error) { console.error("Failed to load MCP network settings:", error); @@ -247,14 +272,34 @@ const MCPNetworkSettings: React.FC = ({ accessToken }) setStoredClientIdHeader(null); }; + const persistDisabledCapabilities = async (token: string) => { + if (unchangedSinceLoad(disabledCapabilities, storedDisabledCapabilities)) return; + if (disabledCapabilities.length > 0) { + await updateConfigFieldSetting(token, "mcp_disabled_capabilities", disabledCapabilities); + setStoredDisabledCapabilities(disabledCapabilities); + return; + } + await deleteConfigFieldSetting(token, "mcp_disabled_capabilities"); + setStoredDisabledCapabilities(null); + }; + + const setCapabilityServed = (capability: DisableableCapability, served: boolean) => { + setDisabledCapabilities( + served + ? disabledCapabilities.filter((name) => name !== capability) + : [...disabledCapabilities.filter((name) => name !== capability), capability], + ); + }; + const handleSave = async () => { if (!accessToken) return; setSaving(true); const [rangeResult] = await Promise.allSettled([persistRanges(accessToken)]); const [clientResult] = await Promise.allSettled([persistAllowedClients(accessToken)]); const [headerResult] = await Promise.allSettled([persistClientIdHeader(accessToken)]); + const [capabilityResult] = await Promise.allSettled([persistDisabledCapabilities(accessToken)]); setSaving(false); - const failures = [rangeResult, clientResult, headerResult].filter( + const failures = [rangeResult, clientResult, headerResult, capabilityResult].filter( (result): result is PromiseRejectedResult => result.status === "rejected", ); if (failures.length === 0) { @@ -452,6 +497,34 @@ const MCPNetworkSettings: React.FC = ({ accessToken })

+
+

Served Capabilities

+

+ Turn off prompts or resources when none of your MCP servers offer them. A capability that is turned off is + left out of the initialize response, its list calls return an empty list without contacting any upstream + server, and its get or read calls are rejected. +

+
+ + + {CAPABILITY_TOGGLES.map(({ capability, label, description }) => ( +
+
+ +

{description}

+
+ setCapabilityServed(capability, served)} + /> +
+ ))} + +