From c6b5f6da7f9bf785f89bcb59e48c3a72d6afd3af Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Tue, 22 Sep 2026 12:04:09 -0700 Subject: [PATCH 01/24] fix(mcp): refresh server catalog for each gateway operation --- .../mcp_server/auth/user_api_key_auth_mcp.py | 2 + .../mcp_server/byok_oauth_endpoints.py | 2 +- .../proxy/_experimental/mcp_server/catalog.py | 94 ++++++++ .../mcp_server/discoverable_endpoints.py | 103 ++++----- .../mcp_server/mcp_server_manager.py | 20 +- .../_experimental/mcp_server/operations.py | 7 + .../mcp_server/rest_endpoints.py | 3 + .../mcp_management_endpoints.py | 7 + .../mcp/litellm_proxy_mcp_handler.py | 4 + .../auth/test_user_api_key_auth_mcp.py | 8 +- .../mcp_server/test_byok_oauth_endpoints.py | 1 + .../mcp_server/test_discoverable_endpoints.py | 72 +++++- .../mcp_server/test_mcp_server_manager.py | 206 ++++++++++++++++++ .../mcp_server/test_rest_endpoints.py | 1 + .../test_mcp_management_endpoints.py | 8 +- .../mcp/test_litellm_proxy_mcp_handler.py | 9 +- 16 files changed, 478 insertions(+), 69 deletions(-) create mode 100644 litellm/proxy/_experimental/mcp_server/catalog.py diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 60dc91a69cc..1f5eb9db66f 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -13,6 +13,7 @@ from typing_extensions import assert_never import litellm from litellm._logging import verbose_logger +from litellm.proxy._experimental.mcp_server.catalog import with_mcp_catalog from litellm.proxy._experimental.mcp_server.oauth_utils import ( get_passthrough_resource_metadata_url, get_passthrough_www_authenticate, @@ -398,6 +399,7 @@ class MCPRequestHandler: LITELLM_MCP_ACCESS_GROUPS_HEADER_NAME = SpecialHeaders.mcp_access_groups.value @staticmethod + @with_mcp_catalog async def process_mcp_request( scope: Scope, ) -> tuple[ diff --git a/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py b/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py index 2c63e0a96d8..4df115727c0 100644 --- a/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py @@ -689,7 +689,7 @@ async def byok_authorize_get( global_mcp_server_manager, ) - registry: Final = global_mcp_server_manager.get_registry() + registry: Final = await global_mcp_server_manager.catalog.list() if server_id in registry: srv: Final = registry[server_id] server_name = srv.server_name or srv.name diff --git a/litellm/proxy/_experimental/mcp_server/catalog.py b/litellm/proxy/_experimental/mcp_server/catalog.py new file mode 100644 index 00000000000..673a1260c09 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/catalog.py @@ -0,0 +1,94 @@ +from __future__ import annotations + +import asyncio +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping +from contextlib import asynccontextmanager +from contextvars import ContextVar +from dataclasses import dataclass +from functools import wraps +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, ParamSpec, TypeVar + +if TYPE_CHECKING: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager + from litellm.types.mcp_server.mcp_server_manager import MCPServer + +_P: Final = ParamSpec("_P") +_R: Final = TypeVar("_R") + + +@dataclass(slots=True) +class _CatalogScope: + servers: Mapping[str, MCPServer] + owner_task_id: int + active: bool = True + + +class TargetCatalog: + def __init__(self, manager: MCPServerManager) -> None: + self._manager = manager + self._reload_lock = asyncio.Lock() + self._scope: ContextVar[_CatalogScope | None] = ContextVar("mcp_catalog_scope", default=None) + + @property + def current(self) -> Mapping[str, MCPServer] | None: + scope: Final = self._scope.get() + return scope.servers if scope is not None and scope.active else None + + async def refresh(self) -> None: + async with self._reload_lock: + await self._manager._reload_servers_from_database() # pyright: ignore[reportPrivateUsage] # existing staged loader + + async def list(self) -> Mapping[str, MCPServer]: + from litellm.proxy.proxy_server import prisma_client # noqa: PLC0415 # runtime proxy dependency + + scope: Final = self._scope.get() + if scope is not None and scope.active and scope.owner_task_id == id(asyncio.current_task()): + return scope.servers + async with self._reload_lock: + if prisma_client is not None: + try: + await self._manager._reload_servers_from_database(reuse_unchanged=True) # pyright: ignore[reportPrivateUsage] # existing staged loader + except Exception as exc: # noqa: BLE001 # never serve an unverified database snapshot + from fastapi import HTTPException # noqa: PLC0415 # optional proxy dependency + + raise HTTPException( + status_code=503, detail="MCP server configuration could not be refreshed" + ) from exc + return MappingProxyType(self._manager.config_mcp_servers | self._manager.registry) + + @asynccontextmanager + async def operation(self) -> AsyncIterator[Mapping[str, MCPServer]]: + existing: Final = self._scope.get() + if existing is not None and existing.active and existing.owner_task_id == id(asyncio.current_task()): + yield existing.servers + return + scope: Final = _CatalogScope(await self.list(), id(asyncio.current_task())) + token: Final = self._scope.set(scope) + try: + yield scope.servers + finally: + scope.active = False + self._scope.reset(token) + + async def resolve(self, lookup: str, client_ip: str | None = None) -> MCPServer | None: + async with self.operation(): + return self._manager.get_mcp_server_by_name( + lookup, client_ip=client_ip + ) or self._manager.get_mcp_server_by_id(lookup, client_ip=client_ip) + + +def with_mcp_catalog(function: Callable[_P, Awaitable[_R]]) -> Callable[_P, Awaitable[_R]]: + @wraps(function) + async def wrapped( + *args: _P.args, + **kwargs: _P.kwargs, # kwargs-ok: ParamSpec preserves each wrapped endpoint keyword contract + ) -> _R: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # manager imports catalog + global_mcp_server_manager, + ) + + async with global_mcp_server_manager.catalog.operation(): + return await function(*args, **kwargs) + + return wrapped diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 64bab0a7832..56c2896ea4f 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -36,6 +36,7 @@ from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( can_store_oauth_credential, oauth_authorization_uses_gateway_credential, ) +from litellm.proxy._experimental.mcp_server.catalog import with_mcp_catalog from litellm.proxy._experimental.mcp_server.faults import ( CallerRejected, CredentialSource, @@ -1875,6 +1876,7 @@ async def register_client_with_server( @router.get("/authorize/mcp-session") +@with_mcp_catalog async def authorize_mcp_session( request: Request, redirect_uri: str, @@ -1900,6 +1902,7 @@ async def authorize_mcp_session( @router.get("/{mcp_server_name}/authorize") @router.get("/authorize") +@with_mcp_catalog async def authorize( request: Request, redirect_uri: str, @@ -1938,9 +1941,11 @@ async def authorize( resource=resource, ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + lookup_name: Final[str | None] = mcp_server_name or client_id client_ip: Final = IPAddressUtils.get_mcp_client_ip(request) - mcp_server = _resolve_mcp_server_by_name_or_id(lookup_name, client_ip) if lookup_name else None + mcp_server = await global_mcp_server_manager.catalog.resolve(lookup_name, client_ip) if lookup_name else None if mcp_server is None and mcp_server_name is None: mcp_server = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip) if mcp_server is None: @@ -1974,6 +1979,7 @@ async def authorize( @router.post("/{mcp_server_name}/token") @router.post("/token") +@with_mcp_catalog async def token_endpoint( request: Request, grant_type: str = Form(...), @@ -2067,6 +2073,7 @@ async def _vendor_credential_state(user_id: str, server_id: str) -> VendorCreden @router.get("/authorize/flow") +@with_mcp_catalog async def authorize_flow(request: Request, flow: str) -> Response: return await describe_connect_flow( request=request, @@ -2078,6 +2085,7 @@ async def authorize_flow(request: Request, flow: str) -> Response: @router.post("/authorize/complete") +@with_mcp_catalog async def authorize_complete( request: Request, flow: str = Form(...), @@ -2435,6 +2443,7 @@ def is_network_error(exc: Exception) -> bool: return isinstance(exc, httpx.TransportError) +@with_mcp_catalog async def _build_oauth_protected_resource_response( request: Request, mcp_server_name: str | None, @@ -2739,12 +2748,7 @@ def _build_oauth_authorization_server_response( *, issuer_path: str | None = None, ) -> dict: - """Build OAuth authorization server metadata response (gateway-as-AS shape). - - Synchronous because the body only does dict construction and synchronous - registry lookups; unlike :func:`_build_oauth_protected_resource_response` - it does not need to await any upstream IO. - """ + """Build OAuth authorization server metadata response (gateway-as-AS shape).""" request_base_url: Final = get_request_base_url(request) client_ip: Final = IPAddressUtils.get_mcp_client_ip(request) explicitly_named: Final = mcp_server_name is not None @@ -2792,6 +2796,7 @@ def _build_oauth_authorization_server_response( # Standard MCP pattern: /.well-known/oauth-authorization-server/mcp/{server_name} @router.get(f"/.well-known/oauth-authorization-server{well_known_root_suffix()}/mcp/{{mcp_server_name}}") +@with_mcp_catalog async def oauth_authorization_server_mcp_standard(request: Request, mcp_server_name: str): """ OAuth authorization server discovery endpoint using standard MCP URL pattern. @@ -2809,6 +2814,7 @@ async def oauth_authorization_server_mcp_standard(request: Request, mcp_server_n # LiteLLM legacy pattern and root endpoint @router.get(f"/.well-known/oauth-authorization-server{well_known_root_suffix()}/{{mcp_server_name}}") @router.get("/.well-known/oauth-authorization-server") +@with_mcp_catalog async def oauth_authorization_server_mcp(request: Request, mcp_server_name: str | None = None): """ OAuth authorization server discovery endpoint. @@ -2882,6 +2888,7 @@ async def jwks_json(request: Request): # Additional legacy pattern support @router.get(f"/.well-known/oauth-authorization-server{well_known_root_suffix()}/{{mcp_server_name}}/mcp") +@with_mcp_catalog async def oauth_authorization_server_legacy(request: Request, mcp_server_name: str): """ OAuth authorization server discovery for legacy /{server_name}/mcp pattern. @@ -2903,46 +2910,44 @@ async def register_client(request: Request, mcp_server_name: str | None = None): data: Final[dict] = {**request_data} client_redirect_uris: Final = client_supplied_redirect_uris(data.get("redirect_uris")) - dummy_return: Final = { - "client_id": mcp_server_name or "dummy_client", - "client_secret": "dummy", - "redirect_uris": client_redirect_uris or [f"{request_base_url}/callback"], - } - client_ip: Final = IPAddressUtils.get_mcp_client_ip(request) - if not mcp_server_name: - # A real DCR request carries redirect_uris (RFC 7591): route it to the aggregate DCR - # endpoint the aggregate authorization-server metadata advertises. A single-server - # deployment registers at /{server}/register instead (its bare-origin discovery - # advertises that), so this does not affect it. A request without redirect_uris is not - # a DCR request, so the legacy single-server-or-dummy fallback is kept for it. - if data.get("redirect_uris"): - return await register_aggregate_client( - request=request, request_body=data, token_exchange_available=token_exchange_available() - ) - resolved: Final = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip) - if resolved: - return await register_client_with_server( - request=request, - mcp_server=resolved, - client_name=data.get("client_name", ""), - grant_types=data.get("grant_types", []), - response_types=data.get("response_types", []), - token_endpoint_auth_method=data.get("token_endpoint_auth_method", ""), - fallback_client_id=resolved.server_name or resolved.name, - client_redirect_uris=client_redirect_uris, - ) - return dummy_return + if not mcp_server_name and data.get("redirect_uris"): + return await register_aggregate_client( + request=request, request_body=data, token_exchange_available=token_exchange_available() + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager - mcp_server: Final = _resolve_mcp_server_by_name_or_id(mcp_server_name, client_ip) - if mcp_server is None: - return dummy_return - return await register_client_with_server( - request=request, - mcp_server=mcp_server, - client_name=data.get("client_name", ""), - grant_types=data.get("grant_types", []), - response_types=data.get("response_types", []), - token_endpoint_auth_method=data.get("token_endpoint_auth_method", ""), - fallback_client_id=mcp_server_name, - client_redirect_uris=client_redirect_uris, - ) + async with global_mcp_server_manager.catalog.operation(): + dummy_return: Final = { + "client_id": mcp_server_name or "dummy_client", + "client_secret": "dummy", + "redirect_uris": client_redirect_uris or [f"{request_base_url}/callback"], + } + client_ip: Final = IPAddressUtils.get_mcp_client_ip(request) + if not mcp_server_name: + resolved: Final = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip) + if resolved: + return await register_client_with_server( + request=request, + mcp_server=resolved, + client_name=data.get("client_name", ""), + grant_types=data.get("grant_types", []), + response_types=data.get("response_types", []), + token_endpoint_auth_method=data.get("token_endpoint_auth_method", ""), + fallback_client_id=resolved.server_name or resolved.name, + client_redirect_uris=client_redirect_uris, + ) + return dummy_return + + mcp_server: Final = _resolve_mcp_server_by_name_or_id(mcp_server_name, client_ip) + if mcp_server is None: + return dummy_return + return await register_client_with_server( + request=request, + mcp_server=mcp_server, + client_name=data.get("client_name", ""), + grant_types=data.get("grant_types", []), + response_types=data.get("response_types", []), + token_endpoint_auth_method=data.get("token_endpoint_auth_method", ""), + fallback_client_id=mcp_server_name, + client_redirect_uris=client_redirect_uris, + ) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 20c114a2f3e..9806bf1be68 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -73,6 +73,7 @@ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPServerAccess, _is_mcp_admitted_user_subject, ) +from litellm.proxy._experimental.mcp_server.catalog import TargetCatalog from litellm.proxy._experimental.mcp_server.contracts import OperationContext from litellm.proxy._experimental.mcp_server.elicitation_handler import ( MCP_ELICITATION_AVAILABLE, @@ -1860,6 +1861,7 @@ class MCPServerManager: self._template_discovery_cache = _DiscoveryCache[ResourceTemplate]( discovery_ttl, discovery_clock, TypeAdapter(tuple[ResourceTemplate, ...]) ) + self.catalog = TargetCatalog(self) self.registry: dict[str, MCPServer] = {} self._openapi_health_probes: Callable[[str], _OpenAPIHealthProbe] = lru_cache(maxsize=128)(_OpenAPIHealthProbe) self.config_mcp_servers: dict[str, MCPServer] = {} @@ -2287,11 +2289,12 @@ class MCPServerManager: e, ) - def get_registry(self) -> dict[str, MCPServer]: + def get_registry(self) -> Mapping[str, MCPServer]: """ Get the registered MCP Servers from the registry and union with the config MCP Servers """ - return self.config_mcp_servers | self.registry + snapshot: Final = self.catalog.current + return snapshot if snapshot is not None else self.config_mcp_servers | self.registry def is_config_declared_server(self, server_id: str) -> bool: """True when server_id was declared in config.yaml (present in the in-memory config map). @@ -6507,14 +6510,15 @@ class MCPServerManager: return None async def reload_servers_from_database(self): + await self.catalog.refresh() + + async def _reload_servers_from_database(self, *, reuse_unchanged: bool = False): """Re-synchronize the in-memory MCP server registry with the database.""" from litellm.proxy.management_endpoints.mcp_management_endpoints import ( get_prisma_client_or_throw, ) verbose_logger.debug("Loading MCP servers from database into registry...") - self._upstream_initialize_instructions_by_server_id.clear() - self._upstream_initialize_instructions_probed_at.clear() # perform authz check to filter the mcp servers user has access to prisma_client: Final = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") @@ -6593,7 +6597,8 @@ class MCPServerManager: # Register OpenAPI tools *after* the final short prefix is assigned # so the tools are stored in the global registry under the same # prefix that lookups will use. - await self._maybe_register_openapi_tools(new_server, initialize_mapping=False) + if not reuse_unchanged or new_server is not previous_registry.get(server_id): + await self._maybe_register_openapi_tools(new_server, initialize_mapping=False) registered_registry[server_id] = new_server if new_server.spec_path: registered_openapi_tools = True @@ -6607,10 +6612,13 @@ class MCPServerManager: dropped_registry_keys: Final = previous_registry.keys() - registered_registry.keys() for registry_key in dropped_registry_keys: + self._cleanup_server_tool_routing_artifacts(previous_registry[registry_key]) self._invalidate_oauth_discovery_state(previous_registry[registry_key].server_id) for server_id in previous_registry.keys() | registered_registry.keys(): if previous_registry.get(server_id) != registered_registry.get(server_id): + self._upstream_initialize_instructions_by_server_id.pop(server_id, None) + self._upstream_initialize_instructions_probed_at.pop(server_id, None) self._invalidate_discovery_lists(server_id) self.registry = registered_registry # A discovery task may have published into ``previous_registry`` while @@ -6834,7 +6842,7 @@ class MCPServerManager: return server return None - def get_filtered_registry(self, client_ip: str | None = None) -> dict[str, MCPServer]: + def get_filtered_registry(self, client_ip: str | None = None) -> Mapping[str, MCPServer]: """ Get registry filtered by client IP access control. diff --git a/litellm/proxy/_experimental/mcp_server/operations.py b/litellm/proxy/_experimental/mcp_server/operations.py index fcee3483e15..05a630c38c0 100644 --- a/litellm/proxy/_experimental/mcp_server/operations.py +++ b/litellm/proxy/_experimental/mcp_server/operations.py @@ -50,6 +50,7 @@ from litellm.proxy._experimental.mcp_server.byok_credential_cache import ( cache_byok_credential, get_cached_byok_credential, ) +from litellm.proxy._experimental.mcp_server.catalog import with_mcp_catalog from litellm.proxy._experimental.mcp_server.contracts import ( AuthorizedToolCall, OperationContext, @@ -614,6 +615,7 @@ def apply_tool_overrides( return tools +@with_mcp_catalog async def _get_allowed_mcp_servers( user_api_key_auth: UserAPIKeyAuth | None, mcp_servers: Sequence[str] | None, @@ -1435,6 +1437,7 @@ async def filter_tools_by_key_team_permissions( ] +@with_mcp_catalog async def _list_mcp_tools( user_api_key_auth: UserAPIKeyAuth | None = None, mcp_auth_header: str | None = None, @@ -1485,6 +1488,7 @@ async def _list_mcp_tools( return AggregateToolListing(tools=[], outcomes={}) +@with_mcp_catalog async def _list_mcp_prompts( user_api_key_auth: UserAPIKeyAuth | None = None, mcp_auth_header: str | None = None, @@ -1526,6 +1530,7 @@ async def _list_mcp_prompts( return managed_prompts +@with_mcp_catalog async def _list_mcp_resources( user_api_key_auth: UserAPIKeyAuth | None = None, mcp_auth_header: str | None = None, @@ -1555,6 +1560,7 @@ async def _list_mcp_resources( return managed_resources +@with_mcp_catalog async def _list_mcp_resource_templates( user_api_key_auth: UserAPIKeyAuth | None = None, mcp_auth_header: str | None = None, @@ -3055,6 +3061,7 @@ class GatewayOperations: @overload async def execute(self, operation: ReadResourceRequest, context: OperationContext) -> ReadResourceResult: ... + @with_mcp_catalog async def execute(self, operation: GatewayOperation, context: OperationContext) -> GatewayResult: match operation: case AuthorizedToolCall(): diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index c2f7bf7d531..1c476dc98fe 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -22,6 +22,7 @@ from litellm.exceptions import ( GuardrailRaisedException, ModifyResponseException, ) +from litellm.proxy._experimental.mcp_server.catalog import with_mcp_catalog from litellm.proxy._experimental.mcp_server.exceptions import ( MCPServerListError, MCPServerURLCredentialsError, @@ -861,6 +862,7 @@ if MCP_AVAILABLE: return await _apply_toolset_scope(user_api_key_dict, toolset.toolset_id) @router.get("/tools/list", dependencies=[Depends(user_api_key_auth)]) + @with_mcp_catalog async def list_tool_rest_api( request: Request, server_id: str | None = Query(None, description="The server id to list tools for"), @@ -1085,6 +1087,7 @@ if MCP_AVAILABLE: } @router.post("/tools/call", dependencies=[Depends(user_api_key_auth)]) + @with_mcp_catalog async def call_tool_rest_api( request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 9ad78876043..f79e6bd69c4 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -45,6 +45,8 @@ from fastapi import ( from fastapi.responses import JSONResponse from typing_extensions import ReadOnly, TypedDict +from litellm.proxy._experimental.mcp_server.catalog import with_mcp_catalog + try: from prisma.errors import RecordNotFoundError, UniqueViolationError except ImportError: @@ -1043,6 +1045,7 @@ if MCP_AVAILABLE: tags=["mcp"], description="MCP registry endpoint. Spec: https://github.com/modelcontextprotocol/registry", ) + @with_mcp_catalog async def get_mcp_registry(request: Request): if not _is_public_registry_enabled(): raise HTTPException( @@ -1166,6 +1169,7 @@ if MCP_AVAILABLE: dependencies=[Depends(user_api_key_auth)], response_model=list[LiteLLM_MCPServerTable], ) + @with_mcp_catalog async def fetch_all_mcp_servers( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), team_id: str | None = Query( @@ -1284,6 +1288,7 @@ if MCP_AVAILABLE: description="Health check for MCP servers", dependencies=[Depends(user_api_key_auth)], ) + @with_mcp_catalog async def health_check_servers( server_ids: list[str] | None = Query( None, @@ -1959,6 +1964,7 @@ if MCP_AVAILABLE: return _redact_mcp_credentials(temp_record) + @with_mcp_catalog async def _mcp_oauth_user_api_key_auth(request: Request) -> UserAPIKeyAuth: """ Auth dependency for MCP OAuth browser-navigation endpoints (/authorize, /token). @@ -2059,6 +2065,7 @@ if MCP_AVAILABLE: request_data=request_data, ) + @with_mcp_catalog async def _get_cached_temporary_mcp_server_or_404( server_id: str, user_api_key_dict: UserAPIKeyAuth, diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 88a2b92c680..c9425828407 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -10,6 +10,7 @@ from openai.types.responses.function_tool_param import FunctionToolParam from litellm._logging import verbose_logger from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.proxy._experimental.mcp_server.catalog import with_mcp_catalog from litellm.proxy._experimental.mcp_server.utils import ( iter_known_server_prefixes, logging_safe_mcp_headers, @@ -178,6 +179,7 @@ class LiteLLM_Proxy_MCP_Handler: ) @staticmethod + @with_mcp_catalog async def routes_through_gateway( tools: Iterable[Mapping[str, object]] | None, served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] = _gateway_served_names, @@ -229,6 +231,7 @@ class LiteLLM_Proxy_MCP_Handler: return user_api_key_auth @staticmethod + @with_mcp_catalog async def _get_mcp_tools_from_manager( user_api_key_auth: "UserAPIKeyAuth | None", mcp_tools_with_litellm_proxy: Iterable[Mapping[str, object]] | None, @@ -680,6 +683,7 @@ class LiteLLM_Proxy_MCP_Handler: return result_text or "Tool executed successfully" @staticmethod + @with_mcp_catalog async def _execute_tool_calls( tool_server_map: dict[str, str], tool_calls: Sequence[object], diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 35e7055bbc0..a18d17ef553 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -7482,9 +7482,11 @@ class TestGatewaySessionAdmission: rpm_limit=rpm_limit, ) ) + prisma = MagicMock() + prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[]) with ( patch("litellm.proxy.auth.auth_checks.get_user_object", get_user_object), - patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.prisma_client", prisma), patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), ): yield get_user_object @@ -9571,10 +9573,12 @@ class TestScopedSessionAdmission: rpm_limit=None, ) ) + prisma = MagicMock() + prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[]) with ( patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), patch("litellm.proxy.auth.auth_checks.get_user_object", get_user_object), - patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.prisma_client", prisma), patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), ): auth_result, *_rest = await MCPRequestHandler.process_mcp_request(scope_dict) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py index a77b4c8d565..327a6370580 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py @@ -632,6 +632,7 @@ async def test_execute_byok_tool_missing_credential_advertises_api_key_flow(monk mcp_operations.byok_credential_cache.flush_cache() server = MCPServer(server_id="byok-discovery", name="byok-discovery", transport=MCPTransport.http, is_byok=True) prisma = MagicMock() + prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[]) prisma.db.litellm_mcpusercredentials.find_unique = AsyncMock(return_value=None) monkeypatch.setattr(proxy_server, "prisma_client", prisma) with pytest.raises(HTTPException) as exc_info: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index c1d5cedeba0..b3e10561c8d 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -738,7 +738,8 @@ async def test_token_endpoint_forwards_code_verifier(): @pytest.mark.asyncio -async def test_register_client_without_mcp_server_name_returns_dummy(): +@pytest.mark.parametrize("server_name", [None, "missing"]) +async def test_register_client_without_mcp_server_name_returns_dummy(server_name): try: from fastapi import Request @@ -761,17 +762,18 @@ async def test_register_client_without_mcp_server_name_returns_dummy(): "litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body", new=AsyncMock(return_value={}), ): - result = await register_client(request=mock_request) + result = await register_client(request=mock_request, mcp_server_name=server_name) assert result == { - "client_id": "dummy_client", + "client_id": server_name or "dummy_client", "client_secret": "dummy", "redirect_uris": ["https://proxy.litellm.example/callback"], } @pytest.mark.asyncio -async def test_register_client_returns_existing_server_credentials(): +@pytest.mark.parametrize("use_root", [False, True]) +async def test_register_client_returns_existing_server_credentials(use_root): try: from fastapi import Request @@ -811,7 +813,9 @@ async def test_register_client_returns_existing_server_credentials(): "litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body", new=AsyncMock(return_value={}), ): - result = await register_client(request=mock_request, mcp_server_name=oauth2_server.server_name) + result = await register_client( + request=mock_request, mcp_server_name=None if use_root else oauth2_server.server_name + ) finally: global_mcp_server_manager.registry.clear() @@ -12347,3 +12351,61 @@ async def test_identity_bound_authorize_unrelated_bearer_uses_browser_session( proxy_server.prisma_client.db.litellm_mcpusercredentials.upsert.assert_not_called() proxy_server.prisma_client.db.litellm_usertable.create.assert_not_called() proxy_server.prisma_client.db.litellm_teamtable.create.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("change", ["create", "update", "delete"]) +async def test_authorize_observes_committed_peer_server_changes(monkeypatch, change): + from datetime import datetime, timedelta, timezone + from types import SimpleNamespace + + from fastapi import HTTPException, Request + + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server import discoverable_endpoints + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.proxy._types import LiteLLM_MCPServerTable + + monkeypatch.setenv("LITELLM_SALT_KEY", "catalog-consistency-test-key") + stamp = datetime.now(timezone.utc) + old_server = _create_id_lookup_oauth2_server() + old_server.updated_at = stamp + row = LiteLLM_MCPServerTable( + server_id=old_server.server_id, + server_name=old_server.server_name, + alias=old_server.alias, + url="https://upstream.example/mcp", + transport="http", + auth_type="oauth2", + authorization_url="https://new-provider.example/authorize", + token_url="https://new-provider.example/token", + scopes=["read"], + credentials={"client_id": "current-client", "client_secret": "current-secret"}, + created_at=stamp, + updated_at=stamp + timedelta(seconds=1), + ) + read_rows = AsyncMock(return_value=[] if change == "delete" else [row]) + prisma = SimpleNamespace(db=SimpleNamespace(litellm_mcpservertable=SimpleNamespace(find_many=read_rows))) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr( + global_mcp_server_manager, "registry", {} if change == "create" else {old_server.server_id: old_server} + ) + monkeypatch.setattr(global_mcp_server_manager, "config_mcp_servers", {}) + request = Request( + {"type": "http", "scheme": "https", "server": ("gateway.example", 443), "path": "/authorize", "headers": []} + ) + + if change == "delete": + with pytest.raises(HTTPException) as exc: + await discoverable_endpoints.authorize( + request, "http://localhost/callback", mcp_server_name=old_server.server_id + ) + assert exc.value.status_code == 404 + else: + response = await discoverable_endpoints.authorize( + request, "http://localhost/callback", mcp_server_name=old_server.server_id + ) + assert response.status_code == 307 + assert response.headers["location"].startswith("https://new-provider.example/authorize?") + assert "client_id=current-client" in response.headers["location"] + read_rows.assert_awaited_once() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 7725aca1948..30297322b0c 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -14516,3 +14516,209 @@ async def test_client_sampling_does_not_fill_explicit_context_from_another_ambie assert captured["client_ip"] is None finally: auth_context_var.reset(token) + + +def _catalog_row(name="initial"): + return LiteLLM_MCPServerTable( + server_id="catalog-server", + server_name=name, + transport="http", + url="https://upstream.example/mcp", + created_at=datetime(2026, 1, 1), + updated_at=datetime(2026, 1, 1 if name == "initial" else 2), + ) + + +def _catalog_database(monkeypatch, read_rows): + from types import SimpleNamespace + + from litellm.proxy import proxy_server + + client = SimpleNamespace(db=SimpleNamespace(litellm_mcpservertable=SimpleNamespace(find_many=read_rows))) + monkeypatch.setattr(proxy_server, "prisma_client", client) + + +@pytest.mark.asyncio +async def test_catalog_pins_one_snapshot_and_next_operation_reads_current_rows(monkeypatch): + read_rows = AsyncMock(return_value=[_catalog_row()]) + _catalog_database(monkeypatch, read_rows) + manager = MCPServerManager() + async with manager.catalog.operation(): + initial = manager.get_mcp_server_by_id("catalog-server") + read_rows.return_value = [_catalog_row("updated")] + async with manager.catalog.operation(): + assert await manager.catalog.resolve("initial") is initial + assert tuple((await manager.catalog.list()).values()) == (initial,) + read_rows.assert_awaited_once() + async with manager.catalog.operation(): + assert manager.get_mcp_server_by_id("catalog-server").name == "updated" + assert await manager.catalog.resolve("initial") is None + assert manager.catalog.current is None + assert read_rows.await_count == 2 + + +@pytest.mark.asyncio +async def test_catalog_database_failure_keeps_registry_but_rejects_operation(monkeypatch): + read_rows = AsyncMock(return_value=[_catalog_row()]) + _catalog_database(monkeypatch, read_rows) + manager = MCPServerManager() + snapshot = await manager.catalog.list() + read_rows.side_effect = RuntimeError("private connection detail") + with pytest.raises(HTTPException) as exc: + async with manager.catalog.operation(): + pytest.fail("An unverified database snapshot must not execute") + assert exc.value.status_code == 503 + assert exc.value.detail == "MCP server configuration could not be refreshed" + assert manager.registry == dict(snapshot) + assert manager.catalog.current is None + + +@pytest.mark.asyncio +async def test_catalog_serializes_background_and_operation_refreshes(monkeypatch): + entered = asyncio.Event() + release = asyncio.Event() + + async def first_read(**kwargs): + entered.set() + await release.wait() + return [_catalog_row()] + + read_rows = AsyncMock(side_effect=first_read) + _catalog_database(monkeypatch, read_rows) + manager = MCPServerManager() + first = asyncio.create_task(manager.reload_servers_from_database()) + await entered.wait() + read_rows.side_effect = None + read_rows.return_value = [_catalog_row("updated")] + second = asyncio.create_task(manager.catalog.list()) + await asyncio.sleep(0) + assert read_rows.await_count == 1 + release.set() + await first + latest = await second + assert latest["catalog-server"].name == "updated" + assert manager.registry["catalog-server"].name == "updated" + + +@pytest.mark.asyncio +async def test_catalog_cancellation_releases_refresh_lock(monkeypatch): + entered = asyncio.Event() + blocked = asyncio.Event() + + async def blocked_read(**kwargs): + entered.set() + await blocked.wait() + return [] + + read_rows = AsyncMock(side_effect=blocked_read) + _catalog_database(monkeypatch, read_rows) + manager = MCPServerManager() + task = asyncio.create_task(manager.catalog.list()) + await entered.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + read_rows.side_effect = None + read_rows.return_value = [_catalog_row()] + snapshot = await asyncio.wait_for(manager.catalog.list(), timeout=1) + assert tuple(snapshot) == ("catalog-server",) + assert manager.catalog.current is None + + +@pytest.mark.asyncio +async def test_catalog_child_operation_does_not_inherit_stale_session_snapshot(monkeypatch): + read_rows = AsyncMock(return_value=[_catalog_row()]) + _catalog_database(monkeypatch, read_rows) + manager = MCPServerManager() + release = asyncio.Event() + + async def child_operation(): + await release.wait() + async with manager.catalog.operation(): + return manager.get_mcp_server_by_id("catalog-server") + + async with manager.catalog.operation(): + child = asyncio.create_task(child_operation()) + read_rows.return_value = [_catalog_row("updated")] + release.set() + assert (await child).name == "updated" + assert read_rows.await_count == 2 + + +@pytest.mark.asyncio +async def test_catalog_config_only_snapshot_cleans_up_after_error(monkeypatch): + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "prisma_client", None) + manager = MCPServerManager() + configured = MCPServer(server_id="configured", name="configured", transport="stdio", command="echo") + manager.config_mcp_servers = {configured.server_id: configured} + async def fail_operation(): + async with manager.catalog.operation(): + assert await manager.catalog.resolve("configured") is configured + raise ValueError("stop operation") + + with pytest.raises(ValueError, match="stop operation"): + await fail_operation() + assert manager.catalog.current is None + assert dict(await manager.catalog.list()) == {"configured": configured} + + +@pytest.mark.asyncio +async def test_catalog_delete_drops_derived_tool_mapping(monkeypatch): + read_rows = AsyncMock(return_value=[_catalog_row()]) + _catalog_database(monkeypatch, read_rows) + manager = MCPServerManager() + await manager.catalog.list() + manager.tool_name_to_mcp_server_name_mapping = {"initial-echo": "initial"} + read_rows.return_value = [] + assert not await manager.catalog.list() + assert manager.tool_name_to_mcp_server_name_mapping == {} + + +@pytest.mark.asyncio +async def test_catalog_unchanged_read_preserves_derived_initialize_instructions(monkeypatch): + read_rows = AsyncMock(return_value=[_catalog_row()]) + _catalog_database(monkeypatch, read_rows) + manager = MCPServerManager() + await manager.catalog.list() + manager._upstream_initialize_instructions_by_server_id = {"catalog-server": "cached instructions"} + manager._upstream_initialize_instructions_probed_at = {"catalog-server": 123.0} + await manager.catalog.list() + assert manager._upstream_initialize_instructions_by_server_id == {"catalog-server": "cached instructions"} + assert manager._upstream_initialize_instructions_probed_at == {"catalog-server": 123.0} + + +@pytest.mark.asyncio +async def test_catalog_reuses_openapi_tools_until_configuration_or_background_refresh(monkeypatch, tmp_path): + from litellm.proxy._experimental.mcp_server import openapi_to_mcp_generator, tool_registry + + registry = tool_registry.MCPToolRegistry() + monkeypatch.setattr(tool_registry, "global_mcp_tool_registry", registry) + spec_path = tmp_path / "catalog.json" + spec_path.write_text(json.dumps({ + "openapi": "3.0.0", "info": {"title": "Catalog", "version": "1"}, + "paths": {"/echo": {"get": {"operationId": "echo"}}}, + })) + row = _catalog_row().model_copy(update={"spec_path": str(spec_path)}) + read_rows = AsyncMock(return_value=[row]) + _catalog_database(monkeypatch, read_rows) + manager = MCPServerManager() + with patch.object( + openapi_to_mcp_generator, "load_openapi_spec_async", + wraps=openapi_to_mcp_generator.load_openapi_spec_async, + ) as load_spec: + await manager.catalog.list() + first_tools = tuple(registry.list_tools()) + assert len(first_tools) == 1 + await manager.catalog.list() + assert load_spec.await_count == 1 + assert [tool.name for tool in registry.list_tools()] == [tool.name for tool in first_tools] + await manager.reload_servers_from_database() + assert load_spec.await_count == 2 + read_rows.return_value = [row.model_copy(update={"updated_at": datetime(2026, 1, 2)})] + await manager.catalog.list() + assert load_spec.await_count == 3 + read_rows.return_value = [] + await manager.catalog.list() + assert registry.list_tools() == [] diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index 233a8cc96ba..90178414fd6 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -1495,6 +1495,7 @@ class TestListToolsRestAPI: mcp_info={"server_name": "stub"}, ) stub_server.available_on_public_internet = True + monkeypatch.setattr(rest_endpoints.global_mcp_server_manager, "registry", {"server-1": stub_server}) mock_transport_ctx = AsyncMock() mock_transport_ctx.__aenter__ = AsyncMock(return_value=(MagicMock(), MagicMock())) diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index 53645e62034..acb248748f5 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -136,7 +136,7 @@ def create_mcp_router_test_client() -> TestClient: def patch_proxy_general_settings(settings: dict): - fake_proxy_server_module = types.SimpleNamespace(general_settings=settings) + fake_proxy_server_module = types.SimpleNamespace(general_settings=settings, prisma_client=None) return patch.dict( sys.modules, {"litellm.proxy.proxy_server": fake_proxy_server_module}, @@ -2552,7 +2552,7 @@ class TestTemporaryMCPSessionEndpoints: expected_auth = generate_mock_user_api_key_auth( user_role=LitellmUserRoles.PROXY_ADMIN, api_key=api_key_in_cookie ) - fake_proxy_server = types.SimpleNamespace(master_key=master_key) + fake_proxy_server = types.SimpleNamespace(master_key=master_key, prisma_client=None) with ( patch.dict(sys.modules, {"litellm.proxy.proxy_server": fake_proxy_server}), @@ -2629,7 +2629,7 @@ class TestTemporaryMCPSessionEndpoints: mock_manager = MagicMock() mock_manager.get_mcp_server_by_id.return_value = non_oauth_server mock_manager.get_mcp_server_by_name.return_value = None - fake_proxy_server = types.SimpleNamespace(master_key=None) + fake_proxy_server = types.SimpleNamespace(master_key=None, prisma_client=None) with ( patch.dict(sys.modules, {"litellm.proxy.proxy_server": fake_proxy_server}), @@ -2681,7 +2681,7 @@ class TestTemporaryMCPSessionEndpoints: mock_manager = MagicMock() mock_manager.get_mcp_server_by_id.return_value = internal_server mock_manager.get_mcp_server_by_name.return_value = None - fake_proxy_server = types.SimpleNamespace(master_key=None) + fake_proxy_server = types.SimpleNamespace(master_key=None, prisma_client=None) with ( patch.dict(sys.modules, {"litellm.proxy.proxy_server": fake_proxy_server}), diff --git a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py index 83537c236a3..6058f0bbd61 100644 --- a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py +++ b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -1,3 +1,4 @@ +from contextlib import nullcontext import subprocess import sys import textwrap @@ -28,10 +29,11 @@ class _DummyMCPResult: def _setup_mcp_call_environment(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: """Patch MCP globals so _execute_tool_calls can run in tests.""" - proxy_module = types.SimpleNamespace(proxy_logging_obj=object()) + proxy_module = types.SimpleNamespace(proxy_logging_obj=object(), prisma_client=None) monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_module) fake_manager = types.SimpleNamespace( + catalog=types.SimpleNamespace(operation=nullcontext), get_registry=MagicMock(return_value={}), call_tool=AsyncMock(return_value=_DummyMCPResult()), # Newer logging path calls this to enrich spend logs metadata @@ -49,7 +51,7 @@ def _setup_proxy_logging(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: """Patch proxy_logging_obj so failure hook can be asserted.""" proxy_logging_obj = MagicMock() proxy_logging_obj.post_call_failure_hook = AsyncMock() - proxy_module = types.SimpleNamespace(proxy_logging_obj=proxy_logging_obj) + proxy_module = types.SimpleNamespace(proxy_logging_obj=proxy_logging_obj, prisma_client=None) monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_module) return proxy_logging_obj.post_call_failure_hook @@ -378,6 +380,7 @@ async def test_execute_tool_calls_logs_failure_via_post_call_failure_hook(monkey post_call_failure_hook = _setup_proxy_logging(monkeypatch) fake_manager = types.SimpleNamespace( + catalog=types.SimpleNamespace(operation=nullcontext), get_registry=MagicMock(return_value={}), call_tool=AsyncMock(side_effect=HTTPException(status_code=500, detail="boom")) ) @@ -506,6 +509,7 @@ async def test_get_mcp_tools_from_manager_enables_list_tools_logging(monkeypatch # Patch manager methods used by _get_mcp_tools_from_manager to avoid needing full UserAPIKeyAuth fields. fake_manager = types.SimpleNamespace( + catalog=types.SimpleNamespace(operation=nullcontext), get_registry=MagicMock(return_value={}), get_allowed_mcp_servers=AsyncMock(return_value=[]), get_mcp_servers_from_ids=MagicMock(return_value=[]), @@ -559,6 +563,7 @@ async def test_get_mcp_tools_from_manager_forwards_request_tags(monkeypatch): mock_get_tools, ) fake_manager = types.SimpleNamespace( + catalog=types.SimpleNamespace(operation=nullcontext), get_registry=MagicMock(return_value={}), get_allowed_mcp_servers=AsyncMock(return_value=[]), get_mcp_servers_from_ids=MagicMock(return_value=[]), From 44a5992cf8422b5fa8d7128a936142edd99f83f5 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Tue, 22 Sep 2026 12:11:14 -0700 Subject: [PATCH 02/24] fix(mcp): reject configuration changes during scoped dispatch --- .../proxy/_experimental/mcp_server/catalog.py | 28 ++++- .../mcp_server/mcp_server_manager.py | 8 +- .../_experimental/mcp_server/operations.py | 3 + .../mcp_server/test_mcp_server_manager.py | 106 ++++++++++++++++++ 4 files changed, 141 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/catalog.py b/litellm/proxy/_experimental/mcp_server/catalog.py index 673a1260c09..a7dabdb0a0e 100644 --- a/litellm/proxy/_experimental/mcp_server/catalog.py +++ b/litellm/proxy/_experimental/mcp_server/catalog.py @@ -37,7 +37,31 @@ class TargetCatalog: async def refresh(self) -> None: async with self._reload_lock: - await self._manager._reload_servers_from_database() # pyright: ignore[reportPrivateUsage] # existing staged loader + await self._reload() + + async def _reload(self, *, reuse_unchanged: bool = False) -> None: + token: Final = self._scope.set(None) + try: + await self._manager._reload_servers_from_database(reuse_unchanged=reuse_unchanged) # pyright: ignore[reportPrivateUsage] # existing staged loader + finally: + self._scope.reset(token) + + def assert_current(self, server: MCPServer) -> None: + snapshot: Final = self.current + if snapshot is None or server.server_id not in snapshot: + return + expected: Final = snapshot[server.server_id] + registered: Final = self._manager.registry.get(server.server_id) or self._manager.config_mcp_servers.get( + server.server_id + ) + if ( + registered is None + or registered.updated_at != expected.updated_at + or server.updated_at != expected.updated_at + ): + from fastapi import HTTPException # noqa: PLC0415 # optional proxy dependency + + raise HTTPException(status_code=503, detail="MCP server configuration changed; retry the operation") async def list(self) -> Mapping[str, MCPServer]: from litellm.proxy.proxy_server import prisma_client # noqa: PLC0415 # runtime proxy dependency @@ -48,7 +72,7 @@ class TargetCatalog: async with self._reload_lock: if prisma_client is not None: try: - await self._manager._reload_servers_from_database(reuse_unchanged=True) # pyright: ignore[reportPrivateUsage] # existing staged loader + await self._reload(reuse_unchanged=True) except Exception as exc: # noqa: BLE001 # never serve an unverified database snapshot from fastapi import HTTPException # noqa: PLC0415 # optional proxy dependency diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 9806bf1be68..0f31efcf83a 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2181,6 +2181,7 @@ class MCPServerManager: incomplete metadata for a server whose OAuth flow the gateway runs itself. """ + self.catalog.assert_current(server) acquisition: Final = self._get_or_start_oauth_discovery_task(server) if acquisition is None: return self._registered_server(server) @@ -2193,11 +2194,13 @@ class MCPServerManager: raise match outcome: case _OAuthDiscoveryResolved(resolved_server): + self.catalog.assert_current(resolved_server) return resolved_server case _OAuthDiscoveryStale(): return await self._rejoin_oauth_metadata_discovery(server, retry_stale=_retry_stale) case _OAuthDiscoveryFailed(timed_out=timed_out): current: Final = self._registered_server(server) + self.catalog.assert_current(current) if current.is_client_forwarded_token: return current server_ref: Final = current.alias or current.server_name or current.name or current.server_id @@ -5544,6 +5547,7 @@ class MCPServerManager: # Registration used add_server_prefix_to_name(base, get_server_prefix(server)), # and tool_name is the bare base name by the time call_tool reaches here, so # rebuilding the key the same way reproduces it exactly + self.catalog.assert_current(server) registry_key: Final = add_server_prefix_to_name(tool_name, get_server_prefix(server)) tool: Final = global_mcp_tool_registry.get_tool(registry_key) if tool is None: @@ -6599,9 +6603,9 @@ class MCPServerManager: # prefix that lookups will use. if not reuse_unchanged or new_server is not previous_registry.get(server_id): await self._maybe_register_openapi_tools(new_server, initialize_mapping=False) + if new_server.spec_path: + registered_openapi_tools = True registered_registry[server_id] = new_server - if new_server.spec_path: - registered_openapi_tools = True except Exception as e: verbose_logger.exception( "Skipping MCP server %s (%s) during DB reload: %s", diff --git a/litellm/proxy/_experimental/mcp_server/operations.py b/litellm/proxy/_experimental/mcp_server/operations.py index 05a630c38c0..0adbe7f6769 100644 --- a/litellm/proxy/_experimental/mcp_server/operations.py +++ b/litellm/proxy/_experimental/mcp_server/operations.py @@ -2552,6 +2552,9 @@ async def _handle_local_mcp_tool(name: str, arguments: dict[str, object]) -> Cal """ import inspect + server: Final = global_mcp_server_manager.server_owning_tool_name_prefix(name) + if server is not None: + global_mcp_server_manager.catalog.assert_current(server) tool: Final = global_mcp_tool_registry.get_tool(name) if not tool: raise HTTPException(status_code=404, detail=f"Tool '{name}' not found") diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 30297322b0c..b031280f975 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -14722,3 +14722,109 @@ async def test_catalog_reuses_openapi_tools_until_configuration_or_background_re read_rows.return_value = [] await manager.catalog.list() assert registry.list_tools() == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("change", ["edit", "delete"]) +async def test_catalog_rejects_configuration_switch_before_client_creation(monkeypatch, change): + read_rows = AsyncMock(return_value=[_catalog_row()]) + _catalog_database(monkeypatch, read_rows) + manager = MCPServerManager() + async with manager.catalog.operation(): + admitted = manager.get_mcp_server_by_id("catalog-server") + read_rows.return_value = [_catalog_row("updated")] if change == "edit" else [] + await asyncio.create_task(manager.reload_servers_from_database()) + with patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPClient") as factory: + with pytest.raises(HTTPException) as exc: + await manager._create_mcp_client(admitted) + assert exc.value.status_code == 503 + assert exc.value.detail == "MCP server configuration changed; retry the operation" + factory.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("dispatch", ["managed", "local"]) +async def test_catalog_changed_openapi_handler_never_dispatches(monkeypatch, tmp_path, respx_mock, dispatch): + from litellm.proxy._experimental.mcp_server import operations, tool_registry + from litellm.proxy._experimental.mcp_server.utils import add_server_prefix_to_name, get_server_prefix + + registry = tool_registry.MCPToolRegistry() + monkeypatch.setattr(tool_registry, "global_mcp_tool_registry", registry) + monkeypatch.setattr(operations, "global_mcp_tool_registry", registry) + spec_path = tmp_path / "catalog-race.json" + spec_path.write_text(json.dumps({ + "openapi": "3.0.0", "info": {"title": "Catalog", "version": "1"}, + "paths": {"/echo": {"get": {"operationId": "echo"}}}, + })) + row = _catalog_row().model_copy(update={"spec_path": str(spec_path)}) + read_rows = AsyncMock(return_value=[row]) + _catalog_database(monkeypatch, read_rows) + manager = MCPServerManager() + monkeypatch.setattr(operations, "global_mcp_server_manager", manager) + destination = respx_mock.get("https://changed.example/echo").respond(200, text="must not execute") + async with manager.catalog.operation(): + admitted = manager.get_mcp_server_by_id("catalog-server") + read_rows.return_value = [row.model_copy(update={ + "url": "https://changed.example", "updated_at": datetime(2026, 1, 2), + })] + await asyncio.create_task(manager.reload_servers_from_database()) + with pytest.raises(HTTPException) as exc: + if dispatch == "managed": + await manager._call_openapi_tool_handler(admitted, "echo", {}) + else: + await operations._handle_local_mcp_tool( + add_server_prefix_to_name("echo", get_server_prefix(admitted)), {} + ) + assert exc.value.status_code == 503 + assert destination.call_count == 0 + + +@pytest.mark.asyncio +async def test_catalog_rejects_old_admission_in_a_new_operation(monkeypatch): + read_rows = AsyncMock(return_value=[_catalog_row()]) + _catalog_database(monkeypatch, read_rows) + manager = MCPServerManager() + original = (await manager.catalog.list())["catalog-server"] + read_rows.return_value = [_catalog_row("updated")] + async with manager.catalog.operation(): + with pytest.raises(HTTPException) as exc: + await manager._create_mcp_client(original) + assert exc.value.detail == "MCP server configuration changed; retry the operation" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("change", [False, True]) +async def test_catalog_deferred_discovery_preserves_admitted_configuration(monkeypatch, change): + row = _catalog_row().model_copy(update={"auth_type": MCPAuth.true_passthrough}) + read_rows = AsyncMock(return_value=[row]) + _catalog_database(monkeypatch, read_rows) + manager = MCPServerManager() + entered = asyncio.Event() + release = asyncio.Event() + metadata = MCPOAuthMetadata( + authorization_url="https://issuer.example/authorize", token_url="https://issuer.example/token", + ) + + async def discover(server): + entered.set() + await release.wait() + return metadata + + with patch.object(manager, "_discover_oauth_metadata_for_server", side_effect=discover): + async with manager.catalog.operation(): + admitted = manager.get_mcp_server_by_id("catalog-server") + resolving = asyncio.create_task(manager.ensure_oauth_metadata_discovered(admitted)) + await asyncio.wait_for(entered.wait(), timeout=1) + if change: + read_rows.return_value = [row.model_copy(update={"updated_at": datetime(2026, 1, 2)})] + await manager.reload_servers_from_database() + release.set() + if change: + with pytest.raises(HTTPException) as exc: + await asyncio.wait_for(resolving, timeout=1) + assert exc.value.detail == "MCP server configuration changed; retry the operation" + else: + resolved = await asyncio.wait_for(resolving, timeout=1) + assert resolved.authorization_url == metadata.authorization_url + assert resolved.token_url == metadata.token_url + assert resolved.updated_at == admitted.updated_at From 74a470410d49633fc5850fa7e7b911f022d68cd8 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Tue, 22 Sep 2026 12:40:45 -0700 Subject: [PATCH 03/24] fix(mcp): refresh shared catalog state for each operation --- .../mcp_server/auth/user_api_key_auth_mcp.py | 2 + .../mcp_server/byok_oauth_endpoints.py | 2 + .../proxy/_experimental/mcp_server/catalog.py | 394 +++++++++++++++ litellm/proxy/_experimental/mcp_server/db.py | 9 + .../mcp_server/discoverable_endpoints.py | 12 +- .../mcp_server/gateway_dcr_flow.py | 2 + .../mcp_server/mcp_server_manager.py | 275 +++-------- .../_experimental/mcp_server/operations.py | 28 +- .../mcp_server/rest_endpoints.py | 3 + .../mcp_server/semantic_tool_filter.py | 2 + .../proxy/_experimental/mcp_server/server.py | 3 + .../_experimental/mcp_server/tool_registry.py | 30 +- .../proxy/_experimental/mcp_server/utils.py | 6 +- .../mcp_security/mcp_security_guardrail.py | 11 +- .../mcp_management_endpoints.py | 268 ++++++----- litellm/proxy/proxy_server.py | 51 +- .../mcp/litellm_proxy_mcp_handler.py | 4 + .../types/mcp_server/mcp_server_manager.py | 2 +- .../auth/test_user_api_key_auth_mcp.py | 2 + .../mcp_server/test_byok_oauth_endpoints.py | 1 + .../mcp_server/test_discoverable_endpoints.py | 10 +- .../mcp_server/test_mcp_server.py | 8 +- .../mcp_server/test_mcp_server_manager.py | 449 +++++++++++++++--- .../mcp_server/test_rest_endpoints.py | 5 +- .../mcp_server/test_short_mcp_tool_prefix.py | 12 +- .../guardrail_hooks/test_mcp_security.py | 27 ++ .../test_mcp_management_endpoints.py | 154 ++++-- .../mcp/test_litellm_proxy_mcp_handler.py | 6 + 28 files changed, 1288 insertions(+), 490 deletions(-) create mode 100644 litellm/proxy/_experimental/mcp_server/catalog.py diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 60dc91a69cc..6b8006dbfc9 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -13,6 +13,7 @@ from typing_extensions import assert_never import litellm from litellm._logging import verbose_logger +from litellm.proxy._experimental.mcp_server.catalog import catalog_operation, global_manager from litellm.proxy._experimental.mcp_server.oauth_utils import ( get_passthrough_resource_metadata_url, get_passthrough_www_authenticate, @@ -398,6 +399,7 @@ class MCPRequestHandler: LITELLM_MCP_ACCESS_GROUPS_HEADER_NAME = SpecialHeaders.mcp_access_groups.value @staticmethod + @catalog_operation(global_manager) async def process_mcp_request( scope: Scope, ) -> tuple[ diff --git a/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py b/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py index 2c63e0a96d8..086dd946efe 100644 --- a/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py @@ -25,6 +25,7 @@ from fastapi import APIRouter, Depends, Form, HTTPException, Request from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse from litellm._logging import verbose_proxy_logger +from litellm.proxy._experimental.mcp_server.catalog import catalog_operation, global_manager from litellm.proxy._experimental.mcp_server.db import store_user_credential from litellm.proxy._experimental.mcp_server.oauth_utils import ( BYOK_RESOURCE_METADATA_PATH, @@ -649,6 +650,7 @@ async def byok_protected_resource_metadata(request: Request) -> JSONResponse: @router.get("/v1/mcp/oauth/authorize", include_in_schema=False) +@catalog_operation(global_manager) async def byok_authorize_get( request: Request, client_id: str | None = None, diff --git a/litellm/proxy/_experimental/mcp_server/catalog.py b/litellm/proxy/_experimental/mcp_server/catalog.py new file mode 100644 index 00000000000..0f5eca285e3 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/catalog.py @@ -0,0 +1,394 @@ +"""Authoritative MCP catalog snapshots shared by legacy lookup adapters.""" + +from __future__ import annotations + +import asyncio +import hashlib +import json +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequence +from contextlib import asynccontextmanager +from contextvars import ContextVar +from dataclasses import dataclass +from functools import wraps +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, ParamSpec, TypeVar + +from litellm._logging import verbose_logger + +if TYPE_CHECKING: + from mcp.types import Tool as SDKTool + from pydantic import BaseModel + + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import AggregateToolListing, ServerOutcome + from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager + from litellm.types.mcp_server.mcp_server_manager import MCPServer + from litellm.types.mcp_server.tool_registry import MCPTool + +_P = ParamSpec("_P") +_R = TypeVar("_R") + + +@dataclass(frozen=True, slots=True) +class CatalogSnapshot: + servers: Mapping[str, MCPServer] + identity: str + tools: Mapping[str, MCPTool] + routing: dict[str, str] + + +def _configuration_identity(server: MCPServer) -> str: + return json.dumps( + server.model_dump( + mode="json", + exclude=frozenset(("short_prefix", "scopes", "authorization_url", "token_url", "registration_url")) + | (frozenset() if server.issuer_is_anchored else frozenset(("issuer",))), + ), + sort_keys=True, + ) + + +def _check_oauth_revision(selected: MCPServer, candidate: MCPServer | None) -> None: + from fastapi import HTTPException + + if candidate is None or _configuration_identity(selected) != _configuration_identity(candidate): + raise HTTPException(status_code=503, detail="OAuth metadata discovery changed repeatedly; retry shortly") + + +def _snapshot(manager: MCPServerManager, database_identity: str) -> CatalogSnapshot: + from litellm.proxy._experimental.mcp_server.tool_registry import global_mcp_tool_registry + + servers: Final = manager.config_mcp_servers | manager.registry + detached: Final = MappingProxyType({key: value.model_copy(deep=True) for key, value in servers.items()}) + serialized: Final = json.dumps( + ( + database_identity, + tuple(sorted(manager.registry)), + tuple((key, _configuration_identity(value)) for key, value in sorted(manager.config_mcp_servers.items())), + ), + sort_keys=True, + ) + return CatalogSnapshot( + detached, + hashlib.sha256(serialized.encode()).hexdigest(), + MappingProxyType(dict(global_mcp_tool_registry.published_tools)), + dict(manager.published_tool_routes), + ) + + +class TargetCatalog: + def __init__(self, manager: MCPServerManager) -> None: + self.manager = manager + self._refresh_lock = asyncio.Lock() + self._database_identity = "" + self._warned_shadowed_config_server_ids: frozenset[str] = frozenset() + self._warned_capturing_config_server_ids: frozenset[str] = frozenset() + self._operation: ContextVar[tuple[CatalogSnapshot, asyncio.Event] | None] = ContextVar( + "mcp_catalog_snapshot", default=None + ) + self._staged_routing: ContextVar[tuple[dict[str, str], asyncio.Event] | None] = ContextVar( + "mcp_catalog_routing", default=None + ) + + def current(self) -> CatalogSnapshot | None: + scoped: Final = self._operation.get() + return scoped[0] if scoped is not None and not scoped[1].is_set() else None + + def routing(self) -> dict[str, str]: + staged: Final = self._staged_routing.get() + if staged is not None and not staged[1].is_set(): + return staged[0] + snapshot: Final = self.current() + return snapshot.routing if snapshot is not None else self.manager.published_tool_routes + + def registry(self) -> dict[str, MCPServer]: + snapshot: Final = self.current() + return ( + dict(snapshot.servers) if snapshot is not None else self.manager.config_mcp_servers | self.manager.registry + ) + + @asynccontextmanager + async def operation(self) -> AsyncIterator[CatalogSnapshot]: + current: Final = self.current() + if current is not None: + yield current + return + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is not None: + from litellm.proxy.proxy_server import should_load_db_object + + if should_load_db_object("mcp"): + await self.reload() + from litellm.proxy._experimental.mcp_server.tool_registry import global_mcp_tool_registry + + snapshot: Final = _snapshot(self.manager, self._database_identity) + closed: Final = asyncio.Event() + token: Final = self._operation.set((snapshot, closed)) + try: + with global_mcp_tool_registry.catalog_scope(snapshot.tools): + yield snapshot + finally: + closed.set() + self._operation.reset(token) + self._retain_discovered_routing(snapshot) + + def _retain_discovered_routing(self, snapshot: CatalogSnapshot) -> None: + from litellm.proxy._experimental.mcp_server.utils import normalize_server_name + + current: Final = self.manager.config_mcp_servers | self.manager.registry + unchanged_owners: Final = frozenset( + owner + for key, server in snapshot.servers.items() + if current.get(key) == server + for owner in self.manager.owned_mapping_values(server) + ) + self.manager.published_tool_routes = self.manager.published_tool_routes | MappingProxyType( + { + name: owner + for name, owner in snapshot.routing.items() + if normalize_server_name(owner) in unchanged_owners + } + ) + + async def resolve(self, identifier: str, client_ip: str | None = None) -> MCPServer | None: + async with self.operation(): + return self.manager.get_mcp_server_by_id(identifier, client_ip) or self.manager.get_mcp_server_by_name( + identifier, client_ip + ) + + async def resolve_oauth_metadata( + self, + server: MCPServer, + resolve: Callable[[MCPServer], Awaitable[MCPServer]], + ) -> MCPServer: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import oauth_endpoints_unresolved + + snapshot: Final = self.current() + selected: Final = snapshot.servers.get(server.server_id) if snapshot is not None else None + if selected is None: + return await resolve(server) + if not oauth_endpoints_unresolved(selected): + return selected + registered: Final = self.manager.registry.get(server.server_id) or self.manager.config_mcp_servers.get( + server.server_id + ) + _check_oauth_revision(selected, registered) + resolved: Final = await resolve(selected) + _check_oauth_revision(selected, resolved) + return resolved + + @staticmethod + async def list( + servers: Sequence[MCPServer], + fetch: Callable[[MCPServer], Awaitable[tuple[list[SDKTool], ServerOutcome]]], + server_key: Callable[[MCPServer], str], + ) -> AggregateToolListing: + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import AggregateToolListing + + tasks: Final = tuple(asyncio.ensure_future(fetch(server)) for server in servers) + try: + results: Final = await asyncio.gather(*tasks) + return AggregateToolListing( + tools=[tool for tools, _ in results for tool in tools], + outcomes={server_key(server): outcome for server, (_, outcome) in zip(servers, results)}, + ) + finally: + for task in tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + async def reload(self) -> None: + from litellm.proxy._experimental.mcp_server.tool_registry import global_mcp_tool_registry + + async with self._refresh_lock: + staged_config: Final = { + key: value.model_copy(deep=True) for key, value in self.manager.config_mcp_servers.items() + } + await self.manager.hydrate_config_servers_dcr_clients(tuple(staged_config.values())) + staged_routing: Final = dict(self.manager.published_tool_routes) + closed: Final = asyncio.Event() + routing_token: Final = self._staged_routing.set((staged_routing, closed)) + try: + with global_mcp_tool_registry.catalog_scope(global_mcp_tool_registry.published_tools) as staged_tools: + await self._reload() + self.manager.config_mcp_servers = staged_config + global_mcp_tool_registry.tools = staged_tools + self.manager.published_tool_routes = staged_routing + finally: + closed.set() + self._staged_routing.reset(routing_token) + + async def _reload(self) -> None: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + carry_forward_resolved_oauth_endpoints, + config_ids_capturing_db_identifiers, + oauth_endpoints_unresolved, + warn_on_server_name_fields, + ) + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + get_prisma_client_or_throw, + ) + + verbose_logger.debug("Loading MCP servers from database into registry...") + + prisma_client: Final = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") + # Load only "active", legacy "approved", and NULL (no approval workflow) rows. + # Pending/rejected servers are excluded at the DB level so we never load them. + from litellm.proxy._experimental.mcp_server.db import LiteLLM_MCPServerTable, get_runtime_mcp_server_rows + + raw_rows: Final[Sequence[BaseModel]] = await get_runtime_mcp_server_rows(prisma_client) + database_identity: Final = hashlib.sha256( + json.dumps( + tuple(sorted(json.dumps(row.model_dump(mode="json"), sort_keys=True, default=str) for row in raw_rows)) + ).encode() + ).hexdigest() + verbose_logger.info("Found %s MCP servers in database", len(raw_rows)) + + previous_registry: Final = self.manager.registry + new_registry: Final[dict[str, MCPServer]] = {} + + # Stage one: build every server. Stage two assigns short prefixes + # against the *full* set so dedup is deterministic regardless of + # iteration order. + for row in raw_rows: + try: + server = LiteLLM_MCPServerTable.model_validate(row.model_dump()) + existing_server = previous_registry.get(server.server_id) + + if ( + existing_server is not None + and existing_server.updated_at is not None + and server.updated_at is not None + and existing_server.updated_at == server.updated_at + and ( + self.manager.oauth_discovery_slot(server.server_id) is not None + or not oauth_endpoints_unresolved(existing_server) + ) + ): + # Re-use existing server instance to avoid re-running build_mcp_server_from_table() + # which can perform network discovery for OAuth2 servers. + new_registry[server.server_id] = existing_server + continue + + warn_on_server_name_fields( + server_id=server.server_id, + alias=getattr(server, "alias", None), + server_name=getattr(server, "server_name", None), + ) + verbose_logger.debug("Building server from DB: %s (%s)", server.server_id, server.server_name) + # raw_rows come straight from the DB, so their global env var + # values (like credentials) are still encrypted here, unlike the + # already-decrypted records add_server/update_server are handed. + # Decrypt them while building the registry entry. + new_server = await self.manager.build_mcp_server_from_table( + server, env_vars_are_encrypted=True, register_oauth_discovery=False + ) + # Carry the cached short_prefix from the previous registry entry + # (if any) so the prefix is stable across reloads. + if existing_server is not None and existing_server.short_prefix: + new_server.short_prefix = existing_server.short_prefix + carry_forward_resolved_oauth_endpoints(new_server=new_server, previous_server=existing_server) + new_registry[server.server_id] = new_server + except Exception as e: + verbose_logger.exception( + "Skipping MCP server %s (%s) during DB reload: %s", + getattr(row, "server_id", None), + getattr(row, "alias", None), + e, + ) + + # Assign short prefixes against the full candidate set without + # publishing the staged registry to concurrent callers. + registered_registry: Final[dict[str, MCPServer]] = {} + for server_id, new_server in new_registry.items(): + try: + if new_server is not previous_registry.get(server_id): + self.manager.assign_unique_short_prefix(new_server, registry=new_registry) + # Register OpenAPI tools *after* the final short prefix is assigned + # so the tools are stored in the global registry under the same + # prefix that lookups will use. + if new_server is not previous_registry.get(server_id): + if previous_server := previous_registry.get(server_id): + self.manager.remove_server_tool_routing(previous_server) + await self.manager.maybe_register_openapi_tools(new_server, initialize_mapping=False) + registered_registry[server_id] = new_server + except Exception as e: + self.manager.remove_server_tool_routing(new_server) + verbose_logger.exception( + "Skipping MCP server %s (%s) during DB reload: %s", + new_server.server_id, + getattr(new_server, "alias", None), + e, + ) + + dropped_registry_keys: Final = previous_registry.keys() - registered_registry.keys() + for registry_key in dropped_registry_keys: + self.manager.remove_server_tool_routing(previous_registry[registry_key]) + self.manager.invalidate_oauth_discovery_state(previous_registry[registry_key].server_id) + + for server_id in previous_registry.keys() | registered_registry.keys(): + if previous_registry.get(server_id) != registered_registry.get(server_id): + self.manager.invalidate_discovery_lists(server_id) + self.manager.invalidate_oauth_discovery_state(server_id) + self._database_identity = database_identity + self.manager.registry = registered_registry + # A discovery task may have published into ``previous_registry`` while + # this replacement was being staged. Reconcile every published entry + # synchronously after the swap so a lost publication cannot also leave + # the replacement unresolved with no retry slot. + registered_servers: Final = tuple(registered_registry.values()) + self.manager.reconcile_oauth_discovery_slots_for_servers(registered_servers) + self.manager.prime_oauth_metadata_discovery_for_servers(registered_servers) + + verbose_logger.debug("MCP registry refreshed (%s servers in registry)", len(registered_registry)) + + # get_registry() is ``config_mcp_servers | registry``, so a database row sharing an id with a + # config.yaml server hides that server everywhere. Only reachable once an operator pins + # ``server_id`` in config.yaml; say so rather than letting the server disappear silently. + shadowed_config_server_ids: Final = frozenset( + self.manager.config_mcp_servers.keys() & registered_registry.keys() + ) + if shadowed_config_server_ids and shadowed_config_server_ids != self._warned_shadowed_config_server_ids: + verbose_logger.warning( + "config.yaml MCP server_id(s) %s are also database-backed MCP servers. The database " + "entry takes precedence, so the config.yaml server is unreachable. Give the config " + "entry a different server_id.", + ", ".join(sorted(shadowed_config_server_ids)), + ) + self._warned_shadowed_config_server_ids = shadowed_config_server_ids + + # The mirror image of the block above: a config server_id that is a database server's name + # answers that server's grants instead, because ids are matched before names. + capturing_config_server_ids: Final = config_ids_capturing_db_identifiers( + self.manager.config_mcp_servers.keys(), registered_registry.values() + ) + if capturing_config_server_ids and capturing_config_server_ids != self._warned_capturing_config_server_ids: + verbose_logger.warning( + "config.yaml MCP server_id(s) %s are the name or alias of a database-backed MCP " + "server. Permission entries naming them resolve to the config.yaml server, not the " + "database one. Give the config entry a different server_id.", + ", ".join(sorted(capturing_config_server_ids)), + ) + self._warned_capturing_config_server_ids = capturing_config_server_ids + + +def catalog_operation( + manager: Callable[[], MCPServerManager], +) -> Callable[[Callable[_P, Awaitable[_R]]], Callable[_P, Awaitable[_R]]]: + def decorate(function: Callable[_P, Awaitable[_R]]) -> Callable[_P, Awaitable[_R]]: + @wraps(function) + async def wrapped(*args: _P.args, **kwargs: _P.kwargs) -> _R: # kwargs-ok: preserves ParamSpec + + async with manager().catalog.operation(): + return await function(*args, **kwargs) + + return wrapped + + return decorate + + +def global_manager() -> MCPServerManager: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + return global_mcp_server_manager diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 30ee8b7a4fc..ee7298fe146 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -645,6 +645,15 @@ async def get_all_mcp_servers( return list(_readable_mcp_servers(mcp_servers)) +async def get_runtime_mcp_server_rows( + prisma_client: PrismaClient, +) -> Sequence["prisma_db_models.LiteLLM_MCPServerTable"]: + where: Final[prisma_db_types.LiteLLM_MCPServerTableWhereInput] = { + "OR": [{"approval_status": None}, {"approval_status": {"in": ["active", "approved"]}}] + } + return await _db_find_mcp_server_rows(prisma_client, where) + + async def get_mcp_server(prisma_client: PrismaClient, server_id: str) -> LiteLLM_MCPServerTable | None: """ Returns the matching mcp server from the db iff exists diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 64bab0a7832..3378b1df521 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -36,6 +36,7 @@ from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( can_store_oauth_credential, oauth_authorization_uses_gateway_credential, ) +from litellm.proxy._experimental.mcp_server.catalog import catalog_operation, global_manager from litellm.proxy._experimental.mcp_server.faults import ( CallerRejected, CredentialSource, @@ -1769,7 +1770,7 @@ async def resolve_ephemeral_dcr_client( def _register_flow_needed_endpoint(mcp_server: MCPServer) -> str | None: """The register flow's deferred-discovery join gate. A DCR bridge with no admin-configured client can only register callers through the upstream's registration endpoint - (``_oauth_endpoints_unresolved`` keeps its discovery slot armed for exactly this shape), so + (``oauth_endpoints_unresolved`` keeps its discovery slot armed for exactly this shape), so the flow must keep joining discovery while registration is still missing instead of silently degrading to the dummy short-circuit. Every other shape only needs the authorization url.""" if mcp_server.is_dcr_bridge and not mcp_server.client_id and mcp_server.effective_registration_url is None: @@ -1900,6 +1901,7 @@ async def authorize_mcp_session( @router.get("/{mcp_server_name}/authorize") @router.get("/authorize") +@catalog_operation(global_manager) async def authorize( request: Request, redirect_uri: str, @@ -1974,6 +1976,7 @@ async def authorize( @router.post("/{mcp_server_name}/token") @router.post("/token") +@catalog_operation(global_manager) async def token_endpoint( request: Request, grant_type: str = Form(...), @@ -2078,6 +2081,7 @@ async def authorize_flow(request: Request, flow: str) -> Response: @router.post("/authorize/complete") +@catalog_operation(global_manager) async def authorize_complete( request: Request, flow: str = Form(...), @@ -2696,6 +2700,7 @@ async def oauth_authorization_server_aggregate(request: Request): # Standard MCP pattern: /.well-known/oauth-protected-resource/mcp/{server_name} # This is the pattern expected by standard MCP clients (mcp-inspector, VSCode Copilot) @router.get(f"/.well-known/oauth-protected-resource{well_known_root_suffix()}/mcp/{{mcp_server_name}}") +@catalog_operation(global_manager) async def oauth_protected_resource_mcp_standard(request: Request, mcp_server_name: str): """ OAuth protected resource discovery endpoint using standard MCP URL pattern. @@ -2716,6 +2721,7 @@ async def oauth_protected_resource_mcp_standard(request: Request, mcp_server_nam # LiteLLM legacy pattern: /.well-known/oauth-protected-resource/{server_name}/mcp # Kept for backward compatibility with existing deployments @router.get(f"/.well-known/oauth-protected-resource{well_known_root_suffix()}/{{mcp_server_name}}/mcp") +@catalog_operation(global_manager) async def oauth_protected_resource_mcp(request: Request, mcp_server_name: str | None = None): """ OAuth protected resource discovery endpoint using LiteLLM legacy URL pattern. @@ -2792,6 +2798,7 @@ def _build_oauth_authorization_server_response( # Standard MCP pattern: /.well-known/oauth-authorization-server/mcp/{server_name} @router.get(f"/.well-known/oauth-authorization-server{well_known_root_suffix()}/mcp/{{mcp_server_name}}") +@catalog_operation(global_manager) async def oauth_authorization_server_mcp_standard(request: Request, mcp_server_name: str): """ OAuth authorization server discovery endpoint using standard MCP URL pattern. @@ -2809,6 +2816,7 @@ async def oauth_authorization_server_mcp_standard(request: Request, mcp_server_n # LiteLLM legacy pattern and root endpoint @router.get(f"/.well-known/oauth-authorization-server{well_known_root_suffix()}/{{mcp_server_name}}") @router.get("/.well-known/oauth-authorization-server") +@catalog_operation(global_manager) async def oauth_authorization_server_mcp(request: Request, mcp_server_name: str | None = None): """ OAuth authorization server discovery endpoint. @@ -2882,6 +2890,7 @@ async def jwks_json(request: Request): # Additional legacy pattern support @router.get(f"/.well-known/oauth-authorization-server{well_known_root_suffix()}/{{mcp_server_name}}/mcp") +@catalog_operation(global_manager) async def oauth_authorization_server_legacy(request: Request, mcp_server_name: str): """ OAuth authorization server discovery for legacy /{server_name}/mcp pattern. @@ -2895,6 +2904,7 @@ async def oauth_authorization_server_legacy(request: Request, mcp_server_name: s @router.post("/{mcp_server_name}/register") @router.post("/register") +@catalog_operation(global_manager) async def register_client(request: Request, mcp_server_name: str | None = None): # Get the correct base URL considering X-Forwarded-* headers request_base_url: Final = get_request_base_url(request) diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index e66504af47a..624b1795a14 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -55,6 +55,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict, assert_never from litellm._logging import verbose_logger from litellm.caching.caching import DualCache +from litellm.proxy._experimental.mcp_server.catalog import catalog_operation, global_manager from litellm.proxy._experimental.mcp_server.oauth_utils import ( TOKEN_NO_CACHE_HEADERS, canonical_resource_uri, @@ -770,6 +771,7 @@ def _open_flow_for( return flow +@catalog_operation(global_manager) async def _flow_target( flow: _ConnectFlow, lookup_server_reachability: LookupServerReachability ) -> tuple[Literal["unscoped", "interactive", "m2m", "stale"], MCPServer | None]: diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 20c114a2f3e..93899ea3a78 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -178,7 +178,6 @@ from litellm.proxy.middleware.per_request_root_path_middleware import ( get_request_root_path, ) from litellm.proxy.utils import PrismaClient, ProxyLogging -from litellm.repositories.table_repositories import MCPServerRepository from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.mcp import ( DEFAULT_SUBJECT_TOKEN_TYPE, @@ -287,7 +286,7 @@ def _requires_oauth_discovery( use_issuer_anchor: bool, server: MCPServer, ) -> bool: - return _has_oauth_discovery_source(server_url, use_issuer_anchor) and _oauth_endpoints_unresolved(server) + return _has_oauth_discovery_source(server_url, use_issuer_anchor) and oauth_endpoints_unresolved(server) _StringList: TypeAlias = list[str] @@ -541,7 +540,7 @@ def _config_identifier_owners( ) -def _config_ids_capturing_db_identifiers( +def config_ids_capturing_db_identifiers( config_server_ids: Container[str], db_servers: Iterable[MCPServer], ) -> frozenset[str]: @@ -722,7 +721,7 @@ def _flow_endpoints_missing( return authorization_url is None or token_url is None -def _oauth_endpoints_unresolved(server: MCPServer) -> bool: +def oauth_endpoints_unresolved(server: MCPServer) -> bool: """``_flow_endpoints_missing`` over a built registry entry, for the reload fast-path check. The flow comes from ``effective_oauth2_flow``, the one column-first, shape-fallback judge every @@ -781,7 +780,7 @@ def _endpoints_corroborate_authorization_url( ) == _normalized_authorize_endpoint(trusted_authorization_url) -def _carry_forward_resolved_oauth_endpoints(new_server: MCPServer, previous_server: MCPServer | None) -> None: +def carry_forward_resolved_oauth_endpoints(new_server: MCPServer, previous_server: MCPServer | None) -> None: """Keep the last known good OAuth endpoints when a rebuild's re-discovery comes back empty. A rebuild wholesale-replaces the registry entry, so without this a transient upstream outage @@ -1427,7 +1426,7 @@ def _obo_retry_applies(server: MCPServer, subject_token: str | None) -> bool: return server.auth_type == MCPAuth.oauth2_token_exchange and bool(subject_token) -def _warn_on_server_name_fields( +def warn_on_server_name_fields( *, server_id: str, alias: str | None, @@ -1860,6 +1859,9 @@ class MCPServerManager: self._template_discovery_cache = _DiscoveryCache[ResourceTemplate]( discovery_ttl, discovery_clock, TypeAdapter(tuple[ResourceTemplate, ...]) ) + from litellm.proxy._experimental.mcp_server.catalog import TargetCatalog + + self.catalog = TargetCatalog(self) self.registry: dict[str, MCPServer] = {} self._openapi_health_probes: Callable[[str], _OpenAPIHealthProbe] = lru_cache(maxsize=128)(_OpenAPIHealthProbe) self.config_mcp_servers: dict[str, MCPServer] = {} @@ -1886,7 +1888,7 @@ class MCPServerManager: # semaphore so an edited limit rebuilds it instead of keeping the old cap # until restart. self._server_call_semaphores: dict[str, tuple[int, asyncio.Semaphore]] = {} - self.tool_name_to_mcp_server_name_mapping: dict[str, str] = {} + self.published_tool_routes: dict[str, str] = {} """ { "gmail_send_email": "zapier_mcp_server", @@ -1900,13 +1902,11 @@ class MCPServerManager: # Last set of config server ids found shadowed by database rows. reload_servers_from_database # runs on the config-reload timer, so this keeps a standing misconfiguration from re-logging # the same warning every interval; a change in the set logs again. - self._warned_shadowed_config_server_ids: frozenset[str] = frozenset() - self._warned_capturing_config_server_ids: frozenset[str] = frozenset() self._oauth_discovery_on_startup = _mcp_oauth_discovery_on_startup_enabled() self._oauth_discovery_generation_counter = 0 self._oauth_discovery_slots: tuple[_OAuthDiscoverySlot, ...] = () - def _oauth_discovery_slot(self, server_id: str) -> _OAuthDiscoverySlot | None: + def oauth_discovery_slot(self, server_id: str) -> _OAuthDiscoverySlot | None: return next((slot for slot in self._oauth_discovery_slots if slot.server_id == server_id), None) def _remove_oauth_discovery_slot(self, server_id: str) -> None: @@ -1919,7 +1919,7 @@ class MCPServerManager: ) def _set_oauth_discovery_deferred(self, server_id: str, discovery_deferred: bool) -> None: - previous: Final = self._oauth_discovery_slot(server_id) + previous: Final = self.oauth_discovery_slot(server_id) self._remove_oauth_discovery_slot(server_id) if previous is not None and previous.task is not None and not previous.task.done(): previous.task.cancel() @@ -1932,8 +1932,8 @@ class MCPServerManager: ) ) - def _invalidate_oauth_discovery_state(self, server_id: str) -> None: - previous: Final = self._oauth_discovery_slot(server_id) + def invalidate_oauth_discovery_state(self, server_id: str) -> None: + previous: Final = self.oauth_discovery_slot(server_id) self._remove_oauth_discovery_slot(server_id) if previous is not None and previous.task is not None and not previous.task.done(): previous.task.cancel() @@ -2008,7 +2008,7 @@ class MCPServerManager: return resolved def _oauth_discovery_slot_is_current(self, server_id: str, generation: int) -> bool: - slot: Final = self._oauth_discovery_slot(server_id) + slot: Final = self.oauth_discovery_slot(server_id) return slot is not None and slot.generation == generation def _expire_temporary_oauth_discovery(self, server_id: str, generation: int) -> None: @@ -2045,7 +2045,7 @@ class MCPServerManager: if not self._oauth_discovery_slot_is_current(server.server_id, generation): return _OAuthDiscoveryStale(server_id=server.server_id) current: Final = self._registered_server(server) - if not _oauth_endpoints_unresolved(current): + if not oauth_endpoints_unresolved(current): published: Final = self._publish_resolved_oauth_server(current, generation) return ( _OAuthDiscoveryResolved(server=published) @@ -2056,7 +2056,7 @@ class MCPServerManager: if not self._oauth_discovery_slot_is_current(server.server_id, generation): return _OAuthDiscoveryStale(server_id=server.server_id) candidate: Final = self._merge_discovered_oauth_metadata(self._registered_server(server), metadata) - if _oauth_endpoints_unresolved(candidate): + if oauth_endpoints_unresolved(candidate): return None published_candidate: Final = self._publish_resolved_oauth_server(candidate, generation) return ( @@ -2103,7 +2103,7 @@ class MCPServerManager: return outcome def _record_oauth_discovery_failure(self, server_id: str, generation: int) -> None: - slot: Final = self._oauth_discovery_slot(server_id) + slot: Final = self.oauth_discovery_slot(server_id) if slot is None or slot.generation != generation: return consecutive_failures: Final = slot.consecutive_failures + 1 @@ -2119,7 +2119,7 @@ class MCPServerManager: self, server: MCPServer, ) -> tuple[asyncio.Task[_OAuthDiscoveryOutcome], int] | None: - slot: Final = self._oauth_discovery_slot(server.server_id) + slot: Final = self.oauth_discovery_slot(server.server_id) if slot is None: return None if slot.task is not None: @@ -2148,19 +2148,24 @@ class MCPServerManager: """ self._get_or_start_oauth_discovery_task(server) - def _prime_oauth_metadata_discovery_for_servers(self, servers: Sequence[MCPServer]) -> None: + def prime_oauth_metadata_discovery_for_servers(self, servers: Sequence[MCPServer]) -> None: for server in servers: self.prime_oauth_metadata_discovery(server) - def _reconcile_oauth_discovery_slots_for_servers(self, servers: Sequence[MCPServer]) -> None: + def reconcile_oauth_discovery_slots_for_servers(self, servers: Sequence[MCPServer]) -> None: """Align retry slots after an atomic registry replacement.""" for server in servers: should_defer = _requires_oauth_discovery(server.url, server.issuer_is_anchored, server) - has_slot = self._oauth_discovery_slot(server.server_id) is not None + has_slot = self.oauth_discovery_slot(server.server_id) is not None if should_defer != has_slot: self._set_oauth_discovery_deferred(server.server_id, should_defer) async def ensure_oauth_metadata_discovered(self, server: MCPServer, *, _retry_stale: bool = True) -> MCPServer: + return await self.catalog.resolve_oauth_metadata( + server, lambda selected: self._ensure_oauth_metadata_discovered(selected, _retry_stale=_retry_stale) + ) + + async def _ensure_oauth_metadata_discovered(self, server: MCPServer, *, _retry_stale: bool = True) -> MCPServer: """Join the bounded discovery task and return the resolved server. Concurrent callers share one task per server. A failed attempt remains @@ -2209,7 +2214,7 @@ class MCPServerManager: if retry_stale: return await self.ensure_oauth_metadata_discovered(server, _retry_stale=False) current: Final = self._registered_server(server) - if not _oauth_endpoints_unresolved(current) or current.is_client_forwarded_token: + if not oauth_endpoints_unresolved(current) or current.is_client_forwarded_token: return current raise HTTPException(status_code=503, detail="OAuth metadata discovery changed repeatedly; retry shortly") @@ -2291,7 +2296,15 @@ class MCPServerManager: """ Get the registered MCP Servers from the registry and union with the config MCP Servers """ - return self.config_mcp_servers | self.registry + return self.catalog.registry() + + @property + def tool_name_to_mcp_server_name_mapping(self) -> dict[str, str]: + return self.catalog.routing() + + @tool_name_to_mcp_server_name_mapping.setter + def tool_name_to_mcp_server_name_mapping(self, mapping: dict[str, str]) -> None: + self.published_tool_routes = mapping def is_config_declared_server(self, server_id: str) -> bool: """True when server_id was declared in config.yaml (present in the in-memory config map). @@ -2378,7 +2391,7 @@ class MCPServerManager: ) assigned_server_ids[server_id] = server_name - _warn_on_server_name_fields( + warn_on_server_name_fields( server_id=server_id, alias=alias, server_name=server_name, @@ -2584,10 +2597,10 @@ class MCPServerManager: token_validation=server_config.get("token_validation", None), oauth_identity_binding=server_config.get("oauth_identity_binding", None), ) - self._assign_unique_short_prefix(new_server) + self.assign_unique_short_prefix(new_server) _warn_legacy_delegate_auth_if_applicable(new_server, source="config") _warn_config_id_jag_server_outruns_sso(new_server) - self._invalidate_discovery_lists(server_id) + self.invalidate_discovery_lists(server_id) self.config_mcp_servers[server_id] = new_server self._set_oauth_discovery_deferred( server_id, @@ -2608,13 +2621,13 @@ class MCPServerManager: "Loaded MCP Servers: %s", json.dumps(_redacted_registry_dump(self.config_mcp_servers), indent=4) ) - await self._hydrate_config_servers_dcr_clients() + await self.hydrate_config_servers_dcr_clients() - self._prime_oauth_metadata_discovery_for_servers(tuple(self.config_mcp_servers.values())) + self.prime_oauth_metadata_discovery_for_servers(tuple(self.config_mcp_servers.values())) self.initialize_tool_name_to_mcp_server_name_mapping() - async def _hydrate_config_servers_dcr_clients(self) -> None: + async def hydrate_config_servers_dcr_clients(self, servers: Sequence[MCPServer] | None = None) -> None: """Overlay each config-declared server's persisted DCR client (from the server-scoped store) onto its in-memory object so token refresh authenticates after a restart. A best-effort no-op when the DB is unreachable at config-load time.""" @@ -2622,7 +2635,7 @@ class MCPServerManager: hydrate_config_server_dcr_client, ) - for server in self.config_mcp_servers.values(): + for server in servers if servers is not None else self.config_mcp_servers.values(): try: if await hydrate_config_server_dcr_client(server): verbose_logger.debug( @@ -2787,17 +2800,18 @@ class MCPServerManager: mappings make ``_get_mcp_server_from_tool_name`` resolve to a prefix that no longer exists in the live registry. """ - from litellm.proxy._experimental.mcp_server.tool_registry import ( - global_mcp_tool_registry, - ) + self.invalidate_discovery_lists(server.server_id) + self.remove_server_tool_routing(server) + + def remove_server_tool_routing(self, server: MCPServer) -> None: + from litellm.proxy._experimental.mcp_server.tool_registry import global_mcp_tool_registry - self._invalidate_discovery_lists(server.server_id) prefix_root: Final = normalize_server_name(get_server_prefix(server)) if server.spec_path and prefix_root: openapi_key_prefix: Final = prefix_root + MCP_TOOL_PREFIX_SEPARATOR global_mcp_tool_registry.unregister_tools_with_prefix(openapi_key_prefix) - owned_normalized: Final = self._owned_mapping_values(server) + owned_normalized: Final = self.owned_mapping_values(server) stale_mapping_keys: Final = tuple( tool_name @@ -2808,13 +2822,13 @@ class MCPServerManager: for key in stale_mapping_keys: del self.tool_name_to_mcp_server_name_mapping[key] - def _owned_mapping_values(self, server: MCPServer) -> frozenset[str]: + def owned_mapping_values(self, server: MCPServer) -> frozenset[str]: return frozenset( normalize_server_name(value) for value in (*iter_known_server_prefixes(server), server.name) if value ) def server_exposes_tool(self, server: MCPServer, tool_name: str) -> bool: - owned: Final = self._owned_mapping_values(server) + owned: Final = self.owned_mapping_values(server) mapped_owners: Final = ( self.tool_name_to_mcp_server_name_mapping.get(spelling) for spelling in iter_known_tool_name_spellings(tool_name, server) @@ -2845,7 +2859,7 @@ class MCPServerManager: if evicted is not None: verbose_logger.debug("Removed MCP Server: %s", mcp_server.server_id or mcp_server.server_name) self._cleanup_server_tool_routing_artifacts(evicted) - self._invalidate_oauth_discovery_state(evicted.server_id) + self.invalidate_oauth_discovery_state(evicted.server_id) else: verbose_logger.warning("Server ID %s not found in registry", mcp_server.server_id) @@ -2936,6 +2950,7 @@ class MCPServerManager: *, credentials_are_encrypted: bool = True, env_vars_are_encrypted: bool | None = None, + register_oauth_discovery: bool = True, ) -> MCPServer: _mcp_info: Final[MCPInfo] = mcp_server.mcp_info or {} env_dict: Final = _deserialize_json_dict(getattr(mcp_server, "env", None)) @@ -3160,13 +3175,14 @@ class MCPServerManager: max_concurrent_requests=getattr(mcp_server, "max_concurrent_requests", None), ) _warn_legacy_delegate_auth_if_applicable(new_server, source="database") - self._set_oauth_discovery_deferred( - new_server.server_id, - _requires_oauth_discovery(server_url, use_issuer_anchor, new_server), - ) + if register_oauth_discovery: + self._set_oauth_discovery_deferred( + new_server.server_id, + _requires_oauth_discovery(server_url, use_issuer_anchor, new_server), + ) return new_server - async def _maybe_register_openapi_tools(self, server: MCPServer, *, initialize_mapping: bool = True): + async def maybe_register_openapi_tools(self, server: MCPServer, *, initialize_mapping: bool = True): """Register OpenAPI tools if the server has a spec_path configured.""" if server.spec_path: verbose_logger.info("Loading OpenAPI spec from %s for server %s", server.spec_path, server.name) @@ -3194,10 +3210,10 @@ class MCPServerManager: # Re-decrypting plaintext would zero the values, so build with # env_vars_are_encrypted=False. new_server: Final = await self.build_mcp_server_from_table(mcp_server, env_vars_are_encrypted=False) - self._assign_unique_short_prefix(new_server) - self._invalidate_discovery_lists(mcp_server.server_id) + self.assign_unique_short_prefix(new_server) + self.invalidate_discovery_lists(mcp_server.server_id) self.registry[mcp_server.server_id] = new_server - await self._maybe_register_openapi_tools(new_server) + await self.maybe_register_openapi_tools(new_server) self.prime_oauth_metadata_discovery(new_server) verbose_logger.debug("Added MCP Server: %s", new_server.name) @@ -3215,7 +3231,7 @@ class MCPServerManager: evicted = self.registry.pop(mcp_server.server_name, None) if evicted is not None: self._cleanup_server_tool_routing_artifacts(evicted) - self._invalidate_oauth_discovery_state(evicted.server_id) + self.invalidate_oauth_discovery_state(evicted.server_id) return try: if mcp_server.server_id in self.registry: @@ -3227,14 +3243,14 @@ class MCPServerManager: existing_prefix: Final = self.registry[mcp_server.server_id].short_prefix if existing_prefix and not new_server.short_prefix: new_server.short_prefix = existing_prefix - _carry_forward_resolved_oauth_endpoints( + carry_forward_resolved_oauth_endpoints( new_server=new_server, previous_server=self.registry[mcp_server.server_id], ) - self._assign_unique_short_prefix(new_server) - self._invalidate_discovery_lists(mcp_server.server_id) + self.assign_unique_short_prefix(new_server) + self.invalidate_discovery_lists(mcp_server.server_id) self.registry[mcp_server.server_id] = new_server - await self._maybe_register_openapi_tools(new_server) + await self.maybe_register_openapi_tools(new_server) self.prime_oauth_metadata_discovery(new_server) verbose_logger.debug("Updated MCP Server: %s", new_server.name) @@ -4466,7 +4482,9 @@ class MCPServerManager: ) raise_classified_list_failure(e, server.name, suppress_challenge=server.is_dcr_bridge) - def _invalidate_discovery_lists(self, server_id: str) -> None: + def invalidate_discovery_lists(self, server_id: str) -> None: + self._upstream_initialize_instructions_by_server_id.pop(server_id, None) + self._upstream_initialize_instructions_probed_at.pop(server_id, None) self._prompt_discovery_cache.invalidate(server_id) self._resource_discovery_cache.invalidate(server_id) self._template_discovery_cache.invalidate(server_id) @@ -5256,7 +5274,7 @@ class MCPServerManager: _SHORT_PREFIX_MAX_REHASH_ATTEMPTS = 1024 - def _assign_unique_short_prefix( + def assign_unique_short_prefix( self, server: MCPServer, registry: dict[str, MCPServer] | None = None, @@ -6146,7 +6164,7 @@ class MCPServerManager: failure is logged, never raised, because the DB write already succeeded and the TTL remains the backstop. """ - self._invalidate_discovery_lists(server_id) + self.invalidate_discovery_lists(server_id) try: await self._per_user_oauth_token_store.invalidate(user_id, server_id) except Exception as exc: # noqa: BLE001 - cache drop is best-effort; TTL is the backstop @@ -6445,7 +6463,7 @@ class MCPServerManager: Note: This now handles prefixed tool names """ for server in self.get_registry().values(): - if self._oauth_discovery_slot(server.server_id) is not None: + if self.oauth_discovery_slot(server.server_id) is not None: continue if server.needs_user_oauth_token: # Skip OAuth2 servers that rely on user-provided tokens @@ -6507,152 +6525,7 @@ class MCPServerManager: return None async def reload_servers_from_database(self): - """Re-synchronize the in-memory MCP server registry with the database.""" - from litellm.proxy.management_endpoints.mcp_management_endpoints import ( - get_prisma_client_or_throw, - ) - - verbose_logger.debug("Loading MCP servers from database into registry...") - self._upstream_initialize_instructions_by_server_id.clear() - self._upstream_initialize_instructions_probed_at.clear() - - # perform authz check to filter the mcp servers user has access to - prisma_client: Final = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") - # Load only "active", legacy "approved", and NULL (no approval workflow) rows. - # Pending/rejected servers are excluded at the DB level so we never load them. - from litellm.proxy._experimental.mcp_server.db import LiteLLM_MCPServerTable - - raw_rows: Final[Sequence[BaseModel]] = await MCPServerRepository(prisma_client).table.find_many( - where={ - "OR": [ - {"approval_status": None}, - {"approval_status": {"in": ["active", "approved"]}}, - ] - } - ) - verbose_logger.info("Found %s MCP servers in database", len(raw_rows)) - - previous_registry: Final = self.registry - new_registry: Final[dict[str, MCPServer]] = {} - - # Stage one: build every server. Stage two assigns short prefixes - # against the *full* set so dedup is deterministic regardless of - # iteration order. - for row in raw_rows: - try: - server = LiteLLM_MCPServerTable.model_validate(row.model_dump()) - existing_server = previous_registry.get(server.server_id) - - if ( - existing_server is not None - and existing_server.updated_at is not None - and server.updated_at is not None - and existing_server.updated_at == server.updated_at - and ( - self._oauth_discovery_slot(server.server_id) is not None - or not _oauth_endpoints_unresolved(existing_server) - ) - ): - # Re-use existing server instance to avoid re-running build_mcp_server_from_table() - # which can perform network discovery for OAuth2 servers. - new_registry[server.server_id] = existing_server - continue - - _warn_on_server_name_fields( - server_id=server.server_id, - alias=getattr(server, "alias", None), - server_name=getattr(server, "server_name", None), - ) - verbose_logger.debug("Building server from DB: %s (%s)", server.server_id, server.server_name) - # raw_rows come straight from the DB, so their global env var - # values (like credentials) are still encrypted here, unlike the - # already-decrypted records add_server/update_server are handed. - # Decrypt them while building the registry entry. - new_server = await self.build_mcp_server_from_table(server, env_vars_are_encrypted=True) - # Carry the cached short_prefix from the previous registry entry - # (if any) so the prefix is stable across reloads. - if existing_server is not None and existing_server.short_prefix: - new_server.short_prefix = existing_server.short_prefix - _carry_forward_resolved_oauth_endpoints(new_server=new_server, previous_server=existing_server) - new_registry[server.server_id] = new_server - except Exception as e: - verbose_logger.exception( - "Skipping MCP server %s (%s) during DB reload: %s", - getattr(row, "server_id", None), - getattr(row, "alias", None), - e, - ) - - # Assign short prefixes against the full candidate set without - # publishing the staged registry to concurrent callers. - registered_registry: Final[dict[str, MCPServer]] = {} - registered_openapi_tools = False - for server_id, new_server in new_registry.items(): - try: - self._assign_unique_short_prefix(new_server, registry=new_registry) - # Register OpenAPI tools *after* the final short prefix is assigned - # so the tools are stored in the global registry under the same - # prefix that lookups will use. - await self._maybe_register_openapi_tools(new_server, initialize_mapping=False) - registered_registry[server_id] = new_server - if new_server.spec_path: - registered_openapi_tools = True - except Exception as e: - verbose_logger.exception( - "Skipping MCP server %s (%s) during DB reload: %s", - new_server.server_id, - getattr(new_server, "alias", None), - e, - ) - - dropped_registry_keys: Final = previous_registry.keys() - registered_registry.keys() - for registry_key in dropped_registry_keys: - self._invalidate_oauth_discovery_state(previous_registry[registry_key].server_id) - - for server_id in previous_registry.keys() | registered_registry.keys(): - if previous_registry.get(server_id) != registered_registry.get(server_id): - self._invalidate_discovery_lists(server_id) - self.registry = registered_registry - # A discovery task may have published into ``previous_registry`` while - # this replacement was being staged. Reconcile every published entry - # synchronously after the swap so a lost publication cannot also leave - # the replacement unresolved with no retry slot. - registered_servers: Final = tuple(registered_registry.values()) - self._reconcile_oauth_discovery_slots_for_servers(registered_servers) - self._prime_oauth_metadata_discovery_for_servers(registered_servers) - if registered_openapi_tools: - self.initialize_tool_name_to_mcp_server_name_mapping() - - verbose_logger.debug("MCP registry refreshed (%s servers in registry)", len(registered_registry)) - - # get_registry() is ``config_mcp_servers | registry``, so a database row sharing an id with a - # config.yaml server hides that server everywhere. Only reachable once an operator pins - # ``server_id`` in config.yaml; say so rather than letting the server disappear silently. - shadowed_config_server_ids: Final = frozenset(self.config_mcp_servers.keys() & registered_registry.keys()) - if shadowed_config_server_ids and shadowed_config_server_ids != self._warned_shadowed_config_server_ids: - verbose_logger.warning( - "config.yaml MCP server_id(s) %s are also database-backed MCP servers. The database " - "entry takes precedence, so the config.yaml server is unreachable. Give the config " - "entry a different server_id.", - ", ".join(sorted(shadowed_config_server_ids)), - ) - self._warned_shadowed_config_server_ids = shadowed_config_server_ids - - # The mirror image of the block above: a config server_id that is a database server's name - # answers that server's grants instead, because ids are matched before names. - capturing_config_server_ids: Final = _config_ids_capturing_db_identifiers( - self.config_mcp_servers.keys(), registered_registry.values() - ) - if capturing_config_server_ids and capturing_config_server_ids != self._warned_capturing_config_server_ids: - verbose_logger.warning( - "config.yaml MCP server_id(s) %s are the name or alias of a database-backed MCP " - "server. Permission entries naming them resolve to the config.yaml server, not the " - "database one. Give the config entry a different server_id.", - ", ".join(sorted(capturing_config_server_ids)), - ) - self._warned_capturing_config_server_ids = capturing_config_server_ids - - await self._hydrate_config_servers_dcr_clients() + await self.catalog.reload() def get_mcp_servers_from_ids(self, server_ids: list[str]) -> list[MCPServer]: servers: Final = [] diff --git a/litellm/proxy/_experimental/mcp_server/operations.py b/litellm/proxy/_experimental/mcp_server/operations.py index fcee3483e15..bbe2987b856 100644 --- a/litellm/proxy/_experimental/mcp_server/operations.py +++ b/litellm/proxy/_experimental/mcp_server/operations.py @@ -1,6 +1,5 @@ """Shared MCP operation policy and dispatch.""" -import asyncio import traceback import types import uuid @@ -50,6 +49,7 @@ from litellm.proxy._experimental.mcp_server.byok_credential_cache import ( cache_byok_credential, get_cached_byok_credential, ) +from litellm.proxy._experimental.mcp_server.catalog import TargetCatalog, catalog_operation from litellm.proxy._experimental.mcp_server.contracts import ( AuthorizedToolCall, OperationContext, @@ -614,6 +614,7 @@ def apply_tool_overrides( return tools +@catalog_operation(lambda: global_mcp_server_manager) async def _get_allowed_mcp_servers( user_api_key_auth: UserAPIKeyAuth | None, mcp_servers: Sequence[str] | None, @@ -930,6 +931,7 @@ def _aggregate_server_key(server: MCPServer) -> str: return get_server_prefix(server) or "unknown" +@catalog_operation(lambda: global_mcp_server_manager) async def _get_tools_from_mcp_servers( user_api_key_auth: UserAPIKeyAuth | None, mcp_auth_header: str | None, @@ -1151,24 +1153,17 @@ async def _get_tools_from_mcp_servers( verbose_logger.exception("Error getting tools from server %s: %s", server.name, e) return [], classify_list_exception(e) - # Fetch tools from all servers in parallel - tasks: Final = [_fetch_and_filter_server_tools(server) for server in allowed_mcp_servers] - results: Final = await asyncio.gather(*tasks) - - # Flatten results into single list - all_tools: Final[list[MCPTool]] = [tool for tools, _ in results for tool in tools] - server_outcomes: Final[dict[str, ServerOutcome]] = { - _aggregate_server_key(server): outcome - for server, (_, outcome) in zip(allowed_mcp_servers, results) - if server is not None - } + listing: Final = await TargetCatalog.list( + allowed_mcp_servers, _fetch_and_filter_server_tools, _aggregate_server_key + ) + all_tools: Final = listing.tools + server_outcomes: Final = listing.outcomes # If logging is enabled, enrich spend_logs_metadata with counts if litellm_logging_obj: per_server_tool_counts: Final[dict[str, int]] = { - _aggregate_server_key(server): len(server_tools) - for server, (server_tools, _) in zip(allowed_mcp_servers, results) - if server is not None + key: outcome.tool_count if isinstance(outcome, ServerListOk) else 0 + for key, outcome in server_outcomes.items() } metadata_dict: Final = litellm_logging_obj.model_call_details.get("metadata") @@ -1435,6 +1430,7 @@ async def filter_tools_by_key_team_permissions( ] +@catalog_operation(lambda: global_mcp_server_manager) async def _list_mcp_tools( user_api_key_auth: UserAPIKeyAuth | None = None, mcp_auth_header: str | None = None, @@ -2271,6 +2267,7 @@ async def fire_mcp_tool_call_failure_logging( @client +@catalog_operation(lambda: global_mcp_server_manager) async def call_mcp_tool( name: str, arguments: dict[str, object] | None = None, @@ -3055,6 +3052,7 @@ class GatewayOperations: @overload async def execute(self, operation: ReadResourceRequest, context: OperationContext) -> ReadResourceResult: ... + @catalog_operation(lambda: global_mcp_server_manager) async def execute(self, operation: GatewayOperation, context: OperationContext) -> GatewayResult: match operation: case AuthorizedToolCall(): diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index c2f7bf7d531..74e530bb12f 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -22,6 +22,7 @@ from litellm.exceptions import ( GuardrailRaisedException, ModifyResponseException, ) +from litellm.proxy._experimental.mcp_server.catalog import catalog_operation, global_manager from litellm.proxy._experimental.mcp_server.exceptions import ( MCPServerListError, MCPServerURLCredentialsError, @@ -861,6 +862,7 @@ if MCP_AVAILABLE: return await _apply_toolset_scope(user_api_key_dict, toolset.toolset_id) @router.get("/tools/list", dependencies=[Depends(user_api_key_auth)]) + @catalog_operation(global_manager) async def list_tool_rest_api( request: Request, server_id: str | None = Query(None, description="The server id to list tools for"), @@ -1085,6 +1087,7 @@ if MCP_AVAILABLE: } @router.post("/tools/call", dependencies=[Depends(user_api_key_auth)]) + @catalog_operation(global_manager) async def call_tool_rest_api( request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), diff --git a/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py b/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py index dcf1b01bc25..0644238071e 100644 --- a/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py +++ b/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py @@ -11,6 +11,7 @@ from typing import TYPE_CHECKING, Any, Final from litellm._logging import verbose_logger from litellm.exceptions import ContextWindowExceededError from litellm.litellm_core_utils.exception_mapping_utils import ExceptionCheckers +from litellm.proxy._experimental.mcp_server.catalog import catalog_operation, global_manager from litellm.proxy._experimental.mcp_server.faults import iter_exception_tree from litellm.proxy._experimental.mcp_server.utils import MCP_TOOL_PREFIX_SEPARATOR @@ -78,6 +79,7 @@ class SemanticMCPToolFilter: self._tool_map: dict[str, object] = {} # MCPTool objects or OpenAI function dicts self._index_sync_lock = asyncio.Lock() + @catalog_operation(global_manager) async def build_router_from_mcp_registry(self) -> None: """Build semantic router from all MCP tools in the registry (no auth checks).""" from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 433b693fcae..f160ce0f4a9 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -1559,6 +1559,9 @@ if MCP_AVAILABLE: ) return user_api_key_auth.model_copy(update={"object_permission": updated_op, "mcp_toolset_id": toolset_id}) + from litellm.proxy._experimental.mcp_server.catalog import catalog_operation + + @catalog_operation(lambda: operations.global_mcp_server_manager) async def _raise_preemptive_401_for_unauthenticated_servers( scope: Scope, mcp_servers: list[str] | None, diff --git a/litellm/proxy/_experimental/mcp_server/tool_registry.py b/litellm/proxy/_experimental/mcp_server/tool_registry.py index e9e28c8a782..2adf5db8f67 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_registry.py +++ b/litellm/proxy/_experimental/mcp_server/tool_registry.py @@ -1,5 +1,8 @@ +import asyncio import json -from collections.abc import Callable +from collections.abc import Callable, Iterator, Mapping +from contextlib import contextmanager +from contextvars import ContextVar from typing import TYPE_CHECKING, Any, Final from litellm._logging import verbose_logger @@ -22,7 +25,30 @@ class MCPToolRegistry: def __init__(self): # Registry to store all registered tools - self.tools: dict[str, MCPTool] = {} + self.published_tools: dict[str, MCPTool] = {} + self._catalog_tools: ContextVar[tuple[dict[str, MCPTool], asyncio.Event] | None] = ContextVar( + "mcp_catalog_tools", default=None + ) + + @property + def tools(self) -> dict[str, MCPTool]: + scoped: Final = self._catalog_tools.get() + return scoped[0] if scoped is not None and not scoped[1].is_set() else self.published_tools + + @tools.setter + def tools(self, tools: dict[str, MCPTool]) -> None: + self.published_tools = tools + + @contextmanager + def catalog_scope(self, tools: Mapping[str, MCPTool]) -> Iterator[dict[str, MCPTool]]: + detached: Final = dict(tools) + closed: Final = asyncio.Event() + token: Final = self._catalog_tools.set((detached, closed)) + try: + yield detached + finally: + closed.set() + self._catalog_tools.reset(token) def register_tool( self, diff --git a/litellm/proxy/_experimental/mcp_server/utils.py b/litellm/proxy/_experimental/mcp_server/utils.py index 6bd080f5216..b06098ef474 100644 --- a/litellm/proxy/_experimental/mcp_server/utils.py +++ b/litellm/proxy/_experimental/mcp_server/utils.py @@ -81,7 +81,7 @@ MCP_TOOL_PREFIX_FORMAT: Final = "{server_name}{separator}{tool_name}" # principle hash to the same three chars; that natural-hash collision # IS a routing-correctness issue (the second registrant would otherwise # have its tools misrouted to the first), so registration goes through -# ``MCPServerManager._assign_unique_short_prefix`` which rehashes with +# ``MCPServerManager.assign_unique_short_prefix`` which rehashes with # a deterministic attempt counter until it finds an unused prefix and # caches the result on ``MCPServer.short_prefix``. A collision is # logged at INFO when it happens. @@ -114,7 +114,7 @@ def compute_short_server_prefix(server_id: str, attempt: int = 0) -> str: and whose remaining characters are drawn from the full base62 alphabet. Pass ``attempt > 0`` to rehash to a different prefix when the natural hash collides with a prefix already assigned to another - server (see ``MCPServerManager._assign_unique_short_prefix``). An + server (see ``MCPServerManager.assign_unique_short_prefix``). An empty ``server_id`` raises ``ValueError`` — short prefixes require a stable identifier to be deterministic. """ @@ -314,7 +314,7 @@ def get_server_prefix(server: object) -> str: When the short-prefix mode is enabled (``LITELLM_USE_SHORT_MCP_TOOL_PREFIX``) a three-character base62 ID is returned. We prefer the cached ``server.short_prefix`` value when set — that field is populated at - registration time by ``MCPServerManager._assign_unique_short_prefix`` + registration time by ``MCPServerManager.assign_unique_short_prefix`` and resolves natural-hash collisions deterministically — and only fall back to the natural hash for ad-hoc / temp-server objects without a cached value. In default mode the historical behaviour is preserved: diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py index 35480f5e397..aedc50e613a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py @@ -46,7 +46,7 @@ class MCPSecurityGuardrail(CustomGuardrail): if self.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) is not True: return data - unregistered: Final = self._find_unregistered_mcp_servers(data) + unregistered: Final = await self._find_unregistered_mcp_servers(data) if not unregistered: return data @@ -90,7 +90,7 @@ class MCPSecurityGuardrail(CustomGuardrail): return server_names @staticmethod - def _find_unregistered_mcp_servers(data: dict) -> set[str]: + async def _find_unregistered_mcp_servers(data: dict) -> set[str]: """Check tools in data against the MCP server registry. Returns set of unregistered server names.""" tools: Final = data.get("tools") if not tools or not isinstance(tools, list): @@ -104,7 +104,8 @@ class MCPSecurityGuardrail(CustomGuardrail): global_mcp_server_manager, ) - registry: Final = global_mcp_server_manager.get_registry() - registered_names: Final = set(registry.keys()) + async with global_mcp_server_manager.catalog.operation(): + registry: Final = global_mcp_server_manager.get_registry() + registered_names: Final = set(registry.keys()) - return requested_servers - registered_names + return requested_servers - registered_names diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 9ad78876043..a1e72158794 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -19,7 +19,8 @@ import functools import importlib import json import os -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import AsyncIterator, Iterable, Mapping, Sequence +from contextlib import asynccontextmanager from dataclasses import dataclass from datetime import datetime, timedelta, timezone from typing import ( @@ -45,6 +46,8 @@ from fastapi import ( from fastapi.responses import JSONResponse from typing_extensions import ReadOnly, TypedDict +from litellm.proxy._experimental.mcp_server.catalog import catalog_operation + try: from prisma.errors import RecordNotFoundError, UniqueViolationError except ImportError: @@ -1061,7 +1064,8 @@ if MCP_AVAILABLE: registry_servers.append({"server": _build_builtin_registry_entry(base_url)}) # Centralized IP-based filtering: external callers only see public servers - registered_servers: Final = list(global_mcp_server_manager.get_filtered_registry(client_ip).values()) + async with global_mcp_server_manager.catalog.operation(): + registered_servers: Final = list(global_mcp_server_manager.get_filtered_registry(client_ip).values()) registered_servers.sort(key=_build_mcp_registry_server_name) @@ -1091,6 +1095,7 @@ if MCP_AVAILABLE: return "view_all" return "restricted" + @catalog_operation(lambda: global_mcp_server_manager) async def _get_team_scoped_mcp_server_list( team_id: str, ) -> list[LiteLLM_MCPServerTable]: @@ -1129,6 +1134,7 @@ if MCP_AVAILABLE: return _redact_mcp_credentials_list(servers) + @catalog_operation(lambda: global_mcp_server_manager) async def _resolve_accessible_mcp_servers( user_api_key_dict: UserAPIKeyAuth, ) -> list[LiteLLM_MCPServerTable]: @@ -1150,6 +1156,7 @@ if MCP_AVAILABLE: aggregated.setdefault(server.server_id, server) return list(aggregated.values()) + @catalog_operation(lambda: global_mcp_server_manager) async def _connected_app_reachable_server_ids(user_api_key_dict: UserAPIKeyAuth) -> frozenset[str]: """Server ids a connected app authorized by this dashboard user is served on the aggregate MCP endpoint, resolved through the one owner of the admitted subject so the page and the @@ -1284,6 +1291,7 @@ if MCP_AVAILABLE: description="Health check for MCP servers", dependencies=[Depends(user_api_key_auth)], ) + @catalog_operation(lambda: global_mcp_server_manager) async def health_check_servers( server_ids: list[str] | None = Query( None, @@ -1592,6 +1600,7 @@ if MCP_AVAILABLE: dependencies=[Depends(user_api_key_auth)], response_model=LiteLLM_MCPServerTable, ) + @catalog_operation(lambda: global_mcp_server_manager) async def fetch_mcp_server( request: Request, server_id: str, @@ -2016,9 +2025,7 @@ if MCP_AVAILABLE: server_id: Final[str] = request.path_params.get("server_id", "") if server_id: - _s = global_mcp_server_manager.get_mcp_server_by_id(server_id) - if not _s: - _s = global_mcp_server_manager.get_mcp_server_by_name(server_id) + _s = await global_mcp_server_manager.catalog.resolve(server_id) if ( _s and getattr(_s, "auth_type", None) == MCPAuth.oauth2 @@ -2064,42 +2071,44 @@ if MCP_AVAILABLE: user_api_key_dict: UserAPIKeyAuth, request: Request | None = None, ) -> MCPServer: - server = await get_cached_temporary_mcp_server(server_id) - resolved_from_temp_cache: Final = server is not None - if server is None: - # Fall back to real DB/config server (e.g. for the user-side OAuth flow - # which calls these endpoints with a real server_id, not a temp session id). - from litellm.proxy.auth.ip_address_utils import IPAddressUtils + async with _oauth_server_operation(server_id, user_api_key_dict, request=request) as server: + return server - client_ip: Final = IPAddressUtils.get_mcp_client_ip(request) if request else None - server = global_mcp_server_manager.get_mcp_server_by_id( - server_id - ) or global_mcp_server_manager.get_mcp_server_by_name(server_id, client_ip=client_ip) - if server is None: - raise HTTPException( - status_code=status.HTTP_404_NOT_FOUND, - detail={"error": f"MCP server {server_id} not found"}, - ) + @asynccontextmanager + async def _oauth_server_operation( + server_id: str, + user_api_key_dict: UserAPIKeyAuth, + request: Request | None = None, + ) -> AsyncIterator[MCPServer]: + temporary: Final = await get_cached_temporary_mcp_server(server_id) + if temporary is not None: + if not _user_has_admin_view(user_api_key_dict): + raise HTTPException(status_code=403, detail={"error": f"Access denied to MCP server {server_id}"}) + yield temporary + return + async with global_mcp_server_manager.catalog.operation(): + yield await _resolve_saved_oauth_server(server_id, user_api_key_dict, request) - # Per-server access policy mirrors `fetch_mcp_server`: admin-view - # callers are unrestricted; non-admins must have the server in their - # allowed-servers set. Temporary cached servers come from the - # admin-only `/server/oauth/session` setup flow and are not exposed - # to non-admins. + @catalog_operation(lambda: global_mcp_server_manager) + async def _resolve_saved_oauth_server( + server_id: str, + user_api_key_dict: UserAPIKeyAuth, + request: Request | None, + ) -> MCPServer: + from litellm.proxy.auth.ip_address_utils import IPAddressUtils + + client_ip: Final = IPAddressUtils.get_mcp_client_ip(request) if request else None + server: Final = global_mcp_server_manager.get_mcp_server_by_id( + server_id + ) or global_mcp_server_manager.get_mcp_server_by_name(server_id, client_ip=client_ip) + if server is None: + raise HTTPException(status_code=404, detail={"error": f"MCP server {server_id} not found"}) if not _user_has_admin_view(user_api_key_dict): - if resolved_from_temp_cache: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail={"error": f"Access denied to MCP server {server_id}"}, - ) - allowed_server_ids: Final[set[str]] = set() - for auth_context in await build_effective_auth_contexts(user_api_key_dict): - allowed_server_ids.update(await global_mcp_server_manager.get_allowed_mcp_servers(auth_context)) - if server.server_id not in allowed_server_ids: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail={"error": f"Access denied to MCP server {server_id}"}, - ) + allowed_ids: Final[set[str]] = set() + for context in await build_effective_auth_contexts(user_api_key_dict): + allowed_ids.update(await global_mcp_server_manager.get_allowed_mcp_servers(context)) + if server.server_id not in allowed_ids: + raise HTTPException(status_code=403, detail={"error": f"Access denied to MCP server {server_id}"}) return server @router.get( @@ -2119,47 +2128,47 @@ if MCP_AVAILABLE: response_type: str | None = None, scope: str | None = None, ): - mcp_server: Final = await _get_cached_temporary_mcp_server_or_404(server_id, user_api_key_dict, request=request) - _raise_if_not_oauth2(mcp_server) - # Use the server's stored client_id when the caller doesn't supply one - stored_or_supplied_client_id: Final = mcp_server.client_id or client_id or "" - ephemeral_dcr_client: Final = ( - await resolve_ephemeral_dcr_client( + async with _oauth_server_operation(server_id, user_api_key_dict, request=request) as mcp_server: + _raise_if_not_oauth2(mcp_server) + # Use the server's stored client_id when the caller doesn't supply one + stored_or_supplied_client_id: Final = mcp_server.client_id or client_id or "" + ephemeral_dcr_client: Final = ( + await resolve_ephemeral_dcr_client( + request=request, + mcp_server=mcp_server, + code_challenge=code_challenge, + code_challenge_method=code_challenge_method, + redirect_uri=redirect_uri, + ) + if not stored_or_supplied_client_id + else None + ) + resolved_client_id: Final = stored_or_supplied_client_id or ( + ephemeral_dcr_client.client_id if ephemeral_dcr_client else "" + ) + if not resolved_client_id: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={ + "error": "missing_client_id", + "message": ( + "No client_id available for this MCP server. " + "Either configure the server with a client_id or supply one in the request." + ), + }, + ) + return await authorize_with_server( request=request, mcp_server=mcp_server, + client_id=resolved_client_id, + redirect_uri=redirect_uri, + state=state, code_challenge=code_challenge, code_challenge_method=code_challenge_method, - redirect_uri=redirect_uri, + response_type=response_type, + scope=scope, + ephemeral_dcr_client=ephemeral_dcr_client, ) - if not stored_or_supplied_client_id - else None - ) - resolved_client_id: Final = stored_or_supplied_client_id or ( - ephemeral_dcr_client.client_id if ephemeral_dcr_client else "" - ) - if not resolved_client_id: - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail={ - "error": "missing_client_id", - "message": ( - "No client_id available for this MCP server. " - "Either configure the server with a client_id or supply one in the request." - ), - }, - ) - return await authorize_with_server( - request=request, - mcp_server=mcp_server, - client_id=resolved_client_id, - redirect_uri=redirect_uri, - state=state, - code_challenge=code_challenge, - code_challenge_method=code_challenge_method, - response_type=response_type, - scope=scope, - ephemeral_dcr_client=ephemeral_dcr_client, - ) @router.post( "/server/oauth/{server_id}/token", @@ -2179,47 +2188,47 @@ if MCP_AVAILABLE: refresh_token: str | None = Form(None), scope: str | None = Form(None), ): - mcp_server: Final = await _get_cached_temporary_mcp_server_or_404(server_id, user_api_key_dict, request=request) - _raise_if_not_oauth2(mcp_server) - # Sealed passthrough codes exist only for the authorization_code grant. A refresh_token - # grant must never open one: the minted client is unrecoverable after the single flow by - # contract, so an expired browser-held token re-runs authorize instead. - sealed_code: Final = ( - redeem_passthrough_authorization_code(code=code, mcp_server=mcp_server, code_verifier=code_verifier) - if grant_type == "authorization_code" - else None - ) - resolved_code: Final = sealed_code.upstream_code if sealed_code else code - # A sealed flow ran the gateway /callback as its upstream redirect (bridge short-circuit - # or plain flow alike), so the exchange must present that binding, not the browser page. - resolved_redirect_uri: Final = f"{get_request_base_url(request)}/callback" if sealed_code else redirect_uri - caller_client_id: Final = sealed_code.client_id if sealed_code else client_id - caller_client_secret: Final = sealed_code.client_secret if sealed_code else client_secret - resolved_client_id: Final = mcp_server.client_id or caller_client_id or "" - if not resolved_client_id: - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail={ - "error": "missing_client_id", - "message": ( - "No client_id available for this MCP server. " - "Either configure the server with a client_id or supply one in the request." - ), - }, + async with _oauth_server_operation(server_id, user_api_key_dict, request=request) as mcp_server: + _raise_if_not_oauth2(mcp_server) + # Sealed passthrough codes exist only for the authorization_code grant. A refresh_token + # grant must never open one: the minted client is unrecoverable after the single flow by + # contract, so an expired browser-held token re-runs authorize instead. + sealed_code: Final = ( + redeem_passthrough_authorization_code(code=code, mcp_server=mcp_server, code_verifier=code_verifier) + if grant_type == "authorization_code" + else None + ) + resolved_code: Final = sealed_code.upstream_code if sealed_code else code + # A sealed flow ran the gateway /callback as its upstream redirect (bridge short-circuit + # or plain flow alike), so the exchange must present that binding, not the browser page. + resolved_redirect_uri: Final = f"{get_request_base_url(request)}/callback" if sealed_code else redirect_uri + caller_client_id: Final = sealed_code.client_id if sealed_code else client_id + caller_client_secret: Final = sealed_code.client_secret if sealed_code else client_secret + resolved_client_id: Final = mcp_server.client_id or caller_client_id or "" + if not resolved_client_id: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={ + "error": "missing_client_id", + "message": ( + "No client_id available for this MCP server. " + "Either configure the server with a client_id or supply one in the request." + ), + }, + ) + return await exchange_token_with_server( + request=request, + mcp_server=mcp_server, + grant_type=grant_type, + code=resolved_code, + redirect_uri=resolved_redirect_uri, + client_id=resolved_client_id, + client_secret=caller_client_secret, + code_verifier=code_verifier, + refresh_token=refresh_token, + scope=scope, + client_token_endpoint_auth_method=sealed_code.token_endpoint_auth_method if sealed_code else None, ) - return await exchange_token_with_server( - request=request, - mcp_server=mcp_server, - grant_type=grant_type, - code=resolved_code, - redirect_uri=resolved_redirect_uri, - client_id=resolved_client_id, - client_secret=caller_client_secret, - code_verifier=code_verifier, - refresh_token=refresh_token, - scope=scope, - client_token_endpoint_auth_method=sealed_code.token_endpoint_auth_method if sealed_code else None, - ) @router.post( "/server/oauth/{server_id}/register", @@ -2231,22 +2240,22 @@ if MCP_AVAILABLE: server_id: str, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): - mcp_server: Final = await _get_cached_temporary_mcp_server_or_404(server_id, user_api_key_dict, request=request) - request_data: Final = await _read_request_body(request=request) - data: Final[dict] = {**request_data} - client_redirect_uris: Final = client_supplied_redirect_uris(data.get("redirect_uris")) + async with _oauth_server_operation(server_id, user_api_key_dict, request=request) as mcp_server: + request_data: Final = await _read_request_body(request=request) + data: Final[dict] = {**request_data} + client_redirect_uris: Final = client_supplied_redirect_uris(data.get("redirect_uris")) - return await register_client_with_server( - request=request, - mcp_server=mcp_server, - client_name=data.get("client_name", ""), - grant_types=data.get("grant_types", []), - response_types=data.get("response_types", []), - token_endpoint_auth_method=data.get("token_endpoint_auth_method", ""), - fallback_client_id=server_id, - persist_credentials=_user_is_full_admin(user_api_key_dict), - client_redirect_uris=client_redirect_uris, - ) + return await register_client_with_server( + request=request, + mcp_server=mcp_server, + client_name=data.get("client_name", ""), + grant_types=data.get("grant_types", []), + response_types=data.get("response_types", []), + token_endpoint_auth_method=data.get("token_endpoint_auth_method", ""), + fallback_client_id=server_id, + persist_credentials=_user_is_full_admin(user_api_key_dict), + client_redirect_uris=client_redirect_uris, + ) @router.delete( "/server/{server_id}", @@ -2597,6 +2606,7 @@ if MCP_AVAILABLE: # ── Per-user MCP env var endpoints ──────────────────────────────────────── + @catalog_operation(lambda: global_mcp_server_manager) async def _authorize_and_fetch_mcp_server( prisma_client, user_api_key_dict: UserAPIKeyAuth, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 18bc02467d1..8e2bb3676d9 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -19677,30 +19677,33 @@ async def _resolve_mcp_csv_tokens(csv_segment: str, client_ip: str | None) -> li all-unmatched server filter falls back to the full ``allowed_mcp_servers`` list and silently broadens the request scope). """ - from litellm.constants import DEFAULT_MCP_NAMESPACE_CSV_MAX_TOKENS - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) + from litellm.proxy._experimental.mcp_server.catalog import global_manager - seen: Final[set] = set() - deduped: Final[list[str]] = [] - for raw in csv_segment.split(","): - token = raw.strip() - if not token or token in seen: - continue - seen.add(token) - deduped.append(token) - if len(deduped) >= DEFAULT_MCP_NAMESPACE_CSV_MAX_TOKENS: - break + async with global_manager().catalog.operation(): + from litellm.constants import DEFAULT_MCP_NAMESPACE_CSV_MAX_TOKENS + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) - resolved: Final[list[str]] = [] - for token in deduped: - if global_mcp_server_manager.get_mcp_server_by_name(token, client_ip=client_ip): - resolved.append(token) - continue - if await _is_mcp_access_group_cached(token): - resolved.append(token) - return resolved + seen: Final[set] = set() + deduped: Final[list[str]] = [] + for raw in csv_segment.split(","): + token = raw.strip() + if not token or token in seen: + continue + seen.add(token) + deduped.append(token) + if len(deduped) >= DEFAULT_MCP_NAMESPACE_CSV_MAX_TOKENS: + break + + resolved: Final[list[str]] = [] + for token in deduped: + if global_mcp_server_manager.get_mcp_server_by_name(token, client_ip=client_ip): + resolved.append(token) + continue + if await _is_mcp_access_group_cached(token): + resolved.append(token) + return resolved async def _is_mcp_access_group_cached(name: str) -> bool: @@ -19755,7 +19758,9 @@ async def dynamic_mcp_route(mcp_server_name: str, request: Request): client_ip: Final = IPAddressUtils.get_mcp_client_ip(request) # 1. Registered MCP server alias - if global_mcp_server_manager.get_mcp_server_by_name(mcp_server_name, client_ip=client_ip): + async with global_mcp_server_manager.catalog.operation(): + server: Final = global_mcp_server_manager.get_mcp_server_by_name(mcp_server_name, client_ip=client_ip) + if server is not None: return await _mcp_forward_as_path(mcp_server_name, request) # 2. Comma-separated list — validate every token resolves to a known diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 88a2b92c680..739b4087f5b 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -10,6 +10,7 @@ from openai.types.responses.function_tool_param import FunctionToolParam from litellm._logging import verbose_logger from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.proxy._experimental.mcp_server.catalog import catalog_operation, global_manager from litellm.proxy._experimental.mcp_server.utils import ( iter_known_server_prefixes, logging_safe_mcp_headers, @@ -111,6 +112,7 @@ async def _toolset_exists(name: str) -> bool: return False +@catalog_operation(global_manager) async def _gateway_served_names( names: Collection[str], servers: Callable[[], Collection[MCPServer]] = _registered_mcp_servers, @@ -229,6 +231,7 @@ class LiteLLM_Proxy_MCP_Handler: return user_api_key_auth @staticmethod + @catalog_operation(global_manager) async def _get_mcp_tools_from_manager( user_api_key_auth: "UserAPIKeyAuth | None", mcp_tools_with_litellm_proxy: Iterable[Mapping[str, object]] | None, @@ -680,6 +683,7 @@ class LiteLLM_Proxy_MCP_Handler: return result_text or "Tool executed successfully" @staticmethod + @catalog_operation(global_manager) async def _execute_tool_calls( tool_server_map: dict[str, str], tool_calls: Sequence[object], diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index cb32299b143..192391c6f44 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -204,7 +204,7 @@ class MCPServer(BaseModel): # None or a value <= 0 means unlimited. max_concurrent_requests: int | None = None # Resolved short-ID tool prefix when LITELLM_USE_SHORT_MCP_TOOL_PREFIX is - # enabled. Set by ``MCPServerManager._assign_unique_short_prefix`` at + # enabled. Set by ``MCPServerManager.assign_unique_short_prefix`` at # registration time so that natural-hash collisions between two # different ``server_id`` values are bumped deterministically. Left # ``None`` in default-prefix mode. diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 35e7055bbc0..8449fb94831 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -7485,6 +7485,7 @@ class TestGatewaySessionAdmission: with ( patch("litellm.proxy.auth.auth_checks.get_user_object", get_user_object), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.should_load_db_object", return_value=False), patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), ): yield get_user_object @@ -9575,6 +9576,7 @@ class TestScopedSessionAdmission: patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), patch("litellm.proxy.auth.auth_checks.get_user_object", get_user_object), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.should_load_db_object", return_value=False), patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), ): auth_result, *_rest = await MCPRequestHandler.process_mcp_request(scope_dict) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py index a77b4c8d565..90b92048b22 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py @@ -631,6 +631,7 @@ async def test_execute_byok_tool_missing_credential_advertises_api_key_flow(monk monkeypatch.setenv("PROXY_BASE_URL", "https://gateway.example.com/proxy") mcp_operations.byok_credential_cache.flush_cache() server = MCPServer(server_id="byok-discovery", name="byok-discovery", transport=MCPTransport.http, is_byok=True) + monkeypatch.setattr(proxy_server, "should_load_db_object", lambda _kind: False) prisma = MagicMock() prisma.db.litellm_mcpusercredentials.find_unique = AsyncMock(return_value=None) monkeypatch.setattr(proxy_server, "prisma_client", prisma) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index c1d5cedeba0..16698042361 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -9060,7 +9060,7 @@ async def test_load_servers_from_config_hydrates_dcr_clients(): ) hydrate_spy = AsyncMock() - with patch.object(global_mcp_server_manager, "_hydrate_config_servers_dcr_clients", new=hydrate_spy): + with patch.object(global_mcp_server_manager, "hydrate_config_servers_dcr_clients", new=hydrate_spy): await global_mcp_server_manager.load_servers_from_config({}) hydrate_spy.assert_awaited_once() @@ -9085,7 +9085,7 @@ async def test_reload_servers_from_database_hydrates_dcr_clients(): "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", return_value=prisma, ), - patch.object(global_mcp_server_manager, "_hydrate_config_servers_dcr_clients", new=hydrate_spy), + patch.object(global_mcp_server_manager, "hydrate_config_servers_dcr_clients", new=hydrate_spy), ): await global_mcp_server_manager.reload_servers_from_database() @@ -9618,7 +9618,7 @@ def test_oauth_endpoints_count_admin_entered_urls_as_resolved(): """A leftover issuer empties the resolved authorize/token fields but must not keep the server on the deferred-discovery retry path when the admin already stored those URLs.""" from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - _oauth_endpoints_unresolved, + oauth_endpoints_unresolved, ) from litellm.types.mcp import MCPAuth, MCPTransport from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -9635,7 +9635,7 @@ def test_oauth_endpoints_count_admin_entered_urls_as_resolved(): configured_authorization_url="https://github.com/login/oauth/authorize", configured_token_url="https://github.com/login/oauth/access_token", ) - assert _oauth_endpoints_unresolved(server) is False + assert oauth_endpoints_unresolved(server) is False @pytest.mark.asyncio @@ -11135,7 +11135,7 @@ def test_discovery_advertises_the_exchange_grant_only_where_the_gateway_can_serv litellm_jwtauth=LiteLLM_JWTAuth(virtual_key_claim_field=virtual_key_claim_field), ) monkeypatch.setattr("litellm.proxy.proxy_server.jwt_handler", handler) - monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"enable_jwt_auth": jwt_auth_enabled}) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"enable_jwt_auth": jwt_auth_enabled, "supported_db_objects": []}) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", object()) monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) exchange_grant = ["urn:ietf:params:oauth:grant-type:token-exchange"] if exchange_servable else [] diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index b715fe67e20..d80c5e84ee8 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -5535,7 +5535,7 @@ class TestMCPServerManagerReload: ): await manager.reload_servers_from_database() - mock_build.assert_awaited_once_with(db_row, env_vars_are_encrypted=True) + mock_build.assert_awaited_once_with(db_row, env_vars_are_encrypted=True, register_oauth_discovery=False) assert manager.registry["server-1"] is rebuilt_server @pytest.mark.asyncio @@ -5587,7 +5587,7 @@ class TestMCPServerManagerReload: "build_mcp_server_from_table", AsyncMock(side_effect=build_server), ), - patch.object(manager, "_maybe_register_openapi_tools", AsyncMock()), + patch.object(manager, "maybe_register_openapi_tools", AsyncMock()), caplog.at_level("ERROR", logger="LiteLLM"), ): await manager.reload_servers_from_database() @@ -5659,7 +5659,7 @@ class TestMCPServerManagerReload: ), patch.object( manager, - "_maybe_register_openapi_tools", + "maybe_register_openapi_tools", AsyncMock(side_effect=register_openapi_tools), ), caplog.at_level("ERROR", logger="LiteLLM"), @@ -9534,7 +9534,7 @@ class TestPreemptive401ModeAware: assert resolved.authorization_url == "https://idp.example.com/authorize" assert resolved.token_url == "https://idp.example.com/token" assert resolved.registration_url == "https://idp.example.com/register" - assert manager._oauth_discovery_slot(server.server_id) is None + assert manager.oauth_discovery_slot(server.server_id) is None assert exc.value.status_code == 401 @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 7725aca1948..3615d9681f0 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -44,7 +44,7 @@ from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( _deserialize_json_dict, _flow_endpoints_missing, _mcp_oauth_discovery_on_startup_enabled, - _oauth_endpoints_unresolved, + oauth_endpoints_unresolved, _deserialize_json_list, _normalize_mcp_server_cost_info, _obo_retry_applies, @@ -694,7 +694,7 @@ class TestMCPServerManager: assert resolved[0].scopes == ["mcp.read"] assert manager.config_mcp_servers[server.server_id] is resolved[0] assert server.authorization_url is None - assert manager._oauth_discovery_slot(server.server_id) is None + assert manager.oauth_discovery_slot(server.server_id) is None @pytest.mark.asyncio async def test_table_oauth_discovery_can_be_deferred_until_first_use(self): @@ -720,7 +720,7 @@ class TestMCPServerManager: server = await manager.build_mcp_server_from_table(row, credentials_are_encrypted=False) discovery.assert_not_awaited() - assert manager._oauth_discovery_slot(server.server_id) is not None + assert manager.oauth_discovery_slot(server.server_id) is not None manager.registry[server.server_id] = server with patch.object(manager, "_descovery_metadata", new=discovery): @@ -777,7 +777,7 @@ class TestMCPServerManager: assert len({id(resolution) for resolution in resolutions}) == 1 assert resolutions[0].authorization_url == "https://idp.example.com/authorize" assert resolutions[0].token_url == "https://idp.example.com/token" - assert manager._oauth_discovery_slot(server.server_id) is None + assert manager.oauth_discovery_slot(server.server_id) is None @pytest.mark.asyncio async def test_lazy_oauth_discovery_timeout_is_bounded(self): @@ -809,7 +809,7 @@ class TestMCPServerManager: assert exc.value.status_code == 503 assert "timed out" in str(exc.value.detail) discovery.assert_awaited_once_with(server) - assert manager._oauth_discovery_slot(server.server_id) is not None + assert manager.oauth_discovery_slot(server.server_id) is not None @pytest.mark.asyncio async def test_cancelling_one_waiter_does_not_cancel_shared_discovery(self): @@ -890,7 +890,7 @@ class TestMCPServerManager: assert replacement.token_url is None assert resolved.authorization_url == "https://idp.example.com/authorize" assert resolved.token_url == "https://idp.example.com/token" - assert manager._oauth_discovery_slot(replacement.server_id) is None + assert manager.oauth_discovery_slot(replacement.server_id) is None def test_registry_swap_reconcile_keeps_slot_for_issuer_anchored_server_without_url(self): manager = MCPServerManager() @@ -907,9 +907,9 @@ class TestMCPServerManager: manager.registry[server.server_id] = server manager._set_oauth_discovery_deferred(server.server_id, True) - manager._reconcile_oauth_discovery_slots_for_servers([server]) + manager.reconcile_oauth_discovery_slots_for_servers([server]) - assert manager._oauth_discovery_slot(server.server_id) is not None + assert manager.oauth_discovery_slot(server.server_id) is not None resolved = server.model_copy( update={ @@ -918,12 +918,12 @@ class TestMCPServerManager: } ) manager.registry[resolved.server_id] = resolved - manager._reconcile_oauth_discovery_slots_for_servers([resolved]) + manager.reconcile_oauth_discovery_slots_for_servers([resolved]) - assert manager._oauth_discovery_slot(server.server_id) is None + assert manager.oauth_discovery_slot(server.server_id) is None def _assert_oauth_discovery_state_removed(self, manager, server_id): - assert manager._oauth_discovery_slot(server_id) is None + assert manager.oauth_discovery_slot(server_id) is None @pytest.mark.asyncio async def test_deactivated_server_clears_lazy_oauth_discovery_state(self): @@ -965,7 +965,7 @@ class TestMCPServerManager: with ( patch( - "litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPServerRepository", + "litellm.proxy._experimental.mcp_server.db.MCPServerRepository", return_value=repository, ), patch( @@ -998,7 +998,7 @@ class TestMCPServerManager: manager.registry[server.server_id] = server previous_registry = manager.registry manager._set_oauth_discovery_deferred(server.server_id, True) - old_generation = manager._oauth_discovery_slot(server.server_id).generation + old_generation = manager.oauth_discovery_slot(server.server_id).generation resolved = server.model_copy( update={ "authorization_url": "https://idp.example.com/authorize", @@ -1017,7 +1017,8 @@ class TestMCPServerManager: raw_row = MagicMock() raw_row.model_dump.return_value = row.model_dump() repository = MagicMock() - repository.table.find_many = AsyncMock(return_value=[raw_row]) + other_row: Final = row.model_copy(update={"server_id": "changed-other", "server_name": "changed_other", "auth_type": MCPAuth.none}) + repository.table.find_many = AsyncMock(return_value=[raw_row, other_row]) async def publish_while_staged(*_args, **_kwargs): assert manager.registry is previous_registry @@ -1025,7 +1026,7 @@ class TestMCPServerManager: with ( patch( - "litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPServerRepository", + "litellm.proxy._experimental.mcp_server.db.MCPServerRepository", return_value=repository, ), patch( @@ -1034,16 +1035,16 @@ class TestMCPServerManager: ), patch.object( manager, - "_maybe_register_openapi_tools", + "maybe_register_openapi_tools", new=AsyncMock(side_effect=publish_while_staged), ), - patch.object(manager, "_prime_oauth_metadata_discovery_for_servers"), + patch.object(manager, "prime_oauth_metadata_discovery_for_servers"), ): await manager.reload_servers_from_database() assert previous_registry[server.server_id] is resolved assert manager.registry[server.server_id] is server - retry_slot = manager._oauth_discovery_slot(server.server_id) + retry_slot = manager.oauth_discovery_slot(server.server_id) assert retry_slot is not None assert retry_slot.generation > old_generation @@ -1143,7 +1144,7 @@ class TestMCPServerManager: assert manager.config_mcp_servers[server.server_id].authorization_url == "https://idp.example.com/authorize" assert manager.config_mcp_servers[server.server_id].token_url is None assert manager.config_mcp_servers[server.server_id].scopes is None - assert manager._oauth_discovery_slot(server.server_id) is not None + assert manager.oauth_discovery_slot(server.server_id) is not None @pytest.mark.asyncio async def test_create_mcp_client_triggers_deferred_oauth_discovery(self): @@ -1370,7 +1371,7 @@ class TestMCPServerManager: manager = MCPServerManager() with ( patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=None)), - patch.object(manager, "_hydrate_config_servers_dcr_clients", new=AsyncMock()), + patch.object(manager, "hydrate_config_servers_dcr_clients", new=AsyncMock()), caplog.at_level(logging.WARNING, logger="LiteLLM"), ): await manager.load_servers_from_config(self._id_jag_config()) @@ -1386,7 +1387,7 @@ class TestMCPServerManager: manager = MCPServerManager() with ( patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=None)), - patch.object(manager, "_hydrate_config_servers_dcr_clients", new=AsyncMock()), + patch.object(manager, "hydrate_config_servers_dcr_clients", new=AsyncMock()), caplog.at_level(logging.WARNING, logger="LiteLLM"), ): await manager.load_servers_from_config(self._id_jag_config()) @@ -1408,7 +1409,7 @@ class TestMCPServerManager: } with ( patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=None)), - patch.object(manager, "_hydrate_config_servers_dcr_clients", new=AsyncMock()), + patch.object(manager, "hydrate_config_servers_dcr_clients", new=AsyncMock()), caplog.at_level(logging.WARNING, logger="LiteLLM"), ): await manager.load_servers_from_config(config) @@ -1422,7 +1423,7 @@ class TestMCPServerManager: manager = MCPServerManager() with ( patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=None)), - patch.object(manager, "_hydrate_config_servers_dcr_clients", new=AsyncMock()), + patch.object(manager, "hydrate_config_servers_dcr_clients", new=AsyncMock()), caplog.at_level(logging.WARNING, logger="LiteLLM"), ): await manager.load_servers_from_config(self._id_jag_config()) @@ -7377,7 +7378,7 @@ class TestMCPServerTimestamps: with ( patch("litellm.proxy._experimental.mcp_server.db.update_mcp_server", new=update_mcp_server_mock), patch( - "litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPServerRepository", + "litellm.proxy._experimental.mcp_server.db.MCPServerRepository", return_value=repo_instance, ), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), @@ -7463,10 +7464,10 @@ class TestMCPServerTimestamps: client_id="cid", client_secret="csec", ) - assert _oauth_endpoints_unresolved(m2m_shaped) is False + assert oauth_endpoints_unresolved(m2m_shaped) is False interactive_unresolved = m2m_shaped.model_copy(update={"client_id": None, "client_secret": None}) - assert _oauth_endpoints_unresolved(interactive_unresolved) is True + assert oauth_endpoints_unresolved(interactive_unresolved) is True def test_dcr_bridge_relay_arm_needs_its_registration_endpoint(self): """A dcr_bridge server with no admin-configured client can only register callers through the @@ -7486,11 +7487,11 @@ class TestMCPServerTimestamps: token_url="https://idp.example.com/token", registration_url=None, ) - assert _oauth_endpoints_unresolved(relay_arm) is True + assert oauth_endpoints_unresolved(relay_arm) is True assert ( - _oauth_endpoints_unresolved(relay_arm.model_copy(update={"registration_url": "https://idp/reg"})) is False + oauth_endpoints_unresolved(relay_arm.model_copy(update={"registration_url": "https://idp/reg"})) is False ) - assert _oauth_endpoints_unresolved(relay_arm.model_copy(update={"client_id": "admin-client"})) is False + assert oauth_endpoints_unresolved(relay_arm.model_copy(update={"client_id": "admin-client"})) is False def test_entra_obo_without_scopes_is_unresolved(self): """entra_obo token exchange fails closed without a scope, and scopes can come from resource @@ -7506,9 +7507,9 @@ class TestMCPServerTimestamps: token_url="https://idp.example.com/token", scopes=None, ) - assert _oauth_endpoints_unresolved(entra) is True - assert _oauth_endpoints_unresolved(entra.model_copy(update={"scopes": ["api://app/.default"]})) is False - assert _oauth_endpoints_unresolved(entra.model_copy(update={"token_exchange_profile": "rfc8693"})) is False + assert oauth_endpoints_unresolved(entra) is True + assert oauth_endpoints_unresolved(entra.model_copy(update={"scopes": ["api://app/.default"]})) is False + assert oauth_endpoints_unresolved(entra.model_copy(update={"token_exchange_profile": "rfc8693"})) is False @pytest.mark.asyncio async def test_reload_fast_path_retries_unresolved_oauth_servers(self): @@ -7553,7 +7554,7 @@ class TestMCPServerTimestamps: build_mock = AsyncMock(return_value=previous_entry) with ( patch( - "litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPServerRepository", + "litellm.proxy._experimental.mcp_server.db.MCPServerRepository", return_value=repo_instance, ), patch( @@ -7646,7 +7647,7 @@ class TestMCPServerTimestamps: def test_carry_forward_skips_when_url_or_auth_type_changed(self): """Stale endpoints from a different upstream or auth mode must not carry forward.""" from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - _carry_forward_resolved_oauth_endpoints, + carry_forward_resolved_oauth_endpoints, ) def make_server(url: str, auth_type: MCPAuth, authorization_url: Optional[str]) -> MCPServer: @@ -7662,19 +7663,19 @@ class TestMCPServerTimestamps: previous = make_server("https://old.example.com/mcp", MCPAuth.oauth2, "https://idp.example.com/authorize") url_changed = make_server("https://new.example.com/mcp", MCPAuth.oauth2, None) - _carry_forward_resolved_oauth_endpoints(new_server=url_changed, previous_server=previous) + carry_forward_resolved_oauth_endpoints(new_server=url_changed, previous_server=previous) assert url_changed.authorization_url is None auth_changed = make_server("https://old.example.com/mcp", MCPAuth.true_passthrough, None) - _carry_forward_resolved_oauth_endpoints(new_server=auth_changed, previous_server=previous) + carry_forward_resolved_oauth_endpoints(new_server=auth_changed, previous_server=previous) assert auth_changed.authorization_url is None same = make_server("https://old.example.com/mcp", MCPAuth.oauth2, None) - _carry_forward_resolved_oauth_endpoints(new_server=same, previous_server=previous) + carry_forward_resolved_oauth_endpoints(new_server=same, previous_server=previous) assert same.authorization_url == "https://idp.example.com/authorize" explicit = make_server("https://old.example.com/mcp", MCPAuth.oauth2, "https://configured.example.com/auth") - _carry_forward_resolved_oauth_endpoints(new_server=explicit, previous_server=previous) + carry_forward_resolved_oauth_endpoints(new_server=explicit, previous_server=previous) assert explicit.authorization_url == "https://configured.example.com/auth" def test_carry_forward_does_not_revive_token_url_across_authorization_url_change(self): @@ -7685,7 +7686,7 @@ class TestMCPServerTimestamps: endpoint recreates the RFC 9700 mix-up, durably, and the discovery gate alone cannot catch it because the stale endpoint comes from the registry, not from discovery.""" from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - _carry_forward_resolved_oauth_endpoints, + carry_forward_resolved_oauth_endpoints, ) previous = MCPServer( @@ -7707,7 +7708,7 @@ class TestMCPServerTimestamps: authorization_url="https://idp-b.example.com/authorize", ) - _carry_forward_resolved_oauth_endpoints(new_server=repointed, previous_server=previous) + carry_forward_resolved_oauth_endpoints(new_server=repointed, previous_server=previous) assert repointed.authorization_url == "https://idp-b.example.com/authorize" assert repointed.token_url is None @@ -7719,7 +7720,7 @@ class TestMCPServerTimestamps: consistent group, and a rebuild that re-pins the same authorize endpoint (formatting aside) keeps carrying the corroborated token endpoint.""" from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - _carry_forward_resolved_oauth_endpoints, + carry_forward_resolved_oauth_endpoints, ) def previous() -> MCPServer: @@ -7742,7 +7743,7 @@ class TestMCPServerTimestamps: auth_type=MCPAuth.oauth2, authorization_url=None, ) - _carry_forward_resolved_oauth_endpoints(new_server=blipped, previous_server=previous()) + carry_forward_resolved_oauth_endpoints(new_server=blipped, previous_server=previous()) assert blipped.authorization_url == "https://idp.example.com/authorize" assert blipped.token_url == "https://idp.example.com/token" assert blipped.registration_url == "https://idp.example.com/register" @@ -7755,7 +7756,7 @@ class TestMCPServerTimestamps: auth_type=MCPAuth.oauth2, authorization_url="https://IDP.example.com:443/authorize/", ) - _carry_forward_resolved_oauth_endpoints(new_server=same_authorize, previous_server=previous()) + carry_forward_resolved_oauth_endpoints(new_server=same_authorize, previous_server=previous()) assert same_authorize.token_url == "https://idp.example.com/token" assert same_authorize.registration_url == "https://idp.example.com/register" @@ -7767,7 +7768,7 @@ class TestMCPServerTimestamps: scopes still carry as last-known-good. Anchoring is keyed on the explicit issuer_is_anchored flag, not on issuer truthiness, so a discovered issuer does not trip this fail-closed branch.""" from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - _carry_forward_resolved_oauth_endpoints, + carry_forward_resolved_oauth_endpoints, ) previous = MCPServer( @@ -7793,7 +7794,7 @@ class TestMCPServerTimestamps: issuer_is_anchored=True, ) - _carry_forward_resolved_oauth_endpoints(new_server=failed_rebuild, previous_server=previous) + carry_forward_resolved_oauth_endpoints(new_server=failed_rebuild, previous_server=previous) assert failed_rebuild.authorization_url is None assert failed_rebuild.token_url is None @@ -7807,7 +7808,7 @@ class TestMCPServerTimestamps: regression the explicit issuer_is_anchored flag prevents: keying fail-closed on issuer truthiness alone would drop the working endpoints the moment the server learned its issuer.""" from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - _carry_forward_resolved_oauth_endpoints, + carry_forward_resolved_oauth_endpoints, ) previous = MCPServer( @@ -7834,7 +7835,7 @@ class TestMCPServerTimestamps: authorization_url=None, ) - _carry_forward_resolved_oauth_endpoints(new_server=blipped_rebuild, previous_server=previous) + carry_forward_resolved_oauth_endpoints(new_server=blipped_rebuild, previous_server=previous) assert blipped_rebuild.authorization_url == "https://idp.example.com/authorize" assert blipped_rebuild.token_url == "https://idp.example.com/token" @@ -11975,7 +11976,7 @@ class TestClientForwardedDiscoveryFailureIsNotFatal: assert resolved is manager.config_mcp_servers[server.server_id] assert resolved.authorization_url is None assert resolved.token_url is None - assert manager._oauth_discovery_slot(server.server_id) is not None + assert manager.oauth_discovery_slot(server.server_id) is not None @pytest.mark.parametrize( "auth_type, serves_the_listing", @@ -12037,7 +12038,7 @@ class TestClientForwardedDiscoveryFailureIsNotFatal: assert resolved.token_url == "https://idp.example.com/token" assert resolved.registration_url == "https://idp.example.com/register" assert manager.config_mcp_servers[server.server_id].authorization_url == "https://idp.example.com/authorize" - assert manager._oauth_discovery_slot(server.server_id) is None + assert manager.oauth_discovery_slot(server.server_id) is None class TestResolveOpenapiToolAuth: @@ -12395,7 +12396,7 @@ class TestConfigServerIdPinning: ) with ( patch( # test-quality-ok: the db reload path has no seam but its own repository - "litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPServerRepository", + "litellm.proxy._experimental.mcp_server.db.MCPServerRepository", return_value=repository, ), patch( # test-quality-ok: same, the prisma client is fetched inside the reload @@ -13316,7 +13317,7 @@ def test_stale_discovery_cannot_overwrite_new_registered_server() -> None: transport=MCPTransport.http, auth_type=MCPAuth.oauth2, ) manager._set_oauth_discovery_deferred(original.server_id, True) - original_slot: Final = manager._oauth_discovery_slot(original.server_id) + original_slot: Final = manager.oauth_discovery_slot(original.server_id) assert original_slot is not None replacement: Final = original.model_copy(update={"url": "https://new.example.com/mcp"}) manager.registry[original.server_id] = replacement @@ -13335,30 +13336,30 @@ async def test_temporary_oauth_discovery_expires_without_more_requests() -> None ) manager._set_oauth_discovery_deferred(server.server_id, True) resolved: Final = await manager.ensure_oauth_metadata_discovered(server) - assert manager._oauth_discovery_slot(server.server_id) is not None + assert manager.oauth_discovery_slot(server.server_id) is not None loop: Final = asyncio.get_running_loop() expired: Final = loop.create_future() with patch.object(loop, "time", return_value=loop.time() + 301): loop.call_later(0, expired.set_result, None) await expired assert resolved.authorization_url == server.authorization_url - assert manager._oauth_discovery_slot(server.server_id) is None + assert manager.oauth_discovery_slot(server.server_id) is None def test_old_temporary_discovery_expiry_preserves_replacement() -> None: manager: Final = MCPServerManager() manager._set_oauth_discovery_deferred("reused-session", True) - old_slot: Final = manager._oauth_discovery_slot("reused-session") + old_slot: Final = manager.oauth_discovery_slot("reused-session") assert old_slot is not None manager._set_oauth_discovery_deferred("reused-session", True) - replacement: Final = manager._oauth_discovery_slot("reused-session") + replacement: Final = manager.oauth_discovery_slot("reused-session") manager._expire_temporary_oauth_discovery("reused-session", old_slot.generation) - assert manager._oauth_discovery_slot("reused-session") is replacement + assert manager.oauth_discovery_slot("reused-session") is replacement assert replacement is not None manager._expire_temporary_oauth_discovery("reused-session", replacement.generation) - assert manager._oauth_discovery_slot("reused-session") is None + assert manager.oauth_discovery_slot("reused-session") is None manager._expire_temporary_oauth_discovery("reused-session", replacement.generation) - assert manager._oauth_discovery_slot("reused-session") is None + assert manager.oauth_discovery_slot("reused-session") is None @pytest.mark.asyncio @@ -13719,12 +13720,12 @@ async def test_discovery_cache_invalidation_during_fetch_does_not_repopulate_old with _mcp_upstream(upstream.respond): task: Final = asyncio.create_task(manager.get_prompts_from_server(_discovery_server(), None)) await asyncio.wait_for(upstream.entered.wait(), timeout=5) - manager._invalidate_discovery_lists("discovery") + manager.invalidate_discovery_lists("discovery") upstream.release.set() assert (await task)[0].name == "discovery-example" assert len(await manager.get_prompts_from_server(_discovery_server(), None)) == 1 assert upstream.initializes == 2 - manager._invalidate_discovery_lists("discovery") + manager.invalidate_discovery_lists("discovery") assert len(await manager.get_prompts_from_server(_discovery_server(), None)) == 1 assert upstream.initializes == 3 @@ -14516,3 +14517,331 @@ async def test_client_sampling_does_not_fill_explicit_context_from_another_ambie assert captured["client_ip"] is None finally: auth_context_var.reset(token) + + +@pytest.mark.asyncio +async def test_catalog_observes_committed_update_and_delete_without_background_reload(): + from datetime import timedelta + + timestamp: Final = datetime.now() + row: Final = LiteLLM_MCPServerTable( + server_id="catalog-server", server_name="catalog_server", alias="catalog_server", + transport=MCPTransport.http, url="https://first.example.com/mcp", updated_at=timestamp, + ) + updated: Final = row.model_copy(update={"url": "https://second.example.com/mcp", "updated_at": timestamp + timedelta(seconds=1)}) + prisma: Final = MagicMock() + prisma.db.litellm_mcpservertable.find_many = AsyncMock(side_effect=([row], [updated], [])) + manager: Final = MCPServerManager() + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.proxy_server.general_settings", {}), + ): + first = await manager.catalog.resolve(row.server_id) + manager.registry[row.server_id].short_prefix = "a12" + second = await manager.catalog.resolve(row.server_id) + deleted = await manager.catalog.resolve(row.server_id) + + assert first is not None and first.url == row.url + assert second is not None and second.url == updated.url + assert second.short_prefix == "a12" + assert deleted is None + assert prisma.db.litellm_mcpservertable.find_many.await_count == 3 + + +@pytest.mark.asyncio +async def test_catalog_lookup_uses_one_snapshot_until_operation_finishes(): + from datetime import timedelta + + row: Final = LiteLLM_MCPServerTable( + server_id="snapshot-server", alias="snapshot_server", transport=MCPTransport.http, + url="https://first.example.com/mcp", updated_at=datetime.now(), + ) + updated: Final = row.model_copy(update={"url": "https://second.example.com/mcp", "updated_at": row.updated_at + timedelta(seconds=1)}) + prisma: Final = MagicMock() + prisma.db.litellm_mcpservertable.find_many = AsyncMock(side_effect=([row], [updated], [updated])) + manager: Final = MCPServerManager() + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.proxy_server.general_settings", {}), + ): + async with manager.catalog.operation() as before: + await manager.reload_servers_from_database() + during = await manager.catalog.resolve(row.server_id) + assert during is not None and during.url == row.url + assert manager.registry[row.server_id].url == updated.url + async with manager.catalog.operation() as after: + current = manager.get_mcp_server_by_id(row.server_id) + assert current is not None and current.url == updated.url + assert before.identity != after.identity + assert prisma.db.litellm_mcpservertable.find_many.await_count == 3 + + +@pytest.mark.asyncio +async def test_catalog_failed_reload_preserves_published_discovery_state(): + manager: Final = MCPServerManager() + server: Final = MCPServer(server_id="healthy", name="healthy", transport=MCPTransport.http) + manager.registry = {server.server_id: server} + manager._upstream_initialize_instructions_by_server_id[server.server_id] = "healthy instructions" + prisma: Final = MagicMock() + prisma.db.litellm_mcpservertable.find_many = AsyncMock(side_effect=RuntimeError("database unavailable")) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + with pytest.raises(RuntimeError, match="database unavailable"): + await manager.reload_servers_from_database() + assert manager.get_mcp_server_by_id(server.server_id) is server + assert manager._upstream_initialize_instructions_by_server_id == {server.server_id: "healthy instructions"} + + +@pytest.mark.asyncio +async def test_catalog_cancellation_retains_state_and_releases_refresh_lock(): + entered: Final = asyncio.Event() + release: Final = asyncio.Event() + + async def blocked_read(**kwargs: object) -> list[LiteLLM_MCPServerTable]: + entered.set() + await release.wait() + return [] + + manager: Final = MCPServerManager() + server: Final = MCPServer(server_id="healthy", name="healthy", transport=MCPTransport.http) + manager.registry = {server.server_id: server} + manager._upstream_initialize_instructions_by_server_id[server.server_id] = "healthy instructions" + prisma: Final = MagicMock() + prisma.db.litellm_mcpservertable.find_many = AsyncMock(side_effect=blocked_read) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + pending = asyncio.create_task(manager.reload_servers_from_database()) + await asyncio.wait_for(entered.wait(), timeout=2) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + assert manager.registry == {server.server_id: server} + assert manager._upstream_initialize_instructions_by_server_id == {server.server_id: "healthy instructions"} + release.set() + await asyncio.wait_for(manager.reload_servers_from_database(), timeout=2) + assert manager.registry == {} + + +@pytest.mark.asyncio +async def test_catalog_background_lookup_after_operation_exit_observes_deletion(): + manager: Final = MCPServerManager() + server: Final = MCPServer(server_id="expired-snapshot", name="expired_snapshot", transport=MCPTransport.http) + manager.registry = {server.server_id: server} + release: Final = asyncio.Event() + + async def lookup_after_exit() -> MCPServer | None: + await release.wait() + return await manager.catalog.resolve(server.server_id) + + with patch("litellm.proxy.proxy_server.prisma_client", None): + async with manager.catalog.operation(): + pending: Final = asyncio.create_task(lookup_after_exit()) + manager.registry = {} + release.set() + assert await pending is None + + +@pytest.mark.asyncio +async def test_catalog_cancelled_openapi_refresh_retains_tools_and_discovery(monkeypatch): + from datetime import timedelta + from litellm.proxy._experimental.mcp_server import tool_registry + + manager: Final = MCPServerManager() + stamp: Final = datetime.now() + server: Final = MCPServer(server_id="staged", name="staged", transport=MCPTransport.http, + url="https://before.example/mcp", spec_path="before.json", updated_at=stamp, auth_type=MCPAuth.oauth2) + manager.registry = {server.server_id: server} + manager._set_oauth_discovery_deferred(server.server_id, True) + original_slot: Final = manager.oauth_discovery_slot(server.server_id) + manager._upstream_initialize_instructions_by_server_id[server.server_id] = "keep instructions" + manager.tool_name_to_mcp_server_name_mapping = {"staged-existing": "staged"} + registry: Final = tool_registry.MCPToolRegistry() + registry.register_tool("staged-existing", "existing", {}, lambda: "existing") + monkeypatch.setattr(tool_registry, "global_mcp_tool_registry", registry) + row: Final = LiteLLM_MCPServerTable(server_id=server.server_id, alias="staged", transport=MCPTransport.http, + url="https://after.example/mcp", spec_path="after.json", updated_at=stamp + timedelta(seconds=1), auth_type=MCPAuth.oauth2) + prisma: Final = MagicMock() + prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[row]) + + async def cancelled_registration(*args: object, **kwargs: object) -> None: + registry.register_tool("staged-new", "new", {}, lambda: "new") + raise asyncio.CancelledError() + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch.object(manager, "maybe_register_openapi_tools", side_effect=cancelled_registration), + patch.object(manager, "invalidate_discovery_lists") as invalidate, + ): + with pytest.raises(asyncio.CancelledError): + await manager.reload_servers_from_database() + invalidate.assert_not_called() + assert manager.registry == {server.server_id: server} + assert [tool.name for tool in registry.list_tools()] == ["staged-existing"] + assert manager.tool_name_to_mcp_server_name_mapping == {"staged-existing": "staged"} + assert manager._upstream_initialize_instructions_by_server_id == {server.server_id: "keep instructions"} + + assert manager.oauth_discovery_slot(server.server_id) is original_slot + + +@pytest.mark.asyncio +async def test_catalog_snapshot_identity_is_independent_of_worker_oauth_discovery(): + row: Final = LiteLLM_MCPServerTable(server_id="identity-server", alias="identity_server", transport=MCPTransport.http, + url="https://upstream.example/mcp", auth_type=MCPAuth.oauth2, updated_at=datetime.now()) + prisma: Final = MagicMock() + prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[row]) + manager: Final = MCPServerManager() + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.proxy_server.general_settings", {}), + patch.object(manager, "prime_oauth_metadata_discovery_for_servers"), + ): + async with manager.catalog.operation() as unresolved: + assert manager.get_mcp_server_by_id(row.server_id).authorization_url is None + resolved: Final = manager.registry[row.server_id].model_copy(update={ + "authorization_url": "https://idp.example/authorize", "token_url": "https://idp.example/token"}) + manager.registry[row.server_id] = resolved + async with manager.catalog.operation() as discovered: + assert manager.get_mcp_server_by_id(row.server_id).authorization_url == resolved.authorization_url + assert unresolved.identity == discovered.identity + + +@pytest.mark.asyncio +async def test_catalog_failed_openapi_row_does_not_publish_partial_handlers(monkeypatch): + from litellm.proxy._experimental.mcp_server import tool_registry + + manager: Final = MCPServerManager() + registry: Final = tool_registry.MCPToolRegistry() + monkeypatch.setattr(tool_registry, "global_mcp_tool_registry", registry) + row: Final = LiteLLM_MCPServerTable(server_id="broken", alias="broken", transport=MCPTransport.http, + url="https://upstream.example/mcp", spec_path="broken.json") + prisma: Final = MagicMock() + prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[row]) + + async def broken_registration(server: MCPServer, **kwargs: object) -> None: + registry.register_tool("broken-partial", "partial", {}, lambda: "must not run") + manager.tool_name_to_mcp_server_name_mapping["broken-partial"] = "broken" + raise ValueError("invalid remaining operation") + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch.object(manager, "maybe_register_openapi_tools", side_effect=broken_registration), + ): + await manager.reload_servers_from_database() + assert manager.registry == {} + assert registry.list_tools() == [] + assert manager.tool_name_to_mcp_server_name_mapping == {} + + +@pytest.mark.asyncio +async def test_catalog_cancelled_config_hydration_preserves_published_credentials(): + manager: Final = MCPServerManager() + server: Final = MCPServer(server_id="config-hydration", name="config_hydration", transport=MCPTransport.http, + client_id="previous-client") + manager.config_mcp_servers = {server.server_id: server} + + async def cancelled_hydration(target: MCPServer) -> bool: + target.client_id = "unpublished-client" + raise asyncio.CancelledError() + + with patch("litellm.proxy._experimental.mcp_server.discoverable_endpoints.hydrate_config_server_dcr_client", side_effect=cancelled_hydration): + with pytest.raises(asyncio.CancelledError): + await manager.reload_servers_from_database() + assert manager.get_mcp_server_by_id(server.server_id).client_id == "previous-client" + + +@pytest.mark.asyncio +async def test_catalog_oauth_resolution_cannot_replace_an_operation_target_after_update(): + manager: Final = MCPServerManager() + server: Final = MCPServer(server_id="pinned-oauth", name="pinned_oauth", transport=MCPTransport.http, + url="https://before.example/mcp", auth_type=MCPAuth.oauth2, + authorization_url="https://before.example/authorize", token_url="https://before.example/token") + manager.registry = {server.server_id: server} + with patch("litellm.proxy.proxy_server.prisma_client", None): + async with manager.catalog.operation(): + selected: Final = manager.get_mcp_server_by_id(server.server_id) + manager.registry[server.server_id] = server.model_copy(update={ + "url": "https://after.example/mcp", "authorization_url": "https://after.example/authorize"}) + resolved: Final = await manager.ensure_oauth_metadata_discovered(selected) + assert resolved.url == "https://before.example/mcp" + assert resolved.authorization_url == "https://before.example/authorize" + assert manager.registry[server.server_id].url == "https://after.example/mcp" + + +@pytest.mark.asyncio +async def test_catalog_unresolved_oauth_snapshot_fails_closed_if_target_was_replaced(): + manager: Final = MCPServerManager() + server: Final = MCPServer(server_id="replaced-oauth", name="replaced_oauth", transport=MCPTransport.http, + url="https://before.example/mcp", auth_type=MCPAuth.oauth2) + manager.registry = {server.server_id: server} + with ( + patch("litellm.proxy.proxy_server.prisma_client", None), + patch.object(manager, "_discover_oauth_metadata_for_server", new_callable=AsyncMock) as discovery, + ): + async with manager.catalog.operation(): + selected: Final = manager.get_mcp_server_by_id(server.server_id) + manager.registry[server.server_id] = server.model_copy(update={ + "url": "https://after.example/mcp", "authorization_url": "https://after.example/authorize", + "token_url": "https://after.example/token"}) + with pytest.raises(HTTPException) as exc: + await manager.ensure_oauth_metadata_discovered(selected) + assert exc.value.status_code == 503 + discovery.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_catalog_unresolved_oauth_snapshot_accepts_discovery_for_the_same_target(): + manager: Final = MCPServerManager() + server: Final = MCPServer(server_id="same-oauth", name="same_oauth", transport=MCPTransport.http, + url="https://upstream.example/mcp", auth_type=MCPAuth.oauth2) + manager.registry = {server.server_id: server} + discovered: Final = server.model_copy(update={"authorization_url": "https://issuer.example/authorize", + "token_url": "https://issuer.example/token", "scopes": ["read"]}) + with ( + patch("litellm.proxy.proxy_server.prisma_client", None), + patch.object(manager, "_ensure_oauth_metadata_discovered", return_value=discovered) as discovery, + ): + async with manager.catalog.operation(): + resolved: Final = await manager.ensure_oauth_metadata_discovered(server) + assert resolved.authorization_url == "https://issuer.example/authorize" + assert resolved.token_url == "https://issuer.example/token" + assert resolved.scopes == ["read"] + discovery.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_catalog_fresh_lookup_does_not_fall_back_to_stale_grants_when_database_fails(): + manager: Final = MCPServerManager() + server: Final = MCPServer(server_id="stale-grant", name="stale_grant", transport=MCPTransport.http, + allow_all_keys=True) + manager.registry = {server.server_id: server} + prisma: Final = MagicMock() + prisma.db.litellm_mcpservertable.find_many = AsyncMock(side_effect=RuntimeError("unavailable")) + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.proxy_server.general_settings", {}), + ): + with pytest.raises(RuntimeError, match="unavailable"): + await manager.catalog.resolve(server.server_id) + assert manager.registry == {server.server_id: server} + + +@pytest.mark.asyncio +async def test_catalog_list_failure_cancels_and_joins_other_upstream_fetches(): + from litellm.proxy._experimental.mcp_server.catalog import TargetCatalog + + started: Final = asyncio.Event() + stopped: Final = asyncio.Event() + pending: Final = MCPServer(server_id="pending", name="pending", transport=MCPTransport.http) + broken: Final = MCPServer(server_id="broken", name="broken", transport=MCPTransport.http) + + async def fetch(server: MCPServer): + if server.server_id == broken.server_id: + await started.wait() + raise RuntimeError("unexpected fetch failure") + started.set() + try: + await asyncio.Event().wait() + finally: + stopped.set() + + with pytest.raises(RuntimeError, match="unexpected fetch failure"): + await TargetCatalog.list((pending, broken), fetch, lambda server: server.server_id) + assert stopped.is_set() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index 233a8cc96ba..676cc31fe05 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -1559,9 +1559,8 @@ class TestListToolsRestAPI: ) monkeypatch.setattr( rest_endpoints.global_mcp_server_manager, - "get_mcp_server_by_id", - lambda server_id: stub_server if server_id == "server-1" else None, - raising=False, + "registry", + {stub_server.server_id: stub_server}, ) request = _build_request(path="/mcp-rest/tools/list", method="GET") diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py index 941e5deee93..7b6007c405e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py @@ -345,7 +345,7 @@ class TestManagerShortPrefix: class TestShortPrefixCollisionResolution: - """``_assign_unique_short_prefix`` must rehash on collision. + """``assign_unique_short_prefix`` must rehash on collision. The dedup path is exercised by forcing two distinct ``server_id`` values to both hash to the same natural prefix via a monkeypatched @@ -355,7 +355,7 @@ class TestShortPrefixCollisionResolution: def test_no_op_when_flag_off(self): manager = MCPServerManager() server = _make_server(server_id="abc") - manager._assign_unique_short_prefix(server) + manager.assign_unique_short_prefix(server) assert server.short_prefix is None def test_assigns_natural_hash_when_no_collision(self, monkeypatch): @@ -364,7 +364,7 @@ class TestShortPrefixCollisionResolution: monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true") manager = MCPServerManager() server = _make_server(server_id="abc") - manager._assign_unique_short_prefix(server) + manager.assign_unique_short_prefix(server) assert server.short_prefix == mcp_utils.compute_short_server_prefix("abc") @@ -394,9 +394,9 @@ class TestShortPrefixCollisionResolution: # Pretend both are already in the registry so dedup sees both. manager.registry[first.server_id] = first - manager._assign_unique_short_prefix(first) + manager.assign_unique_short_prefix(first) manager.registry[second.server_id] = second - manager._assign_unique_short_prefix(second) + manager.assign_unique_short_prefix(second) assert first.short_prefix == "AAA" assert second.short_prefix == "AAB" @@ -408,7 +408,7 @@ class TestShortPrefixCollisionResolution: server = _make_server(server_id="abc") server.short_prefix = "ZZZ" # pretend a previous registration set this - manager._assign_unique_short_prefix(server) + manager.assign_unique_short_prefix(server) assert server.short_prefix == "ZZZ" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_mcp_security.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_mcp_security.py index d57a91d45bf..7ed0596fbe2 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_mcp_security.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_mcp_security.py @@ -205,3 +205,30 @@ class TestInitializeGuardrail: assert isinstance(result, MCPSecurityGuardrail) assert result.on_violation == expected assert result in litellm.callbacks + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_type", ["acompletion", "aresponses"]) +async def test_guardrail_observes_saved_server_creation_and_deletion_on_another_worker(guardrail, call_type): + from unittest.mock import AsyncMock + from litellm.proxy._types import LiteLLM_MCPServerTable + from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager + + manager = MCPServerManager() + row = LiteLLM_MCPServerTable(server_id="peer-server", alias="peer_server", transport="http", + url="https://upstream.example/mcp") + prisma = MagicMock() + prisma.db.litellm_mcpservertable.find_many = AsyncMock(side_effect=([row], [])) + data = {"tools": [{"type": "mcp", "server_url": "litellm_proxy/mcp/peer-server"}], + "guardrails": ["test-mcp-security"]} + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.proxy_server.general_settings", {}), + patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", manager), + ): + result = await guardrail.async_pre_call_hook(UserAPIKeyAuth(), MagicMock(), data, call_type) + assert result == data + with pytest.raises(HTTPException) as exc: + await guardrail.async_pre_call_hook(UserAPIKeyAuth(), MagicMock(), data, call_type) + assert exc.value.status_code == 400 + assert exc.value.detail["unregistered_servers"] == ["peer-server"] diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index 53645e62034..e21dce01119 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -1,3 +1,4 @@ +from contextlib import nullcontext import os import sys import types @@ -2629,6 +2630,7 @@ class TestTemporaryMCPSessionEndpoints: mock_manager = MagicMock() mock_manager.get_mcp_server_by_id.return_value = non_oauth_server mock_manager.get_mcp_server_by_name.return_value = None + mock_manager.catalog.resolve = AsyncMock(return_value=non_oauth_server) fake_proxy_server = types.SimpleNamespace(master_key=None) with ( @@ -2681,6 +2683,7 @@ class TestTemporaryMCPSessionEndpoints: mock_manager = MagicMock() mock_manager.get_mcp_server_by_id.return_value = internal_server mock_manager.get_mcp_server_by_name.return_value = None + mock_manager.catalog.resolve = AsyncMock(return_value=internal_server) fake_proxy_server = types.SimpleNamespace(master_key=None) with ( @@ -2738,6 +2741,63 @@ class TestTemporaryMCPSessionEndpoints: } assert dependency_names == {None, "user_api_key_dict"} + @pytest.mark.asyncio + async def test_authorize_saved_server_on_cold_worker_without_oauth_session(self): + from collections.abc import Mapping + + from starlette.requests import Request + + from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager + from litellm.proxy.management_endpoints.mcp_management_endpoints import mcp_authorize + + row: Final = LiteLLM_MCPServerTable( + server_id="saved-server", + server_name="saved_server", + alias="saved_server", + transport=MCPTransport.http, + url="https://upstream.example.com/mcp", + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + authorization_url="https://upstream.example.com/authorize", + token_url="https://upstream.example.com/token", + approval_status="active", + ) + manager: Final = MCPServerManager() + prisma: Final = MagicMock() + + async def persisted_rows(*, where: Mapping[str, object]) -> list[LiteLLM_MCPServerTable]: + return [] if where.get("approval_status") == "draft" else [row] + + prisma.db.litellm_mcpservertable.find_many = AsyncMock(side_effect=persisted_rows) + prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=None) + prisma.db.litellm_mcpservertable.find_first = AsyncMock(return_value=None) + request: Final = Request( + {"type": "http", "method": "GET", "scheme": "http", "server": ("localhost", 4000), + "path": "/v1/mcp/server/oauth/saved-server/authorize", "headers": [], + "query_string": b"", "client": ("127.0.0.1", 1234)} + ) + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.proxy_server.general_settings", {}), + patch("litellm.proxy.proxy_server.master_key", "sk-unit-test-catalog"), + patch.object(mgmt_endpoints, "global_mcp_server_manager", manager), + ): + response = await mcp_authorize( + request=request, + server_id=row.server_id, + user_api_key_dict=generate_mock_user_api_key_auth(), + client_id="client-id", + redirect_uri="http://localhost:9876/callback", + state="saved-server-test", + code_challenge=None, + code_challenge_method=None, + response_type="code", + scope=None, + ) + + assert response.status_code == 307 + assert response.headers["location"].startswith("https://upstream.example.com/authorize?") + @pytest.mark.asyncio async def test_mcp_authorize_proxies_to_discoverable_endpoint(self): from litellm.proxy.management_endpoints.mcp_management_endpoints import ( @@ -2754,8 +2814,8 @@ class TestTemporaryMCPSessionEndpoints: with ( patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._get_cached_temporary_mcp_server_or_404", - return_value=server, + "litellm.proxy.management_endpoints.mcp_management_endpoints._oauth_server_operation", + return_value=nullcontext(server), ) as get_server, patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.authorize_with_server", @@ -2776,7 +2836,7 @@ class TestTemporaryMCPSessionEndpoints: ) assert result is authorize_response - get_server.assert_awaited_once_with("server-1", admin_auth, request=request) + get_server.assert_called_once_with("server-1", admin_auth, request=request) authorize_mock.assert_awaited_once_with( request=request, mcp_server=server, @@ -2804,8 +2864,8 @@ class TestTemporaryMCPSessionEndpoints: admin_auth = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN) patches = [ patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._get_cached_temporary_mcp_server_or_404", - return_value=server, + "litellm.proxy.management_endpoints.mcp_management_endpoints._oauth_server_operation", + return_value=nullcontext(server), ), patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.authorize_with_server", @@ -2909,8 +2969,8 @@ class TestTemporaryMCPSessionEndpoints: with ( patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._get_cached_temporary_mcp_server_or_404", - return_value=server, + "litellm.proxy.management_endpoints.mcp_management_endpoints._oauth_server_operation", + return_value=nullcontext(server), ), patch( "litellm.proxy._experimental.mcp_server.discoverable_endpoints.mint_ephemeral_dcr_client", @@ -3055,8 +3115,8 @@ class TestTemporaryMCPSessionEndpoints: with ( patch("litellm.proxy.proxy_server.master_key", "sk-lit4581-test-master-key"), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._get_cached_temporary_mcp_server_or_404", - return_value=server, + "litellm.proxy.management_endpoints.mcp_management_endpoints._oauth_server_operation", + return_value=nullcontext(server), ), patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.exchange_token_with_server", @@ -3114,8 +3174,8 @@ class TestTemporaryMCPSessionEndpoints: with ( patch("litellm.proxy.proxy_server.master_key", "sk-lit4581-test-master-key"), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._get_cached_temporary_mcp_server_or_404", - return_value=server, + "litellm.proxy.management_endpoints.mcp_management_endpoints._oauth_server_operation", + return_value=nullcontext(server), ), patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.exchange_token_with_server", @@ -3168,8 +3228,8 @@ class TestTemporaryMCPSessionEndpoints: with ( patch("litellm.proxy.proxy_server.master_key", "sk-lit4581-test-master-key"), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._get_cached_temporary_mcp_server_or_404", - return_value=server, + "litellm.proxy.management_endpoints.mcp_management_endpoints._oauth_server_operation", + return_value=nullcontext(server), ), patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.exchange_token_with_server", @@ -3218,8 +3278,8 @@ class TestTemporaryMCPSessionEndpoints: with ( patch("litellm.proxy.proxy_server.master_key", "sk-lit4581-test-master-key"), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._get_cached_temporary_mcp_server_or_404", - return_value=server, + "litellm.proxy.management_endpoints.mcp_management_endpoints._oauth_server_operation", + return_value=nullcontext(server), ), patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.exchange_token_with_server", @@ -3266,8 +3326,8 @@ class TestTemporaryMCPSessionEndpoints: with ( patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._get_cached_temporary_mcp_server_or_404", - return_value=server, + "litellm.proxy.management_endpoints.mcp_management_endpoints._oauth_server_operation", + return_value=nullcontext(server), ), patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.authorize_with_server", @@ -3306,8 +3366,8 @@ class TestTemporaryMCPSessionEndpoints: with ( patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._get_cached_temporary_mcp_server_or_404", - return_value=server, + "litellm.proxy.management_endpoints.mcp_management_endpoints._oauth_server_operation", + return_value=nullcontext(server), ), patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.exchange_token_with_server", @@ -3351,8 +3411,8 @@ class TestTemporaryMCPSessionEndpoints: with ( patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._get_cached_temporary_mcp_server_or_404", - return_value=server, + "litellm.proxy.management_endpoints.mcp_management_endpoints._oauth_server_operation", + return_value=nullcontext(server), ) as get_server, patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.exchange_token_with_server", @@ -3374,7 +3434,7 @@ class TestTemporaryMCPSessionEndpoints: ) assert result is exchange_response - get_server.assert_awaited_once_with("server-1", admin_auth, request=request) + get_server.assert_called_once_with("server-1", admin_auth, request=request) exchange_mock.assert_awaited_once_with( request=request, mcp_server=server, @@ -3405,8 +3465,8 @@ class TestTemporaryMCPSessionEndpoints: with ( patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._get_cached_temporary_mcp_server_or_404", - return_value=server, + "litellm.proxy.management_endpoints.mcp_management_endpoints._oauth_server_operation", + return_value=nullcontext(server), ) as get_server, patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.exchange_token_with_server", @@ -3428,7 +3488,7 @@ class TestTemporaryMCPSessionEndpoints: ) assert result is exchange_response - get_server.assert_awaited_once_with("server-1", admin_auth, request=request) + get_server.assert_called_once_with("server-1", admin_auth, request=request) exchange_mock.assert_awaited_once_with( request=request, mcp_server=server, @@ -3465,8 +3525,8 @@ class TestTemporaryMCPSessionEndpoints: with ( patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._get_cached_temporary_mcp_server_or_404", - return_value=server, + "litellm.proxy.management_endpoints.mcp_management_endpoints._oauth_server_operation", + return_value=nullcontext(server), ) as get_server, patch( "litellm.proxy.management_endpoints.mcp_management_endpoints._read_request_body", @@ -3484,7 +3544,7 @@ class TestTemporaryMCPSessionEndpoints: ) assert result is register_response - get_server.assert_awaited_once_with("server-1", admin_auth, request=request) + get_server.assert_called_once_with("server-1", admin_auth, request=request) read_body.assert_awaited_once_with(request=request) register_mock.assert_awaited_once_with( request=request, @@ -3528,8 +3588,8 @@ class TestTemporaryMCPSessionEndpoints: with ( patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._get_cached_temporary_mcp_server_or_404", - return_value=server, + "litellm.proxy.management_endpoints.mcp_management_endpoints._oauth_server_operation", + return_value=nullcontext(server), ), patch( "litellm.proxy.management_endpoints.mcp_management_endpoints._read_request_body", @@ -3573,8 +3633,8 @@ class TestTemporaryMCPSessionEndpoints: with ( patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._get_cached_temporary_mcp_server_or_404", - return_value=server, + "litellm.proxy.management_endpoints.mcp_management_endpoints._oauth_server_operation", + return_value=nullcontext(server), ), patch( "litellm.proxy.management_endpoints.mcp_management_endpoints._read_request_body", @@ -7940,3 +8000,33 @@ async def test_config_server_edit_preserves_api_contract_without_creating_rows(r prisma.tx.assert_not_called() assert server.model_dump() == original assert manager.registry == {} + + +@pytest.mark.asyncio +async def test_saved_server_authorize_denial_does_not_dispatch_upstream(): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager + + manager: Final = MCPServerManager() + row: Final = LiteLLM_MCPServerTable(server_id="denied-peer", alias="denied_peer", transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, url="https://upstream.example/mcp", + authorization_url="https://upstream.example/authorize", token_url="https://upstream.example/token") + prisma: Final = MagicMock() + prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[row]) + user: Final = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.INTERNAL_USER) + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.proxy_server.general_settings", {}), + patch.object(mgmt_endpoints, "global_mcp_server_manager", manager), + patch.object(mgmt_endpoints, "get_cached_temporary_mcp_server", AsyncMock(return_value=None)), + patch.object(mgmt_endpoints, "build_effective_auth_contexts", AsyncMock(return_value=(user,))), + patch.object(manager, "get_allowed_mcp_servers", AsyncMock(return_value=[])), + patch.object(mgmt_endpoints, "authorize_with_server", new_callable=AsyncMock) as authorize, + patch.object(mgmt_endpoints, "resolve_ephemeral_dcr_client", new_callable=AsyncMock) as register, + ): + with pytest.raises(HTTPException) as exc: + await mgmt_endpoints.mcp_authorize(request=None, server_id=row.server_id, user_api_key_dict=user, + client_id="client", redirect_uri="http://localhost/callback") + assert exc.value.status_code == 403 + authorize.assert_not_awaited() + register.assert_not_awaited() + assert prisma.db.litellm_mcpservertable.find_many.await_count == 1 diff --git a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py index 83537c236a3..9c6b2404d1e 100644 --- a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py +++ b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -1,3 +1,5 @@ +from contextlib import nullcontext + import subprocess import sys import textwrap @@ -32,6 +34,7 @@ def _setup_mcp_call_environment(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_module) fake_manager = types.SimpleNamespace( + catalog=types.SimpleNamespace(operation=nullcontext), get_registry=MagicMock(return_value={}), call_tool=AsyncMock(return_value=_DummyMCPResult()), # Newer logging path calls this to enrich spend logs metadata @@ -378,6 +381,7 @@ async def test_execute_tool_calls_logs_failure_via_post_call_failure_hook(monkey post_call_failure_hook = _setup_proxy_logging(monkeypatch) fake_manager = types.SimpleNamespace( + catalog=types.SimpleNamespace(operation=nullcontext), get_registry=MagicMock(return_value={}), call_tool=AsyncMock(side_effect=HTTPException(status_code=500, detail="boom")) ) @@ -506,6 +510,7 @@ async def test_get_mcp_tools_from_manager_enables_list_tools_logging(monkeypatch # Patch manager methods used by _get_mcp_tools_from_manager to avoid needing full UserAPIKeyAuth fields. fake_manager = types.SimpleNamespace( + catalog=types.SimpleNamespace(operation=nullcontext), get_registry=MagicMock(return_value={}), get_allowed_mcp_servers=AsyncMock(return_value=[]), get_mcp_servers_from_ids=MagicMock(return_value=[]), @@ -559,6 +564,7 @@ async def test_get_mcp_tools_from_manager_forwards_request_tags(monkeypatch): mock_get_tools, ) fake_manager = types.SimpleNamespace( + catalog=types.SimpleNamespace(operation=nullcontext), get_registry=MagicMock(return_value={}), get_allowed_mcp_servers=AsyncMock(return_value=[]), get_mcp_servers_from_ids=MagicMock(return_value=[]), From 34ed44d9b38d17e714ce5ed46c57e188ce3ce5c9 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Tue, 22 Sep 2026 12:48:42 -0700 Subject: [PATCH 04/24] fix(mcp): refresh catalog before native alias routing --- litellm/proxy/proxy_server.py | 2 + .../proxy/test_dynamic_mcp_route.py | 61 +++++++++++++++++++ 2 files changed, 63 insertions(+) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index efc428fa4c2..91e2a051a27 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -321,6 +321,7 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.openai_like.model_info import MODEL_INFO_REFRESH_SECONDS from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.proxy._experimental.mcp_server.byok_credential_cache import byok_credential_cache +from litellm.proxy._experimental.mcp_server.catalog import with_mcp_catalog from litellm.proxy._lazy_features import attach_lazy_features, reserve_lazy_slot from litellm.proxy._types import * from litellm.proxy.analytics_endpoints.analytics_endpoints import ( @@ -19755,6 +19756,7 @@ async def _is_mcp_access_group_cached(name: str) -> bool: "/{mcp_server_name}/mcp", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"], ) +@with_mcp_catalog async def dynamic_mcp_route(mcp_server_name: str, request: Request): """Handle /{name}/mcp for MCP server aliases, toolsets, MCP access group tags, and comma-separated lists. diff --git a/tests/test_litellm/proxy/test_dynamic_mcp_route.py b/tests/test_litellm/proxy/test_dynamic_mcp_route.py index da7b8e01f46..0f5297d55cb 100644 --- a/tests/test_litellm/proxy/test_dynamic_mcp_route.py +++ b/tests/test_litellm/proxy/test_dynamic_mcp_route.py @@ -611,3 +611,64 @@ def test_aggregate_mcp_route_returns_404_when_mcp_unavailable(): assert response.status_code == 404 assert handler_calls == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("change", ["create", "rename", "delete"]) +@pytest.mark.parametrize("csv", [False, True]) +async def test_dynamic_route_observes_committed_peer_catalog_changes(monkeypatch, change, csv): + from datetime import datetime + from types import SimpleNamespace + + from starlette.responses import Response + + from litellm.caching.caching import DualCache + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server import mcp_server_manager + from litellm.proxy._types import LiteLLM_MCPServerTable + + old_row = LiteLLM_MCPServerTable( + server_id="catalog-route", server_name="previous", alias="previous", transport="http", + url="https://previous.example/mcp", updated_at=datetime(2026, 1, 1), + ) + new_row = old_row.model_copy(update={ + "server_name": "current", "alias": "current", "url": "https://current.example/mcp", + "updated_at": datetime(2026, 1, 2), + }) + + async def find_rows(*, where): + if "mcp_access_groups" in where or change == "delete": + return [] + return [new_row] + + read_rows = AsyncMock(side_effect=find_rows) + prisma = SimpleNamespace(db=SimpleNamespace( + litellm_mcpservertable=SimpleNamespace(find_many=read_rows), + litellm_mcptoolsettable=SimpleNamespace(find_first=AsyncMock(return_value=None)), + )) + manager = mcp_server_manager.MCPServerManager() + if change != "create": + manager.registry = {old_row.server_id: await manager.build_mcp_server_from_table(old_row)} + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache()) + monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager) + name = "previous" if change == "delete" else "current" + segment = f"{name},missing" if csv else name + request = _make_request(f"/{segment}/mcp") + + async def forwarded(path_segment, request): + if change != "delete": + assert manager.get_mcp_server_by_name(path_segment).url == new_row.url + return Response(content=b"forwarded", status_code=200) + + with patch(_FORWARD, new=AsyncMock(side_effect=forwarded)) as forward: + if change == "delete": + with pytest.raises(HTTPException) as exc: + await proxy_server.dynamic_mcp_route(segment, request) + assert exc.value.status_code == 404 + forward.assert_not_awaited() + else: + response = await proxy_server.dynamic_mcp_route(segment, request) + assert response.status_code == 200 + forward.assert_awaited_once_with(name, request) + assert sum("OR" in call.kwargs["where"] for call in read_rows.await_args_list) == 1 From 2690651b3bde294a299b54dcee7a1eb3107e250e Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Tue, 22 Sep 2026 12:53:18 -0700 Subject: [PATCH 05/24] docs(mcp): clarify native route resolution order --- litellm/proxy/proxy_server.py | 2 +- tests/test_litellm/proxy/test_dynamic_mcp_route.py | 3 +-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 91e2a051a27..838749b9933 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -19762,7 +19762,7 @@ async def dynamic_mcp_route(mcp_server_name: str, request: Request): Resolution order: 1. Registered MCP server alias / name - 2. Comma-separated list (short-circuits before any DB call) + 2. Comma-separated list 3. Toolset name (DB lookup, cached) 4. MCP access group tag (DB lookup, cached) """ diff --git a/tests/test_litellm/proxy/test_dynamic_mcp_route.py b/tests/test_litellm/proxy/test_dynamic_mcp_route.py index 0f5297d55cb..5426cd92005 100644 --- a/tests/test_litellm/proxy/test_dynamic_mcp_route.py +++ b/tests/test_litellm/proxy/test_dynamic_mcp_route.py @@ -3,8 +3,7 @@ Tests for the dynamic_mcp_route handler in proxy_server.py. Covers the resolution order: 1. Registered MCP server alias → forwards to /mcp/{name} - 2. Comma-separated list → short-circuits before any DB call; - forwarded to /mcp/{segment} + 2. Comma-separated list → forwards to /mcp/{segment} 3. Toolset name (cached) → sets toolset scope, forwards to /mcp 4. MCP access group tag (cached) → forwards to /mcp/{name} when the group resolves to at least one server From 180acefb3dfad1f71d32439ff3b726fc5b69932c Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Tue, 22 Sep 2026 13:31:03 -0700 Subject: [PATCH 06/24] test(mcp): align streaming fixture and generated API documentation --- .../responses/mcp/test_mcp_streaming_iterator.py | 4 +++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 14 +++++++------- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/tests/test_litellm/responses/mcp/test_mcp_streaming_iterator.py b/tests/test_litellm/responses/mcp/test_mcp_streaming_iterator.py index c2c204e7024..3982081706f 100644 --- a/tests/test_litellm/responses/mcp/test_mcp_streaming_iterator.py +++ b/tests/test_litellm/responses/mcp/test_mcp_streaming_iterator.py @@ -1,5 +1,6 @@ import sys import types +from contextlib import nullcontext from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock @@ -77,6 +78,7 @@ def _mock_mcp_environment(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: """Patch the MCP tool-call plumbing so _execute_tool_calls can run in tests.""" call_tool = AsyncMock(return_value=CallToolResult(content=[TextContent(type="text", text="ok")], isError=False)) fake_manager = types.SimpleNamespace( + catalog=types.SimpleNamespace(operation=nullcontext), get_registry=MagicMock(return_value={}), call_tool=call_tool, _get_mcp_server_from_tool_name=MagicMock(return_value=None), @@ -89,7 +91,7 @@ def _mock_mcp_environment(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: monkeypatch.setitem( sys.modules, "litellm.proxy.proxy_server", - types.SimpleNamespace(proxy_logging_obj=MagicMock()), + types.SimpleNamespace(proxy_logging_obj=MagicMock(), prisma_client=None), ) return call_tool diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 68d3ab364d6..b419189e70f 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -23134,7 +23134,7 @@ export interface paths { * * Resolution order: * 1. Registered MCP server alias / name - * 2. Comma-separated list (short-circuits before any DB call) + * 2. Comma-separated list * 3. Toolset name (DB lookup, cached) * 4. MCP access group tag (DB lookup, cached) */ @@ -23145,7 +23145,7 @@ export interface paths { * * Resolution order: * 1. Registered MCP server alias / name - * 2. Comma-separated list (short-circuits before any DB call) + * 2. Comma-separated list * 3. Toolset name (DB lookup, cached) * 4. MCP access group tag (DB lookup, cached) */ @@ -23156,7 +23156,7 @@ export interface paths { * * Resolution order: * 1. Registered MCP server alias / name - * 2. Comma-separated list (short-circuits before any DB call) + * 2. Comma-separated list * 3. Toolset name (DB lookup, cached) * 4. MCP access group tag (DB lookup, cached) */ @@ -23167,7 +23167,7 @@ export interface paths { * * Resolution order: * 1. Registered MCP server alias / name - * 2. Comma-separated list (short-circuits before any DB call) + * 2. Comma-separated list * 3. Toolset name (DB lookup, cached) * 4. MCP access group tag (DB lookup, cached) */ @@ -23178,7 +23178,7 @@ export interface paths { * * Resolution order: * 1. Registered MCP server alias / name - * 2. Comma-separated list (short-circuits before any DB call) + * 2. Comma-separated list * 3. Toolset name (DB lookup, cached) * 4. MCP access group tag (DB lookup, cached) */ @@ -23189,7 +23189,7 @@ export interface paths { * * Resolution order: * 1. Registered MCP server alias / name - * 2. Comma-separated list (short-circuits before any DB call) + * 2. Comma-separated list * 3. Toolset name (DB lookup, cached) * 4. MCP access group tag (DB lookup, cached) */ @@ -23200,7 +23200,7 @@ export interface paths { * * Resolution order: * 1. Registered MCP server alias / name - * 2. Comma-separated list (short-circuits before any DB call) + * 2. Comma-separated list * 3. Toolset name (DB lookup, cached) * 4. MCP access group tag (DB lookup, cached) */ From ed7cdd151facae5721011b6b3aee8dba7d2e87a0 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Tue, 22 Sep 2026 13:37:09 -0700 Subject: [PATCH 07/24] fix(mcp): preserve discovery published during catalog refresh --- .../proxy/_experimental/mcp_server/catalog.py | 54 +++++++++--- tests/mcp_tests/test_mcp_server.py | 5 +- .../mcp_server/test_mcp_block_recording.py | 3 +- .../mcp_server/test_mcp_server_manager.py | 86 +++++++++++++++++++ .../mcp/test_mcp_streaming_iterator.py | 2 + 5 files changed, 134 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/catalog.py b/litellm/proxy/_experimental/mcp_server/catalog.py index 0f5eca285e3..ed61b4360d9 100644 --- a/litellm/proxy/_experimental/mcp_server/catalog.py +++ b/litellm/proxy/_experimental/mcp_server/catalog.py @@ -133,21 +133,25 @@ class TargetCatalog: self._retain_discovered_routing(snapshot) def _retain_discovered_routing(self, snapshot: CatalogSnapshot) -> None: + self.manager.published_tool_routes = self.manager.published_tool_routes | self._unchanged_routing( + snapshot.servers, snapshot.routing + ) + + def _unchanged_routing( + self, servers: Mapping[str, MCPServer], routing: Mapping[str, str] + ) -> MappingProxyType[str, str]: from litellm.proxy._experimental.mcp_server.utils import normalize_server_name current: Final = self.manager.config_mcp_servers | self.manager.registry unchanged_owners: Final = frozenset( owner - for key, server in snapshot.servers.items() - if current.get(key) == server + for key, server in servers.items() + if (candidate := current.get(key)) is not None + and _configuration_identity(candidate) == _configuration_identity(server) for owner in self.manager.owned_mapping_values(server) ) - self.manager.published_tool_routes = self.manager.published_tool_routes | MappingProxyType( - { - name: owner - for name, owner in snapshot.routing.items() - if normalize_server_name(owner) in unchanged_owners - } + return MappingProxyType( + {name: owner for name, owner in routing.items() if normalize_server_name(owner) in unchanged_owners} ) async def resolve(self, identifier: str, client_ip: str | None = None) -> MCPServer | None: @@ -202,9 +206,18 @@ class TargetCatalog: from litellm.proxy._experimental.mcp_server.tool_registry import global_mcp_tool_registry async with self._refresh_lock: - staged_config: Final = { - key: value.model_copy(deep=True) for key, value in self.manager.config_mcp_servers.items() - } + previous_config: Final = MappingProxyType( + {key: value.model_copy(deep=True) for key, value in self.manager.config_mcp_servers.items()} + ) + previous_servers: Final = previous_config | MappingProxyType( + {key: value.model_copy(deep=True) for key, value in self.manager.registry.items()} + ) + config_identities: Final = MappingProxyType( + {key: _configuration_identity(value) for key, value in previous_config.items()} + ) + staged_config: Final = MappingProxyType( + {key: value.model_copy(deep=True) for key, value in previous_config.items()} + ) await self.manager.hydrate_config_servers_dcr_clients(tuple(staged_config.values())) staged_routing: Final = dict(self.manager.published_tool_routes) closed: Final = asyncio.Event() @@ -212,9 +225,24 @@ class TargetCatalog: try: with global_mcp_tool_registry.catalog_scope(global_mcp_tool_registry.published_tools) as staged_tools: await self._reload() - self.manager.config_mcp_servers = staged_config + concurrent_routes: Final = self._unchanged_routing( + previous_servers, self.manager.published_tool_routes + ) + retained_staged_routes: Final = self._unchanged_routing( + previous_config | self.manager.registry, staged_routing + ) + self.manager.config_mcp_servers = { + key: value.model_copy( + update=staged_config[key].model_dump( + include=frozenset(("client_id", "client_secret", "token_endpoint_auth_method")) + ) + ) + if config_identities.get(key) == _configuration_identity(value) + else value + for key, value in self.manager.config_mcp_servers.items() + } global_mcp_tool_registry.tools = staged_tools - self.manager.published_tool_routes = staged_routing + self.manager.published_tool_routes = concurrent_routes | retained_staged_routes finally: closed.set() self._staged_routing.reset(routing_token) diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index 2b92367f186..cff3d5f78c6 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -2,7 +2,7 @@ import os import pytest from unittest.mock import AsyncMock, MagicMock, patch -from contextlib import asynccontextmanager +from contextlib import asynccontextmanager, nullcontext from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( @@ -909,6 +909,7 @@ async def test_get_tools_from_mcp_servers(): # Create a mock manager mock_manager = AsyncMock() + mock_manager.catalog.operation = nullcontext mock_manager.get_allowed_mcp_servers = AsyncMock( return_value=["server1_id", "server2_id"] ) @@ -937,6 +938,7 @@ async def test_get_tools_from_mcp_servers(): # Test Case 2: Without specific MCP servers # Create a different mock manager for the second test case mock_manager_2 = AsyncMock() + mock_manager_2.catalog.operation = nullcontext mock_manager_2.get_allowed_mcp_servers = AsyncMock( return_value=["server1_id", "server2_id"] ) @@ -984,6 +986,7 @@ async def test_get_tools_from_mcp_servers(): # Test Case 3: With specific MCP servers and access groups # Create a mock manager mock_manager = AsyncMock() + mock_manager.catalog.operation = nullcontext mock_manager.get_allowed_mcp_servers = AsyncMock( return_value=["server1_id", "server2_id", "server3_id"] ) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_block_recording.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_block_recording.py index b8aadef430f..746130fb892 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_block_recording.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_block_recording.py @@ -66,8 +66,7 @@ async def _call_block(logging_obj, order: list, *, user_api_key_auth=mock.sentin proxy_logging_obj = mock.MagicMock() proxy_logging_obj.post_call_failure_hook.side_effect = _record_post_call_failure_hook - fake_proxy_server = types.ModuleType("litellm.proxy.proxy_server") - fake_proxy_server.proxy_logging_obj = proxy_logging_obj # pyright: ignore[reportAttributeAccessIssue] + fake_proxy_server = types.SimpleNamespace(proxy_logging_obj=proxy_logging_obj, prisma_client=None) with mock.patch.dict(sys.modules, {"litellm.proxy.proxy_server": fake_proxy_server}): with contextlib.suppress(HTTPException): diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 3615d9681f0..fc70689b6c9 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -14519,6 +14519,92 @@ async def test_client_sampling_does_not_fill_explicit_context_from_another_ambie auth_context_var.reset(token) +@pytest.mark.asyncio +@pytest.mark.parametrize("anchored", [False, True]) +async def test_catalog_reload_preserves_concurrent_config_discovery_and_routes(anchored): + manager: Final = MCPServerManager() + server: Final = MCPServer(server_id="config-race", name="config_race", transport=MCPTransport.http, + url="https://upstream.example/mcp", auth_type=MCPAuth.oauth2, + issuer="https://issuer.example" if anchored else None, issuer_is_anchored=anchored) + manager.config_mcp_servers = {server.server_id: server} + manager._set_oauth_discovery_deferred(server.server_id, True) + generation: Final = manager.oauth_discovery_slot(server.server_id).generation + resolved: Final = server.model_copy(update={"authorization_url": "https://issuer.example/authorize", + "token_url": "https://issuer.example/token", "scopes": ["read"], "issuer": "https://issuer.example"}) + + async def hydrate(target: MCPServer) -> bool: + target.client_id = "persisted-client" + return True + + async def read_rows(**_kwargs): + assert manager._publish_resolved_oauth_server(resolved, generation) is resolved + manager.published_tool_routes["config_race-search"] = "config_race" + return [] + + prisma: Final = MagicMock() + prisma.db.litellm_mcpservertable.find_many = AsyncMock(side_effect=read_rows) + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy._experimental.mcp_server.discoverable_endpoints.hydrate_config_server_dcr_client", side_effect=hydrate), + ): + await manager.reload_servers_from_database() + current: Final = manager.get_mcp_server_by_id(server.server_id) + assert current.authorization_url == resolved.authorization_url + assert current.token_url == resolved.token_url + assert current.scopes == ["read"] + assert current.issuer == "https://issuer.example" + assert current.client_id == "persisted-client" + assert manager.oauth_discovery_slot(server.server_id) is None + assert manager.published_tool_routes["config_race-search"] == "config_race" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("change", ["credentials", "delete"]) +async def test_catalog_reload_does_not_restore_replaced_config_credentials_or_routes(change): + manager: Final = MCPServerManager() + server: Final = MCPServer(server_id="config-replaced", name="config_replaced", transport=MCPTransport.http, + client_id="previous-client") + manager.config_mcp_servers = {server.server_id: server} + manager.published_tool_routes["config_replaced-search"] = "config_replaced" + + async def read_rows(**_kwargs): + manager.config_mcp_servers = {} if change == "delete" else { + server.server_id: server.model_copy(update={"client_id": "new-client"})} + return [] + + prisma: Final = MagicMock() + prisma.db.litellm_mcpservertable.find_many = AsyncMock(side_effect=read_rows) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + await manager.reload_servers_from_database() + current: Final = manager.get_mcp_server_by_id(server.server_id) + if change == "delete": + assert current is None + else: + assert current.client_id == "new-client" + assert "config_replaced-search" not in manager.published_tool_routes + + +@pytest.mark.asyncio +@pytest.mark.parametrize("change", ["discovery", "update", "delete"]) +async def test_catalog_operation_retains_routes_only_for_same_configured_target(change): + from mcp.types import Tool + + manager: Final = MCPServerManager() + server: Final = MCPServer(server_id="route-race", name="route_race", transport=MCPTransport.http, + url="https://upstream.example/mcp", auth_type=MCPAuth.oauth2) + manager.registry = {server.server_id: server} + with patch("litellm.proxy.proxy_server.prisma_client", None): + async with manager.catalog.operation(): + tools: Final = manager._create_prefixed_tools([Tool(name="search", inputSchema={})], server) + if change == "delete": + manager.registry = {} + else: + manager.registry[server.server_id] = server.model_copy(update=( + {"authorization_url": "https://issuer.example/authorize", "token_url": "https://issuer.example/token"} + if change == "discovery" else {"url": "https://changed.example/mcp"})) + assert (tools[0].name in manager.published_tool_routes) is (change == "discovery") + + @pytest.mark.asyncio async def test_catalog_observes_committed_update_and_delete_without_background_reload(): from datetime import timedelta diff --git a/tests/test_litellm/responses/mcp/test_mcp_streaming_iterator.py b/tests/test_litellm/responses/mcp/test_mcp_streaming_iterator.py index c2c204e7024..7ed5b5e36a9 100644 --- a/tests/test_litellm/responses/mcp/test_mcp_streaming_iterator.py +++ b/tests/test_litellm/responses/mcp/test_mcp_streaming_iterator.py @@ -1,5 +1,6 @@ import sys import types +from contextlib import nullcontext from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock @@ -77,6 +78,7 @@ def _mock_mcp_environment(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: """Patch the MCP tool-call plumbing so _execute_tool_calls can run in tests.""" call_tool = AsyncMock(return_value=CallToolResult(content=[TextContent(type="text", text="ok")], isError=False)) fake_manager = types.SimpleNamespace( + catalog=types.SimpleNamespace(operation=nullcontext), get_registry=MagicMock(return_value={}), call_tool=call_tool, _get_mcp_server_from_tool_name=MagicMock(return_value=None), From 3a4cd4a8aed41026be7d68edb1b06c951bab0842 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Tue, 22 Sep 2026 13:40:32 -0700 Subject: [PATCH 08/24] fix(mcp): coalesce queued catalog reads without a stale window --- .../proxy/_experimental/mcp_server/catalog.py | 31 +++++-- .../mcp_server/test_mcp_server_manager.py | 86 +++++++++++++++++++ 2 files changed, 110 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/catalog.py b/litellm/proxy/_experimental/mcp_server/catalog.py index a7dabdb0a0e..a3e5508e629 100644 --- a/litellm/proxy/_experimental/mcp_server/catalog.py +++ b/litellm/proxy/_experimental/mcp_server/catalog.py @@ -15,6 +15,7 @@ if TYPE_CHECKING: _P: Final = ParamSpec("_P") _R: Final = TypeVar("_R") +_REFRESH_FAILURE: Final = "MCP server configuration could not be refreshed" @dataclass(slots=True) @@ -29,6 +30,9 @@ class TargetCatalog: self._manager = manager self._reload_lock = asyncio.Lock() self._scope: ContextVar[_CatalogScope | None] = ContextVar("mcp_catalog_scope", default=None) + self._arrival_ticket = 0 + self._completed_ticket = 0 + self._snapshot: Mapping[str, MCPServer] | None = None @property def current(self) -> Mapping[str, MCPServer] | None: @@ -41,8 +45,17 @@ class TargetCatalog: async def _reload(self, *, reuse_unchanged: bool = False) -> None: token: Final = self._scope.set(None) + # A shared read must start after each covered operation arrives. + covered_ticket: Final = self._arrival_ticket try: await self._manager._reload_servers_from_database(reuse_unchanged=reuse_unchanged) # pyright: ignore[reportPrivateUsage] # existing staged loader + except Exception: + self._snapshot = None + self._completed_ticket = covered_ticket + raise + else: + self._snapshot = MappingProxyType(self._manager.config_mcp_servers | self._manager.registry) + self._completed_ticket = covered_ticket finally: self._scope.reset(token) @@ -64,22 +77,26 @@ class TargetCatalog: raise HTTPException(status_code=503, detail="MCP server configuration changed; retry the operation") async def list(self) -> Mapping[str, MCPServer]: + from fastapi import HTTPException # noqa: PLC0415 # optional proxy dependency + from litellm.proxy.proxy_server import prisma_client # noqa: PLC0415 # runtime proxy dependency scope: Final = self._scope.get() if scope is not None and scope.active and scope.owner_task_id == id(asyncio.current_task()): return scope.servers + self._arrival_ticket += 1 + arrival_ticket: Final = self._arrival_ticket async with self._reload_lock: - if prisma_client is not None: + if prisma_client is None: + return MappingProxyType(self._manager.config_mcp_servers | self._manager.registry) + if arrival_ticket > self._completed_ticket: try: await self._reload(reuse_unchanged=True) except Exception as exc: # noqa: BLE001 # never serve an unverified database snapshot - from fastapi import HTTPException # noqa: PLC0415 # optional proxy dependency - - raise HTTPException( - status_code=503, detail="MCP server configuration could not be refreshed" - ) from exc - return MappingProxyType(self._manager.config_mcp_servers | self._manager.registry) + raise HTTPException(status_code=503, detail=_REFRESH_FAILURE) from exc + if self._snapshot is None: + raise HTTPException(status_code=503, detail=_REFRESH_FAILURE) + return self._snapshot @asynccontextmanager async def operation(self) -> AsyncIterator[Mapping[str, MCPServer]]: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index b031280f975..b0ef6d53c21 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -14573,6 +14573,92 @@ async def test_catalog_database_failure_keeps_registry_but_rejects_operation(mon assert manager.catalog.current is None +@pytest.mark.asyncio +@pytest.mark.parametrize("failed", [False, True]) +async def test_catalog_coalesces_waiters_only_behind_a_read_started_after_their_arrival(monkeypatch, failed): + started = asyncio.Event() + release = asyncio.Event() + old_row = _catalog_row() + new_row = _catalog_row("updated") + read_rows = AsyncMock(return_value=[old_row]) + _catalog_database(monkeypatch, read_rows) + manager = MCPServerManager() + previous = await manager.catalog.list() + read_rows.reset_mock() + + async def read_current_rows(**kwargs): + first = not started.is_set() + if first: + started.set() + await release.wait() + if failed: + raise RuntimeError("unavailable database") + return [old_row if first else new_row] + + read_rows.side_effect = read_current_rows + first = asyncio.create_task(manager.catalog.list()) + await started.wait() + followers = [asyncio.create_task(manager.catalog.list()) for _ in range(8)] + await asyncio.sleep(0) + release.set() + results = await asyncio.gather(first, *followers, return_exceptions=True) + if failed: + assert all(isinstance(result, HTTPException) and result.status_code == 503 for result in results) + assert manager.registry == dict(previous) + else: + assert results[0]["catalog-server"].name == "initial" + assert all(result["catalog-server"].name == "updated" for result in results[1:]) + assert read_rows.await_count == 2 + + read_rows.side_effect = None + read_rows.return_value = [new_row] + assert (await manager.catalog.list())["catalog-server"].name == "updated" + assert read_rows.await_count == 3 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("interruption", ["background", "cancel_reader", "cancel_waiter"]) +async def test_catalog_queued_refresh_preserves_freshness_through_cancellation(monkeypatch, interruption): + started = asyncio.Event() + release = asyncio.Event() + + async def read_current_rows(**kwargs): + first = not started.is_set() + if first: + started.set() + await release.wait() + return [_catalog_row("initial" if first else "updated")] + + read_rows = AsyncMock(side_effect=read_current_rows) + _catalog_database(monkeypatch, read_rows) + manager = MCPServerManager() + first = asyncio.create_task(manager.catalog.list()) + await started.wait() + background = asyncio.create_task(manager.reload_servers_from_database()) if interruption == "background" else None + waiter = asyncio.create_task(manager.catalog.list()) + survivors = [asyncio.create_task(manager.catalog.list()) for _ in range(4)] + await asyncio.sleep(0) + if interruption == "cancel_reader": + first.cancel() + elif interruption == "cancel_waiter": + waiter.cancel() + release.set() + first_result, waiter_result, *results = await asyncio.gather(first, waiter, *survivors, return_exceptions=True) + if background is not None: + await background + if interruption == "cancel_reader": + assert isinstance(first_result, asyncio.CancelledError) + else: + assert first_result["catalog-server"].name == "initial" + if interruption == "cancel_waiter": + assert isinstance(waiter_result, asyncio.CancelledError) + else: + assert waiter_result["catalog-server"].name == "updated" + assert all(result["catalog-server"].name == "updated" for result in results) + assert read_rows.await_count == 2 + assert manager.catalog.current is None + + @pytest.mark.asyncio async def test_catalog_serializes_background_and_operation_refreshes(monkeypatch): entered = asyncio.Event() From 42821c102e487ca8e702eea6e2db9c3d0f76eeed Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Tue, 22 Sep 2026 13:47:10 -0700 Subject: [PATCH 09/24] fix(mcp): reconcile concurrent route changes when publishing catalog --- .../proxy/_experimental/mcp_server/catalog.py | 25 ++++++++++++--- .../mcp_server/test_mcp_server_manager.py | 32 +++++++++++++++++++ 2 files changed, 53 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/catalog.py b/litellm/proxy/_experimental/mcp_server/catalog.py index ed61b4360d9..f121ab08963 100644 --- a/litellm/proxy/_experimental/mcp_server/catalog.py +++ b/litellm/proxy/_experimental/mcp_server/catalog.py @@ -219,17 +219,34 @@ class TargetCatalog: {key: value.model_copy(deep=True) for key, value in previous_config.items()} ) await self.manager.hydrate_config_servers_dcr_clients(tuple(staged_config.values())) - staged_routing: Final = dict(self.manager.published_tool_routes) + initial_routing: Final = MappingProxyType(dict(self.manager.published_tool_routes)) + staged_routing: Final = dict(initial_routing) closed: Final = asyncio.Event() routing_token: Final = self._staged_routing.set((staged_routing, closed)) try: with global_mcp_tool_registry.catalog_scope(global_mcp_tool_registry.published_tools) as staged_tools: await self._reload() concurrent_routes: Final = self._unchanged_routing( - previous_servers, self.manager.published_tool_routes + previous_servers, + MappingProxyType( + { + name: owner + for name, owner in self.manager.published_tool_routes.items() + if initial_routing.get(name) != owner + and staged_routing.get(name) == initial_routing.get(name) + } + ), ) + removed_routes: Final = initial_routing.keys() - self.manager.published_tool_routes.keys() retained_staged_routes: Final = self._unchanged_routing( - previous_config | self.manager.registry, staged_routing + previous_config | self.manager.registry, + MappingProxyType( + { + name: owner + for name, owner in staged_routing.items() + if name not in removed_routes or owner != initial_routing[name] + } + ), ) self.manager.config_mcp_servers = { key: value.model_copy( @@ -242,7 +259,7 @@ class TargetCatalog: for key, value in self.manager.config_mcp_servers.items() } global_mcp_tool_registry.tools = staged_tools - self.manager.published_tool_routes = concurrent_routes | retained_staged_routes + self.manager.published_tool_routes = retained_staged_routes | concurrent_routes finally: closed.set() self._staged_routing.reset(routing_token) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index fc70689b6c9..7462f948240 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -14605,6 +14605,38 @@ async def test_catalog_operation_retains_routes_only_for_same_configured_target( assert (tools[0].name in manager.published_tool_routes) is (change == "discovery") +@pytest.mark.asyncio +@pytest.mark.parametrize("change", ["concurrent", "staged", "deleted", "staged_after_delete", "staged_with_concurrent"]) +async def test_catalog_reload_keeps_new_route_owner_over_earlier_route(change): + manager: Final = MCPServerManager() + servers: Final = {name: MCPServer(server_id=name, name=name, transport=MCPTransport.http, + client_id="configured-client") for name in ("old_owner", "new_owner", "concurrent_owner")} + manager.config_mcp_servers = servers + manager.published_tool_routes = {"shared-search": "old_owner"} + + async def read_rows(**_kwargs): + if change == "concurrent": + manager.published_tool_routes["shared-search"] = "new_owner" + elif change in ("staged", "staged_after_delete", "staged_with_concurrent"): + if change == "staged_after_delete": + manager.published_tool_routes.clear() + elif change == "staged_with_concurrent": + manager.published_tool_routes["shared-search"] = "concurrent_owner" + manager.tool_name_to_mcp_server_name_mapping["shared-search"] = "new_owner" + else: + manager.published_tool_routes.clear() + return [] + + prisma: Final = MagicMock() + prisma.db.litellm_mcpservertable.find_many = AsyncMock(side_effect=read_rows) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + await manager.reload_servers_from_database() + if change == "deleted": + assert "shared-search" not in manager.published_tool_routes + else: + assert manager.published_tool_routes["shared-search"] == "new_owner" + + @pytest.mark.asyncio async def test_catalog_observes_committed_update_and_delete_without_background_reload(): from datetime import timedelta From 81b6035beed8861a406fac593b189b9f421e59de Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Tue, 22 Sep 2026 14:08:27 -0700 Subject: [PATCH 10/24] test(mcp): provide catalog scope in post-call hook fixtures --- .../responses/mcp/test_litellm_proxy_mcp_handler.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py index ead9f9ac713..cc2062da5ba 100644 --- a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py +++ b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -507,6 +507,7 @@ async def test_execute_tool_calls_applies_post_call_hook_content(monkeypatch): isError=False, ) fake_manager = types.SimpleNamespace( + catalog=types.SimpleNamespace(operation=nullcontext), get_registry=MagicMock(return_value={}), call_tool=AsyncMock(return_value=result), _get_mcp_server_from_tool_name=MagicMock(return_value=None), @@ -546,6 +547,7 @@ async def test_execute_tool_calls_returns_proxy_result_without_logging(monkeypat ) fake_manager = types.SimpleNamespace( + catalog=types.SimpleNamespace(operation=nullcontext), get_registry=MagicMock(return_value={}), call_tool=AsyncMock(return_value=result), _get_mcp_server_from_tool_name=MagicMock(return_value=None), @@ -579,6 +581,7 @@ async def test_execute_tool_calls_passes_logging_details_to_proxy_hook(monkeypat ) fake_manager = types.SimpleNamespace( + catalog=types.SimpleNamespace(operation=nullcontext), get_registry=MagicMock(return_value={}), call_tool=AsyncMock(return_value=result), _get_mcp_server_from_tool_name=MagicMock(return_value=None), @@ -614,6 +617,7 @@ async def test_execute_tool_calls_continues_when_post_call_logging_fails(monkeyp result = CallToolResult(content=[TextContent(type="text", text="ok")], isError=False) fake_manager = types.SimpleNamespace( + catalog=types.SimpleNamespace(operation=nullcontext), get_registry=MagicMock(return_value={}), call_tool=AsyncMock(return_value=result), _get_mcp_server_from_tool_name=MagicMock(return_value=None), From d4c1afdebcba325a12f404e7438ef6e8f5066104 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Tue, 22 Sep 2026 14:11:07 -0700 Subject: [PATCH 11/24] fix(mcp): retain valid live routes and handlers during refresh --- .../proxy/_experimental/mcp_server/catalog.py | 39 +++++- .../mcp_server/test_mcp_server_manager.py | 113 +++++++++++++++++- 2 files changed, 143 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/catalog.py b/litellm/proxy/_experimental/mcp_server/catalog.py index f121ab08963..35d6d266a5d 100644 --- a/litellm/proxy/_experimental/mcp_server/catalog.py +++ b/litellm/proxy/_experimental/mcp_server/catalog.py @@ -209,8 +209,9 @@ class TargetCatalog: previous_config: Final = MappingProxyType( {key: value.model_copy(deep=True) for key, value in self.manager.config_mcp_servers.items()} ) + live_registry: Final = self.manager.registry previous_servers: Final = previous_config | MappingProxyType( - {key: value.model_copy(deep=True) for key, value in self.manager.registry.items()} + {key: value.model_copy(deep=True) for key, value in live_registry.items()} ) config_identities: Final = MappingProxyType( {key: _configuration_identity(value) for key, value in previous_config.items()} @@ -223,17 +224,21 @@ class TargetCatalog: staged_routing: Final = dict(initial_routing) closed: Final = asyncio.Event() routing_token: Final = self._staged_routing.set((staged_routing, closed)) + initial_tools: Final = MappingProxyType(dict(global_mcp_tool_registry.published_tools)) try: - with global_mcp_tool_registry.catalog_scope(global_mcp_tool_registry.published_tools) as staged_tools: + with global_mcp_tool_registry.catalog_scope(initial_tools) as staged_tools: await self._reload() + live_routes: Final = self._unchanged_routing(previous_servers, self.manager.published_tool_routes) concurrent_routes: Final = self._unchanged_routing( - previous_servers, + self.manager.config_mcp_servers | live_registry, MappingProxyType( { name: owner for name, owner in self.manager.published_tool_routes.items() if initial_routing.get(name) != owner - and staged_routing.get(name) == initial_routing.get(name) + and ( + name not in staged_routing or staged_routing.get(name) == initial_routing.get(name) + ) } ), ) @@ -248,6 +253,18 @@ class TargetCatalog: } ), ) + concurrent_tools: Final = MappingProxyType( + { + name: tool + for name, tool in global_mcp_tool_registry.published_tools.items() + if ( + name in live_routes + or name in concurrent_routes + or (name not in initial_routing and name not in self.manager.published_tool_routes) + ) + and initial_tools.get(name) is staged_tools.get(name) + } + ) self.manager.config_mcp_servers = { key: value.model_copy( update=staged_config[key].model_dump( @@ -258,8 +275,18 @@ class TargetCatalog: else value for key, value in self.manager.config_mcp_servers.items() } - global_mcp_tool_registry.tools = staged_tools - self.manager.published_tool_routes = retained_staged_routes | concurrent_routes + global_mcp_tool_registry.tools = ( + MappingProxyType( + { + name: tool + for name, tool in staged_tools.items() + if name in global_mcp_tool_registry.published_tools + or tool is not initial_tools.get(name) + } + ) + | concurrent_tools + ) + self.manager.published_tool_routes = live_routes | retained_staged_routes | concurrent_routes finally: closed.set() self._staged_routing.reset(routing_token) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 7462f948240..6f903aebbf5 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -14607,34 +14607,63 @@ async def test_catalog_operation_retains_routes_only_for_same_configured_target( @pytest.mark.asyncio @pytest.mark.parametrize("change", ["concurrent", "staged", "deleted", "staged_after_delete", "staged_with_concurrent"]) -async def test_catalog_reload_keeps_new_route_owner_over_earlier_route(change): +@pytest.mark.parametrize("mapped", [False, True]) +async def test_catalog_reload_keeps_new_route_owner_over_earlier_route(change, mapped, monkeypatch): + from litellm.proxy._experimental.mcp_server import tool_registry + manager: Final = MCPServerManager() + registry: Final = tool_registry.MCPToolRegistry() + monkeypatch.setattr(tool_registry, "global_mcp_tool_registry", registry) servers: Final = {name: MCPServer(server_id=name, name=name, transport=MCPTransport.http, client_id="configured-client") for name in ("old_owner", "new_owner", "concurrent_owner")} manager.config_mcp_servers = servers - manager.published_tool_routes = {"shared-search": "old_owner"} + manager.published_tool_routes = {"shared-search": "old_owner"} if mapped else {} + + async def original_handler(): + return "original response" + + async def new_handler(): + return "new response" + + async def concurrent_handler(): + return "concurrent response" + + registry.register_tool("shared-search", "Search", {}, original_handler) + original_tool: Final = registry.get_tool("shared-search") async def read_rows(**_kwargs): if change == "concurrent": manager.published_tool_routes["shared-search"] = "new_owner" + registry.published_tools["shared-search"] = original_tool.model_copy(update={"handler": new_handler}) elif change in ("staged", "staged_after_delete", "staged_with_concurrent"): if change == "staged_after_delete": manager.published_tool_routes.clear() + registry.published_tools.clear() elif change == "staged_with_concurrent": manager.published_tool_routes["shared-search"] = "concurrent_owner" + registry.published_tools["shared-search"] = original_tool.model_copy(update={"handler": concurrent_handler}) manager.tool_name_to_mcp_server_name_mapping["shared-search"] = "new_owner" + registry.register_tool("shared-search", "Search", {}, new_handler) else: manager.published_tool_routes.clear() + registry.published_tools.clear() + if not mapped: + manager.published_tool_routes.clear() + manager.tool_name_to_mcp_server_name_mapping.clear() return [] prisma: Final = MagicMock() prisma.db.litellm_mcpservertable.find_many = AsyncMock(side_effect=read_rows) with patch("litellm.proxy.proxy_server.prisma_client", prisma): await manager.reload_servers_from_database() - if change == "deleted": + if change == "deleted" or not mapped: assert "shared-search" not in manager.published_tool_routes else: assert manager.published_tool_routes["shared-search"] == "new_owner" + if change == "deleted": + assert registry.get_tool("shared-search") is None + else: + assert await registry.get_tool("shared-search").handler() == "new response" @pytest.mark.asyncio @@ -14666,6 +14695,84 @@ async def test_catalog_observes_committed_update_and_delete_without_background_r assert prisma.db.litellm_mcpservertable.find_many.await_count == 3 +@pytest.mark.asyncio +async def test_catalog_rebuilt_unchanged_server_keeps_discovered_tool_routes(): + from mcp.types import Tool + + row: Final = LiteLLM_MCPServerTable(server_id="unchanged-routes", alias="unchanged_routes", + url="https://upstream.example/mcp", transport=MCPTransport.http) + manager: Final = MCPServerManager() + prisma: Final = MagicMock() + prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[row]) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + await manager.reload_servers_from_database() + server: Final = manager.get_mcp_server_by_id(row.server_id) + tools: Final = manager._create_prefixed_tools([Tool(name="search", inputSchema={})], server) + await manager.reload_servers_from_database() + assert manager.server_exposes_tool(manager.get_mcp_server_by_id(row.server_id), tools[0].name) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("change", ["same", "changed", "deleted"]) +@pytest.mark.parametrize("openapi", [False, True]) +async def test_catalog_reload_retains_routes_discovered_for_a_late_server(change, openapi, monkeypatch): + from datetime import timedelta + + from mcp.types import Tool + from litellm.proxy._experimental.mcp_server import tool_registry + + row: Final = LiteLLM_MCPServerTable(server_id="late-routes", alias="late_routes", + url="https://upstream.example/mcp", transport=MCPTransport.http, updated_at=datetime.now(), + spec_path="late-openapi.json" if openapi else None) + manager: Final = MCPServerManager() + registry: Final = tool_registry.MCPToolRegistry() + monkeypatch.setattr(tool_registry, "global_mcp_tool_registry", registry) + server: Final = await manager.build_mcp_server_from_table(row) + start: Final = asyncio.Event() + published: Final = asyncio.Event() + + async def handler(): + return "late server response" + + async def publish(): + await start.wait() + manager.registry[server.server_id] = server + manager._create_prefixed_tools([Tool(name="search", inputSchema={})], server) + if openapi: + registry.register_tool("late_routes-search", "Search", {}, handler) + published.set() + + async def read_rows(**_kwargs): + start.set() + await published.wait() + if change == "deleted": + return [] + return [row if change == "same" else row.model_copy(update={ + "url": "https://updated.example/mcp", "spec_path": None, + "updated_at": row.updated_at + timedelta(seconds=1)})] + + prisma: Final = MagicMock() + prisma.db.litellm_mcpservertable.find_many = AsyncMock(side_effect=read_rows) + task: Final = asyncio.create_task(publish()) + try: + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + await manager.reload_servers_from_database() + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + if change == "deleted": + assert manager.get_mcp_server_by_id(row.server_id) is None + assert "late_routes-search" not in manager.published_tool_routes + else: + assert manager.server_exposes_tool(manager.get_mcp_server_by_id(row.server_id), "late_routes-search") is (change == "same") + tool: Final = registry.get_tool("late_routes-search") + if openapi and change == "same": + assert tool is not None + assert await tool.handler() == "late server response" + else: + assert tool is None + + @pytest.mark.asyncio async def test_catalog_lookup_uses_one_snapshot_until_operation_finishes(): from datetime import timedelta From d3edb5839dea5c1cdbb42461766f8d1f1f29a2e4 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Tue, 22 Sep 2026 14:25:58 -0700 Subject: [PATCH 12/24] fix(mcp): keep refreshed OpenAPI operation membership authoritative --- .../proxy/_experimental/mcp_server/catalog.py | 18 +++++++- .../mcp_server/test_mcp_server_manager.py | 43 +++++++++++++++++++ 2 files changed, 60 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_experimental/mcp_server/catalog.py b/litellm/proxy/_experimental/mcp_server/catalog.py index 35d6d266a5d..014d84875b5 100644 --- a/litellm/proxy/_experimental/mcp_server/catalog.py +++ b/litellm/proxy/_experimental/mcp_server/catalog.py @@ -204,6 +204,7 @@ class TargetCatalog: async def reload(self) -> None: from litellm.proxy._experimental.mcp_server.tool_registry import global_mcp_tool_registry + from litellm.proxy._experimental.mcp_server.utils import normalize_server_name async with self._refresh_lock: previous_config: Final = MappingProxyType( @@ -228,7 +229,22 @@ class TargetCatalog: try: with global_mcp_tool_registry.catalog_scope(initial_tools) as staged_tools: await self._reload() - live_routes: Final = self._unchanged_routing(previous_servers, self.manager.published_tool_routes) + refreshed_openapi_owners: Final = frozenset( + owner + for server in self.manager.registry.values() + if server.spec_path and server is not live_registry.get(server.server_id) + for owner in self.manager.owned_mapping_values(server) + ) + live_routes: Final = self._unchanged_routing( + previous_servers, + MappingProxyType( + { + name: owner + for name, owner in self.manager.published_tool_routes.items() + if normalize_server_name(owner) not in refreshed_openapi_owners + } + ), + ) concurrent_routes: Final = self._unchanged_routing( self.manager.config_mcp_servers | live_registry, MappingProxyType( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 6f903aebbf5..b3d80beec6b 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -14773,6 +14773,49 @@ async def test_catalog_reload_retains_routes_discovered_for_a_late_server(change assert tool is None +@pytest.mark.asyncio +@pytest.mark.parametrize("retain_operation", [False, True]) +async def test_catalog_openapi_refresh_does_not_restore_removed_operations(tmp_path, monkeypatch, respx_mock, retain_operation): + from litellm.proxy._experimental.mcp_server import tool_registry + + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + manager: Final = MCPServerManager() + registry: Final = tool_registry.MCPToolRegistry() + monkeypatch.setattr(tool_registry, "global_mcp_tool_registry", registry) + spec_path: Final = tmp_path / "openapi.json" + paths: Final = {"/removed": {"get": {"operationId": "removed"}}, "/retained": {"get": {"operationId": "retained"}}} + spec: Final = {"openapi": "3.0.0", "info": {"title": "Refresh", "version": "1"}, "paths": paths} + spec_path.write_text(json.dumps(spec)) + row: Final = LiteLLM_MCPServerTable(server_id="spec-refresh", alias="spec_refresh", + url="https://upstream.example", transport=MCPTransport.http, spec_path=str(spec_path)) + prisma: Final = MagicMock() + prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[row]) + upstream: Final = respx_mock.get("https://upstream.example/retained").respond(200, json={"value": "retained"}) + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + await manager.reload_servers_from_database() + assert manager.server_exposes_tool(manager.registry[row.server_id], "spec_refresh-removed") + assert registry.get_tool("spec_refresh-removed") is not None + paths.pop("/removed") + if not retain_operation: + paths.clear() + spec_path.write_text(json.dumps(spec)) + for _ in range(2): + await manager.reload_servers_from_database() + for name in ("removed", "spec_refresh-removed"): + assert name not in manager.published_tool_routes + assert not manager.server_exposes_tool(manager.registry[row.server_id], name) + assert registry.get_tool("spec_refresh-removed") is None + retained = registry.get_tool("spec_refresh-retained") + if retain_operation: + assert manager.server_exposes_tool(manager.registry[row.server_id], "spec_refresh-retained") + assert retained is not None + assert json.loads(await retained.handler()) == {"value": "retained"} + else: + assert retained is None + assert manager.published_tool_routes == {} + assert upstream.call_count == (2 if retain_operation else 0) + + @pytest.mark.asyncio async def test_catalog_lookup_uses_one_snapshot_until_operation_finishes(): from datetime import timedelta From f6dffda9214922985f454b2480508640aa19a162 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Tue, 22 Sep 2026 14:29:26 -0700 Subject: [PATCH 13/24] test(mcp): preserve logging fixtures after catalog integration --- .../responses/mcp/test_litellm_proxy_mcp_handler.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py index c7557f610f6..3bd78d0b983 100644 --- a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py +++ b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -507,6 +507,7 @@ async def test_execute_tool_calls_applies_post_call_hook_content(monkeypatch): isError=False, ) fake_manager = types.SimpleNamespace( + catalog=types.SimpleNamespace(operation=nullcontext), get_registry=MagicMock(return_value={}), call_tool=AsyncMock(return_value=result), _get_mcp_server_from_tool_name=MagicMock(return_value=None), @@ -546,6 +547,7 @@ async def test_execute_tool_calls_returns_proxy_result_without_logging(monkeypat ) fake_manager = types.SimpleNamespace( + catalog=types.SimpleNamespace(operation=nullcontext), get_registry=MagicMock(return_value={}), call_tool=AsyncMock(return_value=result), _get_mcp_server_from_tool_name=MagicMock(return_value=None), @@ -579,6 +581,7 @@ async def test_execute_tool_calls_passes_logging_details_to_proxy_hook(monkeypat ) fake_manager = types.SimpleNamespace( + catalog=types.SimpleNamespace(operation=nullcontext), get_registry=MagicMock(return_value={}), call_tool=AsyncMock(return_value=result), _get_mcp_server_from_tool_name=MagicMock(return_value=None), @@ -614,6 +617,7 @@ async def test_execute_tool_calls_continues_when_post_call_logging_fails(monkeyp result = CallToolResult(content=[TextContent(type="text", text="ok")], isError=False) fake_manager = types.SimpleNamespace( + catalog=types.SimpleNamespace(operation=nullcontext), get_registry=MagicMock(return_value={}), call_tool=AsyncMock(return_value=result), _get_mcp_server_from_tool_name=MagicMock(return_value=None), From cb4341b3907af6245a5679bbd8b1cb836bbc77ca Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Tue, 22 Sep 2026 17:32:23 -0700 Subject: [PATCH 14/24] fix(mcp): isolate transport test admission state and exhaust OAuth outcomes --- litellm/proxy/_experimental/mcp_server/mcp_server_manager.py | 4 +++- tests/mcp_tests/test_proxy_mcp_e2e.py | 2 ++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 490190204bc..48325a34ef9 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -48,7 +48,7 @@ from mcp.types import ( ) from mcp.types import Tool as MCPTool from pydantic import AnyUrl, BaseModel, TypeAdapter -from typing_extensions import ReadOnly +from typing_extensions import ReadOnly, assert_never import litellm from litellm._logging import verbose_logger @@ -2213,6 +2213,8 @@ class MCPServerManager: detail=f"OAuth metadata discovery {reason} for MCP server {server_ref!r}", ) + return assert_never(outcome) + async def _rejoin_oauth_metadata_discovery(self, server: MCPServer, *, retry_stale: bool) -> MCPServer: if retry_stale: return await self.ensure_oauth_metadata_discovered(server, _retry_stale=False) diff --git a/tests/mcp_tests/test_proxy_mcp_e2e.py b/tests/mcp_tests/test_proxy_mcp_e2e.py index a730f6c10ee..c23a3db9d27 100644 --- a/tests/mcp_tests/test_proxy_mcp_e2e.py +++ b/tests/mcp_tests/test_proxy_mcp_e2e.py @@ -62,9 +62,11 @@ def _clear_proxy_database_env() -> typing.Iterator[None]: async def _initialize_proxy(config_path: str) -> None: + from litellm.proxy._experimental.mcp_server.catalog import TargetCatalog from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager cleanup_router_config_variables() + global_mcp_server_manager.catalog = TargetCatalog(global_mcp_server_manager) await initialize(config=config_path, debug=True) for server_id, upstream in tuple(global_mcp_server_manager.registry.items()): if upstream.server_name != "math_restricted": From 51892184a36937a1e982f9d6ad973f0279756e90 Mon Sep 17 00:00:00 2001 From: joshua Date: Wed, 23 Sep 2026 05:18:27 +0000 Subject: [PATCH 15/24] refactor(mcp): narrow catalog consistency change to ticket scope Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/auth/admission.py | 97 ---------------- .../mcp_server/auth/user_api_key_auth_mcp.py | 2 +- .../proxy/_experimental/mcp_server/catalog.py | 63 ++--------- .../mcp_server/discoverable_endpoints.py | 2 +- .../_experimental/mcp_server/operations.py | 24 ++-- .../mcp_server/auth/test_admission.py | 89 --------------- .../mcp_server/test_mcp_server_manager.py | 104 ------------------ 7 files changed, 25 insertions(+), 356 deletions(-) delete mode 100644 litellm/proxy/_experimental/mcp_server/auth/admission.py delete mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/auth/test_admission.py diff --git a/litellm/proxy/_experimental/mcp_server/auth/admission.py b/litellm/proxy/_experimental/mcp_server/auth/admission.py deleted file mode 100644 index e5d1167d4d3..00000000000 --- a/litellm/proxy/_experimental/mcp_server/auth/admission.py +++ /dev/null @@ -1,97 +0,0 @@ -from __future__ import annotations - -import math -import os -import time -from collections.abc import Callable, Iterator -from contextlib import contextmanager -from dataclasses import dataclass -from typing import Final - -from fastapi import HTTPException, Request - -from litellm.proxy.auth.login_throttle import source_group -from litellm.proxy.auth.network import TrustedProxyConfig, resolve_client_ip - - -@dataclass(slots=True) -class _Window: - started: float - requests: int = 0 - active: int = 0 - - -def _limit(name: str, default: int) -> int: - value: Final = int(os.environ.get(name, str(default))) - if value < 1: - raise ValueError(f"{name} must be positive") - return value - - -class MCPAdmissionLimiter: - def __init__(self, clock: Callable[[], float] = time.monotonic) -> None: - self._clock = clock - self._client_rpm = _limit("LITELLM_MCP_PUBLIC_RPM", 120) - self._worker_rpm = _limit("LITELLM_MCP_PUBLIC_WORKER_RPM", 600) - self._client_active = _limit("LITELLM_MCP_PUBLIC_MAX_IN_FLIGHT", 64) - self._worker_active = _limit("LITELLM_MCP_PUBLIC_WORKER_MAX_IN_FLIGHT", 128) - self._max_sources = _limit("LITELLM_MCP_PUBLIC_MAX_SOURCES", 4096) - self._sources: dict[str, _Window] = {} - self._worker = _Window(clock()) - - @staticmethod - def _reject(retry_after: int) -> None: - raise HTTPException( - status_code=429, - detail="MCP configuration request limit exceeded; retry later", - headers={"Retry-After": str(max(1, retry_after))}, - ) - - @contextmanager - def admit(self, source: str) -> Iterator[None]: - now: Final = self._clock() - if now - self._worker.started >= 60: - self._worker = _Window(now, active=self._worker.active) - self._sources = { - key: value for key, value in self._sources.items() if value.active or now - value.started < 60 - } - if self._worker.active >= self._worker_active: - self._reject(1) - if self._worker.requests >= self._worker_rpm: - self._reject(math.ceil(60 - (now - self._worker.started))) - previous: Final = self._sources.get(source) - if previous is None and len(self._sources) >= self._max_sources: - self._reject(math.ceil(60 - (now - self._worker.started))) - client: Final = ( - previous - if previous is not None and now - previous.started < 60 - else _Window(now, active=previous.active if previous is not None else 0) - ) - if client.active >= self._client_active: - self._reject(1) - if client.requests >= self._client_rpm: - self._reject(math.ceil(60 - (now - client.started))) - self._sources[source] = client - client.requests += 1 - client.active += 1 - self._worker.requests += 1 - self._worker.active += 1 - try: - yield - finally: - self._sources[source].active -= 1 - self._worker.active -= 1 - - -def admission_source(request: Request) -> str: - from litellm.proxy.proxy_server import general_settings # noqa: PLC0415 # runtime proxy configuration - - settings: Final = general_settings or {} - config: Final = TrustedProxyConfig.model_validate( - { - "use_forwarded_for": settings.get("use_x_forwarded_for", False), - "trusted_proxy_cidrs": settings.get("mcp_trusted_proxy_ranges") or (), - } - ) - client_ip, _ = resolve_client_ip(request, config) - return source_group(client_ip or "unknown") diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 0313bbc6fce..d2284cc901c 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -446,7 +446,7 @@ class MCPRequestHandler: Raises: HTTPException: If headers are invalid or missing required headers """ - async with global_manager().catalog.operation(request=Request(scope)): + async with global_manager().catalog.operation(): headers: Final = MCPRequestHandler._safe_get_headers_from_scope(scope) # Check if there is an explicit LiteLLM API key (primary header) diff --git a/litellm/proxy/_experimental/mcp_server/catalog.py b/litellm/proxy/_experimental/mcp_server/catalog.py index bbbecacccdb..ecf71089731 100644 --- a/litellm/proxy/_experimental/mcp_server/catalog.py +++ b/litellm/proxy/_experimental/mcp_server/catalog.py @@ -10,19 +10,14 @@ from contextlib import asynccontextmanager from contextvars import ContextVar from dataclasses import dataclass, replace from functools import wraps -from inspect import signature from types import MappingProxyType from typing import TYPE_CHECKING, Final, ParamSpec, TypeVar from litellm._logging import verbose_logger if TYPE_CHECKING: - from mcp.types import Tool as SDKTool from pydantic import BaseModel - from starlette.requests import Request - from litellm.proxy._experimental.mcp_server.auth.admission import MCPAdmissionLimiter - from litellm.proxy._experimental.mcp_server.faults.list_outcomes import AggregateToolListing, ServerOutcome from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager from litellm.types.mcp_server.mcp_server_manager import MCPServer from litellm.types.mcp_server.tool_registry import MCPTool @@ -86,7 +81,6 @@ class TargetCatalog: self._arrival_ticket = 0 self._completed_ticket = 0 self._shared_snapshot: CatalogSnapshot | None = None - self._admission: MCPAdmissionLimiter | None = None self._warned_shadowed_config_server_ids: frozenset[str] = frozenset() self._warned_capturing_config_server_ids: frozenset[str] = frozenset() self._operation: ContextVar[tuple[CatalogSnapshot, asyncio.Event, int] | None] = ContextVar( @@ -136,18 +130,8 @@ class TargetCatalog: raise HTTPException(status_code=503, detail="MCP server configuration could not be refreshed") return self._shared_snapshot - async def _acquire_snapshot(self, request: Request | None) -> CatalogSnapshot: - if request is None: - return await self._fresh_snapshot() - from litellm.proxy._experimental.mcp_server.auth.admission import MCPAdmissionLimiter, admission_source - - if self._admission is None: - self._admission = MCPAdmissionLimiter() - with self._admission.admit(admission_source(request)): - return await self._fresh_snapshot() - - async def list(self, *, request: Request | None = None) -> Mapping[str, MCPServer]: - async with self.operation(request=request) as snapshot: + async def list(self) -> Mapping[str, MCPServer]: + async with self.operation() as snapshot: return snapshot.servers def assert_current(self, server: MCPServer) -> None: @@ -168,7 +152,7 @@ class TargetCatalog: raise HTTPException(status_code=503, detail="MCP server configuration changed; retry the operation") @asynccontextmanager - async def operation(self, *, request: Request | None = None) -> AsyncIterator[CatalogSnapshot]: + async def operation(self) -> AsyncIterator[CatalogSnapshot]: current: Final = self.current() scoped: Final = self._operation.get() if current is not None and scoped is not None and scoped[2] == id(asyncio.current_task()): @@ -176,7 +160,7 @@ class TargetCatalog: return from litellm.proxy._experimental.mcp_server.tool_registry import global_mcp_tool_registry - shared: Final = await self._acquire_snapshot(request) + shared: Final = await self._fresh_snapshot() snapshot: Final = replace( shared, servers=MappingProxyType({key: value.model_copy(deep=True) for key, value in shared.servers.items()}), @@ -242,27 +226,6 @@ class TargetCatalog: _check_oauth_revision(selected, resolved) return resolved - @staticmethod - async def aggregate_list( - servers: Sequence[MCPServer], - fetch: Callable[[MCPServer], Awaitable[tuple[list[SDKTool], ServerOutcome]]], - server_key: Callable[[MCPServer], str], - ) -> AggregateToolListing: - from litellm.proxy._experimental.mcp_server.faults.list_outcomes import AggregateToolListing - - tasks: Final = tuple(asyncio.ensure_future(fetch(server)) for server in servers) - try: - results: Final = await asyncio.gather(*tasks) - return AggregateToolListing( - tools=[tool for tools, _ in results for tool in tools], - outcomes={server_key(server): outcome for server, (_, outcome) in zip(servers, results)}, - ) - finally: - for task in tasks: - if not task.done(): - task.cancel() - await asyncio.gather(*tasks, return_exceptions=True) - async def reload(self) -> None: async with self._refresh_lock: await self._publish_refresh() @@ -559,18 +522,6 @@ def global_manager() -> MCPServerManager: return global_mcp_server_manager -def public_catalog_operation(function: Callable[_P, Awaitable[_R]]) -> Callable[_P, Awaitable[_R]]: - parameters: Final = signature(function) - - @wraps(function) - async def wrapped(*args: _P.args, **kwargs: _P.kwargs) -> _R: # kwargs-ok: preserves ParamSpec - from starlette.requests import Request - - arguments: Final = parameters.bind(*args, **kwargs).arguments - request: Final = arguments.get("request") - if not isinstance(request, Request): - raise TypeError("Public MCP operations require a Request") - async with global_manager().catalog.operation(request=request): - return await function(*args, **kwargs) - - return wrapped +public_catalog_operation: Final[Callable[[Callable[_P, Awaitable[_R]]], Callable[_P, Awaitable[_R]]]] = ( + catalog_operation(global_manager) +) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 6299068ae6c..f1e19b17d1e 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -2919,7 +2919,7 @@ async def register_client(request: Request, mcp_server_name: str | None = None): ) from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager - async with global_mcp_server_manager.catalog.operation(request=request): + async with global_mcp_server_manager.catalog.operation(): dummy_return: Final = { "client_id": mcp_server_name or "dummy_client", "client_secret": "dummy", diff --git a/litellm/proxy/_experimental/mcp_server/operations.py b/litellm/proxy/_experimental/mcp_server/operations.py index 4ca704d9afc..9bb28793e1e 100644 --- a/litellm/proxy/_experimental/mcp_server/operations.py +++ b/litellm/proxy/_experimental/mcp_server/operations.py @@ -1,5 +1,6 @@ """Shared MCP operation policy and dispatch.""" +import asyncio import traceback import types import uuid @@ -49,7 +50,7 @@ from litellm.proxy._experimental.mcp_server.byok_credential_cache import ( cache_byok_credential, get_cached_byok_credential, ) -from litellm.proxy._experimental.mcp_server.catalog import TargetCatalog, catalog_operation, global_manager +from litellm.proxy._experimental.mcp_server.catalog import catalog_operation, global_manager from litellm.proxy._experimental.mcp_server.contracts import ( AuthorizedToolCall, OperationContext, @@ -1153,17 +1154,24 @@ async def _get_tools_from_mcp_servers( verbose_logger.exception("Error getting tools from server %s: %s", server.name, e) return [], classify_list_exception(e) - listing: Final = await TargetCatalog.aggregate_list( - allowed_mcp_servers, _fetch_and_filter_server_tools, _aggregate_server_key - ) - all_tools: Final = listing.tools - server_outcomes: Final = listing.outcomes + # Fetch tools from all servers in parallel + tasks: Final = [_fetch_and_filter_server_tools(server) for server in allowed_mcp_servers] + results: Final = await asyncio.gather(*tasks) + + # Flatten results into single list + all_tools: Final[list[MCPTool]] = [tool for tools, _ in results for tool in tools] + server_outcomes: Final[dict[str, ServerOutcome]] = { + _aggregate_server_key(server): outcome + for server, (_, outcome) in zip(allowed_mcp_servers, results) + if server is not None + } # If logging is enabled, enrich spend_logs_metadata with counts if litellm_logging_obj: per_server_tool_counts: Final[dict[str, int]] = { - key: outcome.tool_count if isinstance(outcome, ServerListOk) else 0 - for key, outcome in server_outcomes.items() + _aggregate_server_key(server): len(server_tools) + for server, (server_tools, _) in zip(allowed_mcp_servers, results) + if server is not None } metadata_dict: Final = litellm_logging_obj.model_call_details.get("metadata") diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_admission.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_admission.py deleted file mode 100644 index cc689b4db87..00000000000 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_admission.py +++ /dev/null @@ -1,89 +0,0 @@ -from unittest.mock import Mock - -import pytest -from fastapi import HTTPException, Request - -from litellm.proxy._experimental.mcp_server.auth.admission import MCPAdmissionLimiter, admission_source - - -def test_client_and_worker_rate_budgets_recover_after_window(monkeypatch): - monkeypatch.setenv("LITELLM_MCP_PUBLIC_RPM", "2") - monkeypatch.setenv("LITELLM_MCP_PUBLIC_WORKER_RPM", "3") - clock = Mock(return_value=0.0) - limiter = MCPAdmissionLimiter(clock) - for source in ("a", "a"): - with limiter.admit(source): - pass - with pytest.raises(HTTPException) as client_error, limiter.admit("a"): - pytest.fail("client budget must reject") - assert client_error.value.status_code == 429 - assert client_error.value.headers == {"Retry-After": "60"} - with limiter.admit("b"): - pass - clock.return_value = 59.1 - with pytest.raises(HTTPException) as worker_error, limiter.admit("c"): - pytest.fail("worker budget must reject rotated sources") - assert worker_error.value.headers == {"Retry-After": "1"} - clock.return_value = 60.0 - with limiter.admit("a"), limiter.admit("c"): - pass - - -def test_inflight_limits_release_on_failure_across_window_boundary(monkeypatch): - monkeypatch.setenv("LITELLM_MCP_PUBLIC_MAX_IN_FLIGHT", "1") - monkeypatch.setenv("LITELLM_MCP_PUBLIC_WORKER_MAX_IN_FLIGHT", "2") - clock = Mock(return_value=0.0) - limiter = MCPAdmissionLimiter(clock) - with limiter.admit("a"): - with pytest.raises(HTTPException) as client_error, limiter.admit("a"): - pytest.fail("one client cannot occupy another permit") - assert client_error.value.headers == {"Retry-After": "1"} - with limiter.admit("b"): - clock.return_value = 60.0 - with pytest.raises(HTTPException) as worker_error, limiter.admit("c"): - pytest.fail("rotating sources cannot exceed active work") - assert worker_error.value.status_code == 429 - with pytest.raises(RuntimeError, match="upstream"), limiter.admit("c"): - raise RuntimeError("upstream") - with limiter.admit("a"), limiter.admit("c"): - pass - - -def test_source_capacity_never_evicts_live_budget(monkeypatch): - monkeypatch.setenv("LITELLM_MCP_PUBLIC_MAX_SOURCES", "2") - clock = Mock(return_value=0.0) - limiter = MCPAdmissionLimiter(clock) - with limiter.admit("a"): - with limiter.admit("b"): - pass - for index in range(100): - with pytest.raises(HTTPException) as error, limiter.admit(str(index)): - pytest.fail("source spray must be bounded") - assert error.value.status_code == 429 - clock.return_value = 60.0 - with limiter.admit("c"): - pass - with limiter.admit("a"): - pass - - -@pytest.mark.parametrize("value", ["0", "-1", "invalid"]) -def test_invalid_limits_are_not_silently_disabled(monkeypatch, value): - monkeypatch.setenv("LITELLM_MCP_PUBLIC_RPM", value) - with pytest.raises(ValueError, match=r"must be positive|invalid literal"): - MCPAdmissionLimiter() - - -@pytest.mark.parametrize("peer,trusted,xff,expected", [ - ("198.51.100.1", [], "203.0.113.1", "198.51.100.1"), - ("10.0.0.1", ["10.0.0.0/8"], "203.0.113.1, 198.51.100.1", "198.51.100.1"), - ("2001:db8::1", [], "", "2001:db8::/64"), - (None, [], "203.0.113.1", "unknown"), -]) -def test_source_identity_ignores_untrusted_forwarded_addresses(monkeypatch, peer, trusted, xff, expected): - from litellm.proxy import proxy_server - - monkeypatch.setattr(proxy_server, "general_settings", {"use_x_forwarded_for": True, "mcp_trusted_proxy_ranges": trusted}) - request = Request({"type": "http", "client": (peer, 1234) if peer else None, - "headers": [(b"x-forwarded-for", xff.encode())]}) - assert admission_source(request) == expected diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 270a26cb5ff..ba26d4024a4 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -15489,55 +15489,6 @@ async def test_catalog_fresh_lookup_does_not_fall_back_to_stale_grants_when_data assert manager.registry == {server.server_id: server} -@pytest.mark.asyncio -async def test_catalog_list_failure_cancels_and_joins_other_upstream_fetches(): - from litellm.proxy._experimental.mcp_server.catalog import TargetCatalog - - started: Final = asyncio.Event() - stopped: Final = asyncio.Event() - pending: Final = MCPServer(server_id="pending", name="pending", transport=MCPTransport.http) - broken: Final = MCPServer(server_id="broken", name="broken", transport=MCPTransport.http) - - async def fetch(server: MCPServer): - if server.server_id == broken.server_id: - await started.wait() - raise RuntimeError("unexpected fetch failure") - started.set() - try: - await asyncio.Event().wait() - finally: - stopped.set() - - with pytest.raises(RuntimeError, match="unexpected fetch failure"): - await TargetCatalog.aggregate_list((pending, broken), fetch, lambda server: server.server_id) - assert stopped.is_set() - - -@pytest.mark.asyncio -@pytest.mark.parametrize("headers", [[], [(b"authorization", b"Bearer arbitrary")]]) -async def test_public_catalog_budget_rejects_before_database_read(monkeypatch, headers): - from starlette.requests import Request - from litellm.proxy._experimental.mcp_server import mcp_server_manager - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import oauth_authorization_server_mcp - - monkeypatch.setenv("LITELLM_MCP_PUBLIC_RPM", "2") - read_rows = AsyncMock(return_value=[]) - _catalog_database(monkeypatch, read_rows) - manager = MCPServerManager() - monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager) - request = Request({"type": "http", "method": "GET", "path": "/.well-known/oauth-authorization-server/missing", - "headers": headers, "client": ("198.51.100.1", 1234), "scheme": "https", - "server": ("gateway.example", 443), "query_string": b""}) - for expected in (404, 404, 429): - with pytest.raises(HTTPException) as exc: - await oauth_authorization_server_mcp(request=request, mcp_server_name="missing") - assert exc.value.status_code == expected - assert int(exc.value.headers["Retry-After"]) > 0 - assert read_rows.await_count == 2 - await manager.catalog.list() - assert read_rows.await_count == 3 - - @pytest.mark.asyncio async def test_catalog_cancelled_registration_does_not_publish_partial_handlers(monkeypatch): from litellm.proxy._experimental.mcp_server import tool_registry @@ -15558,58 +15509,3 @@ async def test_catalog_cancelled_registration_does_not_publish_partial_handlers( assert [tool.name for tool in registry.list_tools()] == ["existing"] assert manager.registry == {} - -@pytest.mark.asyncio -@pytest.mark.parametrize("cancel", [False, True]) -async def test_public_catalog_inflight_rejection_precedes_read_and_releases_permit(monkeypatch, cancel): - from starlette.requests import Request - - monkeypatch.setenv("LITELLM_MCP_PUBLIC_WORKER_MAX_IN_FLIGHT", "1") - started = asyncio.Event() - release = asyncio.Event() - - async def blocked_rows(**kwargs): - started.set() - await release.wait() - return [_catalog_row()] - - read_rows = AsyncMock(side_effect=blocked_rows) - _catalog_database(monkeypatch, read_rows) - manager = MCPServerManager() - request = Request({"type": "http", "client": ("198.51.100.1", 1234), "headers": []}) - first = asyncio.create_task(manager.catalog.list(request=request)) - await asyncio.wait_for(started.wait(), 1) - with pytest.raises(HTTPException) as error: - await manager.catalog.list(request=request) - assert error.value.status_code == 429 - assert error.value.headers == {"Retry-After": "1"} - assert read_rows.await_count == 1 - authenticated = asyncio.create_task(manager.catalog.list()) - if cancel: - first.cancel() - release.set() - if cancel: - with pytest.raises(asyncio.CancelledError): - await first - else: - assert "catalog-server" in await first - assert "catalog-server" in await asyncio.wait_for(authenticated, 1) - assert "catalog-server" in await manager.catalog.list(request=request) - assert read_rows.await_count == 3 - - -@pytest.mark.asyncio -async def test_public_catalog_releases_permit_before_operation_and_nested_use(monkeypatch): - from starlette.requests import Request - - monkeypatch.setenv("LITELLM_MCP_PUBLIC_WORKER_MAX_IN_FLIGHT", "1") - read_rows = AsyncMock(return_value=[_catalog_row()]) - _catalog_database(monkeypatch, read_rows) - manager = MCPServerManager() - request = Request({"type": "http", "client": ("198.51.100.1", 1234), "headers": []}) - async with manager.catalog.operation(request=request): - async with manager.catalog.operation(request=request): - assert read_rows.await_count == 1 - other = await asyncio.create_task(manager.catalog.list(request=request)) - assert "catalog-server" in other - assert read_rows.await_count == 2 From c31a3129e38589f5f98366911668e007fbca232f Mon Sep 17 00:00:00 2001 From: joshua Date: Wed, 23 Sep 2026 05:35:23 +0000 Subject: [PATCH 16/24] feat(mcp): gate catalog refresh on a database revision marker Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../migration.sql | 13 +++ .../proxy/_experimental/mcp_server/catalog.py | 24 +++-- litellm/proxy/_experimental/mcp_server/db.py | 15 +++ litellm/repositories/config_repository.py | 3 + .../auth/test_user_api_key_auth_mcp.py | 2 + .../mcp_server/test_byok_oauth_endpoints.py | 1 + .../mcp_server/test_db_credentials.py | 9 +- .../mcp_server/test_discoverable_endpoints.py | 3 +- .../mcp_server/test_mcp_server.py | 6 ++ .../mcp_server/test_mcp_server_manager.py | 94 ++++++++++++++++++- .../mcp_server/test_mcp_sigv4_auth.py | 2 + .../mcp_server/test_oauth2_flow_backfill.py | 1 + .../test_oauth_issuer_stamp_backfill.py | 1 + 13 files changed, 159 insertions(+), 15 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260923000000_add_mcp_catalog_revision_trigger/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260923000000_add_mcp_catalog_revision_trigger/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260923000000_add_mcp_catalog_revision_trigger/migration.sql new file mode 100644 index 00000000000..bc43353063e --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260923000000_add_mcp_catalog_revision_trigger/migration.sql @@ -0,0 +1,13 @@ +CREATE OR REPLACE FUNCTION litellm_bump_mcp_catalog_revision() RETURNS TRIGGER AS $$ +BEGIN + INSERT INTO "LiteLLM_Config" ("param_name", "reload_revision") + VALUES ('mcp_catalog', 1) + ON CONFLICT ("param_name") DO UPDATE SET "reload_revision" = "LiteLLM_Config"."reload_revision" + 1; + RETURN NULL; +END; +$$ LANGUAGE plpgsql; + +DROP TRIGGER IF EXISTS "litellm_mcp_catalog_revision" ON "LiteLLM_MCPServerTable"; +CREATE TRIGGER "litellm_mcp_catalog_revision" +AFTER INSERT OR UPDATE OR DELETE ON "LiteLLM_MCPServerTable" +FOR EACH STATEMENT EXECUTE PROCEDURE litellm_bump_mcp_catalog_revision(); diff --git a/litellm/proxy/_experimental/mcp_server/catalog.py b/litellm/proxy/_experimental/mcp_server/catalog.py index ecf71089731..44cbb3c8203 100644 --- a/litellm/proxy/_experimental/mcp_server/catalog.py +++ b/litellm/proxy/_experimental/mcp_server/catalog.py @@ -81,6 +81,7 @@ class TargetCatalog: self._arrival_ticket = 0 self._completed_ticket = 0 self._shared_snapshot: CatalogSnapshot | None = None + self._applied_revision: int | None = None self._warned_shadowed_config_server_ids: frozenset[str] = frozenset() self._warned_capturing_config_server_ids: frozenset[str] = frozenset() self._operation: ContextVar[tuple[CatalogSnapshot, asyncio.Event, int] | None] = ContextVar( @@ -116,12 +117,17 @@ class TargetCatalog: if not should_load_db_object("mcp"): return _snapshot(self.manager, self._database_identity) + from litellm.proxy._experimental.mcp_server.db import get_mcp_catalog_revision + + revision: Final = await get_mcp_catalog_revision(prisma_client) + if revision is not None and revision == self._applied_revision and self._shared_snapshot is not None: + return self._shared_snapshot self._arrival_ticket += 1 arrival: Final = self._arrival_ticket async with self._refresh_lock: if arrival > self._completed_ticket: try: - await self._publish_refresh(reuse_unchanged=True) + await self._publish_refresh(revision, reuse_unchanged=True) except Exception as exc: raise HTTPException( status_code=503, detail="MCP server configuration could not be refreshed" @@ -227,10 +233,14 @@ class TargetCatalog: return resolved async def reload(self) -> None: - async with self._refresh_lock: - await self._publish_refresh() + from litellm.proxy._experimental.mcp_server.db import get_mcp_catalog_revision + from litellm.proxy.proxy_server import prisma_client - async def _publish_refresh(self, *, reuse_unchanged: bool = False) -> None: + revision: Final = await get_mcp_catalog_revision(prisma_client) if prisma_client is not None else None + async with self._refresh_lock: + await self._publish_refresh(revision) + + async def _publish_refresh(self, revision: int | None, *, reuse_unchanged: bool = False) -> None: covered: Final = self._arrival_ticket token: Final = self._operation.set(None) try: @@ -241,6 +251,7 @@ class TargetCatalog: raise else: self._shared_snapshot = _snapshot(self.manager, self._database_identity) + self._applied_revision = revision self._completed_ticket = covered finally: self._operation.reset(token) @@ -522,6 +533,5 @@ def global_manager() -> MCPServerManager: return global_mcp_server_manager -public_catalog_operation: Final[Callable[[Callable[_P, Awaitable[_R]]], Callable[_P, Awaitable[_R]]]] = ( - catalog_operation(global_manager) -) +def public_catalog_operation(function: Callable[_P, Awaitable[_R]]) -> Callable[_P, Awaitable[_R]]: + return catalog_operation(global_manager)(function) diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index ee7298fe146..96921056a29 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -40,6 +40,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( encrypt_value_helper, ) from litellm.proxy.utils import PrismaClient +from litellm.repositories.config_repository import ConfigRepository from litellm.repositories.object_permission_repository import ObjectPermissionRepository from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import ( @@ -513,6 +514,20 @@ async def _db_find_mcp_server_row( return await _mcp_server_table_actions(prisma_client).find_unique(where={"server_id": server_id}) +MCP_CATALOG_REVISION_PARAM_NAME: Final = "mcp_catalog" + + +async def get_mcp_catalog_revision(prisma_client: PrismaClient) -> int | None: + """The ``LiteLLM_Config`` revision the MCP server table trigger last published, or None when + no write has ever bumped it (or the trigger is not installed), meaning always reload.""" + row: Final = await ConfigRepository(prisma_client).table.find_unique( + where={"param_name": MCP_CATALOG_REVISION_PARAM_NAME} + ) + if row is None: + return None + return int(row.reload_revision or 0) + + async def _db_update_mcp_server_row( prisma_client: PrismaClient, server_id: str, diff --git a/litellm/repositories/config_repository.py b/litellm/repositories/config_repository.py index 8b8280622fd..bbc28936234 100644 --- a/litellm/repositories/config_repository.py +++ b/litellm/repositories/config_repository.py @@ -22,6 +22,9 @@ class _ConfigRow(Protocol): @property def param_value(self) -> object: ... + @property + def reload_revision(self) -> int | None: ... + class _ConfigTable(Protocol): async def find_unique(self, *, where: Mapping[str, str]) -> _ConfigRow | None: ... diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index b7959959688..4fc7b0b6bcc 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -7732,6 +7732,7 @@ class TestGatewaySessionAdmission: ) prisma = MagicMock() prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[]) + prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) with ( patch("litellm.proxy.auth.auth_checks.get_user_object", get_user_object), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), @@ -9824,6 +9825,7 @@ class TestScopedSessionAdmission: ) prisma = MagicMock() prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[]) + prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) with ( patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY), patch("litellm.proxy.auth.auth_checks.get_user_object", get_user_object), diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py index 204b50d0372..eed9caa72ee 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py @@ -634,6 +634,7 @@ async def test_execute_byok_tool_missing_credential_advertises_api_key_flow(monk monkeypatch.setattr(proxy_server, "should_load_db_object", lambda _kind: False) prisma = MagicMock() prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[]) + prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) prisma.db.litellm_mcpusercredentials.find_unique = AsyncMock(return_value=None) monkeypatch.setattr(proxy_server, "prisma_client", prisma) with pytest.raises(HTTPException) as exc_info: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py index cfcff73b857..68625dda3a2 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py @@ -464,7 +464,7 @@ class _MapTable: @pytest.mark.parametrize("quoted", [False, True]) async def test_secret_maps_create_update_round_trip(map_algorithm: str, field: str, quoted: bool) -> None: table: Final = _MapTable(quoted=quoted) - prisma: Final = SimpleNamespace(db=SimpleNamespace(litellm_mcpservertable=table)) + prisma: Final = SimpleNamespace(db=SimpleNamespace(litellm_mcpservertable=table, litellm_config=SimpleNamespace(find_unique=AsyncMock(return_value=None)))) original: Final = {"TOKEN": " sensitive-secret\n", "PREFIX": "v2:gcm:literal", "TEMPLATE": "Bearer ${TOKEN}"} create: Final = NewMCPServerRequest.model_validate({ "server_id": "srv-map", "transport": "http", "url": "https://up.example.com/mcp", field: original, @@ -537,7 +537,7 @@ async def test_secret_map_rotation_migrates_rekeys_and_preserves_corrupt( {"server_id": "encrypted", field: old, other: None}, ) prisma: Final = SimpleNamespace(db=SimpleNamespace( - litellm_mcpservertable=table, litellm_mcpserveroauthclient=SimpleNamespace(find_many=AsyncMock(return_value=[])) + litellm_mcpservertable=table, litellm_mcpserveroauthclient=SimpleNamespace(find_many=AsyncMock(return_value=[])), litellm_config=SimpleNamespace(find_unique=AsyncMock(return_value=None)) )) await rotate_mcp_server_credentials_master_key(prisma, touched_by="test", new_master_key="rotated-map-key") assert table.rows["broken"][field] == corrupt @@ -565,7 +565,7 @@ async def test_bulk_reads_isolate_corrupt_secret_maps(reader, field, map_algorit ] snapshot = [row.model_dump() for row in rows] table = SimpleNamespace(find_many=AsyncMock(return_value=rows)) - prisma = SimpleNamespace(db=SimpleNamespace(litellm_mcpservertable=table)) + prisma = SimpleNamespace(db=SimpleNamespace(litellm_mcpservertable=table, litellm_config=SimpleNamespace(find_unique=AsyncMock(return_value=None)))) result = await reader(prisma, ["broken", "healthy"]) if reader is get_mcp_servers else await reader(prisma) items = result.items if reader is get_mcp_submissions else result assert [row.server_id for row in items] == ["healthy"] @@ -584,7 +584,7 @@ async def test_bulk_reads_do_not_swallow_unrelated_validation_errors(reader): row = _prisma_map_row({"server_id": "invalid", "transport": "unsupported"}) table = SimpleNamespace(find_many=AsyncMock(return_value=[row])) - prisma = SimpleNamespace(db=SimpleNamespace(litellm_mcpservertable=table)) + prisma = SimpleNamespace(db=SimpleNamespace(litellm_mcpservertable=table, litellm_config=SimpleNamespace(find_unique=AsyncMock(return_value=None)))) request = reader(prisma, ["invalid"]) if reader is get_mcp_servers else reader(prisma) with pytest.raises(ValidationError, match="transport"): await request @@ -1502,6 +1502,7 @@ async def test_master_key_rotation_reencrypts_oauth_client_store(monkeypatch): prisma = MagicMock() prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[]) + prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) prisma.db.litellm_mcpserveroauthclient.find_many = AsyncMock( return_value=[SimpleNamespace(server_id="config_faros", credentials=blob_old)] ) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 921f144b468..872e3a63ef5 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -9349,6 +9349,7 @@ async def test_reload_servers_from_database_hydrates_dcr_clients(): prisma = MagicMock() prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[]) + prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) hydrate_spy = AsyncMock() with ( @@ -12652,7 +12653,7 @@ async def test_authorize_observes_committed_peer_server_changes(monkeypatch, cha updated_at=stamp + timedelta(seconds=1), ) read_rows = AsyncMock(return_value=[] if change == "delete" else [row]) - prisma = SimpleNamespace(db=SimpleNamespace(litellm_mcpservertable=SimpleNamespace(find_many=read_rows))) + prisma = SimpleNamespace(db=SimpleNamespace(litellm_mcpservertable=SimpleNamespace(find_many=read_rows), litellm_config=SimpleNamespace(find_unique=AsyncMock(return_value=None)))) monkeypatch.setattr(proxy_server, "prisma_client", prisma) monkeypatch.setattr( global_mcp_server_manager, "registry", {} if change == "create" else {old_server.server_id: old_server} diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index ba9a509ccf5..c019c226736 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -5390,6 +5390,7 @@ class TestMCPServerManagerReload: mock_prisma = MagicMock() mock_prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[db_row]) + mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) with ( patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", @@ -5432,6 +5433,7 @@ class TestMCPServerManagerReload: mock_prisma = MagicMock() mock_prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[db_row]) + mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) with ( patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", @@ -5487,6 +5489,7 @@ class TestMCPServerManagerReload: mock_prisma.db.litellm_mcpservertable.find_many = AsyncMock( return_value=[healthy_row, bad_row, another_healthy_row] ) + mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) with ( patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", @@ -5557,6 +5560,7 @@ class TestMCPServerManagerReload: mock_prisma = MagicMock() mock_prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[healthy_row, bad_openapi_row]) + mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) with ( patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", @@ -8718,6 +8722,7 @@ async def test_get_active_submitted_mcp_server_ids_for_user_queries_active_rows( row.server_id = "submitted-1" prisma_client = MagicMock() prisma_client.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[row]) + prisma_client.db.litellm_config.find_unique = AsyncMock(return_value=None) result = await get_active_submitted_mcp_server_ids_for_user(prisma_client, "submitter-user") @@ -8738,6 +8743,7 @@ async def test_get_active_submitted_mcp_server_ids_for_user_empty_user_id_skips_ prisma_client = MagicMock() prisma_client.db.litellm_mcpservertable.find_many = AsyncMock() + prisma_client.db.litellm_config.find_unique = AsyncMock(return_value=None) assert await get_active_submitted_mcp_server_ids_for_user(prisma_client, "") == [] prisma_client.db.litellm_mcpservertable.find_many.assert_not_awaited() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index ba26d4024a4..eb0e8cba9f8 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -1095,7 +1095,7 @@ class TestMCPServerManager: table = SimpleNamespace( find_many=AsyncMock(return_value=[_row(cached.server_id, corrupted), _row("healthy-sibling", stored)]) ) - prisma = SimpleNamespace(db=SimpleNamespace(litellm_mcpservertable=table)) + prisma = SimpleNamespace(db=SimpleNamespace(litellm_config=SimpleNamespace(find_unique=AsyncMock(return_value=None)), litellm_mcpservertable=table)) monkeypatch.setattr(proxy_server, "prisma_client", prisma) with caplog.at_level(logging.DEBUG, logger="LiteLLM"): @@ -14530,12 +14530,19 @@ def _catalog_row(name="initial"): ) -def _catalog_database(monkeypatch, read_rows): +def _catalog_database(monkeypatch, read_rows, read_revision=None): from types import SimpleNamespace from litellm.proxy import proxy_server - client = SimpleNamespace(db=SimpleNamespace(litellm_mcpservertable=SimpleNamespace(find_many=read_rows))) + if read_revision is None: + read_revision = AsyncMock(return_value=None) + client = SimpleNamespace( + db=SimpleNamespace( + litellm_mcpservertable=SimpleNamespace(find_many=read_rows), + litellm_config=SimpleNamespace(find_unique=read_revision), + ) + ) monkeypatch.setattr(proxy_server, "prisma_client", client) @@ -14786,6 +14793,7 @@ async def test_catalog_reload_preserves_concurrent_config_discovery_and_routes(a prisma: Final = MagicMock() prisma.db.litellm_mcpservertable.find_many = AsyncMock(side_effect=read_rows) + prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) with ( patch("litellm.proxy.proxy_server.prisma_client", prisma), patch("litellm.proxy._experimental.mcp_server.discoverable_endpoints.hydrate_config_server_dcr_client", side_effect=hydrate), @@ -14817,6 +14825,7 @@ async def test_catalog_reload_does_not_restore_replaced_config_credentials_or_ro prisma: Final = MagicMock() prisma.db.litellm_mcpservertable.find_many = AsyncMock(side_effect=read_rows) + prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) with patch("litellm.proxy.proxy_server.prisma_client", prisma): await manager.reload_servers_from_database() current: Final = manager.get_mcp_server_by_id(server.server_id) @@ -14897,6 +14906,7 @@ async def test_catalog_reload_keeps_new_route_owner_over_earlier_route(change, m prisma: Final = MagicMock() prisma.db.litellm_mcpservertable.find_many = AsyncMock(side_effect=read_rows) + prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) with patch("litellm.proxy.proxy_server.prisma_client", prisma): await manager.reload_servers_from_database() if change == "deleted" or not mapped: @@ -14921,6 +14931,7 @@ async def test_catalog_observes_committed_update_and_delete_without_background_r updated: Final = row.model_copy(update={"url": "https://second.example.com/mcp", "updated_at": timestamp + timedelta(seconds=1)}) prisma: Final = MagicMock() prisma.db.litellm_mcpservertable.find_many = AsyncMock(side_effect=([row], [updated], [])) + prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) manager: Final = MCPServerManager() with ( patch("litellm.proxy.proxy_server.prisma_client", prisma), @@ -14947,6 +14958,7 @@ async def test_catalog_rebuilt_unchanged_server_keeps_discovered_tool_routes(): manager: Final = MCPServerManager() prisma: Final = MagicMock() prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[row]) + prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) with patch("litellm.proxy.proxy_server.prisma_client", prisma): await manager.reload_servers_from_database() server: Final = manager.get_mcp_server_by_id(row.server_id) @@ -14996,6 +15008,7 @@ async def test_catalog_reload_retains_routes_discovered_for_a_late_server(change prisma: Final = MagicMock() prisma.db.litellm_mcpservertable.find_many = AsyncMock(side_effect=read_rows) + prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) task: Final = asyncio.create_task(publish()) try: with patch("litellm.proxy.proxy_server.prisma_client", prisma): @@ -15033,6 +15046,7 @@ async def test_catalog_openapi_refresh_does_not_restore_removed_operations(tmp_p url="https://upstream.example", transport=MCPTransport.http, spec_path=str(spec_path)) prisma: Final = MagicMock() prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[row]) + prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) upstream: Final = respx_mock.get("https://upstream.example/retained").respond(200, json={"value": "retained"}) with patch("litellm.proxy.proxy_server.prisma_client", prisma): await manager.reload_servers_from_database() @@ -15070,6 +15084,7 @@ async def test_catalog_lookup_uses_one_snapshot_until_operation_finishes(): updated: Final = row.model_copy(update={"url": "https://second.example.com/mcp", "updated_at": row.updated_at + timedelta(seconds=1)}) prisma: Final = MagicMock() prisma.db.litellm_mcpservertable.find_many = AsyncMock(side_effect=([row], [updated], [updated])) + prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) manager: Final = MCPServerManager() with ( patch("litellm.proxy.proxy_server.prisma_client", prisma), @@ -15095,6 +15110,7 @@ async def test_catalog_failed_reload_preserves_published_discovery_state(): manager._upstream_initialize_instructions_by_server_id[server.server_id] = "healthy instructions" prisma: Final = MagicMock() prisma.db.litellm_mcpservertable.find_many = AsyncMock(side_effect=RuntimeError("database unavailable")) + prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) with patch("litellm.proxy.proxy_server.prisma_client", prisma): with pytest.raises(RuntimeError, match="database unavailable"): await manager.reload_servers_from_database() @@ -15118,6 +15134,7 @@ async def test_catalog_cancellation_retains_state_and_releases_refresh_lock(): manager._upstream_initialize_instructions_by_server_id[server.server_id] = "healthy instructions" prisma: Final = MagicMock() prisma.db.litellm_mcpservertable.find_many = AsyncMock(side_effect=blocked_read) + prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) with patch("litellm.proxy.proxy_server.prisma_client", prisma): pending = asyncio.create_task(manager.reload_servers_from_database()) await asyncio.wait_for(entered.wait(), timeout=2) @@ -15171,6 +15188,7 @@ async def test_catalog_cancelled_openapi_refresh_retains_tools_and_discovery(mon url="https://after.example/mcp", spec_path="after.json", updated_at=stamp + timedelta(seconds=1), auth_type=MCPAuth.oauth2) prisma: Final = MagicMock() prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[row]) + prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) async def cancelled_registration(*args: object, **kwargs: object) -> None: registry.register_tool("staged-new", "new", {}, lambda: "new") @@ -15198,6 +15216,7 @@ async def test_catalog_snapshot_identity_is_independent_of_worker_oauth_discover url="https://upstream.example/mcp", auth_type=MCPAuth.oauth2, updated_at=datetime.now()) prisma: Final = MagicMock() prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[row]) + prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) manager: Final = MCPServerManager() with ( patch("litellm.proxy.proxy_server.prisma_client", prisma), @@ -15225,6 +15244,7 @@ async def test_catalog_failed_openapi_row_does_not_publish_partial_handlers(monk url="https://upstream.example/mcp", spec_path="broken.json") prisma: Final = MagicMock() prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[row]) + prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) async def broken_registration(server: MCPServer, **kwargs: object) -> None: registry.register_tool("broken-partial", "partial", {}, lambda: "must not run") @@ -15478,6 +15498,7 @@ async def test_catalog_fresh_lookup_does_not_fall_back_to_stale_grants_when_data manager.registry = {server.server_id: server} prisma: Final = MagicMock() prisma.db.litellm_mcpservertable.find_many = AsyncMock(side_effect=RuntimeError("unavailable")) + prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) with ( patch("litellm.proxy.proxy_server.prisma_client", prisma), patch("litellm.proxy.proxy_server.general_settings", {}), @@ -15509,3 +15530,70 @@ async def test_catalog_cancelled_registration_does_not_publish_partial_handlers( assert [tool.name for tool in registry.list_tools()] == ["existing"] assert manager.registry == {} + + +def _revision_row(revision): + from types import SimpleNamespace + + return SimpleNamespace(reload_revision=revision) + + +@pytest.mark.asyncio +async def test_catalog_unchanged_revision_skips_table_reload(monkeypatch): + read_rows = AsyncMock(side_effect=([_catalog_row()], [_catalog_row("updated")])) + read_revision = AsyncMock(return_value=_revision_row(7)) + _catalog_database(monkeypatch, read_rows, read_revision) + manager = MCPServerManager() + assert (await manager.catalog.list())["catalog-server"].name == "initial" + assert (await manager.catalog.list())["catalog-server"].name == "initial" + read_rows.assert_awaited_once() + assert read_revision.await_count == 2 + + +@pytest.mark.asyncio +async def test_catalog_changed_revision_reloads_table(monkeypatch): + read_rows = AsyncMock(side_effect=([_catalog_row()], [_catalog_row("updated")])) + read_revision = AsyncMock(side_effect=[_revision_row(7), _revision_row(8)]) + _catalog_database(monkeypatch, read_rows, read_revision) + manager = MCPServerManager() + assert (await manager.catalog.list())["catalog-server"].name == "initial" + assert (await manager.catalog.list())["catalog-server"].name == "updated" + assert read_rows.await_count == 2 + assert read_revision.await_count == 2 + + +@pytest.mark.asyncio +async def test_catalog_absent_revision_row_reloads_table_every_operation(monkeypatch): + read_rows = AsyncMock(side_effect=([_catalog_row()], [_catalog_row("updated")])) + read_revision = AsyncMock(return_value=None) + _catalog_database(monkeypatch, read_rows, read_revision) + manager = MCPServerManager() + assert (await manager.catalog.list())["catalog-server"].name == "initial" + assert (await manager.catalog.list())["catalog-server"].name == "updated" + assert read_rows.await_count == 2 + + +@pytest.mark.asyncio +async def test_catalog_failed_refresh_does_not_apply_the_read_revision(monkeypatch): + read_rows = AsyncMock(side_effect=([_catalog_row()], RuntimeError("unavailable"), [_catalog_row("updated")])) + read_revision = AsyncMock(side_effect=[_revision_row(7), _revision_row(8), _revision_row(8)]) + _catalog_database(monkeypatch, read_rows, read_revision) + manager = MCPServerManager() + assert (await manager.catalog.list())["catalog-server"].name == "initial" + with pytest.raises(HTTPException) as error: + await manager.catalog.list() + assert error.value.status_code == 503 + assert (await manager.catalog.list())["catalog-server"].name == "updated" + assert read_rows.await_count == 3 + + +@pytest.mark.asyncio +async def test_catalog_reload_applies_revision_so_next_operation_skips_read(monkeypatch): + read_rows = AsyncMock(return_value=[_catalog_row()]) + read_revision = AsyncMock(return_value=_revision_row(7)) + _catalog_database(monkeypatch, read_rows, read_revision) + manager = MCPServerManager() + await manager.reload_servers_from_database() + assert (await manager.catalog.list())["catalog-server"].name == "initial" + read_rows.assert_awaited_once() + assert read_revision.await_count == 2 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py index 8cf3bc6fcc7..42f8bc4a9ff 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py @@ -993,6 +993,7 @@ class TestRotateCredentials: mock_prisma = MagicMock() mock_prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[server]) + mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) mock_prisma.db.litellm_mcpservertable.update = AsyncMock() mock_prisma.db.litellm_mcpserveroauthclient.find_many = AsyncMock(return_value=[]) @@ -1041,6 +1042,7 @@ class TestRotateCredentials: mock_prisma = MagicMock() mock_prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[server]) + mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) mock_prisma.db.litellm_mcpservertable.update = AsyncMock() mock_prisma.db.litellm_mcpserveroauthclient.find_many = AsyncMock(return_value=[]) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth2_flow_backfill.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth2_flow_backfill.py index c1239c228aa..ce079878328 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth2_flow_backfill.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth2_flow_backfill.py @@ -134,6 +134,7 @@ def _byok_key_row(server_id): def _mock_prisma(null_rows, token_rows): mock_prisma = MagicMock() mock_prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=null_rows) + mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) mock_prisma.db.litellm_mcpservertable.update_many = AsyncMock(return_value=MagicMock()) mock_prisma.db.litellm_mcpusercredentials.find_many = AsyncMock(return_value=token_rows) return mock_prisma diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_issuer_stamp_backfill.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_issuer_stamp_backfill.py index b6c946b95fa..a7e9ac80a1d 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_issuer_stamp_backfill.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_issuer_stamp_backfill.py @@ -29,6 +29,7 @@ def _row(**overrides): def _prisma(rows): prisma_client = MagicMock() prisma_client.db.litellm_mcpservertable.find_many = AsyncMock(return_value=rows) + prisma_client.db.litellm_config.find_unique = AsyncMock(return_value=None) prisma_client.db.litellm_mcpservertable.update = AsyncMock() return prisma_client From 1453a133694e0206842cf8cab457f5804a544b53 Mon Sep 17 00:00:00 2001 From: joshua Date: Wed, 23 Sep 2026 05:47:43 +0000 Subject: [PATCH 17/24] fix(mcp): refresh waiters that observed a newer catalog revision Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/_experimental/mcp_server/catalog.py | 2 +- .../mcp_server/test_mcp_server_manager.py | 43 +++++++++++++++++++ 2 files changed, 44 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_experimental/mcp_server/catalog.py b/litellm/proxy/_experimental/mcp_server/catalog.py index 44cbb3c8203..a24935d6d54 100644 --- a/litellm/proxy/_experimental/mcp_server/catalog.py +++ b/litellm/proxy/_experimental/mcp_server/catalog.py @@ -125,7 +125,7 @@ class TargetCatalog: self._arrival_ticket += 1 arrival: Final = self._arrival_ticket async with self._refresh_lock: - if arrival > self._completed_ticket: + if arrival > self._completed_ticket or revision != self._applied_revision: try: await self._publish_refresh(revision, reuse_unchanged=True) except Exception as exc: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index eb0e8cba9f8..aead498d408 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -15597,3 +15597,46 @@ async def test_catalog_reload_applies_revision_so_next_operation_skips_read(monk assert (await manager.catalog.list())["catalog-server"].name == "initial" read_rows.assert_awaited_once() assert read_revision.await_count == 2 + + +@pytest.mark.asyncio +async def test_catalog_waiter_that_observed_newer_revision_refreshes(monkeypatch): + entered = asyncio.Event() + release = asyncio.Event() + waiter_saw_new_revision = asyncio.Event() + current_revision = [5] + reads = [0] + + async def read_rows(**_kwargs): + reads[0] += 1 + if not entered.is_set(): + entered.set() + await release.wait() + return [_catalog_row()] + if reads[0] == 2: + return [_catalog_row()] + return [ + _catalog_row(), + _catalog_row("added").model_copy(update={"server_id": "added-server"}), + ] + + async def read_current_revision(**_kwargs): + if current_revision[0] == 6: + waiter_saw_new_revision.set() + return _revision_row(current_revision[0]) + + read_rows_mock = AsyncMock(side_effect=read_rows) + read_revision = AsyncMock(side_effect=read_current_revision) + _catalog_database(monkeypatch, read_rows_mock, read_revision) + manager = MCPServerManager() + first = asyncio.create_task(manager.catalog.list()) + await entered.wait() + middle = asyncio.create_task(manager.catalog.list()) + await asyncio.sleep(0) + current_revision[0] = 6 + last = asyncio.create_task(manager.catalog.list()) + await waiter_saw_new_revision.wait() + release.set() + _, _, servers = await asyncio.gather(first, middle, last) + assert "added-server" in servers + assert reads[0] == 3 From 7aea330dfd73285af225ee0cb2deec7549ecb708 Mon Sep 17 00:00:00 2001 From: joshua Date: Wed, 23 Sep 2026 06:07:42 +0000 Subject: [PATCH 18/24] test(mcp): awaitable catalog revision doubles in prisma mocks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/guardrails/guardrail_hooks/test_mcp_security.py | 1 + .../proxy/management_endpoints/test_mcp_management_endpoints.py | 2 ++ tests/test_litellm/proxy/test_dynamic_mcp_route.py | 1 + 3 files changed, 4 insertions(+) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_mcp_security.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_mcp_security.py index 7ed0596fbe2..ba6d54c20cb 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_mcp_security.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_mcp_security.py @@ -219,6 +219,7 @@ async def test_guardrail_observes_saved_server_creation_and_deletion_on_another_ url="https://upstream.example/mcp") prisma = MagicMock() prisma.db.litellm_mcpservertable.find_many = AsyncMock(side_effect=([row], [])) + prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) data = {"tools": [{"type": "mcp", "server_url": "litellm_proxy/mcp/peer-server"}], "guardrails": ["test-mcp-security"]} with ( diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index 175c8e0d760..13f5184d032 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -2775,6 +2775,7 @@ class TestTemporaryMCPSessionEndpoints: prisma.db.litellm_mcpservertable.find_many = AsyncMock(side_effect=persisted_rows) prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=None) prisma.db.litellm_mcpservertable.find_first = AsyncMock(return_value=None) + prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) request: Final = Request( {"type": "http", "method": "GET", "scheme": "http", "server": ("localhost", 4000), "path": "/v1/mcp/server/oauth/saved-server/authorize", "headers": [], @@ -8016,6 +8017,7 @@ async def test_saved_server_authorize_denial_does_not_dispatch_upstream(): authorization_url="https://upstream.example/authorize", token_url="https://upstream.example/token") prisma: Final = MagicMock() prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[row]) + prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) user: Final = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.INTERNAL_USER) with ( patch("litellm.proxy.proxy_server.prisma_client", prisma), diff --git a/tests/test_litellm/proxy/test_dynamic_mcp_route.py b/tests/test_litellm/proxy/test_dynamic_mcp_route.py index 5426cd92005..c70b737728c 100644 --- a/tests/test_litellm/proxy/test_dynamic_mcp_route.py +++ b/tests/test_litellm/proxy/test_dynamic_mcp_route.py @@ -644,6 +644,7 @@ async def test_dynamic_route_observes_committed_peer_catalog_changes(monkeypatch prisma = SimpleNamespace(db=SimpleNamespace( litellm_mcpservertable=SimpleNamespace(find_many=read_rows), litellm_mcptoolsettable=SimpleNamespace(find_first=AsyncMock(return_value=None)), + litellm_config=SimpleNamespace(find_unique=AsyncMock(return_value=None)), )) manager = mcp_server_manager.MCPServerManager() if change != "create": From 515c1248242e4215a8269ccc69095b08a7fee808 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Sat, 3 Oct 2026 14:25:57 -0700 Subject: [PATCH 19/24] refactor(mcp): satisfy catalog refresh lint and type gates --- .../proxy/_experimental/mcp_server/catalog.py | 82 +++++++++++-------- .../mcp_server/mcp_server_manager.py | 20 +++-- .../_experimental/mcp_server/tool_registry.py | 4 +- .../mcp_management_endpoints.py | 4 +- 4 files changed, 63 insertions(+), 47 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/catalog.py b/litellm/proxy/_experimental/mcp_server/catalog.py index 4b43d97d2e2..1395c4b6b08 100644 --- a/litellm/proxy/_experimental/mcp_server/catalog.py +++ b/litellm/proxy/_experimental/mcp_server/catalog.py @@ -5,11 +5,12 @@ from __future__ import annotations import asyncio import hashlib import json -from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequence +from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence from contextlib import asynccontextmanager from contextvars import ContextVar from dataclasses import dataclass, replace from functools import wraps +from itertools import chain from types import MappingProxyType from typing import TYPE_CHECKING, Final, ParamSpec, TypeVar @@ -158,7 +159,7 @@ class TargetCatalog: raise HTTPException(status_code=503, detail="MCP server configuration changed; retry the operation") @asynccontextmanager - async def operation(self) -> AsyncIterator[CatalogSnapshot]: + async def operation(self) -> AsyncGenerator[CatalogSnapshot]: current: Final = self.current() scoped: Final = self._operation.get() if current is not None and scoped is not None and scoped[2] == id(asyncio.current_task()): @@ -193,13 +194,13 @@ class TargetCatalog: from litellm.proxy._experimental.mcp_server.utils import normalize_server_name current: Final = self.manager.config_mcp_servers | self.manager.registry - unchanged_owners: Final = frozenset( - owner + unchanged: Final = ( + server for key, server in servers.items() if (candidate := current.get(key)) is not None and _configuration_identity(candidate) == _configuration_identity(server) - for owner in self.manager.owned_mapping_values(server) ) + unchanged_owners: Final = frozenset(chain.from_iterable(map(self.manager.owned_mapping_values, unchanged))) return MappingProxyType( {name: owner for name, owner in routing.items() if normalize_server_name(owner) in unchanged_owners} ) @@ -282,11 +283,13 @@ class TargetCatalog: try: with global_mcp_tool_registry.catalog_scope(initial_tools) as staged_tools: await self._reload(reuse_unchanged=reuse_unchanged) - refreshed_openapi_owners: Final = frozenset( - owner + refreshed_openapi: Final = ( + server for server in self.manager.registry.values() if server.spec_path and server is not live_registry.get(server.server_id) - for owner in self.manager.owned_mapping_values(server) + ) + refreshed_openapi_owners: Final = frozenset( + chain.from_iterable(map(self.manager.owned_mapping_values, refreshed_openapi)) ) live_routes: Final = self._unchanged_routing( previous_servers, @@ -357,32 +360,13 @@ class TargetCatalog: closed.set() self._staged_routing.reset(routing_token) - async def _reload(self, *, reuse_unchanged: bool) -> None: + async def _stage_servers(self, rows: Sequence[BaseModel], *, reuse_unchanged: bool) -> dict[str, MCPServer]: + from litellm.proxy._experimental.mcp_server.db import LiteLLM_MCPServerTable from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - _warn_on_shared_identifier_prefixes, carry_forward_resolved_oauth_endpoints, - config_ids_capturing_db_identifiers, oauth_endpoints_unresolved, warn_on_server_name_fields, ) - from litellm.proxy.management_endpoints.mcp_management_endpoints import ( - get_prisma_client_or_throw, - ) - - verbose_logger.debug("Loading MCP servers from database into registry...") - - prisma_client: Final = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") - # Load only "active", legacy "approved", and NULL (no approval workflow) rows. - # Pending/rejected servers are excluded at the DB level so we never load them. - from litellm.proxy._experimental.mcp_server.db import LiteLLM_MCPServerTable, get_runtime_mcp_server_rows - - raw_rows: Final[Sequence[BaseModel]] = await get_runtime_mcp_server_rows(prisma_client) - database_identity: Final = hashlib.sha256( - json.dumps( - tuple(sorted(json.dumps(row.model_dump(mode="json"), sort_keys=True, default=str) for row in raw_rows)) - ).encode() - ).hexdigest() - verbose_logger.info("Found %s MCP servers in database", len(raw_rows)) previous_registry: Final = self.manager.registry new_registry: Final[dict[str, MCPServer]] = {} @@ -390,7 +374,7 @@ class TargetCatalog: # Stage one: build every server. Stage two assigns short prefixes # against the *full* set so dedup is deterministic regardless of # iteration order. - for row in raw_rows: + for row in rows: try: server = LiteLLM_MCPServerTable.model_validate(row.model_dump()) existing_server = previous_registry.get(server.server_id) @@ -416,7 +400,7 @@ class TargetCatalog: alias=getattr(server, "alias", None), server_name=getattr(server, "server_name", None), ) - self.manager._warn_if_newly_blocked_stdio(server, existing_server) + self.manager.warn_if_newly_blocked_stdio(server, existing_server) verbose_logger.debug("Building server from DB: %s (%s)", server.server_id, server.server_name) # raw_rows come straight from the DB, so their global env var # values (like credentials) are still encrypted here, unlike the @@ -439,6 +423,35 @@ class TargetCatalog: e, ) + return new_registry + + async def _reload(self, *, reuse_unchanged: bool) -> None: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + config_ids_capturing_db_identifiers, + warn_on_shared_identifier_prefixes, + ) + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + get_prisma_client_or_throw, + ) + + verbose_logger.debug("Loading MCP servers from database into registry...") + + prisma_client: Final = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") + # Load only "active", legacy "approved", and NULL (no approval workflow) rows. + # Pending/rejected servers are excluded at the DB level so we never load them. + from litellm.proxy._experimental.mcp_server.db import get_runtime_mcp_server_rows + + raw_rows: Final[Sequence[BaseModel]] = await get_runtime_mcp_server_rows(prisma_client) + database_identity: Final = hashlib.sha256( + json.dumps( + tuple(sorted(json.dumps(row.model_dump(mode="json"), sort_keys=True, default=str) for row in raw_rows)) + ).encode() + ).hexdigest() + verbose_logger.info("Found %s MCP servers in database", len(raw_rows)) + + previous_registry: Final = self.manager.registry + new_registry: Final = await self._stage_servers(raw_rows, reuse_unchanged=reuse_unchanged) + # Assign short prefixes against the full candidate set without # publishing the staged registry to concurrent callers. registered_registry: Final[dict[str, MCPServer]] = {} @@ -470,14 +483,13 @@ class TargetCatalog: for server_id in previous_registry.keys() | registered_registry.keys(): if previous_registry.get(server_id) != registered_registry.get(server_id): - self.manager._invalidate_server_definition_caches(server_id) + self.manager.invalidate_server_definition_caches(server_id) self.manager.invalidate_oauth_discovery_state(server_id) self._database_identity = database_identity self.manager.registry = registered_registry if not reuse_unchanged: - self.manager._upstream_initialize_instructions_by_server_id.clear() - self.manager._upstream_initialize_instructions_probed_at.clear() - _warn_on_shared_identifier_prefixes(registered_registry.values()) + self.manager.clear_initialize_instructions() + warn_on_shared_identifier_prefixes(registered_registry.values()) # A discovery task may have published into ``previous_registry`` while # this replacement was being staged. Reconcile every published entry # synchronously after the swap so a lost publication cannot also leave diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index d75cb269b39..e043708dd7a 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1458,7 +1458,7 @@ def warn_on_server_name_fields( _warn("server_name", server_name) -def _warn_on_shared_identifier_prefixes(servers: Iterable[MCPServer]) -> None: +def warn_on_shared_identifier_prefixes(servers: Iterable[MCPServer]) -> None: """Warn once per identifier that several servers share. ``get_server_prefix`` resolves alias first, so two servers sharing a @@ -2644,7 +2644,7 @@ class MCPServerManager: self.assign_unique_short_prefix(new_server) _warn_legacy_delegate_auth_if_applicable(new_server, source="config") _warn_config_id_jag_server_outruns_sso(new_server) - self._invalidate_server_definition_caches(server_id) + self.invalidate_server_definition_caches(server_id) self.config_mcp_servers[server_id] = new_server self._set_oauth_discovery_deferred( server_id, @@ -2844,7 +2844,7 @@ class MCPServerManager: mappings make ``_get_mcp_server_from_tool_name`` resolve to a prefix that no longer exists in the live registry. """ - self._invalidate_server_definition_caches(server.server_id) + self.invalidate_server_definition_caches(server.server_id) self.remove_server_tool_routing(server) def remove_server_tool_routing(self, server: MCPServer) -> None: @@ -3257,10 +3257,10 @@ class MCPServerManager: # `credentials` field is the only one still encrypted here). # Re-decrypting plaintext would zero the values, so build with # env_vars_are_encrypted=False. - self._warn_if_newly_blocked_stdio(mcp_server, None) + self.warn_if_newly_blocked_stdio(mcp_server, None) new_server: Final = await self.build_mcp_server_from_table(mcp_server, env_vars_are_encrypted=False) self.assign_unique_short_prefix(new_server) - self._invalidate_server_definition_caches(mcp_server.server_id) + self.invalidate_server_definition_caches(mcp_server.server_id) self.registry[mcp_server.server_id] = new_server await self.maybe_register_openapi_tools(new_server) self.prime_oauth_metadata_discovery(new_server) @@ -3297,7 +3297,7 @@ class MCPServerManager: previous_server=self.registry[mcp_server.server_id], ) self.assign_unique_short_prefix(new_server) - self._invalidate_server_definition_caches(mcp_server.server_id) + self.invalidate_server_definition_caches(mcp_server.server_id) self.registry[mcp_server.server_id] = new_server await self.maybe_register_openapi_tools(new_server) self.prime_oauth_metadata_discovery(new_server) @@ -4343,6 +4343,10 @@ class MCPServerManager: ) raise_classified_list_failure(e, server.name, suppress_challenge=server.is_dcr_bridge) + def clear_initialize_instructions(self) -> None: + self._upstream_initialize_instructions_by_server_id.clear() + self._upstream_initialize_instructions_probed_at.clear() + def invalidate_discovery_lists(self, server_id: str) -> None: self._upstream_initialize_instructions_by_server_id.pop(server_id, None) self._upstream_initialize_instructions_probed_at.pop(server_id, None) @@ -4350,7 +4354,7 @@ class MCPServerManager: self._resource_discovery_cache.invalidate(server_id) self._template_discovery_cache.invalidate(server_id) - def _invalidate_server_definition_caches(self, server_id: str) -> None: + def invalidate_server_definition_caches(self, server_id: str) -> None: from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( # noqa: PLC0415 # lazy: discoverable_endpoints lazily imports this module's manager singleton invalidate_oauth_metadata_cache, ) @@ -4389,7 +4393,7 @@ class MCPServerManager: return server.server_id, hashlib.sha256(material.encode()).hexdigest() @staticmethod - def _warn_if_newly_blocked_stdio(row: LiteLLM_MCPServerTable, previous: MCPServer | None) -> None: + def warn_if_newly_blocked_stdio(row: LiteLLM_MCPServerTable, previous: MCPServer | None) -> None: if previous is None or previous.transport != row.transport: warn_if_mcp_stdio_blocked(row.alias or row.server_name, row.transport) diff --git a/litellm/proxy/_experimental/mcp_server/tool_registry.py b/litellm/proxy/_experimental/mcp_server/tool_registry.py index 2adf5db8f67..830ea7293f5 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_registry.py +++ b/litellm/proxy/_experimental/mcp_server/tool_registry.py @@ -1,6 +1,6 @@ import asyncio import json -from collections.abc import Callable, Iterator, Mapping +from collections.abc import Callable, Generator, Mapping from contextlib import contextmanager from contextvars import ContextVar from typing import TYPE_CHECKING, Any, Final @@ -40,7 +40,7 @@ class MCPToolRegistry: self.published_tools = tools @contextmanager - def catalog_scope(self, tools: Mapping[str, MCPTool]) -> Iterator[dict[str, MCPTool]]: + def catalog_scope(self, tools: Mapping[str, MCPTool]) -> Generator[dict[str, MCPTool]]: detached: Final = dict(tools) closed: Final = asyncio.Event() token: Final = self._catalog_tools.set((detached, closed)) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 4ab82b375fc..1b48ad67ff8 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -19,7 +19,7 @@ import functools import importlib import json import os -from collections.abc import AsyncIterator, Iterable, Mapping, Sequence +from collections.abc import AsyncGenerator, Iterable, Mapping, Sequence from contextlib import asynccontextmanager from dataclasses import dataclass from datetime import datetime, timedelta, timezone @@ -2227,7 +2227,7 @@ if MCP_AVAILABLE: server_id: str, user_api_key_dict: UserAPIKeyAuth, request: Request | None = None, - ) -> AsyncIterator[MCPServer]: + ) -> AsyncGenerator[MCPServer]: if await get_cached_temporary_mcp_server(server_id) is not None: yield await _get_cached_temporary_mcp_server_or_404(server_id, user_api_key_dict, request) return From 4ab35b91d1dcedf252651893b450c9fe994200d7 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Sat, 3 Oct 2026 16:48:46 -0700 Subject: [PATCH 20/24] test(mcp): provide catalog scope in toolset fixtures --- tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py index 4f48cea2468..5dfb61b52ef 100644 --- a/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py +++ b/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -1322,6 +1322,7 @@ async def test_responses_discovery_logs_sanitized_caller_headers(monkeypatch: py def _toolset_gateway_manager(toolset_id: str, server_id: str) -> types.SimpleNamespace: return types.SimpleNamespace( + catalog=types.SimpleNamespace(operation=nullcontext), get_registry=MagicMock(return_value={}), get_allowed_mcp_servers=AsyncMock(return_value=[]), get_mcp_servers_from_ids=MagicMock(return_value=[]), From ed3c941c1268205ce6eaafd2a63fb5a2d4ac9334 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Sat, 3 Oct 2026 17:04:48 -0700 Subject: [PATCH 21/24] test(mcp): exercise temporary OAuth through catalog operations --- .../test_mcp_management_endpoints.py | 33 ++++++++----------- 1 file changed, 13 insertions(+), 20 deletions(-) diff --git a/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py index 0364c77f519..1e0919f22f1 100644 --- a/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -10604,7 +10604,7 @@ class TestMCPServerResolutionCharacterization: effects: Final = _ResolutionEffects() from litellm.proxy.management_endpoints.mcp_management_endpoints import ( _cache_temporary_mcp_server_in_redis, - _get_cached_temporary_mcp_server_or_404, + _oauth_server_operation, _TemporaryMCPServerEntry, ) @@ -10689,11 +10689,8 @@ class TestMCPServerResolutionCharacterization: ): if expected_status in (403, 404): with pytest.raises(HTTPException) as exc_info: - await _get_cached_temporary_mcp_server_or_404( - server_id, - auth, - request=_make_mock_request(), - ) + async with _oauth_server_operation(server_id, auth, request=_make_mock_request()): + pytest.fail("Denied OAuth resolution must not enter the operation") expected_detail: Final = ( {"error": f"MCP server {server_id} not found"} @@ -10707,20 +10704,16 @@ class TestMCPServerResolutionCharacterization: effects.assert_no_writes() assert httpx_mock.calls.call_count == 0, f"{source}/{caller}: no upstream HTTP" else: - resolved: Final = await _get_cached_temporary_mcp_server_or_404( - server_id, - auth, - request=_make_mock_request(), - ) - expected_alias: Final = ( - db_server.alias - if source == "temp_draft" - else "LIT3974 OAuth" - if source == "config" - else temp_server.alias - ) - assert resolved.server_id == server_id, f"{source}/{caller}: resolved OAuth server" - assert resolved.alias == expected_alias, f"{source}/{caller}: resolved OAuth display name" + async with _oauth_server_operation(server_id, auth, request=_make_mock_request()) as resolved: + expected_alias: Final = ( + db_server.alias + if source == "temp_draft" + else "LIT3974 OAuth" + if source == "config" + else temp_server.alias + ) + assert resolved.server_id == server_id, f"{source}/{caller}: resolved OAuth server" + assert resolved.alias == expected_alias, f"{source}/{caller}: resolved OAuth display name" finally: mgmt_endpoints.litellm.cache = original_cache From 886aadefef49bc41e286dcd6e888d75ed0f6211d Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Sat, 3 Oct 2026 17:21:35 -0700 Subject: [PATCH 22/24] fix(mcp): share active catalog snapshots across discovery tasks --- .../proxy/_experimental/mcp_server/catalog.py | 7 ++-- .../_experimental/mcp_server/conftest.py | 5 +++ .../mcp_server/test_operations.py | 41 +++++++++++++++++++ 3 files changed, 49 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/catalog.py b/litellm/proxy/_experimental/mcp_server/catalog.py index 1395c4b6b08..606ceb22444 100644 --- a/litellm/proxy/_experimental/mcp_server/catalog.py +++ b/litellm/proxy/_experimental/mcp_server/catalog.py @@ -85,7 +85,7 @@ class TargetCatalog: self._applied_revision: int | None = None self._warned_shadowed_config_server_ids: frozenset[str] = frozenset() self._warned_capturing_config_server_ids: frozenset[str] = frozenset() - self._operation: ContextVar[tuple[CatalogSnapshot, asyncio.Event, int] | None] = ContextVar( + self._operation: ContextVar[tuple[CatalogSnapshot, asyncio.Event] | None] = ContextVar( "mcp_catalog_snapshot", default=None ) self._staged_routing: ContextVar[tuple[dict[str, str], asyncio.Event] | None] = ContextVar( @@ -161,8 +161,7 @@ class TargetCatalog: @asynccontextmanager async def operation(self) -> AsyncGenerator[CatalogSnapshot]: current: Final = self.current() - scoped: Final = self._operation.get() - if current is not None and scoped is not None and scoped[2] == id(asyncio.current_task()): + if current is not None: yield current return from litellm.proxy._experimental.mcp_server.tool_registry import global_mcp_tool_registry @@ -174,7 +173,7 @@ class TargetCatalog: routing=dict(shared.routing), ) closed: Final = asyncio.Event() - token: Final = self._operation.set((snapshot, closed, id(asyncio.current_task()))) + token: Final = self._operation.set((snapshot, closed)) try: with global_mcp_tool_registry.catalog_scope(snapshot.tools): yield snapshot diff --git a/tests/unit/proxy/_experimental/mcp_server/conftest.py b/tests/unit/proxy/_experimental/mcp_server/conftest.py index 51cab559797..742c96b1fe8 100644 --- a/tests/unit/proxy/_experimental/mcp_server/conftest.py +++ b/tests/unit/proxy/_experimental/mcp_server/conftest.py @@ -81,10 +81,14 @@ def config_only_mcp_manager_factory(): @pytest.fixture(autouse=True) def _hermetic_mcp_server_registry(): + from litellm.proxy._experimental.mcp_server.catalog import TargetCatalog + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) + saved_catalog = global_mcp_server_manager.catalog + global_mcp_server_manager.catalog = TargetCatalog(global_mcp_server_manager) saved_registry = dict(global_mcp_server_manager.registry) saved_config_servers = dict(global_mcp_server_manager.config_mcp_servers) saved_tool_mapping = dict(global_mcp_server_manager.tool_name_to_mcp_server_name_mapping) @@ -103,6 +107,7 @@ def _hermetic_mcp_server_registry(): global_mcp_server_manager.tool_name_to_mcp_server_name_mapping.clear() global_mcp_server_manager.tool_name_to_mcp_server_name_mapping.update(saved_tool_mapping) global_mcp_server_manager._oauth_discovery_slots = saved_oauth_slots + global_mcp_server_manager.catalog = saved_catalog @pytest.fixture(autouse=True) diff --git a/tests/unit/proxy/_experimental/mcp_server/test_operations.py b/tests/unit/proxy/_experimental/mcp_server/test_operations.py index b16b27ac919..06211f9a0d2 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_operations.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_operations.py @@ -665,3 +665,44 @@ async def test_tools_listing_preserves_explicit_spend_log_policy(log_enabled): ) assert result.tools == [] assert listing.await_args.kwargs["log_list_tools_to_spendlogs"] is log_enabled + + +@pytest.mark.asyncio +async def test_discovery_keeps_one_catalog_revision_across_concurrent_listings(monkeypatch): + from mcp.types import ( + DiscoverRequest, ListToolsResult, ListPromptsResult, ListResourcesResult, + ListResourceTemplatesResult, + ) + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server import operations + from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager + + manager = MCPServerManager() + original = MCPServer(server_id="catalog-server", name="before", transport=MCPTransport.http) + updated = original.model_copy(update={"name": "after"}) + manager.registry = {original.server_id: original} + monkeypatch.setattr(proxy_server, "prisma_client", None) + monkeypatch.setattr(operations, "global_mcp_server_manager", manager) + observed = [] + responses = (ListToolsResult(tools=[]), ListPromptsResult(prompts=[]), + ListResourcesResult(resources=[]), ListResourceTemplatesResult(resource_templates=[])) + + def listing(index): + async def run(*args, **kwargs): + async with manager.catalog.operation(): + observed.append(manager.get_mcp_server_by_id(original.server_id).name) + manager.registry = {updated.server_id: updated} + return responses[index] + return run + + with ( + patch.object(operations, "_execute_handle_list_tools", side_effect=listing(0)), + patch.object(operations, "_execute_list_prompts", side_effect=listing(1)), + patch.object(operations, "_execute_list_resources", side_effect=listing(2)), + patch.object(operations, "_execute_list_resource_templates", side_effect=listing(3)), + ): + result = await GatewayOperations().execute(DiscoverRequest(), prepare_context(UserAPIKeyAuth(user_id="scoped"))) + assert result.capabilities.model_dump(exclude_none=True) == {} + assert observed == ["before"] * 4 + async with manager.catalog.operation(): + assert manager.get_mcp_server_by_id(original.server_id).name == "after" From 57dfed6e39e73324e5a57fd25a8e65e6ca92cf0d Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Sat, 3 Oct 2026 17:46:31 -0700 Subject: [PATCH 23/24] fix(mcp): retain discovered routes across cached catalog operations --- .../proxy/_experimental/mcp_server/catalog.py | 10 +++-- .../mcp_server/test_mcp_server_manager.py | 41 +++++++++++++++++++ .../test_mcp_server_tool_calls_and_headers.py | 4 +- 3 files changed, 49 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/catalog.py b/litellm/proxy/_experimental/mcp_server/catalog.py index 606ceb22444..f319d0530da 100644 --- a/litellm/proxy/_experimental/mcp_server/catalog.py +++ b/litellm/proxy/_experimental/mcp_server/catalog.py @@ -167,10 +167,11 @@ class TargetCatalog: from litellm.proxy._experimental.mcp_server.tool_registry import global_mcp_tool_registry shared: Final = await self._fresh_snapshot() + initial_routing: Final = self._unchanged_routing(shared.servers, self.manager.published_tool_routes) snapshot: Final = replace( shared, servers=MappingProxyType({key: value.model_copy(deep=True) for key, value in shared.servers.items()}), - routing=dict(shared.routing), + routing=dict(initial_routing), ) closed: Final = asyncio.Event() token: Final = self._operation.set((snapshot, closed)) @@ -180,11 +181,12 @@ class TargetCatalog: finally: closed.set() self._operation.reset(token) - self._retain_discovered_routing(snapshot) + self._retain_discovered_routing(snapshot, initial_routing) - def _retain_discovered_routing(self, snapshot: CatalogSnapshot) -> None: + def _retain_discovered_routing(self, snapshot: CatalogSnapshot, initial_routing: Mapping[str, str]) -> None: self.manager.published_tool_routes = self.manager.published_tool_routes | self._unchanged_routing( - snapshot.servers, snapshot.routing + snapshot.servers, + {name: owner for name, owner in snapshot.routing.items() if initial_routing.get(name) != owner}, ) def _unchanged_routing( 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 8e7f694853b..f3299a03686 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 @@ -16990,3 +16990,44 @@ async def test_upstream_preparation_honors_case_sensitive_extra_command(monkeypa client: Final = await MCPServerManager()._create_mcp_client(server) assert client.stdio_config is not None assert client.stdio_config["command"] == "/opt/tools/CustomRunner" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("overlap", [False, True]) +async def test_catalog_cached_revision_retains_newly_discovered_tool_routes(monkeypatch, overlap): + from mcp.types import Tool + + read_rows = AsyncMock(return_value=[]) + _catalog_database(monkeypatch, read_rows, AsyncMock(return_value=_revision_row(7))) + manager = MCPServerManager() + first = MCPServer(server_id="first", name="first", transport=MCPTransport.http) + second = MCPServer(server_id="second", name="second", transport=MCPTransport.http) + manager.config_mcp_servers = {server.server_id: server for server in (first, second)} + manager.published_tool_routes = {"search": "first"} + ready = asyncio.Event() + release = asyncio.Event() + + async def read_catalog(): + async with manager.catalog.operation(): + assert manager._get_mcp_server_from_tool_name("search").server_id == "first" + ready.set() + await release.wait() + + reader = asyncio.create_task(read_catalog()) if overlap else None + try: + if reader is not None: + await asyncio.wait_for(ready.wait(), 2) + async with manager.catalog.operation(): + manager._create_prefixed_tools([Tool(name="search", input_schema={})], second) + assert manager.published_tool_routes["search"] == "second" + release.set() + if reader is not None: + await asyncio.wait_for(reader, 2) + async with manager.catalog.operation(): + assert manager._get_mcp_server_from_tool_name("search").server_id == "second" + assert manager.published_tool_routes["search"] == "second" + read_rows.assert_awaited_once() + finally: + if reader is not None: + reader.cancel() + await asyncio.gather(reader, return_exceptions=True) 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 6c6c1717777..036a1c80042 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 @@ -7735,8 +7735,8 @@ async def test_execute_mcp_tool_rest_server_id_injects_requested_server_credenti ), patch.object( mcp_operations.global_mcp_server_manager, - "get_registry", - return_value={ + "registry", + { requested_server.server_id: requested_server, collision_server.server_id: collision_server, }, From b63edce2dcb6abd17eaa41748593a05fae181f59 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Sat, 3 Oct 2026 18:28:08 -0700 Subject: [PATCH 24/24] fix(mcp): preserve local tool ownership across discovery --- .../mcp_server/mcp_server_manager.py | 14 +++ .../mcp_server/test_mcp_server_manager.py | 86 +++++++++++++++++++ .../test_mcp_server_tool_calls_and_headers.py | 1 + 3 files changed, 101 insertions(+) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index e043708dd7a..c1b892605b3 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -5292,6 +5292,8 @@ class MCPServerManager: Returns: List of tools with prefixed names """ + from litellm.proxy._experimental.mcp_server.tool_registry import global_mcp_tool_registry + prefixed_tools: Final = [] prefix: Final = get_server_prefix(server) @@ -5304,6 +5306,11 @@ class MCPServerManager: # short ID) so call_tool can resolve regardless of which form a # caller / cached client is using. for spelling in iter_known_tool_name_spellings(original_name, server): + namespace_owner = self.server_owning_tool_name_prefix(spelling) + if namespace_owner is not None and namespace_owner.server_id != server.server_id: + continue + if namespace_owner is None and global_mcp_tool_registry.get_tool(spelling) is not None: + continue self.tool_name_to_mcp_server_name_mapping[spelling] = prefix verbose_logger.info("Successfully fetched %s tools from server %s", len(prefixed_tools), server.name) @@ -6466,6 +6473,13 @@ class MCPServerManager: Returns: MCPServer if found, None otherwise """ + from litellm.proxy._experimental.mcp_server.tool_registry import global_mcp_tool_registry + + # Local handlers belong to their registered namespace, even if a native + # discovery result or an older cached route claims the same spelling. + if global_mcp_tool_registry.get_tool(tool_name) is not None: + return self.server_owning_tool_name_prefix(tool_name) + registry_servers: Final = list(self.get_registry().values()) prefix_to_server: Final = self._known_prefix_to_server() 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 f3299a03686..ee78c926fc3 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 @@ -17031,3 +17031,89 @@ async def test_catalog_cached_revision_retains_newly_discovered_tool_routes(monk if reader is not None: reader.cancel() await asyncio.gather(reader, return_exceptions=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("poisoned_route", [False, True]) +async def test_cached_discovery_cannot_authorize_another_servers_local_handler(monkeypatch, poisoned_route): + from datetime import datetime + + from fastapi import HTTPException + from mcp.types import Tool + + from litellm.proxy._experimental.mcp_server import operations + from litellm.proxy._experimental.mcp_server.tool_registry import global_mcp_tool_registry + + _catalog_database(monkeypatch, AsyncMock(return_value=[]), AsyncMock(return_value=_revision_row(7))) + manager = MCPServerManager() + allowed = MCPServer(server_id="allowed", name="allowed", transport=MCPTransport.http) + private = MCPServer(server_id="private", name="private", transport=MCPTransport.http) + manager.config_mcp_servers = {server.server_id: server for server in (allowed, private)} + manager.published_tool_routes = {"getsecret": "private", "private-getsecret": "private"} + handler = AsyncMock(return_value="private result") + monkeypatch.setattr(global_mcp_tool_registry, "published_tools", {}) + global_mcp_tool_registry.register_tool("private-getsecret", "Private tool", {}, handler) + monkeypatch.setattr(operations, "global_mcp_server_manager", manager) + check = AsyncMock(return_value={}) + monkeypatch.setattr(manager, "pre_call_tool_check", check) + async with manager.catalog.operation(): + manager._create_prefixed_tools([Tool(name="private-getsecret", input_schema={})], allowed) + if poisoned_route: + manager.published_tool_routes["private-getsecret"] = "allowed" + async with manager.catalog.operation(): + with pytest.raises(HTTPException) as denied: + await operations._execute_mcp_tool( + name="private-getsecret", arguments={}, allowed_mcp_servers=[allowed], start_time=datetime.now() + ) + assert denied.value.status_code == 403 + handler.assert_not_awaited() + check.assert_not_awaited() + result = await operations._execute_mcp_tool( + name="private-getsecret", arguments={}, allowed_mcp_servers=[private], start_time=datetime.now() + ) + assert result.is_error is False + handler.assert_awaited_once_with() + assert check.await_args.kwargs["server"].server_id == private.server_id + assert manager._get_mcp_server_from_tool_name("allowed-private-getsecret").server_id == allowed.server_id + + +@pytest.mark.parametrize("local_handler", [False, True]) +def test_discovery_preserves_registered_tool_namespaces(monkeypatch, local_handler): + from mcp.types import Tool + + from litellm.proxy._experimental.mcp_server.tool_registry import global_mcp_tool_registry + + manager = MCPServerManager() + allowed = MCPServer(server_id="allowed", name="allowed", transport=MCPTransport.http) + private = MCPServer(server_id="private", name="private", alias="private-alias", transport=MCPTransport.http) + manager.registry = {server.server_id: server for server in (allowed, private)} + manager.published_tool_routes = {"private-getsecret": "private", "private-alias-getsecret": "private"} + monkeypatch.setattr(global_mcp_tool_registry, "published_tools", {}) + if local_handler: + global_mcp_tool_registry.register_tool("orphan", "Orphan local tool", {}, AsyncMock()) + tools = [Tool(name=name, input_schema={}) for name in ("private-getsecret", "private-alias-getsecret", "orphan")] + listed = manager._create_prefixed_tools(tools, allowed) + assert [tool.name for tool in listed] == ["allowed-" + tool.name for tool in tools] + assert manager.published_tool_routes["private-getsecret"] == "private" + assert manager.published_tool_routes["private-alias-getsecret"] == "private" + assert manager._get_mcp_server_from_tool_name("allowed-private-getsecret").server_id == "allowed" + assert ("orphan" in manager.published_tool_routes) is not local_handler + + +@pytest.mark.asyncio +async def test_cached_catalog_excludes_routes_for_servers_outside_its_snapshot(monkeypatch): + read_rows = AsyncMock(return_value=[]) + _catalog_database(monkeypatch, read_rows, AsyncMock(return_value=_revision_row(7))) + manager = MCPServerManager() + pinned = MCPServer(server_id="pinned", name="pinned", transport=MCPTransport.http) + manager.config_mcp_servers = {pinned.server_id: pinned} + async with manager.catalog.operation(): + assert manager.get_mcp_server_by_id("pinned") is not None + published = MCPServer(server_id="published", name="published", transport=MCPTransport.http) + manager.registry[published.server_id] = published + manager.published_tool_routes = {"search": "published", "known": "pinned"} + async with manager.catalog.operation(): + assert manager.get_mcp_server_by_id("published") is None + assert "search" not in manager.tool_name_to_mcp_server_name_mapping + assert manager._get_mcp_server_from_tool_name("known").server_id == "pinned" + read_rows.assert_awaited_once() 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 036a1c80042..fe2a2a1efdc 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 @@ -8632,6 +8632,7 @@ async def test_stateful_mcp_tool_call_uses_current_requests_otel_destinations(_m server = MCPServer( server_id="otel-context-test", name="otelcontext", + server_name="otelcontext", transport=MCPTransport.http, allow_all_keys=True, )