mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
feat(mcp): cache upstream tool lists with an admin-set lifetime (#51160)
This commit is contained in:
parent
e3b32534dc
commit
40981a5d89
21 changed files with 1365 additions and 156 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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 = [
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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
|
|
@ -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())),
|
||||
):
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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"]}
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue