diff --git a/litellm/mcp_client/client.py b/litellm/mcp_client/client.py index 0461f1f06ff..224e14c95c8 100644 --- a/litellm/mcp_client/client.py +++ b/litellm/mcp_client/client.py @@ -30,6 +30,8 @@ from mcp.shared._stream_protocols import ReadStream, WriteStream from mcp.shared.message import SessionMessage from typing_extensions import ReadOnly, TypedDict, assert_never +from litellm.mcp_client.output_schema_memory import OutputSchemaMemory + _TransportStreams: TypeAlias = tuple[ ReadStream[SessionMessage | Exception], WriteStream[SessionMessage], @@ -485,6 +487,39 @@ class _UnvalidatedClientSession(ClientSession): return None +class _RemembersOutputSchemas(ClientSession): + """Preloads the SDK's output schemas from earlier listings, so the output check needs no list of its own.""" + + _seeded: frozenset[str] = frozenset() + _on_listed: Callable[[ListToolsResult], None] | None = None + + def bind(self, remembered: Sequence[MCPTool], on_listed: Callable[[ListToolsResult], None]) -> None: + self._on_listed = on_listed + self._seeded = frozenset(tool.name for tool in remembered) + if remembered: + self._absorb_tool_listing( # pyright: ignore[reportPrivateUsage] # the SDK has no public way to preload schemas + ListToolsResult(tools=list(remembered)), complete=False + ) + + async def list_tools(self, *, params: PaginatedRequestParams | None = None) -> ListToolsResult: + result: Final = await super().list_tools(params=params) + if self._on_listed is not None and (params is None or params.cursor is None): + self._on_listed(result) + return result + + async def validate_tool_result(self, name: str, result: MCPCallToolResult) -> None: + try: + await super().validate_tool_result(name, result) + except RuntimeError: + if name not in self._seeded: + raise + self._seeded = frozenset() + relisted: Final = await self.list_tools() + if all(tool.name != name for tool in relisted.tools): + return + await super().validate_tool_result(name, result) + + class MCPClient: """ MCP Client supporting: @@ -514,7 +549,11 @@ class MCPClient: protocol_version: MCPUpstreamProtocol = "auto", skip_tool_output_validation: bool = False, upstream_transport: httpx2.AsyncBaseTransport | None = None, - ): + output_schema_memory: OutputSchemaMemory | None = None, + output_schema_cache_ttl_seconds: float = 0.0, + ) -> None: + self._output_schema_memory: Final = output_schema_memory + self._output_schema_cache_ttl_seconds: Final = output_schema_cache_ttl_seconds self.protocol_version: MCPUpstreamProtocol = TypeAdapter[MCPUpstreamProtocol]( MCPUpstreamProtocol ).validate_python(protocol_version) @@ -574,6 +613,17 @@ class MCPClient: async def discovery_auth_fingerprint(self) -> str: return self._hash_discovery_auth(await self.prepare_request_auth(), self._aws_identity()) + async def _output_schema_key(self) -> str | None: + if ( + self._output_schema_memory is None + or self._output_schema_cache_ttl_seconds <= 0 + or self.skip_tool_output_validation + or self.protocol_version == "2026-07-28" + or self.transport_type == MCPTransport.stdio + ): + return None + return await self.discovery_auth_fingerprint() + async def subscription_auth_fingerprint(self) -> str: from litellm.proxy.mcp_server.outbound_credentials.client_credentials import ( ClientCredentialsBearerAuth, @@ -767,6 +817,7 @@ class MCPClient: user input (elicitation), or send log messages. """ in_flight_error: BaseException | None = None + schema_key: Final = await self._output_schema_key() with anyio.CancelScope() as cleanup_scope: if isinstance(http_client, _MCPHTTPClient): http_client.cleanup_scope = cleanup_scope @@ -788,7 +839,11 @@ class MCPClient: # The SDK drops a response stream that ends without a JSON-RPC reply, so nothing else # ever fails the request. session_type: Final = ( - _UnvalidatedClientSession if self.skip_tool_output_validation else ClientSession + _UnvalidatedClientSession + if self.skip_tool_output_validation + else ClientSession + if schema_key is None + else _RemembersOutputSchemas ) session_ctx: Final = session_type( read_stream, @@ -808,6 +863,21 @@ class MCPClient: self._last_initialize_instructions = ( instructions.strip() or None if isinstance(instructions, str) else None ) + if ( + isinstance(session_ctx, _RemembersOutputSchemas) + and schema_key is not None + and self._output_schema_memory is not None + ): + memory: Final = self._output_schema_memory + + def remember_listing(listed: ListToolsResult) -> None: + memory.remember( + schema_key, + listed.tools, + ttl_seconds=self._output_schema_cache_ttl_seconds, + ) + + session_ctx.bind(memory.recall(schema_key), remember_listing) result: Final = await operation(session) except BaseException as operation_error: in_flight_error = operation_error diff --git a/litellm/mcp_client/output_schema_memory.py b/litellm/mcp_client/output_schema_memory.py new file mode 100644 index 00000000000..b6d5191220d --- /dev/null +++ b/litellm/mcp_client/output_schema_memory.py @@ -0,0 +1,49 @@ +from __future__ import annotations + +import time +from collections.abc import Callable, Iterable +from dataclasses import dataclass +from typing import Final + +from mcp.types import Tool + + +@dataclass(frozen=True, slots=True) +class _Remembered: + tools: tuple[Tool, ...] + expires_at: float + + +def _schema_only(tool: Tool) -> Tool: + return Tool(name=tool.name, input_schema={"type": "object"}, output_schema=tool.output_schema) + + +class OutputSchemaMemory: + def __init__( + self, + *, + max_entries: int, + clock: Callable[[], float] = time.monotonic, + ) -> None: + self._max_entries: Final = max_entries + self._clock: Final = clock + self._entries: dict[str, _Remembered] = {} # mutable-ok: bounded process-local store + + def recall(self, key: str) -> tuple[Tool, ...]: + entry: Final = self._entries.get(key) + if entry is None: + return () + if entry.expires_at <= self._clock(): + self._entries.pop(key, None) + return () + return entry.tools + + def remember(self, key: str, tools: Iterable[Tool], *, ttl_seconds: float) -> None: + kept: Final = {tool.name: _schema_only(tool) for tool in tools} + self._entries.pop(key, None) + while len(self._entries) >= self._max_entries: + self._entries.pop(next(iter(self._entries))) + self._entries[key] = _Remembered(tools=tuple(kept.values()), expires_at=self._clock() + ttl_seconds) + + def forget(self, key: str) -> None: + self._entries.pop(key, None) diff --git a/litellm/proxy/mcp_server/mcp_server_manager.py b/litellm/proxy/mcp_server/mcp_server_manager.py index 013336e1d5d..eeebf1a294c 100644 --- a/litellm/proxy/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/mcp_server/mcp_server_manager.py @@ -83,6 +83,7 @@ from litellm.mcp_client.client import ( strip_auth_scheme, to_basic_credentials, ) +from litellm.mcp_client.output_schema_memory import OutputSchemaMemory from litellm.proxy._types import ( LiteLLM_MCPServerTable, MCPAuthType, @@ -289,6 +290,7 @@ _USER_ENV_VARS_CACHE_MAX_SIZE: Final = 4096 # cap to prevent unbounded growth _ListedToolsByCaller: TypeAlias = Mapping[str | None, Mapping[str, MCPTool]] _LISTED_TOOLS_CALLERS_PER_SERVER: Final = 256 +_OUTPUT_SCHEMA_MEMORY_MAX_ENTRIES: Final = 256 _UNTYPED_VALUE: Final = TypeAdapter(object) @@ -2024,7 +2026,7 @@ class MCPServerManager: per_user_token_cache: MCPPerUserTokenCache | None = None, discovery_clock: Callable[[], float] = time.monotonic, upstream_transport: httpx2.AsyncBaseTransport | None = None, - ): + ) -> None: self._per_user_oauth_token_store = per_user_oauth_token_store or LazyPerUserOAuthTokenStore( self.get_mcp_server_by_id ) @@ -2084,6 +2086,7 @@ class MCPServerManager: """ self._listed_tools_by_server_id: dict[str, _ListedToolsByCaller] = {} # mutable-ok: refreshed per tools/list self._listed_tools_generations: dict[str, int] = {} # mutable-ok: bumped per server save + self._output_schema_memory: Final = OutputSchemaMemory(max_entries=_OUTPUT_SCHEMA_MEMORY_MAX_ENTRIES) self._upstream_initialize_instructions_by_server_id: dict[str, str] = {} # Per-server monotonic timestamp of last upstream prefetch attempt (success, # empty result, or failure). Used to throttle re-probes for servers that do @@ -4421,6 +4424,7 @@ class MCPServerManager: ), upstream_transport=self._upstream_transport, for_tool_listing=for_tool_listing, + output_schema_memory=self._output_schema_memory, ) _create_mcp_client = create_mcp_client diff --git a/litellm/proxy/mcp_server/upstream.py b/litellm/proxy/mcp_server/upstream.py index 3dd245947db..c510aac818d 100644 --- a/litellm/proxy/mcp_server/upstream.py +++ b/litellm/proxy/mcp_server/upstream.py @@ -8,6 +8,7 @@ from fastapi import HTTPException from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_NPM_CACHE_DIR, MCP_STDIO_ALLOWED_COMMANDS from litellm.mcp_client.client import MCPClient, MCPSigV4Auth +from litellm.mcp_client.output_schema_memory import OutputSchemaMemory from litellm.proxy._types import MCPTransport, UserAPIKeyAuth from litellm.proxy.mcp_server.legacy_callbacks import ElicitationCallback, SamplingCallback from litellm.proxy.mcp_server.mcp_debug import record_auth_resolution @@ -209,6 +210,7 @@ async def prepare_upstream_client( elicitation_callback: ElicitationCallback | None = None, upstream_transport: httpx2.AsyncBaseTransport | None = None, for_tool_listing: bool = False, + output_schema_memory: OutputSchemaMemory | None = None, ) -> MCPClient: transport: Final[MCPTransportType] = server.transport or MCPTransport.sse server_spec: Final = None if transport == MCPTransport.stdio else to_server_spec_fail_closed(server) @@ -232,6 +234,8 @@ async def prepare_upstream_client( transport_type=transport, protocol_version=protocol_version, skip_tool_output_validation=server.skip_tool_output_validation, + output_schema_memory=output_schema_memory, + output_schema_cache_ttl_seconds=server.output_schema_cache_ttl_seconds, auth_type=server.auth_type, auth_value=auth_value, timeout=timeout, @@ -267,6 +271,8 @@ async def prepare_upstream_client( transport_type=transport, protocol_version=protocol_version, skip_tool_output_validation=server.skip_tool_output_validation, + output_schema_memory=output_schema_memory, + output_schema_cache_ttl_seconds=server.output_schema_cache_ttl_seconds, upstream_transport=upstream_transport, auth_type=server.auth_type, timeout=timeout, @@ -311,6 +317,8 @@ async def prepare_upstream_client( transport_type=transport, protocol_version=protocol_version, skip_tool_output_validation=server.skip_tool_output_validation, + output_schema_memory=output_schema_memory, + output_schema_cache_ttl_seconds=server.output_schema_cache_ttl_seconds, upstream_transport=upstream_transport, auth_type=server.auth_type, auth_value=auth_value, diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 1d3d5911eb1..649c8af4ae1 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -1,4 +1,5 @@ import json +import sys from collections.abc import Mapping from datetime import datetime from typing import Annotated, Any, Final, Literal @@ -25,6 +26,15 @@ 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 "output_schema_cache_ttl_seconds" in value: + ttl: Final = value["output_schema_cache_ttl_seconds"] + if ( + isinstance(ttl, bool) + or not isinstance(ttl, (int, float)) + or ttl > sys.float_info.max + or not 0 <= ttl < float("inf") + ): + raise ValueError("output_schema_cache_ttl_seconds must be a finite number of seconds, 0 or more") return value @@ -303,6 +313,10 @@ class MCPServer(LiteLLMBaseModel): def skip_tool_output_validation(self) -> bool: return (self.mcp_info or {}).get("skip_tool_output_validation") is True + @property + def output_schema_cache_ttl_seconds(self) -> float: + return float((self.mcp_info or {}).get("output_schema_cache_ttl_seconds", 0)) + @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: diff --git a/tests/_support/mcp_count_tool_upstream.py b/tests/_support/mcp_count_tool_upstream.py index f1aa61c7dba..7e666eca548 100644 --- a/tests/_support/mcp_count_tool_upstream.py +++ b/tests/_support/mcp_count_tool_upstream.py @@ -9,9 +9,11 @@ COUNT_OUTPUT_SCHEMA: Final = {"type": "object", "properties": {"n": {"type": "in def count_tool_upstream( - structured: Mapping[str, object], methods_seen: "queue.SimpleQueue[str]" + structured: Mapping[str, object], + methods_seen: "queue.SimpleQueue[str]", + output_schema: Mapping[str, object] = COUNT_OUTPUT_SCHEMA, ) -> Callable[[httpx2.Request], httpx2.Response]: - """An MCP upstream with one tool, `count`, whose declared output schema wants an integer `n`.""" + """An MCP upstream with one tool, `count`, whose declared output schema defaults to wanting an integer `n`.""" def respond(request: httpx2.Request) -> httpx2.Response: if request.method != "POST": @@ -22,12 +24,12 @@ def count_tool_upstream( return httpx2.Response(202) results: Final[Mapping[str, Mapping[str, object]]] = { "initialize": { - "protocolVersion": payload["params"].get("protocolVersion"), + "protocolVersion": payload.get("params", {}).get("protocolVersion"), "capabilities": {"tools": {}}, "serverInfo": {"name": "count", "version": "1"}, }, "tools/list": { - "tools": [{"name": "count", "inputSchema": {"type": "object"}, "outputSchema": COUNT_OUTPUT_SCHEMA}] + "tools": [{"name": "count", "inputSchema": {"type": "object"}, "outputSchema": output_schema}] }, "tools/call": {"content": [{"type": "text", "text": "ok"}], "structuredContent": structured}, } diff --git a/tests/unit/mcp_client/test_mcp_client.py b/tests/unit/mcp_client/test_mcp_client.py index 866a29b70e9..b6b69eaf08a 100644 --- a/tests/unit/mcp_client/test_mcp_client.py +++ b/tests/unit/mcp_client/test_mcp_client.py @@ -5,7 +5,7 @@ import json import os import queue import sys -from collections.abc import AsyncIterator, Callable +from collections.abc import AsyncIterator, Callable, Mapping from contextlib import AbstractAsyncContextManager, asynccontextmanager from pathlib import Path from types import ModuleType @@ -75,6 +75,7 @@ from litellm.proxy.mcp_server.outbound_credentials.httpx_auth import StaticHeade from litellm.types.mcp import MCPAuth, MCPStdioConfig, MCPTransport from litellm.types.mcp_server.mcp_server_manager import MCPServer from tests._support.clock_loops import AutojumpClockLoop, ManualClockLoop +from litellm.mcp_client.output_schema_memory import OutputSchemaMemory from tests._support.mcp_count_tool_upstream import count_tool_upstream, methods_received _JSONRPC_MESSAGE_ADAPTER: Final = TypeAdapter(JSONRPCMessage) @@ -5304,3 +5305,171 @@ async def test_a_skipped_validation_passes_the_upstreams_result_through_unchange assert isinstance(result, CallToolResult) assert result.structured_content == {"n": "not-an-integer"} + + +def _remembering_client( + upstream: Callable[[httpx2.Request], httpx2.Response], + memory: OutputSchemaMemory, + *, + tenant: str = "alice", + ttl_seconds: float = 60.0, +) -> "_MockTransportClient": + return _MockTransportClient( + upstream, + server_url="https://upstream.example/mcp", + transport_type=MCPTransport.http, + extra_headers={"x-tenant": tenant}, + output_schema_memory=memory, + output_schema_cache_ttl_seconds=ttl_seconds, + ) + + +def _new_memory() -> OutputSchemaMemory: + return OutputSchemaMemory(max_entries=8) + + +@pytest.mark.asyncio +async def test_a_later_call_reuses_the_schemas_an_earlier_call_listed() -> None: + memory: Final = _new_memory() + first_seen: Final[queue.SimpleQueue[str]] = queue.SimpleQueue() + second_seen: Final[queue.SimpleQueue[str]] = queue.SimpleQueue() + + first: Final = await _remembering_client(count_tool_upstream({"n": 1}, first_seen), memory).call_tool( + CallToolRequestParams(name="count", arguments={}) + ) + second: Final = await _remembering_client(count_tool_upstream({"n": 2}, second_seen), memory).call_tool( + CallToolRequestParams(name="count", arguments={}) + ) + + assert isinstance(first, CallToolResult) and isinstance(second, CallToolResult) + assert not first.is_error and not second.is_error + assert "tools/list" in methods_received(first_seen) + assert "tools/list" not in methods_received(second_seen) + + +@pytest.mark.asyncio +async def test_remembered_schemas_still_reject_invalid_output() -> None: + memory: Final = _new_memory() + await _remembering_client(count_tool_upstream({"n": 1}, queue.SimpleQueue()), memory).call_tool( + CallToolRequestParams(name="count", arguments={}) + ) + + result: Final = await _remembering_client( + count_tool_upstream({"n": "not-an-integer"}, queue.SimpleQueue()), memory + ).call_tool(CallToolRequestParams(name="count", arguments={})) + + assert isinstance(result, CallToolResult) + assert result.is_error is True + + +@pytest.mark.asyncio +async def test_a_schema_the_upstream_changed_does_not_reject_output_that_matches_the_new_schema() -> None: + memory: Final = _new_memory() + await _remembering_client(count_tool_upstream({"n": 1}, queue.SimpleQueue()), memory).call_tool( + CallToolRequestParams(name="count", arguments={}) + ) + renamed_schema: Final = {"type": "object", "properties": {"label": {"type": "string"}}, "required": ["label"]} + seen: Final[queue.SimpleQueue[str]] = queue.SimpleQueue() + + result: Final = await _remembering_client( + count_tool_upstream({"label": "x"}, seen, output_schema=renamed_schema), memory + ).call_tool(CallToolRequestParams(name="count", arguments={})) + + assert isinstance(result, CallToolResult) + assert result.is_error is False + assert "tools/list" in methods_received(seen) + + +@pytest.mark.asyncio +async def test_callers_with_different_credentials_do_not_share_remembered_schemas() -> None: + memory: Final = _new_memory() + await _remembering_client(count_tool_upstream({"n": 1}, queue.SimpleQueue()), memory, tenant="alice").call_tool( + CallToolRequestParams(name="count", arguments={}) + ) + seen: Final[queue.SimpleQueue[str]] = queue.SimpleQueue() + + await _remembering_client(count_tool_upstream({"n": 1}, seen), memory, tenant="bob").call_tool( + CallToolRequestParams(name="count", arguments={}) + ) + + assert "tools/list" in methods_received(seen) + + +@pytest.mark.asyncio +async def test_a_zero_ttl_lists_tools_on_every_call_and_enforces_a_tightened_schema() -> None: + memory: Final = _new_memory() + await _remembering_client(count_tool_upstream({"n": 1}, queue.SimpleQueue()), memory, ttl_seconds=0).call_tool( + CallToolRequestParams(name="count", arguments={}) + ) + tightened_schema: Final = {"type": "object", "properties": {"b": {"type": "string"}}, "required": ["b"]} + seen: Final[queue.SimpleQueue[str]] = queue.SimpleQueue() + + result: Final = await _remembering_client( + count_tool_upstream({"a": "x"}, seen, output_schema=tightened_schema), memory, ttl_seconds=0 + ).call_tool(CallToolRequestParams(name="count", arguments={})) + + assert isinstance(result, CallToolResult) + assert result.is_error is True + assert "tools/list" in methods_received(seen) + + +def _paginated_upstream(structured: Mapping[str, object], *, first_page_tool: str = "first") -> Callable[[httpx2.Request], httpx2.Response]: + schema: Final = {"type": "object", "properties": {"n": {"type": "integer"}}, "required": ["n"]} + + def respond(request: httpx2.Request) -> httpx2.Response: + if request.method != "POST": + return httpx2.Response(200) + payload: Final = json.loads(request.content) + if "id" not in payload: + return httpx2.Response(202) + on_second_page: Final = (payload.get("params") or {}).get("cursor") == "page-2" + results: Final = { + "initialize": { + "protocolVersion": payload.get("params", {}).get("protocolVersion"), + "capabilities": {"tools": {}}, + "serverInfo": {"name": "paged", "version": "1"}, + }, + "tools/list": { + "tools": [ + { + "name": "second" if on_second_page else first_page_tool, + "inputSchema": {"type": "object"}, + "outputSchema": schema, + } + ], + **({} if on_second_page else {"nextCursor": "page-2"}), + }, + "tools/call": {"content": [{"type": "text", "text": "ok"}], "structuredContent": structured}, + } + return httpx2.Response( + 200, json={"jsonrpc": "2.0", "id": payload["id"], "result": results.get(payload["method"], {})} + ) + + return respond + + +@pytest.mark.asyncio +async def test_a_tool_beyond_the_first_listing_page_is_not_validated_when_remembered_schemas_are_in_use() -> None: + memory: Final = _new_memory() + await _remembering_client(_paginated_upstream({"n": 1}), memory).list_tools() + + result: Final = await _remembering_client(_paginated_upstream({"n": "not-an-integer"}), memory).call_tool( + CallToolRequestParams(name="second", arguments={}) + ) + + assert isinstance(result, CallToolResult) + assert result.is_error is False + + +@pytest.mark.asyncio +async def test_a_remembered_tool_that_left_the_first_page_is_not_validated_against_its_old_schema() -> None: + memory: Final = _new_memory() + await _remembering_client(_paginated_upstream({"n": 1}), memory).list_tools() + first_page_only: Final = _paginated_upstream({"n": "not-an-integer"}, first_page_tool="other") + + result: Final = await _remembering_client(first_page_only, memory).call_tool( + CallToolRequestParams(name="first", arguments={}) + ) + + assert isinstance(result, CallToolResult) + assert result.is_error is False diff --git a/tests/unit/mcp_client/test_output_schema_memory.py b/tests/unit/mcp_client/test_output_schema_memory.py new file mode 100644 index 00000000000..ce78f9fbfac --- /dev/null +++ b/tests/unit/mcp_client/test_output_schema_memory.py @@ -0,0 +1,72 @@ +from typing import Final + +from mcp.types import Tool + +from litellm.mcp_client.output_schema_memory import OutputSchemaMemory + + +class _Clock: + def __init__(self) -> None: + self.now = 0.0 + + def __call__(self) -> float: + return self.now + + +def _tool(name: str) -> Tool: + return Tool(name=name, input_schema={"type": "object"}, output_schema={"type": "object", "title": name}) + + +def _names(memory: OutputSchemaMemory, key: str) -> tuple[str, ...]: + return tuple(sorted(tool.name for tool in memory.recall(key))) + + +def test_a_remembered_listing_is_recalled_until_it_expires() -> None: + clock: Final = _Clock() + memory: Final = OutputSchemaMemory(max_entries=4, clock=clock) + memory.remember("upstream", [_tool("a")], ttl_seconds=60) + + clock.now = 59 + assert _names(memory, "upstream") == ("a",) + clock.now = 60 + assert _names(memory, "upstream") == () + + +def test_a_remembered_tool_keeps_its_output_schema() -> None: + memory: Final = OutputSchemaMemory(max_entries=4) + memory.remember("upstream", [_tool("a")], ttl_seconds=60) + + assert memory.recall("upstream")[0].output_schema == {"type": "object", "title": "a"} + + +def test_a_new_listing_replaces_the_remembered_one_and_restarts_its_expiry() -> None: + clock: Final = _Clock() + memory: Final = OutputSchemaMemory(max_entries=4, clock=clock) + memory.remember("upstream", [_tool("a")], ttl_seconds=60) + clock.now = 50 + memory.remember("upstream", [_tool("b")], ttl_seconds=60) + assert _names(memory, "upstream") == ("b",) + + clock.now = 109 + assert _names(memory, "upstream") == ("b",) + clock.now = 110 + assert _names(memory, "upstream") == () + + +def test_the_oldest_upstream_is_dropped_when_the_memory_is_full() -> None: + memory: Final = OutputSchemaMemory(max_entries=2) + memory.remember("one", [_tool("a")], ttl_seconds=60) + memory.remember("two", [_tool("a")], ttl_seconds=60) + memory.remember("three", [_tool("a")], ttl_seconds=60) + + assert _names(memory, "one") == () + assert _names(memory, "two") == ("a",) + assert _names(memory, "three") == ("a",) + + +def test_forgetting_an_upstream_drops_its_listing() -> None: + memory: Final = OutputSchemaMemory(max_entries=4) + memory.remember("upstream", [_tool("a")], ttl_seconds=60) + memory.forget("upstream") + + assert _names(memory, "upstream") == () diff --git a/tests/unit/proxy/mcp_server/test_upstream.py b/tests/unit/proxy/mcp_server/test_upstream.py index 7c88c679972..b558fb4e864 100644 --- a/tests/unit/proxy/mcp_server/test_upstream.py +++ b/tests/unit/proxy/mcp_server/test_upstream.py @@ -190,3 +190,25 @@ async def test_upstream_client_ignores_the_setting_where_the_call_already_lists_ ) assert client.skip_tool_output_validation is expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("setting", "expected"), [(None, 0.0), (0, 0.0), (45, 45.0)]) +async def test_upstream_client_takes_its_output_schema_cache_ttl_from_the_server( + setting: float | None, expected: float +) -> None: + server: Final = MCPServer.model_validate( + { + "server_id": "schema-cache", + "name": "schema_cache", + "url": _HTTP_URL, + "transport": MCPTransport.http, + "mcp_info": {} if setting is None else {"output_schema_cache_ttl_seconds": setting}, + } + ) + + client: Final = await prepare_upstream_client( + server, provider=UpstreamCredentialProvider(), root_path="/mcp", protocol_version="auto" + ) + + assert client._output_schema_cache_ttl_seconds == expected diff --git a/tests/unit/proxy/test__types.py b/tests/unit/proxy/test__types.py index 40e66f1da65..6a5e04afc0a 100644 --- a/tests/unit/proxy/test__types.py +++ b/tests/unit/proxy/test__types.py @@ -553,6 +553,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", ["60", -1, True, None, float("inf"), float("nan"), 10**400]) +def test_mcp_metadata_rejects_an_invalid_output_schema_cache_ttl( + request_model: type[NewMCPServerRequest] | type[UpdateMCPServerRequest], value: object +) -> None: + payload: Final = { + "server_id": "test", + "transport": "http", + "url": "https://example.com/mcp", + "mcp_info": {"output_schema_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, 30, 2.5]) +def test_mcp_metadata_accepts_a_non_negative_output_schema_cache_ttl( + request_model: type[NewMCPServerRequest] | type[UpdateMCPServerRequest], value: float +) -> None: + parsed: Final = request_model.model_validate( + { + "server_id": "test", + "transport": "http", + "url": "https://example.com/mcp", + "mcp_info": {"output_schema_cache_ttl_seconds": value}, + } + ) + + assert parsed.mcp_info == {"output_schema_cache_ttl_seconds": value} + + MCP_SERVER_REQUESTS = (NewMCPServerRequest, UpdateMCPServerRequest) STDIO_SERVER_FIELDS = {"server_id": "stdio-1", "transport": "stdio", "command": "python", "args": ["server.py"]}