mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-18 04:44:37 +00:00
feat(mcp): add mcp_disabled_capabilities to stop advertising prompts and resources (#46563)
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Yassin Kortam
GitHub
parent
82977e2465
commit
92fb55522a
13 changed files with 862 additions and 132 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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.",
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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] },
|
||||
|
|
|
|||
|
|
@ -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<AllowedClientDialogProps> = ({ 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<MCPNetworkSettingsProps> = ({ accessToken }) => {
|
||||
const [loading, setLoading] = useState(true);
|
||||
const [saving, setSaving] = useState(false);
|
||||
|
|
@ -168,9 +185,12 @@ const MCPNetworkSettings: React.FC<MCPNetworkSettingsProps> = ({ accessToken })
|
|||
const [storedRanges, setStoredRanges] = useState<string[] | null>(null);
|
||||
const [storedClients, setStoredClients] = useState<StoredAllowlist>(ABSENT);
|
||||
const [storedClientIdHeader, setStoredClientIdHeader] = useState<string | null>(null);
|
||||
const [disabledCapabilities, setDisabledCapabilities] = useState<DisableableCapability[]>([]);
|
||||
const [storedDisabledCapabilities, setStoredDisabledCapabilities] = useState<DisableableCapability[] | null>(null);
|
||||
const [currentIp, setCurrentIp] = useState<string | null>(null);
|
||||
const [rangeDraft, setRangeDraft] = useState("");
|
||||
const [clientDraft, setClientDraft] = useState<ClientDraft | null>(null);
|
||||
const capabilityToggleId = useId();
|
||||
|
||||
useEffect(() => {
|
||||
loadSettings();
|
||||
|
|
@ -196,6 +216,11 @@ const MCPNetworkSettings: React.FC<MCPNetworkSettingsProps> = ({ 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<MCPNetworkSettingsProps> = ({ 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<MCPNetworkSettingsProps> = ({ accessToken })
|
|||
</p>
|
||||
</Card>
|
||||
|
||||
<div>
|
||||
<p className="text-lg font-semibold">Served Capabilities</p>
|
||||
<p className="mt-1 text-sm text-muted-foreground">
|
||||
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.
|
||||
</p>
|
||||
</div>
|
||||
|
||||
<Card className="space-y-4 p-6">
|
||||
{CAPABILITY_TOGGLES.map(({ capability, label, description }) => (
|
||||
<div key={capability} className="flex items-center justify-between gap-4">
|
||||
<div>
|
||||
<Label htmlFor={`${capabilityToggleId}-${capability}`} className="text-sm font-medium">
|
||||
{label}
|
||||
</Label>
|
||||
<p className="text-xs text-muted-foreground">{description}</p>
|
||||
</div>
|
||||
<Switch
|
||||
id={`${capabilityToggleId}-${capability}`}
|
||||
size="sm"
|
||||
checked={!disabledCapabilities.includes(capability)}
|
||||
onCheckedChange={(served) => setCapabilityServed(capability, served)}
|
||||
/>
|
||||
</div>
|
||||
))}
|
||||
</Card>
|
||||
|
||||
<div className="flex justify-end">
|
||||
<Button onClick={handleSave} disabled={saving}>
|
||||
<Save />
|
||||
|
|
|
|||
5
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
5
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -30036,6 +30036,11 @@ export interface components {
|
|||
* @description Request header whose value names the calling MCP client application (for example 'x-mcp-client') for callers that did not authenticate with a JWT, used only while mcp_allowed_clients is set. The client picks this value itself, so it is a policy control rather than a security boundary; prefer litellm_jwtauth.mcp_client_id_jwt_field where callers use JWTs.
|
||||
*/
|
||||
mcp_client_id_header?: string | null;
|
||||
/**
|
||||
* Mcp Disabled Capabilities
|
||||
* @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_disabled_capabilities?: ("prompts" | "resources")[] | null;
|
||||
/**
|
||||
* Mcp Internal Ip Ranges
|
||||
* @description Custom CIDR ranges that define internal/private networks for MCP access control. When set, only these ranges are treated as internal. Defaults to RFC 1918 private ranges (10.0.0.0/8, 172.16.0.0/12, 192.168.0.0/16, 127.0.0.0/8).
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue