mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-25 05:54:53 +00:00
Merge 5e519482ee into c80f1c2ec8
This commit is contained in:
commit
d7cd26970f
16 changed files with 879 additions and 40 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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 ##################
|
||||
|
|
|
|||
55
litellm/proxy/_experimental/mcp_server/skills.py
Normal file
55
litellm/proxy/_experimental/mcp_server/skills.py
Normal file
|
|
@ -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"
|
||||
|
|
@ -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/")]
|
||||
|
|
|
|||
|
|
@ -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": "",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue