diff --git a/litellm/constants.py b/litellm/constants.py index da95b97036a..c4bedf827c0 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -250,6 +250,7 @@ AGENT_HEALTH_CHECK_TIMEOUT_SECONDS: Final = get_env_float("LITELLM_AGENT_HEALTH_ AGENT_HEALTH_CHECK_GATHER_TIMEOUT_SECONDS: Final = get_env_float("LITELLM_AGENT_HEALTH_CHECK_GATHER_TIMEOUT", 30.0) MCP_UPSTREAM_SESSION_SHUTDOWN_TIMEOUT_SECONDS: Final = 5.0 MCP_TOOL_LISTING_MAX_PAGES: Final = 1000 +MCP_SKILL_OWNER_LOOKUP_MAX_PAGES: Final = 100 MCP_GATEWAY_SESSION_ID_PREFIX_LENGTH: Final = 8 MCP_BYOK_CREDENTIAL_CACHE_TTL_SECONDS: Final = 60 MCP_BYOK_CREDENTIAL_CACHE_MAX_SIZE: Final = 4096 diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index ed6d89397a0..c6a7a64b910 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -50,6 +50,7 @@ class _SessionCallbackMap(TypedDict): from mcp.types import ( CONNECTION_CLOSED, + INVALID_PARAMS, METHOD_NOT_FOUND, REQUEST_TIMEOUT, CacheableResult, @@ -108,6 +109,14 @@ from litellm.proxy._experimental.mcp_server.result_conversion import ( aggregate_freshness, error_text_result, ) +from litellm.proxy._experimental.mcp_server.skills import ( + SKILLS_EXTENSION, + GetSkillRequest, + GetSkillRequestParams, + GetSkillResult, + ListSkillsRequest, + ListSkillsResult, +) from litellm.types.llms.custom_http import VerifyTypes from litellm.types.mcp import ( MCP_LEGACY_VERSIONS, @@ -719,6 +728,7 @@ class MCPClient: ) if self._elicitation_callback is not None else None, + extensions={SKILLS_EXTENSION: {}}, ), ) ), @@ -776,6 +786,7 @@ class MCPClient: sampling_callback=session_kwargs["sampling_callback"], elicitation_callback=session_kwargs["elicitation_callback"], logging_callback=session_kwargs["logging_callback"], + extensions={SKILLS_EXTENSION: {}}, ) session: Final = await session_ctx.__aenter__() try: @@ -1009,6 +1020,8 @@ class MCPClient: async def fetch(session: ClientSession) -> "CatalogListResult": capabilities: Final = session.server_capabilities + if isinstance(request, ListSkillsRequest): + return await self._list_skills_page(session, params) empty: Final = ( ListPromptsResult(prompts=[]) if isinstance(request, ListPromptsRequest) @@ -1041,6 +1054,24 @@ class MCPClient: with anyio.fail_after(max(self.timeout, MCP_TOOL_LISTING_TIMEOUT)): return await self.run_with_session(fetch, quiet_on_error=True) + @staticmethod + async def _list_skills_page(session: ClientSession, params: PaginatedRequestParams) -> ListSkillsResult: + capabilities: Final = session.server_capabilities + if capabilities is not None and SKILLS_EXTENSION in (capabilities.extensions or {}): + return await session.send_request(ListSkillsRequest(params=params), ListSkillsResult) + if params.cursor is not None: + raise MCPError(code=INVALID_PARAMS, message="Upstream catalog became unavailable; start a fresh listing") + return ListSkillsResult(skills=()) + + async def get_skill(self, params: GetSkillRequestParams) -> GetSkillResult: + async def fetch(session: ClientSession) -> GetSkillResult: + capabilities: Final = session.server_capabilities + if capabilities is None or SKILLS_EXTENSION not in (capabilities.extensions or {}): + raise MCPError(code=INVALID_PARAMS, message=f"No skill is served at {params.uri}") + return await session.send_request(GetSkillRequest(params=params), GetSkillResult) + + return await self.run_with_session(fetch, quiet_on_error=True) + async def list_tools_page(self, params: PaginatedRequestParams) -> ListToolsResult: with anyio.fail_after(max(self.timeout, MCP_TOOL_LISTING_TIMEOUT)): return await self.run_with_session(lambda session: session.list_tools(params=params), quiet_on_error=True) diff --git a/litellm/proxy/_experimental/mcp_server/capabilities.py b/litellm/proxy/_experimental/mcp_server/capabilities.py index 92c8004ff6f..cffc1d7497b 100644 --- a/litellm/proxy/_experimental/mcp_server/capabilities.py +++ b/litellm/proxy/_experimental/mcp_server/capabilities.py @@ -11,6 +11,7 @@ from mcp_types.methods import CLIENT_REQUESTS from mcp_types.version import HANDSHAKE_PROTOCOL_VERSIONS, LATEST_HANDSHAKE_VERSION from pydantic import TypeAdapter +from litellm.proxy._experimental.mcp_server.skills import SKILLS_EXTENSION from litellm.types.mcp import ( MCP_LEGACY_VERSIONS, MCPAdvertisedVersion, @@ -60,7 +61,7 @@ REVISION_SUPPORT: Final[Mapping[str, RevisionSupport]] = MappingProxyType( results=frozenset({"complete"}) if version.value in HANDSHAKE_PROTOCOL_VERSIONS else frozenset({"complete", "input_required"}), - extensions=frozenset(), + extensions=frozenset({SKILLS_EXTENSION}) if version.value == "2026-07-28" else frozenset(), completed=version.value in HANDSHAKE_PROTOCOL_VERSIONS or version.value == "2026-07-28", ) for version in MCPSpecVersion diff --git a/litellm/proxy/_experimental/mcp_server/catalog.py b/litellm/proxy/_experimental/mcp_server/catalog.py index d0383be61a8..7df1fb54289 100644 --- a/litellm/proxy/_experimental/mcp_server/catalog.py +++ b/litellm/proxy/_experimental/mcp_server/catalog.py @@ -1259,6 +1259,7 @@ def combine_optional_catalog( ) from litellm.proxy._experimental.mcp_server.faults.list_outcomes import SERVER_OUTCOMES_META_KEY + from litellm.proxy._experimental.mcp_server.skills import ListSkillsRequest, ListSkillsResult outcomes: Final = (meta or {}).get(SERVER_OUTCOMES_META_KEY) incomplete: Final = isinstance(outcomes, dict) and any( @@ -1283,6 +1284,14 @@ def combine_optional_catalog( cache_scope="private", _meta=dict(meta) if meta is not None else None, ) + if isinstance(request, ListSkillsRequest): + return ListSkillsResult( + skills=tuple(chain.from_iterable(page.skills for page in pages if isinstance(page, ListSkillsResult))), + next_cursor=next_cursor, + ttl_ms=ttl_ms, + cache_scope="private", + _meta=dict(meta) if meta is not None else None, + ) return ListResourceTemplatesResult( resource_templates=list( chain.from_iterable( diff --git a/litellm/proxy/_experimental/mcp_server/contracts.py b/litellm/proxy/_experimental/mcp_server/contracts.py index 5c3fbde67a6..0b0ab710e4d 100644 --- a/litellm/proxy/_experimental/mcp_server/contracts.py +++ b/litellm/proxy/_experimental/mcp_server/contracts.py @@ -26,11 +26,13 @@ if TYPE_CHECKING: ListToolsResult, ) + from litellm.proxy._experimental.mcp_server.skills import ListSkillsRequest, ListSkillsResult + CatalogListRequest: TypeAlias = ( - ListToolsRequest | ListPromptsRequest | ListResourcesRequest | ListResourceTemplatesRequest + ListToolsRequest | ListPromptsRequest | ListResourcesRequest | ListResourceTemplatesRequest | ListSkillsRequest ) CatalogListResult: TypeAlias = ( - ListToolsResult | ListPromptsResult | ListResourcesResult | ListResourceTemplatesResult + ListToolsResult | ListPromptsResult | ListResourcesResult | ListResourceTemplatesResult | ListSkillsResult ) from litellm.proxy._experimental.mcp_server.server_resolution import ResolvedMCPServer @@ -70,6 +72,7 @@ class OperationContext: mcp_proxy_mode: bool = False wire_compat: WireCompat = WireCompat.LEGACY protocol_version: str | None = None + client_extensions: frozenset[str] = frozenset() incoming_bearer_token: SecretStr | None | EllipsisType = field(default=..., repr=False) def __post_init__(self) -> None: diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 96cf6ee11aa..aa49116a2c0 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -37,10 +37,10 @@ import httpx import httpx2 from fastapi import HTTPException from httpx import HTTPStatusError -from mcp import ReadResourceResult, Resource +from mcp import MCPError, ReadResourceResult, Resource from mcp.shared.auth import OAuthMetadata -from mcp.types import CallToolRequestParams as MCPCallToolRequestParams from mcp.types import ( + INVALID_PARAMS, CallToolResult, GetPromptRequestParams, GetPromptResult, @@ -57,6 +57,7 @@ from mcp.types import ( Prompt, ResourceTemplate, ) +from mcp.types import CallToolRequestParams as MCPCallToolRequestParams from mcp.types import Tool as MCPTool from pydantic import AnyUrl, Field, SecretStr, TypeAdapter, ValidationError from typing_extensions import ReadOnly, assert_never @@ -157,6 +158,12 @@ from litellm.proxy._experimental.mcp_server.result_conversion import ( from litellm.proxy._experimental.mcp_server.sampling_handler import ( MCP_SAMPLING_AVAILABLE, ) +from litellm.proxy._experimental.mcp_server.skills import ( + GetSkillRequestParams, + GetSkillResult, + ListSkillsRequest, + ListSkillsResult, +) from litellm.proxy._experimental.mcp_server.stdio_gate import ( MCP_STDIO_DISABLED_MESSAGE, is_mcp_stdio_blocked, @@ -4773,31 +4780,17 @@ class MCPServerManager: return ListResourcesResult(resources=[]) case ListResourceTemplatesRequest(): return ListResourceTemplatesResult(resource_templates=[]) + case ListSkillsRequest(): + return ListSkillsResult(skills=()) case _: raise RuntimeError("Unexpected catalog request type") - headers: Final = ( - dict( - chain( - extra_headers.items() if extra_headers else (), - server.static_headers.items() if server.static_headers else (), - ) - ) - or None - ) - stdio_env: Final = self.build_stdio_env(server, raw_headers) - subject_token: Final = self._obo_subject_token(server, raw_headers, user_api_key_auth) - client: Final = await self.create_mcp_client( - server=server, - mcp_auth_header=mcp_auth_header, - extra_headers=headers, - stdio_env=stdio_env, - subject_token=subject_token, - user_api_key_auth=user_api_key_auth, - raw_headers=raw_headers, - client_ip=client_ip, + client: Final = await self._create_catalog_client( + server, user_api_key_auth, mcp_auth_header, extra_headers, raw_headers, client_ip ) page: Final = await client.list_page(request) match page: + case ListSkillsResult(): + return page case ListPromptsResult(): return page.model_copy(update={"prompts": self._create_prefixed_prompts(page.prompts, server)}) case ListResourcesResult(): @@ -4811,6 +4804,55 @@ class MCPServerManager: case _: raise RuntimeError("Unexpected catalog result type") + async def get_skill_from_server( + self, + server: MCPServer, + params: GetSkillRequestParams, + user_api_key_auth: UserAPIKeyAuth | None, + *, + mcp_auth_header: str | Mapping[str, str] | None = None, + extra_headers: Mapping[str, str] | None = None, + raw_headers: Mapping[str, str] | None = None, + client_ip: str | None = None, + ) -> GetSkillResult: + if self._skip_blocked_stdio_listing(server, "skill"): + raise MCPError(code=INVALID_PARAMS, message=f"No skill is served at {params.uri}") + client: Final = await self._create_catalog_client( + server, user_api_key_auth, mcp_auth_header, extra_headers, raw_headers, client_ip + ) + return await client.get_skill(params) + + async def _create_catalog_client( + self, + server: MCPServer, + user_api_key_auth: UserAPIKeyAuth | None, + mcp_auth_header: str | Mapping[str, str] | None, + extra_headers: Mapping[str, str] | None, + raw_headers: Mapping[str, str] | None, + client_ip: str | None, + ) -> MCPClient: + headers: Final = ( + dict( + chain( + extra_headers.items() if extra_headers else (), + server.static_headers.items() if server.static_headers else (), + ) + ) + or None + ) + return await self.create_mcp_client( + server=server, + mcp_auth_header=mcp_auth_header + if mcp_auth_header is None or isinstance(mcp_auth_header, str) + else dict(mcp_auth_header), + extra_headers=headers, + stdio_env=self.build_stdio_env(server, raw_headers), + subject_token=self._obo_subject_token(server, raw_headers, user_api_key_auth), + user_api_key_auth=user_api_key_auth, + raw_headers=raw_headers, + client_ip=client_ip, + ) + async def get_prompts_from_server( self, server: MCPServer, diff --git a/litellm/proxy/_experimental/mcp_server/operations.py b/litellm/proxy/_experimental/mcp_server/operations.py index 90f01368127..cbfbf081546 100644 --- a/litellm/proxy/_experimental/mcp_server/operations.py +++ b/litellm/proxy/_experimental/mcp_server/operations.py @@ -53,6 +53,7 @@ from typing_extensions import ReadOnly, TypedDict, assert_never from litellm._logging import verbose_logger from litellm.constants import ( MAXIMUM_TRACEBACK_LINES_TO_LOG, + MCP_SKILL_OWNER_LOOKUP_MAX_PAGES, MCP_TOOL_LISTING_TIMEOUT, ) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -130,6 +131,15 @@ from litellm.proxy._experimental.mcp_server.result_conversion import ( handler_outcome, to_call_tool_result, ) +from litellm.proxy._experimental.mcp_server.skills import ( + SKILLS_EXTENSION, + GetSkillRequest, + GetSkillRequestParams, + GetSkillResult, + ListSkillsRequest, + ListSkillsResult, + SkillResource, +) from litellm.proxy._experimental.mcp_server.tool_registry import ( global_mcp_tool_registry, ) @@ -2799,14 +2809,26 @@ async def mcp_read_resource( detail="User not allowed to read this resource.", ) - if len(allowed_mcp_servers) != 1: + server: Final = ( + allowed_mcp_servers[0] + if len(allowed_mcp_servers) == 1 + else await _skill_owner( + str(url), + allowed_mcp_servers, + user_api_key_auth=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + client_ip=client_ip, + ) + ) + if server is None: raise HTTPException( status_code=400, detail=("Multiple MCP servers configured; read_resource currently supports exactly one allowed server."), ) - server: Final = allowed_mcp_servers[0] - server_auth_header, extra_headers = _prepare_mcp_server_headers( server=server, mcp_server_auth_headers=mcp_server_auth_headers, @@ -2814,6 +2836,7 @@ async def mcp_read_resource( oauth2_headers=oauth2_headers, raw_headers=raw_headers, user_api_key_auth=user_api_key_auth, + scope_servers=allowed_mcp_servers, ) await _enforce_mcp_server_rate_limit(user_api_key_auth, server) @@ -2831,6 +2854,73 @@ async def mcp_read_resource( ) +async def _skill_owner( + uri: str, + servers: Sequence[MCPServer], + *, + user_api_key_auth: UserAPIKeyAuth | None, + mcp_auth_header: 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, +) -> MCPServer | None: + server_auth_headers: Final = ( + {name: dict(headers) for name, headers in mcp_server_auth_headers.items()} + if mcp_server_auth_headers is not None + else None + ) + oauth_headers: Final = dict(oauth2_headers) if oauth2_headers is not None else None + forwarded_headers: Final = dict(raw_headers) if raw_headers is not None else None + + async def serves(server: MCPServer) -> bool | ProxyRateLimitError: + auth_header, extra_headers = _prepare_mcp_server_headers( + server=server, + mcp_server_auth_headers=server_auth_headers, + mcp_auth_header=mcp_auth_header, + oauth2_headers=oauth_headers, + raw_headers=forwarded_headers, + user_api_key_auth=user_api_key_auth, + scope_servers=list(servers), + ) + + async def scan(cursor: str | None, remaining: int) -> bool | ProxyRateLimitError: + rejection: Final = await _mcp_server_rate_limit_rejection(server, user_api_key_auth) + if rejection is not None: + return rejection + try: + page: Final = await global_mcp_server_manager.get_optional_catalog_page( + server, + ListSkillsRequest(params=PaginatedRequestParams(cursor=cursor)), + user_api_key_auth, + mcp_auth_header=auth_header, + extra_headers=extra_headers, + raw_headers=raw_headers, + client_ip=client_ip, + ) + except Exception as error: # noqa: BLE001 # an unreachable upstream cannot claim the resource + verbose_logger.debug("MCP skill owner lookup failed for %s: %s", server.name, error) + return False + if not isinstance(page, ListSkillsResult): + return False + if any(_skill_lists_uri(skill_entry.uri, skill_entry.resources, uri) for skill_entry in page.skills): + return True + return page.next_cursor is not None and remaining > 1 and await scan(page.next_cursor, remaining - 1) + + return await scan(None, MCP_SKILL_OWNER_LOOKUP_MAX_PAGES) + + owned: Final = await asyncio.gather(*(serves(server) for server in servers)) + owners: Final = tuple(server for server, owns in zip(servers, owned) if owns is True) + rejections: Final = tuple(owns for owns in owned if isinstance(owns, ProxyRateLimitError)) + if not owners and rejections: + raise rejections[0] + return owners[0] if len(owners) == 1 else None + + +def _skill_lists_uri(skill_uri: str, resources: Sequence[SkillResource], uri: str) -> bool: + return skill_uri == uri or any(resource.uri == uri for resource in resources) + + def _get_standard_logging_mcp_tool_call( name: str, arguments: dict[str, object], @@ -3224,6 +3314,94 @@ async def _execute_list_prompts( return ListPromptsResult(prompts=[]) +async def _execute_list_skills(context: OperationContext, params: PaginatedRequestParams) -> ListSkillsResult: + from mcp.types import INVALID_PARAMS, INVALID_REQUEST + + from litellm.proxy._experimental.mcp_server.catalog import list_gateway_catalog + from litellm.proxy._experimental.mcp_server.server_resolution import MCPServerTargetCatalog + + if context.mcp_proxy_mode: + _reject_mcp_proxy_operation() + if "resources" in disabled_capabilities(): + return ListSkillsResult(skills=()) + try: + catalog: Final = MCPServerTargetCatalog(global_mcp_server_manager, listing=list_gateway_catalog) + result: Final = await catalog.list(context, ListSkillsRequest(params=params)) + except MCPError: + raise + except HTTPException as error: + raise MCPError(code=INVALID_REQUEST, message=_http_detail_message(error.detail)) from error + except Exception as error: + if params.cursor is not None: + raise MCPError( + code=INVALID_PARAMS, message="Catalog continuation unavailable; start a fresh listing" + ) from error + return ListSkillsResult(skills=()) + assert isinstance(result, ListSkillsResult) + return result + + +async def _execute_get_skill(context: OperationContext, params: GetSkillRequestParams) -> GetSkillResult: + from mcp.types import INVALID_PARAMS, INVALID_REQUEST + + 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, + mcp_servers, + mcp_server_auth_headers, + oauth2_headers, + raw_headers, + client_ip, + ) = context.legacy_auth() + allowed: Final = await _get_allowed_mcp_servers( + user_api_key_auth=user_api_key_auth, mcp_servers=mcp_servers, client_ip=client_ip + ) + if mcp_servers and not allowed: + try: + await raise_denied_scoped_mcp_access( + requested_names=mcp_servers, user_api_key_auth=user_api_key_auth, client_ip=client_ip + ) + except HTTPException as error: + raise MCPError(code=INVALID_REQUEST, message=_http_detail_message(error.detail)) from error + + owner: Final = await _skill_owner( + params.uri, + allowed, + user_api_key_auth=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + client_ip=client_ip, + ) + if owner is None: + raise MCPError(code=INVALID_PARAMS, message=f"No skill is served at {params.uri}") + auth_header, extra_headers = _prepare_mcp_server_headers( + server=owner, + mcp_server_auth_headers=mcp_server_auth_headers, + mcp_auth_header=mcp_auth_header, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + user_api_key_auth=user_api_key_auth, + scope_servers=allowed, + ) + await _enforce_mcp_server_rate_limit(user_api_key_auth, owner) + result: Final = await global_mcp_server_manager.get_skill_from_server( + owner, + params, + user_api_key_auth, + mcp_auth_header=auth_header, + extra_headers=extra_headers, + raw_headers=raw_headers, + client_ip=client_ip, + ) + return result.model_copy(update={"cache_scope": "private"}) + + async def _execute_get_prompt( context: OperationContext, params: GetPromptRequestParams, host_progress_callback: ProgressCallback | None = None ) -> GetPromptResult | InputRequiredResult: @@ -3404,6 +3582,8 @@ GatewayOperation: TypeAlias = ( | ListResourcesRequest | ListResourceTemplatesRequest | ReadResourceRequest + | ListSkillsRequest + | GetSkillRequest ) GatewayResult: TypeAlias = ( DiscoverResult @@ -3415,6 +3595,8 @@ GatewayResult: TypeAlias = ( | ListResourcesResult | ListResourceTemplatesResult | ReadResourceResult + | ListSkillsResult + | GetSkillResult ) @@ -3470,6 +3652,12 @@ class GatewayOperations: self, operation: ReadResourceRequest, context: OperationContext ) -> ReadResourceResult | InputRequiredResult: ... + @overload + async def execute(self, operation: ListSkillsRequest, context: OperationContext) -> ListSkillsResult: ... + + @overload + async def execute(self, operation: GetSkillRequest, context: OperationContext) -> GetSkillResult: ... + @catalog_operation(lambda: global_mcp_server_manager) async def execute(self, operation: GatewayOperation, context: OperationContext) -> GatewayResult: if not isinstance(operation, (CallToolRequest, GetPromptRequest, ReadResourceRequest)): @@ -3513,6 +3701,11 @@ class GatewayOperations: else ( *(() if "prompts" in disabled else (ListPromptsRequest(),)), *(() if "resources" in disabled else (ListResourcesRequest(), ListResourceTemplatesRequest())), + *( + (ListSkillsRequest(),) + if "resources" not in disabled and SKILLS_EXTENSION in context.client_extensions + else () + ), ) ) memo_token: Final = _mcp_server_admission_memo.set({}) @@ -3547,12 +3740,16 @@ class GatewayOperations: await asyncio.gather(*tasks, return_exceptions=True) finally: _mcp_server_admission_memo.reset(memo_token) + listed: Final[tuple[object, ...]] = tuple(results) # pyright: ignore[reportUnknownArgumentType] # gather over mixed tasks erases element types + skills_listed: Final = any(isinstance(result, ListSkillsResult) and result.skills for result in listed) return build_discovery( configured=configured_versions(), revision=context.protocol_version or "2025-11-25", transport=MCPTransport.http, authorized_operations=enabled_gateway_operations(disabled), upstream_versions=upstream_versions, + client_extensions=context.client_extensions, + upstream_extensions=frozenset({SKILLS_EXTENSION}) if skills_listed else frozenset(), capabilities=ServerCapabilities( tools=notifications.tools or ( @@ -3569,13 +3766,15 @@ class GatewayOperations: resources=notifications.resources or ( ResourcesCapability() - if any( + if skills_listed + or any( (isinstance(result, ListResourcesResult) and result.resources) or (isinstance(result, ListResourceTemplatesResult) and result.resource_templates) for result in results ) else None ), + extensions={SKILLS_EXTENSION: {}} if skills_listed else None, ), ) case AuthorizedToolCall(): @@ -3621,5 +3820,9 @@ class GatewayOperations: ) case ReadResourceRequest(params=params): return await _execute_read_resource(context, params, self._host_progress_callback) + case ListSkillsRequest(params=params): + return await _execute_list_skills(context, params or PaginatedRequestParams()) + case GetSkillRequest(params=params): + return await _execute_get_skill(context, params) case _: return assert_never(operation) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index b3426089375..3375d49558e 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -14,6 +14,7 @@ import time import types from collections import Counter from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterable, Mapping, Sequence, Set +from dataclasses import replace from typing import TYPE_CHECKING, Final, NoReturn, Protocol import httpx @@ -79,6 +80,13 @@ from litellm.proxy._experimental.mcp_server.scope_challenge import ( finish_scope_response, is_tool_call_request, ) +from litellm.proxy._experimental.mcp_server.skills import ( + GetSkillRequest, + GetSkillRequestParams, + GetSkillResult, + ListSkillsRequest, + ListSkillsResult, +) from litellm.proxy._experimental.mcp_server.ui_session_utils import ( ActingUser, GrantedToolsetIds, @@ -1013,9 +1021,26 @@ if MCP_AVAILABLE: ReadResourceRequest(params=params), context ) - async def discover(ctx: ServerRequestContext, params: RequestParams) -> DiscoverResult: + async def list_skills(ctx: ServerRequestContext, params: PaginatedRequestParams) -> ListSkillsResult: async with _legacy_operation_context(ctx, trace=False) as context: - return await operations.GatewayOperations().execute(DiscoverRequest(params=params), context) + return await operations.GatewayOperations().execute(ListSkillsRequest(params=params), context) + + async def get_skill(ctx: ServerRequestContext, params: GetSkillRequestParams) -> GetSkillResult: + async with _legacy_operation_context(ctx, trace=False) as context: + return await operations.GatewayOperations().execute(GetSkillRequest(params=params), context) + + async def discover(ctx: ServerRequestContext, params: RequestParams) -> DiscoverResult: + client_capabilities: Final = ctx.session.client_capabilities + async with _legacy_operation_context(ctx, trace=False) as context: + return await operations.GatewayOperations().execute( + DiscoverRequest(params=params), + replace( + context, + client_extensions=frozenset( + client_capabilities.extensions or {} if client_capabilities is not None else {} + ), + ), + ) from mcp.types import ( EmptyResult, @@ -1107,6 +1132,8 @@ if MCP_AVAILABLE: server.add_request_handler("resources/list", PaginatedRequestParams, list_resources) server.add_request_handler("resources/templates/list", PaginatedRequestParams, list_resource_templates) server.add_request_handler("resources/read", ReadResourceRequestParams, read_resource) + server.add_request_handler("skills/list", PaginatedRequestParams, list_skills) # pyright: ignore[reportUnknownMemberType] # SDK Server.add_request_handler is untyped + server.add_request_handler("skills/get", GetSkillRequestParams, get_skill) # pyright: ignore[reportUnknownMemberType] # SDK Server.add_request_handler is untyped ######################################################## ############ End of MCP Server Routes ################## diff --git a/litellm/proxy/_experimental/mcp_server/skills.py b/litellm/proxy/_experimental/mcp_server/skills.py new file mode 100644 index 00000000000..6b538fa87b9 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/skills.py @@ -0,0 +1,55 @@ +from collections.abc import Mapping +from typing import Final, Literal + +from mcp.types import CacheableResult, PaginatedRequest, PaginatedResult, Request, RequestParams, ResultType +from pydantic import BaseModel, ConfigDict, JsonValue +from pydantic.alias_generators import to_camel + +SKILLS_EXTENSION: Final = "io.modelcontextprotocol/skills" + + +class SkillResource(BaseModel): + model_config = ConfigDict(alias_generator=to_camel, populate_by_name=True, frozen=True, extra="allow") + + uri: str + + +class Skill(BaseModel): + model_config = ConfigDict(alias_generator=to_camel, populate_by_name=True, frozen=True, extra="allow") + + uri: str + frontmatter: Mapping[str, JsonValue] + resources: tuple[SkillResource, ...] = () + + +class ListSkillsRequest(PaginatedRequest[Literal["skills/list"]]): + model_config = ConfigDict(frozen=True) + + method: Literal["skills/list"] = "skills/list" + + +class ListSkillsResult(PaginatedResult, CacheableResult): + model_config = ConfigDict(frozen=True) + + skills: tuple[Skill, ...] + result_type: ResultType = "complete" + + +class GetSkillRequestParams(RequestParams): + model_config = ConfigDict(frozen=True) + + uri: str + + +class GetSkillRequest(Request[GetSkillRequestParams, Literal["skills/get"]]): + model_config = ConfigDict(frozen=True) + + method: Literal["skills/get"] = "skills/get" + params: GetSkillRequestParams + + +class GetSkillResult(CacheableResult): + model_config = ConfigDict(frozen=True) + + skill: Skill + result_type: ResultType = "complete" diff --git a/tests/unit/experimental_mcp_client/test_mcp_client.py b/tests/unit/experimental_mcp_client/test_mcp_client.py index 6ea76a0a8a7..3623bc41eaf 100644 --- a/tests/unit/experimental_mcp_client/test_mcp_client.py +++ b/tests/unit/experimental_mcp_client/test_mcp_client.py @@ -4987,6 +4987,7 @@ async def test_persistent_session_releases_the_current_requests_scope_response(s session.close() await asyncio.wait_for(session.wait_closed(), 5) + @pytest.mark.asyncio @pytest.mark.parametrize("legacy_callback", [False, True]) async def test_tool_progress_preserves_message_and_legacy_callback(legacy_callback: bool) -> None: @@ -5105,3 +5106,73 @@ async def test_notification_callback_keeps_stream_errors_separate_from_protocol_ await _receive_session_message(failure, stream_error=stream_error, on_notification=receiver) assert stream_error.result() is first receiver.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("advertised", [True, False]) +async def test_skills_requests_reach_only_an_upstream_that_advertises_the_extension(advertised: bool) -> None: + from mcp.types import PaginatedRequestParams + + from litellm.proxy._experimental.mcp_server.skills import ( + SKILLS_EXTENSION, + GetSkillRequestParams, + ListSkillsRequest, + ListSkillsResult, + ) + + skill: Final = { + "uri": "skill://billing/refunds/SKILL.md", + "frontmatter": {"name": "refunds"}, + "resources": [{"uri": "skill://billing/refunds/SKILL.md", "digest": "sha256:abc", "size": 3}], + } + observed: Final[list[tuple[str, dict[str, object] | None]]] = [] # mutable-ok: records upstream calls + + def respond(request: httpx2.Request) -> httpx2.Response: + if request.method != "POST": + return httpx2.Response(405) + try: + payload: Final = JSONRPCRequest.model_validate_json(request.content) + except ValidationError: + return httpx2.Response(202) + observed.append((payload.method, payload.params)) + if payload.method == "initialize": + return httpx2.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload.id, + "result": { + "protocolVersion": LATEST_HANDSHAKE_VERSION, + "capabilities": {"resources": {}, "extensions": {SKILLS_EXTENSION: {}} if advertised else {}}, + "serverInfo": {"name": "skills", "version": "1"}, + }, + }, + ) + result: Final = ( + {"skills": [skill], "nextCursor": "upstream-next"} if payload.method == "skills/list" else {"skill": skill} + ) + return httpx2.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": result}) + + client: Final = _MockTransportClient(respond, server_url="https://upstream.example.com/mcp") + page: Final = await client.list_page(ListSkillsRequest(params=PaginatedRequestParams(cursor=None))) + assert isinstance(page, ListSkillsResult) + initialize: Final = next(params for method, params in observed if method == "initialize") + assert initialize is not None + assert SKILLS_EXTENSION in initialize["capabilities"]["extensions"] + if advertised: + fetched = await client.get_skill(GetSkillRequestParams(uri=skill["uri"])) + assert page.model_dump(mode="json", by_alias=True, exclude_none=True)["skills"] == [skill] + assert page.next_cursor == "upstream-next" + assert fetched.skill.model_dump(mode="json", by_alias=True, exclude_none=True) == skill + assert [method for method, _ in observed if method.startswith("skills/")] == ["skills/list", "skills/get"] + get_params: Final = next(params for method, params in observed if method == "skills/get") + assert get_params is not None and get_params["uri"] == skill["uri"] + return + with pytest.raises(MCPError) as missing: + await client.get_skill(GetSkillRequestParams(uri=skill["uri"])) + with pytest.raises(MCPError, match="start a fresh listing") as stale: + await client.list_page(ListSkillsRequest(params=PaginatedRequestParams(cursor="upstream-next"))) + assert stale.value.error.code == -32602 + assert page.skills == () and page.next_cursor is None + assert missing.value.error.code == -32602 + assert not [method for method, _ in observed if method.startswith("skills/")] diff --git a/tests/unit/proxy/_experimental/mcp_server/conftest.py b/tests/unit/proxy/_experimental/mcp_server/conftest.py index de7ed4d45e8..94f85b3f6c1 100644 --- a/tests/unit/proxy/_experimental/mcp_server/conftest.py +++ b/tests/unit/proxy/_experimental/mcp_server/conftest.py @@ -133,7 +133,7 @@ def _mcp_request_ctx(): from mcp.server.context import ServerRequestContext kwargs = { - "session": SimpleNamespace(), + "session": SimpleNamespace(client_capabilities=None), "lifespan_context": {}, "protocol_version": "2025-06-18", "method": "", diff --git a/tests/unit/proxy/_experimental/mcp_server/test_catalog.py b/tests/unit/proxy/_experimental/mcp_server/test_catalog.py index b383078c0ed..7d09ef60086 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_catalog.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_catalog.py @@ -671,3 +671,50 @@ def test_partial_optional_catalog_never_advertises_freshness() -> None: assert result.ttl_ms == 0 assert result.cache_scope == "private" + +@pytest.mark.asyncio +async def test_skills_catalog_fans_out_to_allowed_servers_and_continues_each_cursor(monkeypatch: pytest.MonkeyPatch): + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server import operations + from litellm.proxy._experimental.mcp_server.skills import ListSkillsRequest, ListSkillsResult, Skill + + monkeypatch.setenv("LITELLM_SALT_KEY", "skills-catalog-test") + monkeypatch.setattr(proxy_server, "prisma_client", None) + billing: Final = MCPServer(server_id="billing", name="billing", transport=MCPTransport.http) + support: Final = MCPServer(server_id="support", name="support", transport=MCPTransport.http) + hidden: Final = MCPServer(server_id="hidden", name="hidden", transport=MCPTransport.http) + monkeypatch.setattr( + operations.global_mcp_server_manager, "registry", {s.server_id: s for s in (billing, support, hidden)} + ) + monkeypatch.setattr(operations, "_get_allowed_mcp_servers", AsyncMock(return_value=[billing, support])) + + def skill(uri: str) -> Skill: + return Skill.model_validate( + {"uri": uri, "frontmatter": {"name": uri}, "resources": [{"uri": uri, "digest": "sha256:0", "size": 1}]} + ) + + pages: Final = { + ("billing", None): ListSkillsResult(skills=(skill("skill://billing/a/SKILL.md"),), next_cursor="billing-2"), + ("billing", "billing-2"): ListSkillsResult(skills=(skill("skill://billing/b/SKILL.md"),)), + ("support", None): ListSkillsResult(skills=(skill("skill://support/a/SKILL.md"),)), + } + + async def fetch( + context: object, request: object, server: MCPServer, allowed: object, cursor: str | None + ) -> ListSkillsResult: + assert isinstance(request, ListSkillsRequest) + return pages[(server.server_id, cursor)] + + monkeypatch.setattr(catalog, "fetch_optional_catalog_page", AsyncMock(side_effect=fetch)) + context: Final = operations.prepare_context() + first: Final = await catalog.list_gateway_catalog(context, ListSkillsRequest()) + assert isinstance(first, ListSkillsResult) + assert sorted(entry.uri for entry in first.skills) == ["skill://billing/a/SKILL.md", "skill://support/a/SKILL.md"] + assert first.model_dump(mode="json", by_alias=True)["skills"][0]["resources"][0]["digest"] == "sha256:0" + assert first.next_cursor + second: Final = await catalog.list_gateway_catalog( + context, ListSkillsRequest(params=PaginatedRequestParams(cursor=first.next_cursor)) + ) + assert isinstance(second, ListSkillsResult) + assert [entry.uri for entry in second.skills] == ["skill://billing/b/SKILL.md"] + assert second.next_cursor is None diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_proxy_mode.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_proxy_mode.py index 7dadfc6493a..906cad5ce2e 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_proxy_mode.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_proxy_mode.py @@ -63,6 +63,8 @@ async def test_proxy_rejects_non_tool_protocol_operations() -> None: from mcp.server.context import ServerRequestContext from mcp.types import GetPromptRequestParams, PaginatedRequestParams, ReadResourceRequestParams + from litellm.proxy._experimental.mcp_server.skills import GetSkillRequestParams + ctx = ServerRequestContext( session=SimpleNamespace(), lifespan_context={}, @@ -80,6 +82,10 @@ async def test_proxy_rejects_non_tool_protocol_operations() -> None: await server.list_resource_templates(ctx, PaginatedRequestParams()) with pytest.raises(MCPError): await server.read_resource(ctx, ReadResourceRequestParams(uri="https://example.com/resource")) + with pytest.raises(MCPError, match="unavailable on /mcp/proxy"): + await server.list_skills(ctx, PaginatedRequestParams()) + with pytest.raises(MCPError, match="unavailable on /mcp/proxy"): + await server.get_skill(ctx, GetSkillRequestParams(uri="skill://billing/refunds/SKILL.md")) class FailureRecorder(CustomLogger): diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py index d47b2f73a48..9fc66c763cd 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -20427,3 +20427,30 @@ async def test_legacy_upstream_elicitation_rejects_modern_downstream_without_con finally: active_mcp_request_ctx_var.reset(request_token) legacy_server.active_mcp_session_var.reset(session_token) + + +@pytest.mark.asyncio +async def test_a_blocked_stdio_server_serves_no_skills(monkeypatch: pytest.MonkeyPatch) -> None: + from mcp import MCPError + from mcp.types import INVALID_PARAMS, PaginatedRequestParams + + from litellm.proxy._experimental.mcp_server.skills import GetSkillRequestParams, ListSkillsRequest, ListSkillsResult + + monkeypatch.delenv("LITELLM_ENABLE_MCP_STDIO", raising=False) + manager: Final = MCPServerManager() + server: Final = MCPServer( + server_id="stdio-skills", + name="stdio_skills", + transport=MCPTransport.stdio, + command="python", + args=["server.py"], + ) + first: Final = await manager.get_optional_catalog_page(server, ListSkillsRequest(), None) + assert first == ListSkillsResult(skills=()) + with pytest.raises(RuntimeError, match="Upstream catalog is unavailable"): + await manager.get_optional_catalog_page( + server, ListSkillsRequest(params=PaginatedRequestParams(cursor="next")), None + ) + with pytest.raises(MCPError, match="No skill is served") as missing: + await manager.get_skill_from_server(server, GetSkillRequestParams(uri="skill://stdio/SKILL.md"), None) + assert missing.value.error.code == INVALID_PARAMS diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py index 48106b1de5f..fdc46e95887 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py @@ -1160,6 +1160,7 @@ async def test_mcp_read_resource_success(): oauth2_headers=None, raw_headers=None, user_api_key_auth=user_api_key_auth, + scope_servers=[server], ) mock_manager.read_resource_from_server.assert_awaited_once_with( server=server, @@ -11212,8 +11213,11 @@ def test_protocol_header_respects_configured_advertisement(revision, rejected): @pytest.mark.asyncio -async def test_discovery_adapter_preserves_authenticated_context(_mcp_request_ctx): - from mcp.types import DiscoverResult, RequestParams, ServerCapabilities +@pytest.mark.parametrize("declared", (frozenset(), frozenset({"io.modelcontextprotocol/skills"}))) +async def test_discovery_adapter_preserves_authenticated_context(_mcp_request_ctx, declared: frozenset[str]): + from types import SimpleNamespace + + from mcp.types import ClientCapabilities, DiscoverResult, RequestParams, ServerCapabilities from litellm.proxy._experimental.mcp_server import server expected = DiscoverResult(supported_versions=["2025-11-25"], capabilities=ServerCapabilities()) @@ -11227,11 +11231,19 @@ async def test_discovery_adapter_preserves_authenticated_context(_mcp_request_ct ), patch.object(server.operations.GatewayOperations, "execute", dispatched), ): - result = await server.discover(_mcp_request_ctx(), RequestParams()) + result = await server.discover( + _mcp_request_ctx( + session=SimpleNamespace( + client_capabilities=ClientCapabilities(extensions={name: {} for name in declared}) + ) + ), + RequestParams(), + ) assert result is expected context = dispatched.await_args.args[1] assert context.user_api_key_auth.user_id == "discover-caller" assert context.mcp_servers == ("allowed",) + assert context.client_extensions == declared @pytest.mark.asyncio diff --git a/tests/unit/proxy/_experimental/mcp_server/test_operations.py b/tests/unit/proxy/_experimental/mcp_server/test_operations.py index 3957f486c4c..c9fcb75b85a 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_operations.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_operations.py @@ -1,4 +1,5 @@ import asyncio +import base64 from collections.abc import Callable, Sequence from typing import Final, Literal from unittest.mock import AsyncMock, patch @@ -1335,10 +1336,16 @@ async def test_discovery_skips_disabled_capability_listings(disabled: str, monke ResourceTemplate, ) + from dataclasses import replace + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server.skills import SKILLS_EXTENSION, ListSkillsResult, Skill monkeypatch.setitem(proxy_server.general_settings, "mcp_disabled_capabilities", [disabled]) - context: Final = prepare_context(UserAPIKeyAuth(user_id="scoped"), mcp_servers=["authorized"]) + context: Final = replace( + prepare_context(UserAPIKeyAuth(user_id="scoped"), mcp_servers=["authorized"]), + client_extensions=frozenset({SKILLS_EXTENSION}), + ) tools: Final = AsyncMock(return_value=ListToolsResult(tools=[])) prompts: Final = AsyncMock(return_value=ListPromptsResult(prompts=[Prompt(name="allowed")])) resources: Final = AsyncMock( @@ -1349,7 +1356,11 @@ async def test_discovery_skips_disabled_capability_listings(disabled: str, monke resource_templates=[ResourceTemplate(name="allowed", uri_template="test://{id}")] ) ) + skills: Final = AsyncMock( + return_value=ListSkillsResult(skills=(Skill(uri="skill://allowed/SKILL.md", frontmatter={"name": "allowed"}),)) + ) with ( + patch.object(operations, "_execute_list_skills", skills), patch.object(operations, "_execute_handle_list_tools", tools), patch.object(operations, "_execute_list_prompts", prompts), patch.object(operations, "_execute_list_resources", resources), @@ -1358,6 +1369,7 @@ async def test_discovery_skips_disabled_capability_listings(disabled: str, monke result = await GatewayOperations().execute(DiscoverRequest(), context) assert (result.capabilities.prompts is None) == (disabled == "prompts") assert (result.capabilities.resources is None) == (disabled == "resources") + assert skills.await_count == (0 if disabled == "resources" else 1) if disabled == "prompts": prompts.assert_not_awaited() else: @@ -1950,8 +1962,13 @@ async def test_guardrail_context_never_recovers_a_previous_request_bearer( async def test_discovery_optional_probes_use_remaining_listing_budget(delay, monkeypatch): from types import SimpleNamespace from mcp.types import ( - DiscoverRequest, ListToolsResult, ListPromptsResult, ListResourcesResult, - ListResourceTemplatesResult, ToolsCapability, Tool, + DiscoverRequest, + ListToolsResult, + ListPromptsResult, + ListResourcesResult, + ListResourceTemplatesResult, + ToolsCapability, + Tool, ) from litellm.proxy import proxy_server from litellm.proxy._experimental.mcp_server import subscriptions @@ -1982,9 +1999,15 @@ async def test_discovery_optional_probes_use_remaining_listing_budget(delay, mon patch.object(operations, "_execute_handle_list_tools", side_effect=tools), patch.object(operations, "_execute_list_prompts", return_value=ListPromptsResult(prompts=[])), patch.object(operations, "_execute_list_resources", return_value=ListResourcesResult(resources=[])), - patch.object(operations, "_execute_list_resource_templates", return_value=ListResourceTemplatesResult(resource_templates=[])), + patch.object( + operations, + "_execute_list_resource_templates", + return_value=ListResourceTemplatesResult(resource_templates=[]), + ), patch.object(operations, "_mcp_server_rate_limit_rejection", return_value=None), - patch.object(subscriptions, "_targets", side_effect=blocked if delay == "target" else None, return_value=(target,)) as targets, + patch.object( + subscriptions, "_targets", side_effect=blocked if delay == "target" else None, return_value=(target,) + ) as targets, ): context = prepare_context(UserAPIKeyAuth(user_id="scoped"), protocol_version="2026-07-28") result = await asyncio.wait_for(GatewayOperations().execute(DiscoverRequest(), context), 0.5) @@ -1996,3 +2019,284 @@ async def test_discovery_optional_probes_use_remaining_listing_budget(delay, mon else: assert entered.is_set() assert cleaned.is_set() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("client_declares", [True, False]) +@pytest.mark.parametrize("skills_listed", [True, False]) +async def test_discovery_advertises_skills_only_to_a_declaring_client_with_reachable_skills( + client_declares: bool, skills_listed: bool +) -> None: + from dataclasses import replace + + from mcp.types import ( + DiscoverRequest, + ListPromptsResult, + ListResourcesResult, + ListResourceTemplatesResult, + ListToolsResult, + ) + + from litellm.proxy._experimental.mcp_server.skills import SKILLS_EXTENSION, ListSkillsResult, Skill + + listed: Final = ListSkillsResult( + skills=(Skill(uri="skill://billing/refunds/SKILL.md", frontmatter={"name": "refunds"}),) + if skills_listed + else () + ) + context: Final = replace( + prepare_context(UserAPIKeyAuth(user_id="scoped"), protocol_version="2026-07-28"), + client_extensions=frozenset({SKILLS_EXTENSION}) if client_declares else frozenset(), + ) + with ( + patch.object(operations, "configured_versions", lambda: ("2025-11-25", "2026-07-28")), + patch.object(operations, "_execute_handle_list_tools", AsyncMock(return_value=ListToolsResult(tools=[]))), + patch.object(operations, "_execute_list_prompts", AsyncMock(return_value=ListPromptsResult(prompts=[]))), + patch.object(operations, "_execute_list_resources", AsyncMock(return_value=ListResourcesResult(resources=[]))), + patch.object( + operations, + "_execute_list_resource_templates", + AsyncMock(return_value=ListResourceTemplatesResult(resource_templates=[])), + ), + patch.object(operations, "_execute_list_skills", AsyncMock(return_value=listed)) as skills, + ): + result = await GatewayOperations().execute(DiscoverRequest(), context) + assert skills.await_count == (1 if client_declares else 0) + assert (result.capabilities.extensions or {}) == ( + {SKILLS_EXTENSION: {}} if client_declares and skills_listed else {} + ) + assert (result.capabilities.resources is not None) == (client_declares and skills_listed) + + +def _skills_rig(monkeypatch: pytest.MonkeyPatch, allowed: Sequence[MCPServer]) -> tuple[AsyncMock, AsyncMock]: + from mcp.types import PaginatedRequestParams + + from litellm.proxy._experimental.mcp_server.skills import ( + GetSkillRequestParams, + GetSkillResult, + ListSkillsRequest, + ListSkillsResult, + Skill, + ) + + owned: Final = { + "billing": Skill.model_validate( + { + "uri": "skill://billing/refunds/SKILL.md", + "frontmatter": {"name": "refunds"}, + "resources": [{"uri": "skill://billing/refunds/policy.md", "digest": "sha256:1", "size": 4}], + } + ), + "support": Skill(uri="skill://support/refunds/SKILL.md", frontmatter={"name": "support"}), + } + + async def page( + server: MCPServer, request: ListSkillsRequest, *_args: object, **_kwargs: object + ) -> ListSkillsResult: + assert request.params == PaginatedRequestParams(cursor=None) + return ListSkillsResult(skills=(owned[server.server_id],)) + + async def get( + server: MCPServer, params: GetSkillRequestParams, *_args: object, **_kwargs: object + ) -> GetSkillResult: + return GetSkillResult(skill=owned[server.server_id], cache_scope="public", ttl_ms=5000) + + listing: Final = AsyncMock(side_effect=page) + fetch: Final = AsyncMock(side_effect=get) + monkeypatch.setattr(operations, "_get_allowed_mcp_servers", AsyncMock(return_value=list(allowed))) + monkeypatch.setattr(operations, "_enforce_mcp_server_rate_limit", AsyncMock()) + monkeypatch.setattr(operations.global_mcp_server_manager, "get_optional_catalog_page", listing) + monkeypatch.setattr(operations.global_mcp_server_manager, "get_skill_from_server", fetch) + return listing, fetch + + +def _skill_server(name: str) -> MCPServer: + return MCPServer(server_id=name, name=name, transport=MCPTransport.http, url=f"https://{name}.example/mcp") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "allowed_names, uri, owner", + [ + (("billing", "support"), "skill://billing/refunds/SKILL.md", "billing"), + (("billing", "support"), "skill://support/refunds/SKILL.md", "support"), + (("billing", "support"), "skill://nowhere/SKILL.md", None), + (("support",), "skill://billing/refunds/SKILL.md", None), + ], +) +async def test_get_skill_routes_to_the_allowed_owner_or_reports_invalid_params( + monkeypatch: pytest.MonkeyPatch, allowed_names: tuple[str, ...], uri: str, owner: str | None +) -> None: + from mcp import MCPError + + from litellm.proxy._experimental.mcp_server.skills import GetSkillRequest, GetSkillRequestParams + + allowed: Final = tuple(_skill_server(name) for name in allowed_names) + listing, fetch = _skills_rig(monkeypatch, allowed) + request: Final = GetSkillRequest(params=GetSkillRequestParams(uri=uri)) + context: Final = prepare_context(UserAPIKeyAuth(user_id="scoped")) + if owner is None: + with pytest.raises(MCPError) as missing: + await GatewayOperations().execute(request, context) + assert missing.value.error.code == -32602 + assert fetch.await_count == 0 + assert {call.args[0].server_id for call in listing.await_args_list} == set(allowed_names) + return + result = await GatewayOperations().execute(request, context) + assert [call.args[0].server_id for call in fetch.await_args_list] == [owner] + assert result.skill.uri == uri + assert result.cache_scope == "private" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("uri, owner", [("skill://billing/refunds/policy.md", "billing"), ("skill://none.md", None)]) +async def test_resource_read_across_servers_reaches_only_the_skill_owner_and_preserves_bytes( + monkeypatch: pytest.MonkeyPatch, uri: str, owner: str | None +) -> None: + from fastapi import HTTPException + from mcp.types import BlobResourceContents, ReadResourceResult + from pydantic import AnyUrl + + payload: Final = base64.b64encode(bytes(range(256))).decode() + upstream: Final = ReadResourceResult(contents=[BlobResourceContents(uri=uri, blob=payload)]) + _skills_rig(monkeypatch, (_skill_server("billing"), _skill_server("support"))) + read: Final = AsyncMock(return_value=upstream) + monkeypatch.setattr(operations.global_mcp_server_manager, "read_resource_from_server", read) + if owner is None: + with pytest.raises(HTTPException) as unowned: + await operations.mcp_read_resource(AnyUrl(uri), UserAPIKeyAuth(user_id="scoped")) + assert unowned.value.status_code == 400 + assert read.await_count == 0 + return + result: Final = await operations.mcp_read_resource(AnyUrl(uri), UserAPIKeyAuth(user_id="scoped")) + assert [call.kwargs["server"].server_id for call in read.await_args_list] == [owner] + assert isinstance(result, ReadResourceResult) + content: Final = result.contents[0] + assert isinstance(content, BlobResourceContents) + assert base64.b64decode(content.blob) == bytes(range(256)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "uri, owner", [("skill://billing/refunds/SKILL.md", "billing"), ("skill://support/refunds/SKILL.md", None)] +) +async def test_skill_owner_lookup_admits_each_server_through_its_rate_limit( + monkeypatch: pytest.MonkeyPatch, uri: str, owner: str | None +) -> None: + from litellm.proxy._experimental.mcp_server.skills import GetSkillRequest, GetSkillRequestParams + + listing, fetch = _skills_rig(monkeypatch, (_skill_server("billing"), _skill_server("support"))) + + async def limit(_auth: UserAPIKeyAuth | None, server: MCPServer) -> None: + if server.server_id == "support": + raise ProxyRateLimitError(detail="support RPM exceeded") + + monkeypatch.setattr(operations, "_enforce_mcp_server_rate_limit", AsyncMock(side_effect=limit)) + request: Final = GetSkillRequest(params=GetSkillRequestParams(uri=uri)) + context: Final = prepare_context(UserAPIKeyAuth(user_id="scoped")) + if owner is None: + with pytest.raises(ProxyRateLimitError): + await GatewayOperations().execute(request, context) + assert fetch.await_count == 0 + else: + result: Final = await GatewayOperations().execute(request, context) + assert result.skill.uri == uri + assert [call.args[0].server_id for call in fetch.await_args_list] == [owner] + assert [call.args[0].server_id for call in listing.await_args_list] == ["billing"] + + +@pytest.mark.asyncio +async def test_resource_read_across_servers_withholds_a_caller_bearer_another_server_would_consume( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from mcp.types import ReadResourceResult, TextResourceContents + from pydantic import AnyUrl + + def forwarded(name: str) -> MCPServer: + return MCPServer( + server_id=name, + name=name, + transport=MCPTransport.http, + url=f"https://{name}.example/mcp", + auth_type=MCPAuth.true_passthrough, + ) + + uri: Final = "skill://billing/refunds/policy.md" + _skills_rig(monkeypatch, (forwarded("billing"), forwarded("support"))) + read: Final = AsyncMock(return_value=ReadResourceResult(contents=[TextResourceContents(uri=uri, text="policy")])) + monkeypatch.setattr(operations.global_mcp_server_manager, "read_resource_from_server", read) + await operations.mcp_read_resource( + AnyUrl(uri), + UserAPIKeyAuth(user_id="scoped"), + oauth2_headers={"Authorization": "Bearer caller-token"}, + raw_headers={"authorization": "Bearer caller-token"}, + ) + assert [call.kwargs["server"].server_id for call in read.await_args_list] == ["billing"] + sent: Final = read.await_args_list[0].kwargs + assert sent["mcp_auth_header"] is None + assert "authorization" not in {name.lower() for name in sent["extra_headers"] or {}} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cursor", [None, "continuation-state"]) +async def test_skill_listing_failure_is_empty_first_and_needs_a_fresh_listing_on_continuation( + monkeypatch: pytest.MonkeyPatch, cursor: str | None +) -> None: + from mcp import MCPError + from mcp.types import INVALID_PARAMS, PaginatedRequestParams + + from litellm.proxy._experimental.mcp_server.server_resolution import MCPServerTargetCatalog + from litellm.proxy._experimental.mcp_server.skills import ListSkillsRequest, ListSkillsResult + + failure: Final = RuntimeError("catalog temporarily unavailable") + monkeypatch.setattr(MCPServerTargetCatalog, "list", AsyncMock(side_effect=failure)) + request: Final = ListSkillsRequest(params=PaginatedRequestParams(cursor=cursor)) + if cursor is None: + assert await GatewayOperations().execute(request, prepare_context()) == ListSkillsResult(skills=()) + return + with pytest.raises(MCPError, match="start a fresh listing") as raised: + await GatewayOperations().execute(request, prepare_context()) + assert raised.value.error.code == INVALID_PARAMS + assert raised.value.__cause__ is failure + + +@pytest.mark.asyncio +async def test_skill_listing_access_denial_surfaces_as_invalid_request(monkeypatch: pytest.MonkeyPatch) -> None: + from fastapi import HTTPException + from mcp import MCPError + from mcp.types import INVALID_REQUEST, PaginatedRequestParams + + from litellm.proxy._experimental.mcp_server.server_resolution import MCPServerTargetCatalog + from litellm.proxy._experimental.mcp_server.skills import ListSkillsRequest + + denial: Final = HTTPException(status_code=403, detail={"error": "server billing is not allowed"}) + monkeypatch.setattr(MCPServerTargetCatalog, "list", AsyncMock(side_effect=denial)) + with pytest.raises(MCPError, match="server billing is not allowed") as raised: + await GatewayOperations().execute(ListSkillsRequest(params=PaginatedRequestParams()), prepare_context()) + assert raised.value.error.code == INVALID_REQUEST + + +@pytest.mark.asyncio +async def test_disabling_resources_turns_off_skills_without_asking_upstreams(monkeypatch: pytest.MonkeyPatch) -> None: + from mcp import MCPError + from mcp.types import METHOD_NOT_FOUND + + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server.skills import ( + GetSkillRequest, + GetSkillRequestParams, + ListSkillsRequest, + ) + + monkeypatch.setitem(proxy_server.general_settings, "mcp_disabled_capabilities", ["resources"]) + listing, fetch = _skills_rig(monkeypatch, (_skill_server("billing"),)) + context: Final = prepare_context(UserAPIKeyAuth(user_id="scoped")) + listed: Final = await GatewayOperations().execute(ListSkillsRequest(), context) + with pytest.raises(MCPError) as rejected: + await GatewayOperations().execute( + GetSkillRequest(params=GetSkillRequestParams(uri="skill://billing/refunds/SKILL.md")), context + ) + assert listed.skills == () + assert rejected.value.error.code == METHOD_NOT_FOUND + assert listing.await_count == 0 + assert fetch.await_count == 0