feat(mcp): cache upstream tool lists with an admin-set lifetime (#51160)

This commit is contained in:
shrey 2026-11-12 08:48:04 -08:00 • committed by GitHub
parent e3b32534dc
commit 40981a5d89
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
21 changed files with 1365 additions and 156 deletions

View file

@ -101,7 +101,7 @@ from litellm.constants import (
MCP_TOOL_LISTING_TIMEOUT,
)
from litellm.llms.custom_httpx.http_handler import get_ssl_configuration
from litellm.mcp_client.tools import list_tools_with_pagination
from litellm.mcp_client.tools import list_tools_result_with_pagination, list_tools_with_pagination
from litellm.proxy.mcp_server.mcp_context import get_active_mcp_request_ctx
from litellm.proxy.mcp_server.mcp_debug import capture_upstream_error_response
from litellm.proxy.mcp_server.result_conversion import (
@ -1071,6 +1071,9 @@ class MCPClient:
return await self.run_with_session(lambda session: session.list_tools(params=params), quiet_on_error=True)
async def list_tools(self, raise_on_error: bool = False) -> list[MCPTool]:
return (await self.list_tools_result(raise_on_error=raise_on_error)).tools
async def list_tools_result(self, raise_on_error: bool = False) -> ListToolsResult:
"""List available tools from the server.
Args:
@ -1085,16 +1088,20 @@ class MCPClient:
try:
# A per-server timeout above the global default extends the whole-walk deadline
listing_deadline: Final = max(self.timeout, MCP_TOOL_LISTING_TIMEOUT)
tools: Final = await self.run_with_session(
partial(list_tools_with_pagination, listing_deadline=listing_deadline, require_complete=raise_on_error),
listing: Final = await self.run_with_session(
partial(
list_tools_result_with_pagination,
listing_deadline=listing_deadline,
require_complete=raise_on_error,
),
quiet_on_error=raise_on_error,
)
tool_count: Final = len(tools)
tool_names: Final = tuple(tool.name for tool in tools)
tool_count: Final = len(listing.tools)
tool_names: Final = tuple(tool.name for tool in listing.tools)
verbose_logger.info(
"MCP client listed %s tools from %s: %s", tool_count, self.server_url or "stdio", tool_names
)
return tools
return listing
except asyncio.CancelledError:
verbose_logger.warning("MCP client list_tools was cancelled")
raise
@ -1124,7 +1131,7 @@ class MCPClient:
if raise_on_error:
raise
# Return empty list instead of raising to allow graceful degradation
return []
return ListToolsResult(tools=[])
@staticmethod
def error_tool_result(exc: Exception) -> MCPCallToolResult:

View file

@ -1,9 +1,11 @@
import json
import time
from collections.abc import Sequence
from typing import Final, Literal
import anyio
from mcp import ClientSession, MCPError
from mcp.types import INTERNAL_ERROR, PaginatedRequestParams
from mcp.types import INTERNAL_ERROR, ListToolsResult, PaginatedRequestParams
from mcp.types import CallToolRequestParams as MCPCallToolRequestParams
from mcp.types import CallToolResult as MCPCallToolResult
from mcp.types import Tool as MCPTool
@ -18,6 +20,7 @@ from litellm.constants import (
MCP_TOOL_LISTING_MAX_PAGES,
MCP_TOOL_LISTING_TIMEOUT,
)
from litellm.proxy.mcp_server.result_conversion import age_freshness, aggregate_freshness
from litellm.types.llms.anthropic import AnthropicMessagesTool
from litellm.types.utils import ChatCompletionMessageToolCall
@ -101,9 +104,21 @@ def transform_mcp_tool_to_anthropic_tool(mcp_tool: MCPTool) -> AnthropicMessages
)
def _complete_listing(tools: Sequence[MCPTool], pages: Sequence[tuple[ListToolsResult, float]]) -> ListToolsResult:
now: Final = time.monotonic()
freshness: Final = aggregate_freshness(tuple(age_freshness(page, now - received) for page, received in pages))
return ListToolsResult(tools=list(tools), ttl_ms=freshness.ttl_ms, cache_scope=freshness.cache_scope)
async def list_tools_with_pagination(
session: ClientSession, listing_deadline: float | None = None, *, require_complete: bool = False
) -> list[MCPTool]: # mutable-ok: list return contract
return (await list_tools_result_with_pagination(session, listing_deadline, require_complete=require_complete)).tools
async def list_tools_result_with_pagination(
session: ClientSession, listing_deadline: float | None = None, *, require_complete: bool = False
) -> ListToolsResult:
"""Collect tools from every tools/list page by following nextCursor.
Stops and returns the tools collected so far when the upstream repeats a
@ -111,8 +126,10 @@ async def list_tools_with_pagination(
buggy or slow upstream yields a partial catalog instead of an error.
listing_deadline overrides the default whole-walk deadline; callers with a
per-server timeout above the global default pass it through here.
Only a complete walk carries the upstream ttlMs, so a partial catalog is never cached.
"""
tools: Final[list[MCPTool]] = [] # mutable-ok: accumulates each page's tools
pages: Final[list[tuple[ListToolsResult, float]]] = [] # mutable-ok: bounded pagination evidence
seen_cursors: Final[set[str]] = set() # mutable-ok: guards against cursor loops
cursor: str | None = None # rebind-ok: advances to each page's nextCursor
# The per-request session read timeout restarts on every page, so a multi-page
@ -131,10 +148,11 @@ async def list_tools_with_pagination(
else await session.list_tools(params=PaginatedRequestParams(cursor=cursor))
)
tools.extend(result.tools)
pages.append((result, time.monotonic()))
next_cursor = result.next_cursor
if not isinstance(next_cursor, str) or not next_cursor:
return tools
return _complete_listing(tools, pages)
if next_cursor in seen_cursors:
verbose_logger.warning(
"MCP server repeated a tools/list cursor while listing tools; returning %s tools collected so far",
@ -144,7 +162,7 @@ async def list_tools_with_pagination(
raise MCPError(
code=INTERNAL_ERROR, message="Upstream tool discovery is incomplete: repeated cursor"
)
return tools
return ListToolsResult(tools=tools)
seen_cursors.add(next_cursor)
cursor = next_cursor
@ -155,7 +173,7 @@ async def list_tools_with_pagination(
)
if require_complete:
raise MCPError(code=INTERNAL_ERROR, message="Upstream tool discovery is incomplete: pagination limit")
return tools
return ListToolsResult(tools=tools)
verbose_logger.warning(
"MCP server tools/list pagination exceeded the %s second listing deadline; returning %s tools collected so far",
@ -164,7 +182,7 @@ async def list_tools_with_pagination(
)
if require_complete:
raise MCPError(code=INTERNAL_ERROR, message="Upstream tool discovery is incomplete: listing deadline")
return tools
return ListToolsResult(tools=tools)
async def load_mcp_tools(

View file

@ -15,11 +15,12 @@ from dataclasses import dataclass, replace
from functools import partial, wraps
from itertools import chain
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Generic, ParamSpec, TypeAlias, TypeVar, cast
from typing import TYPE_CHECKING, Final, Generic, ParamSpec, Protocol, TypeAlias, TypeVar, cast
from mcp.shared.exceptions import MCPError
from mcp.types import CacheableResult
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter
from typing_extensions import Self
from litellm._logging import verbose_logger
from litellm.caching.in_memory_cache import InMemoryCache
@ -849,6 +850,7 @@ async def get_filtered_server_tools(
params: PaginatedRequestParams | None = None,
record_listing: bool = False,
listing_updates: ExitStack | None = None,
use_tool_list_cache: bool = True,
) -> tuple[ListToolsResult, ServerOutcome]:
from mcp.types import ListToolsResult
@ -893,6 +895,7 @@ async def get_filtered_server_tools(
proxy_logging_obj=proxy_logging_obj,
catalog_auth_header=catalog_auth_header,
record_listing=False,
use_tool_list_cache=use_tool_list_cache,
)
)
else:
@ -910,6 +913,7 @@ async def get_filtered_server_tools(
catalog_auth_header=catalog_auth_header,
record_listing=False,
listing_updates=listing_updates,
use_tool_list_cache=use_tool_list_cache,
)
tools: Final = page.tools
filtered_tools = filter_tools_by_allowed_tools(tools, server)
@ -1001,6 +1005,7 @@ async def aggregate_gateway_tools(
*,
record_listing: bool = False,
enforce_rate_limits: bool = True,
use_tool_list_cache: bool = True,
) -> AggregateToolListing:
import time
@ -1050,6 +1055,7 @@ async def aggregate_gateway_tools(
params=PaginatedRequestParams(cursor=cursor),
record_listing=record_listing,
listing_updates=listing_updates,
use_tool_list_cache=use_tool_list_cache,
)
if cursor is not None and outcome.tag != "ok":
from mcp.shared.exceptions import MCPError
@ -1337,29 +1343,55 @@ def combine_optional_catalog(
)
_DiscoveryPage = TypeVar("_DiscoveryPage", bound=CacheableResult)
class _Expiring(Protocol):
@property
def ttl_ms(self) -> int: ...
def model_copy(self, *, update: Mapping[str, object] | None = None, deep: bool = False) -> Self: ...
_DiscoveryPage = TypeVar("_DiscoveryPage", bound=_Expiring)
_DiscoveryKey: TypeAlias = tuple[str, str | None]
_DISCOVERY_CACHE_LIMIT: Final = 1024
_DISCOVERY_ENTRY: Final = TypeAdapter(tuple[float, bytes])
class _DiscoveryCache(Generic[_DiscoveryPage]):
def __init__(self, ttl: float, clock: Callable[[], float], adapter: TypeAdapter[_DiscoveryPage]) -> None:
def __init__(
self,
ttl: float,
clock: Callable[[], float],
adapter: TypeAdapter[_DiscoveryPage],
max_entry_kb: int = 64,
storable: Callable[[_DiscoveryPage], bool] | None = None,
) -> None:
self._ttl = ttl
self._clock = clock
self._adapter = adapter
self._entries = InMemoryCache(max_size_in_memory=_DISCOVERY_CACHE_LIMIT, max_size_per_item=64, clock=clock)
self._storable = storable
self._entries = InMemoryCache(
max_size_in_memory=_DISCOVERY_CACHE_LIMIT, max_size_per_item=max_entry_kb, clock=clock
)
self._pending: dict[_DiscoveryKey, asyncio.Task[_DiscoveryPage]] = {}
self._waiters: dict[asyncio.Task[_DiscoveryPage], int] = {} # mutable-ok: constant-time waiter accounting
self._oversized_server_ids: frozenset[str] = frozenset()
self.invalidations = 0
def invalidate(self, server_id: str) -> None:
def _entry_keys(self, server_id: str) -> tuple[str, ...]:
prefix: Final = f"[{json.dumps(server_id)},"
keys: Final = cast( # cast-ok: private cache contains only JSON string keys
"tuple[str, ...]", tuple(self._entries.cache_dict)
)
for entry_key in keys:
if entry_key.startswith(prefix):
self._entries.delete_cache(entry_key)
return tuple(entry_key for entry_key in keys if entry_key.startswith(prefix))
def holds(self, server_id: str) -> bool:
return bool(self._entry_keys(server_id))
def invalidate(self, server_id: str) -> None:
self.invalidations += 1
self._oversized_server_ids = self._oversized_server_ids - {server_id}
for entry_key in self._entry_keys(server_id):
self._entries.delete_cache(entry_key)
for key in tuple(self._pending):
if key[0] == server_id:
self._pending.pop(key)
@ -1369,7 +1401,9 @@ class _DiscoveryCache(Generic[_DiscoveryPage]):
if not task.cancelled():
task.exception()
async def get(self, key: _DiscoveryKey, fetch: Callable[[], Awaitable[_DiscoveryPage]]) -> _DiscoveryPage:
async def get(
self, key: _DiscoveryKey, fetch: Callable[[], Awaitable[_DiscoveryPage]], *, capped: bool = True
) -> _DiscoveryPage:
if self._ttl <= 0:
return await fetch()
entry: Final[object] = self._entries.get_cache(json.dumps(key))
@ -1384,7 +1418,7 @@ class _DiscoveryCache(Generic[_DiscoveryPage]):
return await self._await_fetch(key, pending)
if len(self._pending) >= _DISCOVERY_CACHE_LIMIT:
return await fetch()
task: Final = asyncio.create_task(self._fetch(key, fetch))
task: Final = asyncio.create_task(self._fetch(key, fetch, capped))
self._pending[key] = task
task.add_done_callback(self._observe_completion)
return await self._await_fetch(key, task)
@ -1404,17 +1438,34 @@ class _DiscoveryCache(Generic[_DiscoveryPage]):
if not task.done():
task.cancel()
async def _fetch(self, key: _DiscoveryKey, fetch: Callable[[], Awaitable[_DiscoveryPage]]) -> _DiscoveryPage:
async def _fetch(
self, key: _DiscoveryKey, fetch: Callable[[], Awaitable[_DiscoveryPage]], capped: bool
) -> _DiscoveryPage:
try:
items: Final = await fetch()
ttl: Final = min(self._ttl, items.ttl_ms / 1000)
if ttl > 0 and self._pending.get(key) is asyncio.current_task():
self._entries.set_cache(
json.dumps(key),
_DISCOVERY_ENTRY.dump_json((self._clock() + ttl, self._adapter.dump_json(items))),
ttl=ttl,
)
if self._pending.get(key) is asyncio.current_task():
self.keep(key, items, capped=capped)
return items
finally:
if self._pending.get(key) is asyncio.current_task():
self._pending.pop(key)
def keep(self, key: _DiscoveryKey, items: _DiscoveryPage, *, capped: bool = True) -> None:
declared: Final = items.ttl_ms / 1000
ttl: Final = min(self._ttl, declared) if capped else declared
if ttl <= 0 or (self._storable is not None and not self._storable(items)):
return
entry: Final = _DISCOVERY_ENTRY.dump_json((self._clock() + ttl, self._adapter.dump_json(items)))
if self._entries.check_value_size(entry):
self._entries.set_cache(json.dumps(key), entry, ttl=ttl)
return
server_id: Final = key[0]
if server_id in self._oversized_server_ids:
return
self._oversized_server_ids = self._oversized_server_ids | {server_id}
verbose_logger.warning(
"MCP server %s returned a %s KB listing, above the %s KB cache entry limit. It is fetched on every request",
server_id,
len(entry) // 1024,
self._entries.max_size_per_item,
)

View file

@ -26,7 +26,7 @@ from collections.abc import (
)
from contextlib import AbstractAsyncContextManager, ExitStack, asynccontextmanager, nullcontext
from dataclasses import dataclass, replace
from functools import lru_cache
from functools import lru_cache, partial
from itertools import chain, groupby
from types import EllipsisType, MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeAlias, TypedDict, cast
@ -58,8 +58,8 @@ from mcp.types import (
ResourceTemplate,
)
from mcp.types import Tool as MCPTool
from pydantic import AnyUrl, Field, SecretStr, TypeAdapter, ValidationError
from typing_extensions import ReadOnly, assert_never
from pydantic import AnyUrl, BaseModel, ConfigDict, Field, SecretStr, TypeAdapter, ValidationError
from typing_extensions import ReadOnly, Self, assert_never
import litellm
from litellm._internal_context import with_service_target
@ -236,6 +236,7 @@ from litellm.types.mcp import (
without_header,
)
from litellm.types.mcp_server.mcp_server_manager import (
MAX_TOOLS_LIST_CACHE_TTL_SECONDS,
MCPInfo,
MCPOAuthMetadata,
MCPServer,
@ -1896,6 +1897,49 @@ def _mcp_discovery_cache_ttl() -> float:
return ttl
_TOOL_LIST_ENTRY_LIMIT_KB: Final = 1024
class _UpstreamToolListing(BaseModel):
model_config = ConfigDict(frozen=True)
page: ListToolsResult
ttl_ms: int
upstream_instructions: str | None
@classmethod
def fetched_by(cls, client: MCPClient, page: ListToolsResult) -> Self:
declared: Final[object] = getattr(client, "_last_initialize_instructions", None)
return cls(page=page, ttl_ms=page.ttl_ms, upstream_instructions=declared if isinstance(declared, str) else None)
def is_whole(self) -> bool:
return not self.page.next_cursor
def page_to_serve(self) -> ListToolsResult:
if self.ttl_ms == self.page.ttl_ms:
return self.page
return self.page.model_copy(update={"ttl_ms": self.ttl_ms})
def _tool_page_cursor(params: PaginatedRequestParams | None) -> tuple[str | None] | None:
return None if params is None else (params.cursor,)
def _mcp_tools_list_cache_ttl() -> int | None:
raw: Final = os.environ.get("LITELLM_MCP_TOOLS_LIST_CACHE_TTL")
if raw is None:
return None
try:
ttl: Final = int(raw)
except ValueError:
verbose_logger.warning("Invalid LITELLM_MCP_TOOLS_LIST_CACHE_TTL; ignoring it")
return None
if not 0 <= ttl <= MAX_TOOLS_LIST_CACHE_TTL_SECONDS:
verbose_logger.warning("Invalid LITELLM_MCP_TOOLS_LIST_CACHE_TTL; ignoring it")
return None
return ttl
def _listing_task_was_cancelled(error: asyncio.CancelledError) -> bool:
"""The MCP client raises its own CancelledError(SESSION_CLEANUP_TIMEOUT_MESSAGE) when a session's
teardown outlives its budget while the listing task keeps running. That one is the server's fault,
@ -2045,6 +2089,15 @@ class MCPServerManager:
self._template_discovery_cache = _DiscoveryCache[ListResourceTemplatesResult](
discovery_ttl, discovery_clock, TypeAdapter(ListResourceTemplatesResult)
)
self._tool_discovery_cache = _DiscoveryCache[_UpstreamToolListing](
discovery_ttl,
discovery_clock,
TypeAdapter(_UpstreamToolListing),
max_entry_kb=_TOOL_LIST_ENTRY_LIMIT_KB,
storable=_UpstreamToolListing.is_whole,
)
self._list_caching_enabled = discovery_ttl > 0
self._tools_list_cache_ttl = _mcp_tools_list_cache_ttl()
from litellm.proxy.mcp_server.catalog import CatalogSnapshots
self.catalog = CatalogSnapshots(self)
@ -4440,6 +4493,7 @@ class MCPServerManager:
*,
catalog_auth_header: str | dict[str, str] | None | EllipsisType = ...,
record_listing: bool = False,
use_tool_list_cache: bool = True,
) -> list[MCPTool]:
result: Final = await self.get_tools_page(
server=server,
@ -4453,6 +4507,7 @@ class MCPServerManager:
proxy_logging_obj=proxy_logging_obj,
catalog_auth_header=catalog_auth_header,
record_listing=record_listing,
use_tool_list_cache=use_tool_list_cache,
)
return result.tools
@ -4474,6 +4529,7 @@ class MCPServerManager:
catalog_auth_header: str | dict[str, str] | None | EllipsisType = ...,
record_listing: bool = False,
listing_updates: ExitStack | None = None,
use_tool_list_cache: bool = True,
) -> ListToolsResult:
from litellm.proxy.mcp_server.tool_registry import (
global_mcp_tool_registry,
@ -4521,6 +4577,7 @@ class MCPServerManager:
# upstream API key or a user's OAuth token (MCPClient._get_auth_headers
# applies extra_headers after writing Authorization from auth_value, so
# an injected JWT would otherwise clobber the per-user token).
headers_before_signing: Final = extra_headers
if user_api_key_auth is not None and not server.spec_path:
from litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer import (
get_mcp_jwt_signer,
@ -4595,14 +4652,22 @@ class MCPServerManager:
tools=[t.model_copy(update={"name": registered_names[t.name]}) for t in guarded_openapi]
)
else:
page: Final = (
await retry_upstream_rate_limit(
lambda: self._fetch_tools_page_with_timeout(client, server, params),
server.name,
deadline_seconds=MCP_TOOL_LISTING_TIMEOUT,
)
if params is not None
else ListToolsResult(tools=await self._fetch_tools_with_timeout(client, server.name))
page: Final = await self._list_upstream_tools(
server,
client,
params,
partial(
self._discovery_key,
server,
user_api_key_auth,
mcp_auth_header,
extra_headers,
stdio_env,
subject_token,
raw_headers=raw_headers,
page_cursor=_tool_page_cursor(params),
),
use_cache=use_tool_list_cache and extra_headers is headers_before_signing,
)
tools = page.tools
self._remember_upstream_initialize_instructions(server, client)
@ -4663,6 +4728,92 @@ class MCPServerManager:
)
raise_classified_list_failure(e, server.name, suppress_challenge=server.is_dcr_bridge)
def _admin_tools_list_ttl(self, server: MCPServer) -> int | None:
if not self._list_caching_enabled:
return None
per_server: Final = server.tools_list_cache_ttl_seconds
return self._tools_list_cache_ttl if per_server is None else per_server
def _tools_fetch(
self, client: MCPClient, server: MCPServer, params: PaginatedRequestParams | None
) -> Callable[[], Awaitable[ListToolsResult]]:
if params is None:
return partial(self._fetch_tools_with_timeout, client, server.name)
return partial(
retry_upstream_rate_limit,
partial(self._fetch_tools_page_with_timeout, client, server, params),
server.name,
deadline_seconds=MCP_TOOL_LISTING_TIMEOUT,
)
async def _list_upstream_tools(
self,
server: MCPServer,
client: MCPClient,
params: PaginatedRequestParams | None,
key_for_credential: Callable[[str], _DiscoveryKey],
*,
use_cache: bool,
) -> ListToolsResult:
fetch_page: Final = self._tools_fetch(client, server, params)
async def fetch() -> _UpstreamToolListing:
return _UpstreamToolListing.fetched_by(client, await fetch_page())
listing: Final = await self._cached_or_upstream_tools(
server, client, params, key_for_credential, fetch, use_cache
)
self._remember_instructions_if_missing(server, listing.upstream_instructions)
return listing.page_to_serve()
def _remember_instructions_if_missing(self, server: MCPServer, seen_by_this_caller: str | None) -> None:
if seen_by_this_caller and not self._upstream_initialize_instructions_by_server_id.get(server.server_id):
self._upstream_initialize_instructions_by_server_id[server.server_id] = seen_by_this_caller
async def _cached_or_upstream_tools(
self,
server: MCPServer,
client: MCPClient,
params: PaginatedRequestParams | None,
key_for_credential: Callable[[str], _DiscoveryKey],
fetch: Callable[[], Awaitable[_UpstreamToolListing]],
use_cache: bool,
) -> _UpstreamToolListing:
if not use_cache or (params is not None and params.cursor is not None):
return await fetch()
async def cache_key() -> _DiscoveryKey:
return key_for_credential(await client.discovery_auth_fingerprint())
admin_ttl: Final = self._admin_tools_list_ttl(server)
if admin_ttl is None:
return await self._list_tools_for_upstream_ttl(server, cache_key, fetch)
async def fetch_with_admin_ttl() -> _UpstreamToolListing:
listing: Final = await fetch()
return listing.model_copy(update={"ttl_ms": admin_ttl * 1000}) if listing.is_whole() else listing
if admin_ttl <= 0:
return await fetch_with_admin_ttl()
return await self._tool_discovery_cache.get(await cache_key(), fetch_with_admin_ttl, capped=False)
async def _list_tools_for_upstream_ttl(
self,
server: MCPServer,
cache_key: Callable[[], Awaitable[_DiscoveryKey]],
fetch: Callable[[], Awaitable[_UpstreamToolListing]],
) -> _UpstreamToolListing:
if self._tool_discovery_cache.holds(server.server_id):
return await self._tool_discovery_cache.get(await cache_key(), fetch)
invalidations_before: Final = self._tool_discovery_cache.invalidations
listing: Final = await fetch()
if not self._list_caching_enabled or listing.ttl_ms <= 0 or not listing.is_whole():
return listing
key: Final = await cache_key()
if self._tool_discovery_cache.invalidations == invalidations_before:
self._tool_discovery_cache.keep(key, listing)
return listing
def clear_initialize_instructions(self) -> None:
self._upstream_initialize_instructions_by_server_id.clear()
self._upstream_initialize_instructions_probed_at.clear()
@ -4673,6 +4824,7 @@ class MCPServerManager:
self._prompt_discovery_cache.invalidate(server_id)
self._resource_discovery_cache.invalidate(server_id)
self._template_discovery_cache.invalidate(server_id)
self._tool_discovery_cache.invalidate(server_id)
def invalidate_server_definition_caches(self, server_id: str) -> None:
from litellm.proxy.mcp_server.discoverable_endpoints import ( # noqa: PLC0415 # lazy: discoverable_endpoints lazily imports this module's manager singleton
@ -4774,6 +4926,7 @@ class MCPServerManager:
subject_token: str | None,
credential_fingerprint: str | None = None,
raw_headers: Mapping[str, str] | None = None,
page_cursor: tuple[str | None] | None = None,
) -> _DiscoveryKey:
identity: Final = (
user_api_key_auth.model_dump(
@ -4800,6 +4953,7 @@ class MCPServerManager:
stdio_env,
subject_token,
credential_fingerprint,
page_cursor,
),
sort_keys=True,
separators=(",", ":"),
@ -5644,7 +5798,7 @@ class MCPServerManager:
self,
client: MCPClient,
server_name: str,
) -> list[MCPTool]:
) -> ListToolsResult:
"""
Fetch tools from MCP client with timeout and error handling.
@ -5668,17 +5822,17 @@ class MCPServerManager:
server_name: Name of the server for logging
Returns:
List of tools from the server
The tools from the server with the upstream freshness of the complete walk
"""
try:
with anyio.fail_after(MCP_TOOL_LISTING_TIMEOUT):
tools: Final = await retry_upstream_rate_limit(
lambda: client.list_tools(raise_on_error=True),
listing: Final = await retry_upstream_rate_limit(
lambda: client.list_tools_result(raise_on_error=True),
server_name,
deadline_seconds=MCP_TOOL_LISTING_TIMEOUT,
)
verbose_logger.debug("Tools from %s: %s", server_name, tools)
return tools
verbose_logger.debug("Tools from %s: %s", server_name, listing.tools)
return listing
except TimeoutError as e:
verbose_logger.warning("Timeout while listing tools from %s", server_name)
raise MCPServerListError(ServerListFault(tag="timeout"), server_name) from e
@ -6131,6 +6285,7 @@ class MCPServerManager:
dict(catalog_auth_header) if isinstance(catalog_auth_header, Mapping) else catalog_auth_header
),
record_listing=False,
use_tool_list_cache=False,
)
selected: Final = tuple(tool for tool in upstream if tool.name == name)
scan: Final = await scan_tool_descriptions(

View file

@ -1218,6 +1218,7 @@ async def find_mcp_proxy_tool(
context=context,
allowed_mcp_servers=allowed_mcp_servers,
prefetched_oauth_creds=prefetched_oauth_creds,
use_tool_list_cache=False,
)
return next((tool for tool in page.tools if mcp_proxy_tool_id(tool) == tool_id), None)
@ -1255,6 +1256,7 @@ async def _get_tools_from_mcp_servers(
record_listing: bool = False,
request_trace_id: str | None = None,
enforce_rate_limits: bool = True,
use_tool_list_cache: bool = True,
) -> AggregateToolListing:
"""
Helper method to fetch tools from MCP servers based on server filtering criteria.
@ -1384,6 +1386,7 @@ async def _get_tools_from_mcp_servers(
allowed_mcp_servers=allowed_mcp_servers,
prefetched_oauth_creds=_prefetched_oauth_creds,
record_listing=record_listing,
use_tool_list_cache=use_tool_list_cache,
)
return page.tools, outcome
@ -1419,6 +1422,7 @@ async def _get_tools_from_mcp_servers(
_prefetched_oauth_creds,
record_listing=record_listing,
enforce_rate_limits=enforce_rate_limits,
use_tool_list_cache=use_tool_list_cache,
)
all_tools: Final = aggregated.tools
server_outcomes: Final = aggregated.outcomes
@ -1710,6 +1714,7 @@ async def _list_mcp_tools(
protocol_version: str | None = None,
*,
record_listing: bool = False,
use_tool_list_cache: bool = True,
) -> AggregateToolListing:
"""
List all available MCP tools.
@ -1745,6 +1750,7 @@ async def _list_mcp_tools(
record_listing=record_listing,
params=params,
protocol_version=protocol_version,
use_tool_list_cache=use_tool_list_cache,
)
verbose_logger.debug("Successfully fetched %s tools from managed MCP servers", len(listing.tools))
return listing
@ -2073,6 +2079,7 @@ async def _list_tools_before_first_call(
client_ip=client_ip,
record_listing=False,
enforce_rate_limits=False,
use_tool_list_cache=False,
)
except Exception as e: # noqa: BLE001 # best effort: resolution below answers as it did before
verbose_logger.debug("MCP tools/call: listing %s before its first call failed: %s", server.name, e)

View file

@ -748,6 +748,7 @@ if MCP_AVAILABLE:
proxy_logging_obj: "ProxyLogging | None",
*,
record_listing: bool,
use_tool_list_cache: bool,
) -> list[MCPTool]:
return list(
await global_mcp_server_manager.get_tools_from_server(
@ -760,6 +761,7 @@ if MCP_AVAILABLE:
user_api_key_auth=user_api_key_auth,
proxy_logging_obj=proxy_logging_obj,
record_listing=record_listing,
use_tool_list_cache=use_tool_list_cache,
)
)
@ -793,6 +795,7 @@ if MCP_AVAILABLE:
client_ip,
proxy_logging_obj,
record_listing=False,
use_tool_list_cache=apply_tool_filters,
)
server_filtered: Final = filter_tools_by_allowed_tools(tools, server) if apply_tool_filters else tools
@ -839,6 +842,7 @@ if MCP_AVAILABLE:
IPAddressUtils.get_mcp_client_ip(request),
None,
record_listing=False,
use_tool_list_cache=False,
)
scan: Final = await scan_tool_descriptions(
apply_description_overrides(upstream, server),

View file

@ -588,6 +588,7 @@ async def handle_mcp_tool_search(
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
use_tool_list_cache=False,
)
results: Final = await search_mcp_tools(query, mcp_listing.tools, top_k, settings, ranker)
if isinstance(results, EmbeddingFailed):
@ -628,6 +629,7 @@ async def handle_mcp_proxy_tool(
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
mcp_proxy_mode=True,
use_tool_list_cache=False,
)
llm_router: Final = proxy_server.llm_router
proxy_logging_obj: Final = proxy_server.proxy_logging_obj

View file

@ -18,6 +18,11 @@ from litellm.types.mcp import (
validate_mcp_protocol_transport,
)
MAX_TOOLS_LIST_CACHE_TTL_SECONDS: Final = 365 * 24 * 60 * 60
_TOOLS_LIST_CACHE_TTL_SECONDS: Final = TypeAdapter[int](
Annotated[int, Field(ge=0, le=MAX_TOOLS_LIST_CACHE_TTL_SECONDS)]
)
# MCPInfo now allows arbitrary additional fields for custom metadata
def _validate_mcp_protocol_metadata(value: dict[str, object]) -> dict[str, object]:
@ -25,6 +30,8 @@ def _validate_mcp_protocol_metadata(value: dict[str, object]) -> dict[str, objec
TypeAdapter[MCPUpstreamProtocol](MCPUpstreamProtocol).validate_python(value["protocol_version"])
if "skip_tool_output_validation" in value:
TypeAdapter[bool](bool).validate_python(value["skip_tool_output_validation"], strict=True)
if "tools_list_cache_ttl_seconds" in value:
_TOOLS_LIST_CACHE_TTL_SECONDS.validate_python(value["tools_list_cache_ttl_seconds"], strict=True)
return value
@ -303,6 +310,11 @@ class MCPServer(LiteLLMBaseModel):
def skip_tool_output_validation(self) -> bool:
return (self.mcp_info or {}).get("skip_tool_output_validation") is True
@property
def tools_list_cache_ttl_seconds(self) -> int | None:
configured: Final[object] = (self.mcp_info or {}).get("tools_list_cache_ttl_seconds")
return configured if isinstance(configured, int) else None
@model_validator(mode="after")
def resolve_protocol_version(self) -> Self:
if "protocol_version" not in self.model_fields_set and self.mcp_info is not None:

View file

@ -1,4 +1,8 @@
import json
from functools import partial
from itertools import count
from types import SimpleNamespace
from typing import Final
from unittest.mock import AsyncMock, MagicMock
import pytest
@ -16,6 +20,7 @@ from litellm.mcp_client.tools import (
_normalize_mcp_input_schema,
call_mcp_tool,
call_openai_tool,
list_tools_result_with_pagination,
list_tools_with_pagination,
load_mcp_tools,
transform_mcp_tool_to_anthropic_tool,
@ -169,6 +174,25 @@ async def test_pagination_walk_stops_on_repeated_cursor(mock_session):
assert mock_session.list_tools.call_count == 2
@pytest.mark.asyncio()
@pytest.mark.parametrize("last_cursor", [None, "page-1"])
async def test_pagination_walk_keeps_the_shortest_ttl_only_for_a_complete_walk(
mock_session: MagicMock, last_cursor: str | None, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr("litellm.mcp_client.tools.time", SimpleNamespace(monotonic=partial(next, count(100.0, 0.5))))
mock_session.list_tools.side_effect = [
ListToolsResult(tools=[MCPTool(name="tool_0", inputSchema={})], nextCursor="page-1", ttl_ms=30000),
ListToolsResult(tools=[MCPTool(name="tool_1", inputSchema={})], nextCursor=last_cursor, ttl_ms=20000),
]
result: Final = await list_tools_result_with_pagination(mock_session)
assert [tool.name for tool in result.tools] == ["tool_0", "tool_1"]
if last_cursor is None:
assert result.ttl_ms == 19500
else:
assert result.ttl_ms == 0
@pytest.mark.asyncio()
async def test_pagination_walk_treats_empty_cursor_as_terminal(mock_session):
mock_session.list_tools.side_effect = [

View file

@ -208,12 +208,12 @@ async def test_cancelled_fetch_propagates_instead_of_masquerading_as_a_healthy_e
client = MagicMock()
started = asyncio.Event()
async def list_tools(raise_on_error: bool = False) -> list[object]:
async def list_tools_result(raise_on_error: bool = False) -> tuple[object, ...]:
started.set()
await asyncio.Event().wait()
return []
return ()
client.list_tools = list_tools
client.list_tools_result = list_tools_result
listing = asyncio.create_task(manager._fetch_tools_with_timeout(client, "cancelled_srv"))
await started.wait()
listing.cancel(cancel_message)
@ -235,7 +235,7 @@ async def test_client_side_cancelled_error_is_a_classified_fault_not_a_cancellat
manager = MCPServerManager()
client = MagicMock()
client.list_tools = AsyncMock(side_effect=asyncio.CancelledError(SESSION_CLEANUP_TIMEOUT_MESSAGE))
client.list_tools_result = AsyncMock(side_effect=asyncio.CancelledError(SESSION_CLEANUP_TIMEOUT_MESSAGE))
with pytest.raises(MCPServerListError) as exc_info:
await manager._fetch_tools_with_timeout(client, "slow_teardown_srv")

View file

@ -186,9 +186,7 @@ class TestMCPClientUnitTests:
},
)
]
mock_result = MagicMock()
mock_result.tools = mock_tools
mock_session_instance.list_tools.return_value = mock_result
mock_session_instance.list_tools.return_value = ListToolsResult(tools=mock_tools)
client = MCPClient("http://example.com")
result = await client.list_tools()

View file

@ -10,6 +10,8 @@ from unittest.mock import AsyncMock, MagicMock
import httpx
import pytest
from mcp.types import ListToolsResult
from mcp.types import Tool as MCPTool
if sys.version_info < (3, 11):
from exceptiongroup import ExceptionGroup
@ -115,7 +117,7 @@ async def test_fetch_tools_from_passthrough_raises_on_upstream_401():
)
mock_client = MagicMock()
mock_client.list_tools = AsyncMock(side_effect=upstream_error)
mock_client.list_tools_result = AsyncMock(side_effect=upstream_error)
with pytest.raises(MCPUpstreamAuthError) as exc_info:
await manager._fetch_tools_with_timeout(mock_client, passthrough_server.name)
@ -125,7 +127,7 @@ async def test_fetch_tools_from_passthrough_raises_on_upstream_401():
'Bearer resource_metadata="https://upstream"'
)
assert exc_info.value.server_name == "sample_docs"
mock_client.list_tools.assert_awaited_with(raise_on_error=True)
mock_client.list_tools_result.assert_awaited_with(raise_on_error=True)
@pytest.mark.asyncio
@ -150,7 +152,7 @@ async def test_fetch_tools_from_delegated_oauth2_raises_on_upstream_401():
)
mock_client = MagicMock()
mock_client.list_tools = AsyncMock(side_effect=upstream_error)
mock_client.list_tools_result = AsyncMock(side_effect=upstream_error)
with pytest.raises(MCPUpstreamAuthError) as exc_info:
await manager._fetch_tools_with_timeout(mock_client, delegated_server.name)
@ -160,7 +162,7 @@ async def test_fetch_tools_from_delegated_oauth2_raises_on_upstream_401():
'Bearer resource_metadata="https://upstream"'
)
assert exc_info.value.server_name == "delegated_docs"
mock_client.list_tools.assert_awaited_with(raise_on_error=True)
mock_client.list_tools_result.assert_awaited_with(raise_on_error=True)
@pytest.mark.asyncio
@ -189,14 +191,14 @@ async def test_fetch_tools_from_client_credentials_oauth2_surfaces_upstream_401(
)
mock_client = MagicMock()
mock_client.list_tools = AsyncMock(side_effect=upstream_error)
mock_client.list_tools_result = AsyncMock(side_effect=upstream_error)
with pytest.raises(MCPUpstreamAuthError) as exc_info:
await manager._fetch_tools_with_timeout(mock_client, m2m_server.name)
assert exc_info.value.status_code == 401
assert exc_info.value.server_name == "m2m_docs"
mock_client.list_tools.assert_awaited_with(raise_on_error=True)
mock_client.list_tools_result.assert_awaited_with(raise_on_error=True)
@pytest.mark.asyncio
@ -212,13 +214,12 @@ async def test_fetch_tools_from_passthrough_returns_tools_on_success():
oauth_passthrough=True,
)
tool = MagicMock()
tool.name = "list_documents"
tool: Final = MCPTool(name="list_documents", inputSchema={"type": "object"})
mock_client = MagicMock()
mock_client.list_tools = AsyncMock(return_value=[tool])
mock_client.list_tools_result = AsyncMock(return_value=ListToolsResult(tools=[tool]))
tools = await manager._fetch_tools_with_timeout(mock_client, passthrough_server.name)
assert tools == [tool]
listing: Final = await manager._fetch_tools_with_timeout(mock_client, passthrough_server.name)
assert listing.tools == [tool]
@pytest.mark.parametrize("status_code", [401, 403])
@ -308,7 +309,7 @@ async def test_fetch_tools_from_gateway_managed_surfaces_upstream_401():
"401", request=response.request, response=response
)
mock_client = MagicMock()
mock_client.list_tools = AsyncMock(side_effect=upstream_error)
mock_client.list_tools_result = AsyncMock(side_effect=upstream_error)
with pytest.raises(MCPUpstreamAuthError) as exc_info:
await manager._fetch_tools_with_timeout(mock_client, oauth2_server.name)
@ -316,7 +317,7 @@ async def test_fetch_tools_from_gateway_managed_surfaces_upstream_401():
assert exc_info.value.status_code == 401
assert exc_info.value.www_authenticate is None
assert exc_info.value.server_name == "keycloak_whoami"
mock_client.list_tools.assert_awaited_with(raise_on_error=True)
mock_client.list_tools_result.assert_awaited_with(raise_on_error=True)
def _http_server(server_id: str, name: str, **kwargs) -> MCPServer:
@ -462,7 +463,7 @@ async def test_fetch_tools_logs_upstream_request_details_on_500(caplog):
)
response = httpx.Response(500, request=request)
mock_client = MagicMock()
mock_client.list_tools = AsyncMock(
mock_client.list_tools_result = AsyncMock(
side_effect=httpx.HTTPStatusError("500", request=request, response=response)
)

View file

@ -2,12 +2,14 @@ import asyncio
import json
import logging
from datetime import datetime
from types import SimpleNamespace
from typing import Final
from unittest.mock import AsyncMock, patch
import pytest
from fastapi import HTTPException
from mcp.shared.exceptions import MCPError
from mcp.types import CallToolResult, TextContent
from mcp.types import CallToolResult, ListToolsResult, TextContent
from mcp.types import Tool as MCPTool
from pydantic import AnyUrl
@ -29,6 +31,10 @@ from litellm.types.mcp_server.mcp_server_manager import MCPServer
AUTH = UserAPIKeyAuth(api_key="key")
def _listing_client() -> SimpleNamespace:
return SimpleNamespace(discovery_auth_fingerprint=AsyncMock(return_value="test-credential-hash"))
@pytest.fixture
def proxy_mode():
token = _mcp_proxy_mode.set(True)
@ -165,8 +171,8 @@ async def test_proxy_call_tool_on_a_never_listed_tool_hands_the_pre_hook_no_list
patch.dict(manager.registry, {server.server_id: server}),
patch.dict(manager.tool_name_to_mcp_server_name_mapping),
patch.object(mcp_operations, "_get_allowed_mcp_servers", AsyncMock(return_value=[server])),
patch.object(manager, "create_mcp_client", AsyncMock(return_value=object())),
patch.object(manager, "_fetch_tools_with_timeout", AsyncMock(return_value=upstream)),
patch.object(manager, "create_mcp_client", AsyncMock(return_value=_listing_client())),
patch.object(manager, "_fetch_tools_with_timeout", AsyncMock(return_value=ListToolsResult(tools=upstream))),
patch.object(manager, "pre_call_tool_check", pre_call_tool_check),
patch.object(manager, "_call_regular_mcp_tool", call_regular_mcp_tool),
):
@ -204,14 +210,15 @@ async def test_proxy_schema_and_call_resolve_tool_id_without_waiting_on_unrelate
slow_teardown = asyncio.Event()
served_as = with_mcp_proxy_identity(MCPTool(name="proxy-fast-echo", inputSchema={}), fast.server_id)
async def fetch_tools(client: object, server_name: str) -> list[MCPTool]:
async def fetch_tools(client: object, server_name: str) -> ListToolsResult:
if server_name == slow.name:
try:
await slow_release.wait()
except asyncio.CancelledError:
await slow_teardown.wait()
raise
return [MCPTool(name="echo", description="Echo text back", inputSchema={"type": "object"})]
echo: Final = MCPTool(name="echo", description="Echo text back", inputSchema={"type": "object"})
return ListToolsResult(tools=[echo])
async def call_regular_mcp_tool(*, tasks: list[asyncio.Task[object]], **_: object) -> CallToolResult:
await asyncio.gather(*tasks)
@ -222,7 +229,7 @@ async def test_proxy_schema_and_call_resolve_tool_id_without_waiting_on_unrelate
patch.dict(manager.registry, {fast.server_id: fast, slow.server_id: slow}),
patch.dict(manager.tool_name_to_mcp_server_name_mapping),
patch.object(mcp_operations, "_get_allowed_mcp_servers", AsyncMock(return_value=[slow, fast])),
patch.object(manager, "create_mcp_client", AsyncMock(return_value=object())),
patch.object(manager, "create_mcp_client", AsyncMock(return_value=_listing_client())),
patch.object(manager, "_fetch_tools_with_timeout", fetch_tools),
patch.object(manager, "pre_call_tool_check", AsyncMock(return_value={})),
patch.object(manager, "_call_regular_mcp_tool", call_regular_mcp_tool),

View file

@ -69,7 +69,8 @@ async def test_mcp_server_manager_https_server():
# Create a mock MCPClient
mock_client = AsyncMock()
mock_client.list_tools = AsyncMock(return_value=mock_tools)
mock_client.list_tools_result = AsyncMock(return_value=ListToolsResult(tools=mock_tools))
mock_client.discovery_auth_fingerprint = AsyncMock(return_value="test-credential-hash")
mock_client.call_tool = AsyncMock(return_value=mock_result)
# Mock the MCPClient constructor
@ -133,7 +134,7 @@ async def test_mcp_server_manager_https_server():
assert result.content[0].text == "Email sent successfully"
# Verify client methods were called
mock_client.list_tools.assert_called()
mock_client.list_tools_result.assert_called()
mock_client.call_tool.assert_called_once()
@ -176,7 +177,8 @@ async def test_mcp_http_transport_list_tools_mock():
# Create a mock MCPClient that returns our test tools
mock_client = AsyncMock()
mock_client.list_tools = AsyncMock(return_value=mock_tools)
mock_client.list_tools_result = AsyncMock(return_value=ListToolsResult(tools=mock_tools))
mock_client.discovery_auth_fingerprint = AsyncMock(return_value="test-credential-hash")
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock(return_value=None)
@ -218,7 +220,7 @@ async def test_mcp_http_transport_list_tools_mock():
assert tools[1].name == f"{expected_prefix}-calendar_create_event"
# Verify client methods were called
mock_client.list_tools.assert_called()
mock_client.list_tools_result.assert_called()
# Verify tool mapping was updated
expected_prefix = "test_http_server"
@ -983,6 +985,7 @@ async def test_get_tools_from_mcp_servers():
proxy_logging_obj=None,
catalog_auth_header=None,
record_listing=True,
use_tool_list_cache: bool = True,
):
if server.server_id == "server1_id":
return [mock_tool_1]
@ -1106,9 +1109,10 @@ async def test_list_tools_only_returns_allowed_servers(monkeypatch):
mock_client = AsyncMock()
# Return tools based on server URL
if kwargs.get("server_url") == "https://server-a.com/mcp":
mock_client.list_tools = AsyncMock(return_value=mock_tools_a)
mock_client.list_tools_result = AsyncMock(return_value=ListToolsResult(tools=mock_tools_a))
else:
mock_client.list_tools = AsyncMock(return_value=mock_tools_b)
mock_client.list_tools_result = AsyncMock(return_value=ListToolsResult(tools=mock_tools_b))
mock_client.discovery_auth_fingerprint = AsyncMock(return_value="test-credential-hash")
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock(return_value=None)
return mock_client
@ -1403,7 +1407,8 @@ async def test_mcp_server_manager_alias_tool_prefixing():
# Mock MCPClient
mock_client = AsyncMock()
mock_client.list_tools = AsyncMock(return_value=mock_tools)
mock_client.list_tools_result = AsyncMock(return_value=ListToolsResult(tools=mock_tools))
mock_client.discovery_auth_fingerprint = AsyncMock(return_value="test-credential-hash")
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock(return_value=None)
@ -1463,7 +1468,8 @@ async def test_mcp_server_manager_server_name_tool_prefixing():
# Mock MCPClient
mock_client = AsyncMock()
mock_client.list_tools = AsyncMock(return_value=mock_tools)
mock_client.list_tools_result = AsyncMock(return_value=ListToolsResult(tools=mock_tools))
mock_client.discovery_auth_fingerprint = AsyncMock(return_value="test-credential-hash")
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock(return_value=None)
@ -1523,7 +1529,8 @@ async def test_mcp_server_manager_server_id_tool_prefixing():
# Mock MCPClient
mock_client = AsyncMock()
mock_client.list_tools = AsyncMock(return_value=mock_tools)
mock_client.list_tools_result = AsyncMock(return_value=ListToolsResult(tools=mock_tools))
mock_client.discovery_auth_fingerprint = AsyncMock(return_value="test-credential-hash")
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock(return_value=None)
@ -2022,6 +2029,7 @@ async def test_get_tools_for_single_server():
user_api_key_auth=None,
proxy_logging_obj=ANY,
record_listing=False,
use_tool_list_cache=True,
)
# Verify the result

File diff suppressed because it is too large Load diff

View file

@ -94,6 +94,12 @@ def cleanup_mcp_global_state():
yield
def _listing_client() -> AsyncMock:
client: Final = AsyncMock()
client.discovery_auth_fingerprint = AsyncMock(return_value="test-credential-hash")
return client
def _call_tool_params(name, arguments=None):
from mcp.types import CallToolRequestParams
@ -7980,13 +7986,13 @@ def _worker_that_never_listed(server: MCPServer, upstream_tools: tuple[str, ...]
return CallToolResult(content=[TextContent(type="text", text="ok")], is_error=False)
async def fake_fetch_tools(client, server_name):
return [MCPTool(name=tool_name, inputSchema={}) for tool_name in upstream_tools]
return ListToolsResult(tools=[MCPTool(name=tool_name, inputSchema={}) for tool_name in upstream_tools])
with (
patch.object( # test-quality-ok: the upstream MCP session is the boundary; a real one needs an initialize handshake over a live server
mcp_operations.global_mcp_server_manager,
"create_mcp_client",
new=AsyncMock(return_value=MagicMock()),
new=AsyncMock(return_value=_listing_client()),
) as create_client,
patch.object( # test-quality-ok: same boundary, this is the tools/list answer the upstream would give
mcp_operations.global_mcp_server_manager,
@ -8832,7 +8838,7 @@ async def test_execute_mcp_tool_implicit_listing_before_the_first_call_hands_hoo
server = _never_listed_passthrough_server()
manager.registry[server.server_id] = server
manager._listed_tools_by_server_id.pop(server.server_id, None)
upstream = AsyncMock()
upstream: Final = _listing_client()
upstream.call_tool.return_value = CallToolResult(content=[TextContent(type="text", text="ok")], isError=False)
proxy_logging = _mock_mcp_proxy_logging()
proxy_logging.has_pre_call_guardrails.return_value = False
@ -8842,7 +8848,9 @@ async def test_execute_mcp_tool_implicit_listing_before_the_first_call_hands_hoo
proxy_logging.pre_call_hook = AsyncMock(return_value={})
proxy_logging.during_call_hook = AsyncMock(return_value=None)
fetch_tools = AsyncMock(
return_value=[MCPTool(name="add", description="Adds. FLAGWORD", inputSchema={"type": "object"})]
return_value=ListToolsResult(
tools=[MCPTool(name="add", description="Adds. FLAGWORD", inputSchema={"type": "object"})]
)
)
with (
@ -8889,11 +8897,13 @@ async def test_fetch_pinnable_tool_catalog_records_no_listed_catalog_for_the_adm
request.client.host = "10.1.2.3"
request.headers = {"x-litellm-api-key": "sk-admin"}
fetch_tools = AsyncMock(
return_value=[MCPTool(name="add", description="Upstream wording", inputSchema={"type": "object"})]
return_value=ListToolsResult(
tools=[MCPTool(name="add", description="Upstream wording", inputSchema={"type": "object"})]
)
)
with (
patch.object(manager, "create_mcp_client", new=AsyncMock(return_value=MagicMock())),
patch.object(manager, "create_mcp_client", new=AsyncMock(return_value=_listing_client())),
patch.object(manager, "_fetch_tools_with_timeout", new=fetch_tools),
patch("litellm.proxy.proxy_server.proxy_logging_obj", ProxyLogging(user_api_key_cache=DualCache())),
):

View file

@ -21,7 +21,7 @@ from typing import Any, Final
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from mcp.types import Tool
from mcp.types import ListToolsResult, Tool
from pydantic import JsonValue
import litellm
@ -69,6 +69,10 @@ SAMPLE_TOOLS = _make_tools(
)
def _listing_client() -> SimpleNamespace:
return SimpleNamespace(discovery_auth_fingerprint=AsyncMock(return_value="test-credential-hash"))
@pytest.mark.parametrize(
("schema", "arguments", "error"),
(
@ -1417,8 +1421,8 @@ async def test_mcp_tool_search_leaves_the_listed_tools_slot_empty(monkeypatch: p
with (
patch.dict(manager.tool_name_to_mcp_server_name_mapping),
patch.object(mcp_operations, "_get_allowed_mcp_servers", AsyncMock(return_value=[server])),
patch.object(manager, "create_mcp_client", AsyncMock(return_value=object())),
patch.object(manager, "_fetch_tools_with_timeout", AsyncMock(return_value=upstream)),
patch.object(manager, "create_mcp_client", AsyncMock(return_value=_listing_client())),
patch.object(manager, "_fetch_tools_with_timeout", AsyncMock(return_value=ListToolsResult(tools=upstream))),
):
try:
result = await handle_mcp_tool_search(query="echo", top_k=1, user_api_key_dict=user)

View file

@ -1,11 +1,12 @@
import asyncio
from collections.abc import Callable, Sequence
from types import SimpleNamespace
from typing import Final, Literal
from unittest.mock import AsyncMock, patch
import pytest
from mcp.server.context import ServerRequestContext
from mcp.types import GetPromptRequest, GetPromptRequestParams, GetPromptResult
from mcp.types import GetPromptRequest, GetPromptRequestParams, GetPromptResult, ListToolsResult
from mcp.types import Tool as MCPTool
import litellm
@ -45,6 +46,10 @@ async def _served_catalog_tool() -> str:
return "ok"
def _listing_client() -> SimpleNamespace:
return SimpleNamespace(discovery_auth_fingerprint=AsyncMock(return_value="test-credential-hash"))
@pytest.mark.parametrize(
("settings", "mcp_servers", "toolset_id", "auth_state", "raises"),
[
@ -1602,8 +1607,8 @@ async def test_list_mcp_tools_records_the_catalog_only_when_asked(
upstream = [MCPTool(name="echo", description="Echo text back", inputSchema={"type": "object"})]
with (
patch.object(operations, "_get_allowed_mcp_servers", AsyncMock(return_value=[server])),
patch.object(manager, "create_mcp_client", AsyncMock(return_value=object())),
patch.object(manager, "_fetch_tools_with_timeout", AsyncMock(return_value=upstream)),
patch.object(manager, "create_mcp_client", AsyncMock(return_value=_listing_client())),
patch.object(manager, "_fetch_tools_with_timeout", AsyncMock(return_value=ListToolsResult(tools=upstream))),
patch.dict(manager.tool_name_to_mcp_server_name_mapping),
):
try:

View file

@ -554,6 +554,39 @@ def test_mcp_metadata_accepts_a_boolean_skip_tool_output_validation(
assert parsed.mcp_info == {"skip_tool_output_validation": value}
@pytest.mark.parametrize("request_model", (NewMCPServerRequest, UpdateMCPServerRequest))
@pytest.mark.parametrize("value", [-1, 1.5, "60", True, None, 31536001, 10**400])
def test_mcp_metadata_rejects_a_tools_list_cache_ttl_that_is_not_whole_seconds_within_a_year(
request_model: type[NewMCPServerRequest] | type[UpdateMCPServerRequest], value: object
) -> None:
payload: Final = {
"server_id": "test",
"transport": "http",
"url": "https://example.com/mcp",
"mcp_info": {"tools_list_cache_ttl_seconds": value},
}
with pytest.raises(ValidationError):
request_model.model_validate(payload)
@pytest.mark.parametrize("request_model", (NewMCPServerRequest, UpdateMCPServerRequest))
@pytest.mark.parametrize("value", [0, 60, 31536000])
def test_mcp_metadata_accepts_a_tools_list_cache_ttl_in_whole_seconds(
request_model: type[NewMCPServerRequest] | type[UpdateMCPServerRequest], value: int
) -> None:
parsed: Final = request_model.model_validate(
{
"server_id": "test",
"transport": "http",
"url": "https://example.com/mcp",
"mcp_info": {"tools_list_cache_ttl_seconds": value},
}
)
assert parsed.mcp_info == {"tools_list_cache_ttl_seconds": value}
MCP_SERVER_REQUESTS = (NewMCPServerRequest, UpdateMCPServerRequest)
STDIO_SERVER_FIELDS = {"server_id": "stdio-1", "transport": "stdio", "command": "python", "args": ["server.py"]}

View file

@ -1097,7 +1097,10 @@ async def test_aresponses_mcp_call_forwards_listed_admission_authorization(
monkeypatch.setattr(manager, "registry", {server.server_id: server})
mcp_client: Final = AsyncMock()
mcp_client.list_tools.return_value = [MCPTool(name="whoami", inputSchema={"type": "object", "properties": {}})]
mcp_client.discovery_auth_fingerprint.return_value = "test-credential-hash"
mcp_client.list_tools_result.return_value = ListToolsResult(
tools=[MCPTool(name="whoami", inputSchema={"type": "object", "properties": {}})]
)
mcp_client.call_tool.return_value = CallToolResult(
content=[TextContent(type="text", text="ok")],
isError=False,

View file

@ -12,7 +12,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import HTTPException
from mcp.types import CallToolResult, TextContent
from mcp.types import CallToolResult, ListToolsResult, TextContent
from mcp.types import Tool as MCPTool
from openai.types.responses.tool_param import Mcp
@ -78,6 +78,10 @@ def _setup_proxy_logging(monkeypatch: pytest.MonkeyPatch) -> AsyncMock:
return proxy_logging_obj.post_call_failure_hook
def _listing_client() -> types.SimpleNamespace:
return types.SimpleNamespace(discovery_auth_fingerprint=AsyncMock(return_value="test-credential-hash"))
def test_deduplicate_mcp_tools_single_allowed_server():
tools = [{"name": "search"}, {"name": "search"}] # duplicate on purpose
@ -1592,8 +1596,8 @@ async def test_bridge_listing_leaves_the_callers_catalog_unchanged(
with (
patch.dict(manager.tool_name_to_mcp_server_name_mapping),
patch.object(mcp_operations, "_get_allowed_mcp_servers", AsyncMock(return_value=[server])),
patch.object(manager, "create_mcp_client", AsyncMock(return_value=object())),
patch.object(manager, "_fetch_tools_with_timeout", AsyncMock(return_value=upstream)),
patch.object(manager, "create_mcp_client", AsyncMock(return_value=_listing_client())),
patch.object(manager, "_fetch_tools_with_timeout", AsyncMock(return_value=ListToolsResult(tools=upstream))),
):
try:
if real_listing:
@ -1665,9 +1669,10 @@ async def test_concurrent_bridge_calls_use_their_own_served_metadata(monkeypatch
MCPTool(name="status", description="Read status", inputSchema={"type": "object"}),
]
client: Final = AsyncMock()
client.discovery_auth_fingerprint = AsyncMock(return_value="test-credential-hash")
client.call_tool.return_value = CallToolResult(content=[TextContent(type="text", text="ok")])
manager.create_mcp_client = AsyncMock(return_value=client)
manager._fetch_tools_with_timeout = AsyncMock(return_value=upstream)
manager._fetch_tools_with_timeout = AsyncMock(return_value=ListToolsResult(tools=upstream))
guardrail: Final = _BridgeMetadataGuardrail()
logger: Final = ProxyLogging(user_api_key_cache=DualCache())
monkeypatch.setattr(litellm, "callbacks", [guardrail])
@ -1784,13 +1789,13 @@ async def _listing_through_the_gateway(
agent_permission: LiteLLM_ObjectPermissionTable | None = None,
servers: Sequence[MCPServer] = GATEWAY_SERVERS,
) -> tuple[list[MCPTool], list[str]]:
async def upstream_listing(_client: object, server_name: str) -> list[MCPTool]:
return list(UPSTREAM_TOOLS[server_name])
async def upstream_listing(_client: object, server_name: str) -> ListToolsResult:
return ListToolsResult(tools=list(UPSTREAM_TOOLS[server_name]))
manager: Final = MCPServerManager()
manager.registry = {server.server_id: server for server in servers}
manager.config_mcp_servers = dict(manager.registry)
manager.create_mcp_client = AsyncMock(return_value=object())
manager.create_mcp_client = AsyncMock(return_value=_listing_client())
manager._fetch_tools_with_timeout = AsyncMock(side_effect=upstream_listing)
tables: Final = types.SimpleNamespace(litellm_mcptoolsettable=toolset_table or _toolset_table())
monkeypatch.setattr("litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager", manager)