From ad1f3c5e91b8c7fcb2bcb067e05bafb3dd4bd1be Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Mon, 19 Oct 2026 16:01:23 -0700 Subject: [PATCH] fix(mcp): retry upstream 429 on tools/list and error when every server fails (#46649) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 7 + .../proxy/_experimental/mcp_server/catalog.py | 28 +- .../mcp_server/faults/list_outcomes.py | 121 +++++++- .../mcp_server/mcp_server_manager.py | 13 +- .../_experimental/mcp_server/operations.py | 22 +- .../_experimental/mcp_server/subscriptions.py | 12 +- .../mcp_server/faults/test_list_outcomes.py | 285 ++++++++++++++++++ .../mcp_server/test_mcp_server_manager.py | 150 ++++++++- .../test_mcp_server_tool_calls_and_headers.py | 201 +++++++++++- .../mcp_server/test_operations.py | 39 +++ .../mcp_server/test_subscriptions.py | 50 ++- 11 files changed, 903 insertions(+), 25 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index c7c1576bc20..2d54a873b89 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -252,6 +252,13 @@ MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS: Final = get_env_int("MCP_PER_USER_TOKE # MCP timeout defaults (seconds). Override via env vars for slow/custom MCP servers. MCP_CLIENT_TIMEOUT: Final = get_env_float("LITELLM_MCP_CLIENT_TIMEOUT", 60.0) MCP_TOOL_LISTING_TIMEOUT: Final = get_env_float("LITELLM_MCP_TOOL_LISTING_TIMEOUT", 30.0) +MCP_UPSTREAM_RATE_LIMIT_MAX_RETRIES: Final = get_env_int("LITELLM_MCP_UPSTREAM_RATE_LIMIT_MAX_RETRIES", 3) +MCP_UPSTREAM_RATE_LIMIT_BASE_DELAY_SECONDS: Final = get_env_float( + "LITELLM_MCP_UPSTREAM_RATE_LIMIT_BASE_DELAY_SECONDS", 0.5 +) +MCP_UPSTREAM_RATE_LIMIT_MAX_DELAY_SECONDS: Final = get_env_float( + "LITELLM_MCP_UPSTREAM_RATE_LIMIT_MAX_DELAY_SECONDS", 5.0 +) MCP_METADATA_TIMEOUT: Final = get_env_float("LITELLM_MCP_METADATA_TIMEOUT", 10.0) MCP_HEALTH_CHECK_TIMEOUT: Final = get_env_float("LITELLM_MCP_HEALTH_CHECK_TIMEOUT", 10.0) AGENT_HEALTH_CHECK_TIMEOUT_SECONDS: Final = get_env_float("LITELLM_AGENT_HEALTH_CHECK_TIMEOUT", 5.0) diff --git a/litellm/proxy/_experimental/mcp_server/catalog.py b/litellm/proxy/_experimental/mcp_server/catalog.py index d0383be61a8..71eb2112bd5 100644 --- a/litellm/proxy/_experimental/mcp_server/catalog.py +++ b/litellm/proxy/_experimental/mcp_server/catalog.py @@ -17,6 +17,7 @@ from itertools import chain from types import MappingProxyType from typing import TYPE_CHECKING, Final, Generic, ParamSpec, TypeAlias, TypeVar, cast +from mcp.shared.exceptions import MCPError from mcp.types import CacheableResult from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter @@ -1055,12 +1056,20 @@ async def aggregate_gateway_tools( ) +class ToolListingFailedOutright(MCPError): + pass + + async def list_gateway_tools( context: OperationContext, params: PaginatedRequestParams, *, log_list_tools_to_spendlogs: bool = True ) -> ListToolsResult: - from mcp.types import ListToolsResult + from mcp.types import INTERNAL_ERROR, ListToolsResult - from litellm.proxy._experimental.mcp_server.faults.list_outcomes import SERVER_OUTCOMES_META_KEY, outcome_wire_value + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( + SERVER_OUTCOMES_META_KEY, + listing_failed_outright, + outcome_wire_value, + ) from litellm.proxy._experimental.mcp_server.operations import _list_mcp_tools caller, auth, servers, server_headers, oauth_headers, headers, client_ip = context.legacy_auth() @@ -1078,16 +1087,21 @@ async def list_gateway_tools( record_listing=True, list_tools_log_source="mcp_protocol", ) + wire_outcomes: Final = ( + {key: outcome_wire_value(outcome) for key, outcome in listing.outcomes.items()} if listing.outcomes else None + ) + if params.cursor is None and listing_failed_outright(listing.outcomes): + raise ToolListingFailedOutright( + code=INTERNAL_ERROR, + message="Every MCP server in scope failed to list tools", + data={SERVER_OUTCOMES_META_KEY: wire_outcomes}, + ) return ListToolsResult( tools=listing.tools, next_cursor=listing.next_cursor, ttl_ms=listing.ttl_ms, cache_scope="private", - _meta={ - SERVER_OUTCOMES_META_KEY: {key: outcome_wire_value(outcome) for key, outcome in listing.outcomes.items()} - } - if listing.outcomes - else None, + _meta={SERVER_OUTCOMES_META_KEY: wire_outcomes} if wire_outcomes else None, ) diff --git a/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py b/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py index b94ed1ee251..84673c8c738 100644 --- a/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py +++ b/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py @@ -10,17 +10,27 @@ becomes an outcome, never a second failure. from __future__ import annotations +import email.utils +import random import re -from collections.abc import Iterator -from typing import Final, Literal, NamedTuple, NoReturn, TypeAlias +from collections.abc import Awaitable, Callable, Iterator, Mapping +from datetime import datetime, timezone +from typing import Final, Literal, NamedTuple, NoReturn, TypeAlias, TypeVar from urllib.request import parse_http_list, parse_keqv_list +import anyio import httpx import httpx2 from mcp.types import Tool as MCPTool from pydantic import ConfigDict, TypeAdapter from typing_extensions import assert_never +from litellm._logging import verbose_logger +from litellm.constants import ( + MCP_UPSTREAM_RATE_LIMIT_BASE_DELAY_SECONDS, + MCP_UPSTREAM_RATE_LIMIT_MAX_DELAY_SECONDS, + MCP_UPSTREAM_RATE_LIMIT_MAX_RETRIES, +) from litellm.proxy._experimental.mcp_server.exceptions import ( MCPServerListError, MCPUpstreamAuthError, @@ -176,6 +186,113 @@ def classify_list_exception(exc: BaseException) -> ServerListFault: return ServerListFault(tag="internal") +class UpstreamRateLimit(NamedTuple): + retry_after_seconds: float | None + + +def _parse_retry_after(value: str | None) -> float | None: + if value is None: + return None + stripped: Final = value.strip() + if re.fullmatch(r"\d+", stripped): + return float(stripped) + try: + parsed: Final = email.utils.parsedate_to_datetime(stripped) + except (TypeError, ValueError): + return None + aware: Final = parsed if parsed.tzinfo is not None else parsed.replace(tzinfo=timezone.utc) + return max(0.0, (aware - datetime.now(timezone.utc)).total_seconds()) + + +def upstream_rate_limit(exc: BaseException) -> UpstreamRateLimit | None: + """The upstream 429 signal for the tools/list retry loop: an auth challenge wins over rate + limiting so a 401/403 anywhere in the tree keeps its own channel, then the first 429 response + in deliberate order carries its Retry-After hint. Anything else is not retryable here.""" + if upstream_auth_challenge(exc) is not None: + return None + for response in _iter_upstream_responses(exc): + if response.status_code == 429: + return UpstreamRateLimit(_parse_retry_after(response.headers.get("retry-after"))) + return None + + +def rate_limit_retry_delay( + attempt: int, + retry_after_seconds: float | None, + *, + base_delay: float, + max_delay: float, + jitter: float, +) -> float: + if retry_after_seconds is not None: + return min(retry_after_seconds, max_delay) + return min(max_delay, base_delay * 2**attempt * (1 + jitter)) + + +T: Final = TypeVar("T") + + +async def retry_upstream_rate_limit( + fetch: Callable[[], Awaitable[T]], + server_name: str, + *, + max_retries: int = MCP_UPSTREAM_RATE_LIMIT_MAX_RETRIES, + base_delay: float = MCP_UPSTREAM_RATE_LIMIT_BASE_DELAY_SECONDS, + max_delay: float = MCP_UPSTREAM_RATE_LIMIT_MAX_DELAY_SECONDS, + sleep: Callable[[float], Awaitable[None]] = anyio.sleep, + jitter: Callable[[], float] = random.random, + deadline_seconds: float | None = None, + clock: Callable[[], float] = anyio.current_time, +) -> T: + attempts: Final = max(0, max_retries) + deadline: Final = None if deadline_seconds is None else clock() + deadline_seconds + last_rate_limit: BaseException | None = None # rebind-ok: each 429 replaces the cause a retry deadline re-raises + for attempt in range(attempts + 1): + try: + if attempt > 0 and deadline is not None: + try: + with anyio.fail_after(max(0.0, deadline - clock())): + return await fetch() + except TimeoutError as timeout: + if last_rate_limit is None: + raise + raise last_rate_limit from timeout + return await fetch() + except Exception as exc: + limit = upstream_rate_limit(exc) + if limit is not None: + last_rate_limit = exc + if limit is None or attempt == attempts: + raise + delay = rate_limit_retry_delay( + attempt, + limit.retry_after_seconds, + base_delay=base_delay, + max_delay=max_delay, + jitter=jitter(), + ) + if deadline is not None and clock() + delay >= deadline: + raise + verbose_logger.warning( + "MCP server %s answered 429 to tools/list; retry %d/%d in %.2fs", + server_name, + attempt + 1, + max_retries, + delay, + ) + await sleep(delay) + raise AssertionError("unreachable") + + +def listing_failed_outright(outcomes: Mapping[str, ServerOutcome]) -> bool: + """True only when every server in scope contributed a non-auth fault: an auth outcome means a + caller can fix the listing by authenticating, so the aggregate stays a 200 with outcomes.""" + return bool(outcomes) and all( + isinstance(outcome, ServerListFault) and outcome.tag not in ("auth_required", "forbidden") + for outcome in outcomes.values() + ) + + def outcome_wire_value(outcome: ServerOutcome) -> dict[str, object]: """The client-visible form of one outcome, for the tools/list result ``_meta`` and the REST response: category plus status code only, never upstream prose or URLs.""" diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 96cf6ee11aa..aba5161d56d 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -104,6 +104,7 @@ from litellm.proxy._experimental.mcp_server.exceptions import ( from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( ServerListFault, raise_classified_list_failure, + retry_upstream_rate_limit, upstream_auth_challenge, upstream_insufficient_scope, ) @@ -4531,7 +4532,11 @@ class MCPServerManager: ) else: page: Final = ( - await self._fetch_tools_page_with_timeout(client, server, params) + 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)) ) @@ -5600,7 +5605,11 @@ class MCPServerManager: """ try: with anyio.fail_after(MCP_TOOL_LISTING_TIMEOUT): - tools: Final = await client.list_tools(raise_on_error=True) + tools: Final = await retry_upstream_rate_limit( + lambda: client.list_tools(raise_on_error=True), + server_name, + deadline_seconds=MCP_TOOL_LISTING_TIMEOUT, + ) verbose_logger.debug("Tools from %s: %s", server_name, tools) return tools except TimeoutError as e: diff --git a/litellm/proxy/_experimental/mcp_server/operations.py b/litellm/proxy/_experimental/mcp_server/operations.py index 90f01368127..c45faf38f51 100644 --- a/litellm/proxy/_experimental/mcp_server/operations.py +++ b/litellm/proxy/_experimental/mcp_server/operations.py @@ -3041,6 +3041,19 @@ async def _execute_handle_list_tools( return ListToolsResult(tools=[]) +async def _discover_tools_listing( + context: OperationContext, host_progress_callback: ProgressCallback | None +) -> ListToolsResult: + from litellm.proxy._experimental.mcp_server.catalog import ToolListingFailedOutright + + try: + return await _execute_handle_list_tools( + context, PaginatedRequestParams(), host_progress_callback, log_list_tools_to_spendlogs=False + ) + except ToolListingFailedOutright: + return ListToolsResult(tools=[]) + + async def _execute_mcp_server_tool_call( context: OperationContext, params: CallToolRequestParams, host_progress_callback: ProgressCallback | None = None ) -> CallToolResult | InputRequiredResult: @@ -3517,14 +3530,7 @@ class GatewayOperations: ) memo_token: Final = _mcp_server_admission_memo.set({}) tasks: Final = ( - asyncio.create_task( - _execute_handle_list_tools( - context, - PaginatedRequestParams(), - self._host_progress_callback, - log_list_tools_to_spendlogs=False, - ) - ), + asyncio.create_task(_discover_tools_listing(context, self._host_progress_callback)), *(asyncio.create_task(self.execute(listing, context)) for listing in listings), ) try: diff --git a/litellm/proxy/_experimental/mcp_server/subscriptions.py b/litellm/proxy/_experimental/mcp_server/subscriptions.py index db622ea3cbf..e9de5692d83 100644 --- a/litellm/proxy/_experimental/mcp_server/subscriptions.py +++ b/litellm/proxy/_experimental/mcp_server/subscriptions.py @@ -50,6 +50,9 @@ from mcp.types import ( from litellm.constants import MCP_TOOL_LISTING_MAX_PAGES, MCP_TOOL_LISTING_TIMEOUT from litellm.experimental_mcp_client.client import MCPClient from litellm.proxy._experimental.mcp_server import operations +from litellm.proxy._experimental.mcp_server.catalog import ( + ToolListingFailedOutright, +) from litellm.proxy._experimental.mcp_server.catalog import ( _configuration_identity as configuration_identity, # pyright: ignore[reportPrivateUsage] # reuse catalog-owned configuration identity without treating discovered metadata as configuration ) @@ -283,9 +286,12 @@ async def _listing(event: ServerEvent, context: OperationContext) -> str | None: cursor: str | None = None # rebind-ok: advance the signed gateway cursor for each catalog page with anyio.fail_after(MCP_TOOL_LISTING_TIMEOUT): for _ in range(MCP_TOOL_LISTING_MAX_PAGES): - result = await gateway.execute( - operation.model_copy(update={"params": PaginatedRequestParams(cursor=cursor)}), context - ) + try: + result = await gateway.execute( + operation.model_copy(update={"params": PaginatedRequestParams(cursor=cursor)}), context + ) + except ToolListingFailedOutright: + return digest.hexdigest() items = ( result.tools if isinstance(result, ListToolsResult) diff --git a/tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py b/tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py index ee33c41afb6..f6902c778e5 100644 --- a/tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py +++ b/tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py @@ -326,3 +326,288 @@ def test_managed_scope_step_up_only_accepts_unambiguous_bearer_challenges(status wrapped = RuntimeError("transport failure") wrapped.__cause__ = cause assert upstream_insufficient_scope(wrapped) == expected + + +def _rate_limited_error(retry_after: str | None = None) -> httpx.HTTPStatusError: + request: Final = httpx.Request("POST", "https://mcp.example.com/mcp") + headers: Final = {"retry-after": retry_after} if retry_after is not None else {} + return httpx.HTTPStatusError( + "rate limited", + request=request, + response=httpx.Response(429, headers=headers, request=request), + ) + + +@pytest.mark.asyncio +async def test_retry_upstream_rate_limit_backs_off_then_reraises() -> None: + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import retry_upstream_rate_limit + + sleeps: Final[list[float]] = [] + calls: Final[list[httpx.HTTPStatusError]] = [] + failure: Final = _rate_limited_error() + + async def record(delay: float) -> None: + sleeps.append(delay) + + async def always_limited() -> None: + calls.append(failure) + raise failure + + expected_error: Final = pytest.raises(httpx.HTTPStatusError) + with expected_error as exc_info: + await retry_upstream_rate_limit( + always_limited, + "srv", + max_retries=3, + base_delay=0.5, + max_delay=5.0, + sleep=record, + jitter=lambda: 0.0, + ) + + assert exc_info.value is failure + assert sleeps == [0.5, 1.0, 2.0] + assert len(calls) == 4 + + +@pytest.mark.asyncio +async def test_retry_upstream_rate_limit_honors_retry_after_up_to_cap() -> None: + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import retry_upstream_rate_limit + + sleeps: Final[list[float]] = [] + + async def record(delay: float) -> None: + sleeps.append(delay) + + async def always_limited() -> None: + raise _rate_limited_error("120") + + with pytest.raises(httpx.HTTPStatusError): + await retry_upstream_rate_limit( + always_limited, + "srv", + max_retries=3, + base_delay=0.5, + max_delay=5.0, + sleep=record, + jitter=lambda: 0.0, + ) + + assert sleeps == [5.0, 5.0, 5.0] + + +@pytest.mark.asyncio +async def test_retry_upstream_rate_limit_parses_http_date_retry_after() -> None: + import email.utils + import time + + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import retry_upstream_rate_limit + + sleeps: Final[list[float]] = [] + attempts: Final[list[None]] = [] + + async def record(delay: float) -> None: + sleeps.append(delay) + + async def limited() -> str: + attempts.append(None) + if len(attempts) == 1: + raise _rate_limited_error(email.utils.formatdate(time.time() + 3, usegmt=True)) + raise _rate_limited_error("abc") + + with pytest.raises(httpx.HTTPStatusError): + await retry_upstream_rate_limit( + limited, + "srv", + max_retries=3, + base_delay=0.5, + max_delay=5.0, + sleep=record, + jitter=lambda: 0.0, + ) + + assert 0 < sleeps[0] <= 3.5 + assert sleeps[1:] == [1.0, 2.0] + + +@pytest.mark.asyncio +async def test_retry_upstream_rate_limit_passes_through_non_429() -> None: + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import retry_upstream_rate_limit + + sleeps: Final[list[float]] = [] + calls = 0 + failure: Final = ValueError("broken") + + async def record(delay: float) -> None: + sleeps.append(delay) + + async def broken() -> None: + nonlocal calls + calls += 1 + raise failure + + expected_error: Final = pytest.raises(ValueError, match="broken") + with expected_error as exc_info: + await retry_upstream_rate_limit( + broken, "srv", max_retries=3, base_delay=0.5, max_delay=5.0, sleep=record, jitter=lambda: 0.0 + ) + + assert exc_info.value is failure + assert calls == 1 + assert sleeps == [] + + +@pytest.mark.asyncio +async def test_retry_upstream_rate_limit_negative_max_retries_is_a_single_attempt() -> None: + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import retry_upstream_rate_limit + + sleeps: Final[list[float]] = [] + calls = 0 + failure: Final = _rate_limited_error("0") + + async def record(delay: float) -> None: + sleeps.append(delay) + + async def always_limited() -> None: + nonlocal calls + calls += 1 + raise failure + + expected_error: Final = pytest.raises(httpx.HTTPStatusError) + with expected_error as exc_info: + await retry_upstream_rate_limit( + always_limited, + "srv", + max_retries=-1, + base_delay=0.5, + max_delay=5.0, + sleep=record, + jitter=lambda: 0.0, + ) + + assert exc_info.value is failure + assert calls == 1 + assert sleeps == [] + + +@pytest.mark.asyncio +async def test_retry_deadline_skips_sleep_that_would_pass_deadline() -> None: + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import retry_upstream_rate_limit + + now: Final[list[float]] = [100.0] + sleeps: Final[list[float]] = [] + calls = 0 + failure: Final = _rate_limited_error() + + async def record(delay: float) -> None: + sleeps.append(delay) + + async def always_limited() -> None: + nonlocal calls + calls += 1 + raise failure + + expected_error: Final = pytest.raises(httpx.HTTPStatusError) + with expected_error as exc_info: + await retry_upstream_rate_limit( + always_limited, + "srv", + max_retries=3, + base_delay=2.0, + max_delay=5.0, + sleep=record, + jitter=lambda: 0.0, + deadline_seconds=1.0, + clock=lambda: now[0], + ) + + assert exc_info.value is failure + assert calls == 1 + assert sleeps == [] + + +@pytest.mark.asyncio +async def test_retry_attempt_is_interrupted_at_deadline() -> None: + import anyio + + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import retry_upstream_rate_limit + + calls = 0 + failure: Final = _rate_limited_error("0") + + async def fetch() -> None: + nonlocal calls + calls += 1 + if calls == 1: + raise failure + await anyio.sleep(5) + + expected_error: Final = pytest.raises(httpx.HTTPStatusError) + with anyio.fail_after(5): + with expected_error as exc_info: + await retry_upstream_rate_limit( + fetch, "srv", max_retries=3, base_delay=0.0, max_delay=5.0, deadline_seconds=0.2 + ) + + assert exc_info.value is failure + + +@pytest.mark.asyncio +async def test_retry_deadline_does_not_bound_first_attempt() -> None: + import anyio + + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import retry_upstream_rate_limit + + async def slow() -> str: + await anyio.sleep(0.15) + return "tools" + + result: Final = await retry_upstream_rate_limit( + slow, "srv", max_retries=3, base_delay=0.5, max_delay=5.0, deadline_seconds=0.05 + ) + + assert result == "tools" + + +@pytest.mark.asyncio +async def test_retry_upstream_rate_limit_never_retries_cancellation() -> None: + import asyncio + + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import retry_upstream_rate_limit + + sleeps: Final[list[float]] = [] + calls = 0 + + async def record(delay: float) -> None: + sleeps.append(delay) + + async def cancelled() -> None: + nonlocal calls + calls += 1 + raise asyncio.CancelledError() + + with pytest.raises(asyncio.CancelledError): + await retry_upstream_rate_limit( + cancelled, "srv", max_retries=3, base_delay=0.5, max_delay=5.0, sleep=record, jitter=lambda: 0.0 + ) + + assert calls == 1 + assert sleeps == [] + + +def test_listing_failed_outright_only_when_every_server_faulted_non_auth() -> None: + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( + ServerListOk, + listing_failed_outright, + ) + + assert not listing_failed_outright({}) + assert listing_failed_outright({"a": ServerListFault(tag="upstream_error", status_code=429)}) + assert listing_failed_outright( + {"a": ServerListFault(tag="timeout"), "b": ServerListFault(tag="rate_limited", status_code=429)} + ) + assert not listing_failed_outright( + {"a": ServerListFault(tag="upstream_error", status_code=500), "b": ServerListOk(tool_count=2)} + ) + assert not listing_failed_outright({"a": ServerListFault(tag="auth_required", status_code=401)}) + assert not listing_failed_outright({"a": ServerListFault(tag="forbidden", status_code=403)}) diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py index d47b2f73a48..d314291ef74 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -7,6 +7,7 @@ import logging import os import sys import time +import anyio from collections.abc import AsyncIterator, Awaitable, Callable, Iterator from datetime import datetime from pathlib import Path @@ -37,6 +38,8 @@ from mcp.types import ( CallToolRequestParams as MCPCallToolRequestParams, CallToolResult, GetPromptResult, + ListToolsResult, + PaginatedRequestParams, Prompt, ResourceTemplate, TextResourceContents, @@ -44,7 +47,7 @@ from mcp.types import ( from mcp.types import Tool as MCPTool from pydantic import AnyUrl, TypeAdapter -from litellm.constants import MCP_METADATA_TIMEOUT +from litellm.constants import MCP_METADATA_TIMEOUT, MCP_UPSTREAM_RATE_LIMIT_MAX_RETRIES from litellm.proxy._experimental.mcp_server import discoverable_endpoints from litellm.proxy._experimental.mcp_server.legacy_callbacks import ElicitationCallback from litellm.proxy._experimental.mcp_server.mcp_context import active_mcp_request_ctx_var @@ -20427,3 +20430,148 @@ async def test_legacy_upstream_elicitation_rejects_modern_downstream_without_con finally: active_mcp_request_ctx_var.reset(request_token) legacy_server.active_mcp_session_var.reset(session_token) + + +def _upstream_http_status_error(status_code: int, headers: dict[str, str] | None = None) -> httpx.HTTPStatusError: + request: Final = httpx.Request("POST", "https://upstream.example/mcp") + return httpx.HTTPStatusError( + "upstream rejected", + request=request, + response=httpx.Response(status_code, headers=headers or {}, request=request), + ) + + +class TestMCPUpstreamRateLimitRetry: + """An upstream 429 on tools/list is retried inside the listing deadline instead of folding + into a silent empty outcome on the first attempt.""" + + @pytest.mark.asyncio + async def test_fetch_tools_with_timeout_retries_429_then_returns_tools(self) -> None: + manager: Final = MCPServerManager() + tool: Final = MCPTool(name="ping", inputSchema={"type": "object"}) + client: Final = MagicMock() + client.list_tools = AsyncMock( + side_effect=[ + _upstream_http_status_error(429, {"Retry-After": "0"}), + _upstream_http_status_error(429, {"Retry-After": "0"}), + [tool], + ] + ) + + tools: Final = await manager._fetch_tools_with_timeout(client, "flaky-server") + + assert tools == [tool] + assert client.list_tools.await_count == 3 + + @pytest.mark.asyncio + async def test_fetch_tools_with_timeout_gives_up_after_max_retries(self) -> None: + manager: Final = MCPServerManager() + client: Final = MagicMock() + client.list_tools = AsyncMock(side_effect=_upstream_http_status_error(429, {"Retry-After": "0"})) + + expected_error: Final = pytest.raises(MCPServerListError) + with expected_error as exc_info: + await manager._fetch_tools_with_timeout(client, "saturated-server") + + assert exc_info.value.fault == ServerListFault(tag="upstream_error", status_code=429) + assert client.list_tools.await_count == 1 + MCP_UPSTREAM_RATE_LIMIT_MAX_RETRIES + + @pytest.mark.asyncio + async def test_fetch_tools_with_timeout_does_not_retry_auth_error(self) -> None: + manager: Final = MCPServerManager() + client: Final = MagicMock() + client.list_tools = AsyncMock(side_effect=_upstream_http_status_error(401)) + + with pytest.raises(MCPUpstreamAuthError): + await manager._fetch_tools_with_timeout(client, "srv") + + assert client.list_tools.await_count == 1 + + @pytest.mark.asyncio + async def test_fetch_tools_with_timeout_does_not_retry_server_error(self) -> None: + manager: Final = MCPServerManager() + client: Final = MagicMock() + client.list_tools = AsyncMock(side_effect=_upstream_http_status_error(500)) + + expected_error: Final = pytest.raises(MCPServerListError) + with expected_error as exc_info: + await manager._fetch_tools_with_timeout(client, "srv") + + assert exc_info.value.fault == ServerListFault(tag="upstream_error", status_code=500) + assert client.list_tools.await_count == 1 + + @pytest.mark.asyncio + async def test_fetch_tools_with_timeout_reports_429_when_backoff_would_pass_listing_timeout( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.MCP_TOOL_LISTING_TIMEOUT", 0.5 + ) + manager: Final = MCPServerManager() + client: Final = MagicMock() + client.list_tools = AsyncMock(side_effect=_upstream_http_status_error(429, {"Retry-After": "5"})) + + with anyio.fail_after(3): + expected_error: Final = pytest.raises(MCPServerListError) + with expected_error as exc_info: + await manager._fetch_tools_with_timeout(client, "srv") + + assert exc_info.value.fault == ServerListFault(tag="upstream_error", status_code=429) + assert client.list_tools.await_count == 1 + + @pytest.mark.asyncio + async def test_get_tools_page_retries_429_then_returns_page(self) -> None: + tool: Final = MCPTool(name="list_notes", inputSchema={"type": "object"}) + manager: Final = MCPServerManager() + server: Final = MCPServer(server_id="notes", name="notes", transport=MCPTransport.http) + manager.registry = {server.server_id: server} + manager.create_mcp_client = AsyncMock( + return_value=SimpleNamespace( + _last_initialize_instructions=None, + list_tools_page=AsyncMock( + side_effect=[ + _upstream_http_status_error(429, {"Retry-After": "0"}), + ListToolsResult(tools=[tool]), + ] + ), + ) + ) + + result: Final = await manager.get_tools_page(server, params=PaginatedRequestParams()) + + assert [t.name for t in result.tools] == ["notes-list_notes"] + client: Final = manager.create_mcp_client.return_value + assert client.list_tools_page.await_count == 2 + + @pytest.mark.asyncio + async def test_get_tools_page_retry_attempt_is_bounded_by_listing_timeout( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.MCP_TOOL_LISTING_TIMEOUT", 0.05 + ) + manager: Final = MCPServerManager() + server: Final = MCPServer(server_id="notes", name="notes", transport=MCPTransport.http) + manager.registry = {server.server_id: server} + first_seen: Final = anyio.Event() + never: Final = anyio.Event() + error: Final = _upstream_http_status_error(429, {"Retry-After": "0"}) + + async def _list_tools_page(params: PaginatedRequestParams) -> ListToolsResult: + if not first_seen.is_set(): + first_seen.set() + raise error + await never.wait() + raise AssertionError("unreachable") + + list_tools_page: Final = AsyncMock(side_effect=_list_tools_page) + client: Final = SimpleNamespace(_last_initialize_instructions=None, list_tools_page=list_tools_page) + manager.create_mcp_client = AsyncMock(return_value=client) + + with anyio.fail_after(5): + expected_error: Final = pytest.raises(MCPServerListError) + with expected_error as exc_info: + await manager.get_tools_page(server, params=PaginatedRequestParams()) + + assert exc_info.value.fault == ServerListFault(tag="upstream_error", status_code=429) + assert list_tools_page.await_count == 2 diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py index 48106b1de5f..66d66271f71 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py @@ -6,19 +6,22 @@ import json import os from datetime import datetime, timedelta from types import SimpleNamespace -from typing import Final, Literal +from typing import Callable, Final, Literal from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest from fastapi import HTTPException from mcp import ReadResourceResult, Resource +from mcp.server.context import ServerRequestContext from mcp.server.models import InitializationOptions from mcp.types import ( INVALID_REQUEST, METHOD_NOT_FOUND, BlobResourceContents, CallToolResult, + ListToolsResult, + PaginatedRequestParams, Prompt, ResourceTemplate, TextContent, @@ -32,6 +35,7 @@ from starlette.types import Message, Receive, Scope, Send from litellm.proxy._experimental.mcp_server.mcp_context import active_mcp_request_ctx_var from litellm.proxy._experimental.mcp_server.mcp_server_manager import ListedToolsCaller +from litellm.proxy._experimental.mcp_server.faults.list_outcomes import AggregateToolListing from litellm.proxy._types import ( LiteLLM_MCPServerTable, MCPTransport, @@ -5876,6 +5880,76 @@ async def test_get_tools_from_mcp_servers_logs_list_tools_to_spendlogs_when_enab assert spend_meta["per_server_list_outcomes"] == {"server_a": {"status": "ok", "tool_count": 1}} +@pytest.mark.asyncio +async def test_get_tools_from_mcp_servers_all_failed_listing_never_runs_proxy_failure_hooks() -> None: + from litellm.proxy._types import UserAPIKeyAuth + + servers: Final = ( + MCPServer( + server_id="server_a", + name="server_a", + alias="server_a", + server_name="server_a", + transport=MCPTransport.http, + ), + MCPServer( + server_id="server_b", + name="server_b", + alias="server_b", + server_name="server_b", + transport=MCPTransport.http, + ), + ) + user_auth: Final = UserAPIKeyAuth(api_key="test-key", user_id="test-user") + logging_obj: Final = MagicMock() + logging_obj.model_call_details = {"metadata": {"spend_logs_metadata": {}}} + logging_obj.async_success_handler = AsyncMock() + proxy_logging_obj: Final = MagicMock() + proxy_logging_obj.post_call_failure_hook = AsyncMock() + manager: Final = MagicMock() + manager.get_tools_from_server = AsyncMock(side_effect=RuntimeError("server unavailable")) + + with ( + patch( + "litellm.proxy._experimental.mcp_server.operations._get_allowed_mcp_servers", + new=AsyncMock(return_value=servers), + ), + patch( + "litellm.proxy._experimental.mcp_server.operations._prepare_mcp_server_headers", + return_value=(None, None), + ), + patch("litellm.proxy._experimental.mcp_server.operations.global_mcp_server_manager", manager), + patch("litellm.proxy._experimental.mcp_server.operations.filter_tools_by_allowed_tools", return_value=[]), + patch( + "litellm.proxy._experimental.mcp_server.operations.filter_tools_by_key_team_permissions", + new=AsyncMock(return_value=[]), + ), + patch( + "litellm.proxy._experimental.mcp_server.operations.function_setup", + return_value=(logging_obj, None), + ), + patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj), + ): + listing: Final = await mcp_operations._get_tools_from_mcp_servers( + user_api_key_auth=user_auth, + mcp_auth_header=None, + mcp_servers=["server_a", "server_b"], + mcp_server_auth_headers=None, + params=None, + log_list_tools_to_spendlogs=True, + enforce_rate_limits=False, + ) + + assert listing.tools == [] + logging_obj.async_success_handler.assert_awaited_once() + spend_meta: Final = logging_obj.model_call_details["metadata"]["spend_logs_metadata"] + assert spend_meta["per_server_list_outcomes"] == { + "server_a": {"status": "internal"}, + "server_b": {"status": "internal"}, + } + proxy_logging_obj.post_call_failure_hook.assert_not_awaited() + + @pytest.mark.asyncio async def test_get_tools_from_mcp_servers_returns_tools_when_success_logging_fails(): """ @@ -10071,6 +10145,131 @@ async def test_handle_list_tools_attaches_outcome_meta(_mcp_request_ctx): assert outcomes_meta["broken"] == {"status": "unreachable"} +async def _list_tools_with_listing( + _mcp_request_ctx: Callable[..., ServerRequestContext], + listing: AggregateToolListing, + params: PaginatedRequestParams, +) -> ListToolsResult: + try: + from litellm.proxy._experimental.mcp_server.server import handle_list_tools + except ImportError: + pytest.skip("MCP server not available") + + with ( + patch( + "litellm.proxy._experimental.mcp_server.server.get_or_extract_auth_context", + new=AsyncMock(return_value=(None, None, None, None, None, None, None)), + ), + patch( + "litellm.proxy._experimental.mcp_server.operations._list_mcp_tools", + new=AsyncMock(return_value=listing), + ), + ): + return await handle_list_tools(_mcp_request_ctx(), params) + + +@pytest.mark.asyncio +async def test_handle_list_tools_errors_when_every_server_failed( + _mcp_request_ctx: Callable[..., ServerRequestContext], +) -> None: + from mcp.shared.exceptions import MCPError + from mcp.types import INTERNAL_ERROR + + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( + SERVER_OUTCOMES_META_KEY, + AggregateToolListing, + ServerListFault, + ) + + listing: Final = AggregateToolListing( + tools=[], + outcomes={ + "a": ServerListFault(tag="upstream_error", status_code=429), + "b": ServerListFault(tag="timeout"), + }, + ) + + expected_error: Final = pytest.raises(MCPError) + with expected_error as exc_info: + await _list_tools_with_listing(_mcp_request_ctx, listing, _paged_params()) + + assert exc_info.value.error.code == INTERNAL_ERROR + assert exc_info.value.error.data[SERVER_OUTCOMES_META_KEY] == { + "a": {"status": "upstream_error", "http_status": 429}, + "b": {"status": "timeout"}, + } + + +@pytest.mark.asyncio +async def test_handle_list_tools_auth_fault_keeps_degraded_result( + _mcp_request_ctx: Callable[..., ServerRequestContext], +) -> None: + from mcp.types import ListToolsResult + + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( + SERVER_OUTCOMES_META_KEY, + AggregateToolListing, + ServerListFault, + ) + + listing: Final = AggregateToolListing( + tools=[], + outcomes={"a": ServerListFault(tag="auth_required", status_code=401), "b": ServerListFault(tag="timeout")}, + ) + + result: Final = await _list_tools_with_listing(_mcp_request_ctx, listing, _paged_params()) + + assert isinstance(result, ListToolsResult) + assert result.tools == [] + wire: Final = result.model_dump(by_alias=True) + assert wire["_meta"][SERVER_OUTCOMES_META_KEY]["a"] == {"status": "auth_required", "http_status": 401} + + +@pytest.mark.asyncio +async def test_handle_list_tools_all_faults_on_continuation_stays_degraded( + _mcp_request_ctx: Callable[..., ServerRequestContext], +) -> None: + from mcp.types import ListToolsResult, PaginatedRequestParams + + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( + SERVER_OUTCOMES_META_KEY, + AggregateToolListing, + ServerListFault, + ) + + listing: Final = AggregateToolListing( + tools=[], + outcomes={"a": ServerListFault(tag="upstream_error", status_code=429)}, + ) + + result: Final = await _list_tools_with_listing( + _mcp_request_ctx, listing, PaginatedRequestParams(cursor="cursor-1") + ) + + assert isinstance(result, ListToolsResult) + assert result.tools == [] + wire: Final = result.model_dump(by_alias=True) + assert wire["_meta"][SERVER_OUTCOMES_META_KEY]["a"] == {"status": "upstream_error", "http_status": 429} + + +@pytest.mark.asyncio +async def test_handle_list_tools_empty_outcomes_is_a_clean_empty_result( + _mcp_request_ctx: Callable[..., ServerRequestContext], +) -> None: + from mcp.types import ListToolsResult + + from litellm.proxy._experimental.mcp_server.faults.list_outcomes import AggregateToolListing + + listing: Final = AggregateToolListing(tools=[], outcomes={}) + + result: Final = await _list_tools_with_listing(_mcp_request_ctx, listing, _paged_params()) + + assert isinstance(result, ListToolsResult) + assert result.tools == [] + wire: Final = result.model_dump(by_alias=True) + assert wire.get("_meta") is None + + def _make_oauth2_server( alias: str, *, diff --git a/tests/unit/proxy/_experimental/mcp_server/test_operations.py b/tests/unit/proxy/_experimental/mcp_server/test_operations.py index 3957f486c4c..6f7437c90dc 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_operations.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_operations.py @@ -1263,6 +1263,45 @@ async def test_discovery_preserves_caller_scope_and_proxy_restrictions(): assert listed.await_args.args[0].mcp_servers == ("only-this",) +@pytest.mark.asyncio +async def test_discovery_outright_tools_failure_still_advertises_other_capabilities() -> None: + from mcp.types import DiscoverRequest, ListPromptsResult, ListToolsResult, Prompt + + from litellm.proxy._experimental.mcp_server import operations + from litellm.proxy._experimental.mcp_server.catalog import ToolListingFailedOutright + + tools: Final = AsyncMock(side_effect=ToolListingFailedOutright(code=-32603, message="every server failed")) + prompts: Final = AsyncMock(return_value=ListPromptsResult(prompts=[Prompt(name="p")])) + context: Final = prepare_context(UserAPIKeyAuth(user_id="scoped"), mcp_servers=["broken"]) + with ( + patch.object(operations, "_execute_handle_list_tools", tools), + patch.object(operations, "_execute_list_prompts", prompts), + patch.object(operations, "_execute_list_resources", AsyncMock(return_value=ListToolsResult(tools=[]))), + patch.object(operations, "_execute_list_resource_templates", AsyncMock(return_value=ListToolsResult(tools=[]))), + ): + result: Final = await GatewayOperations().execute(DiscoverRequest(), context) + assert result.capabilities.tools is None + assert result.capabilities.prompts is not None + + +@pytest.mark.asyncio +async def test_discovery_propagates_non_outright_mcp_errors_from_tools_listing() -> None: + from mcp import MCPError + from mcp.types import INVALID_PARAMS, DiscoverRequest + + from litellm.proxy._experimental.mcp_server import operations + + tools: Final = AsyncMock(side_effect=MCPError(code=INVALID_PARAMS, message="bad cursor")) + context: Final = prepare_context( + UserAPIKeyAuth(user_id="scoped"), mcp_servers=["broken"], mcp_proxy_mode=True, protocol_version="2025-06-18" + ) + with patch.object(operations, "_execute_handle_list_tools", tools): + expected_error: Final = pytest.raises(MCPError) + with expected_error as exc_info: + await GatewayOperations().execute(DiscoverRequest(), context) + assert exc_info.value.error.code == INVALID_PARAMS + + @pytest.mark.asyncio async def test_discovery_denial_cannot_advertise_tools(): from mcp.types import DiscoverRequest diff --git a/tests/unit/proxy/_experimental/mcp_server/test_subscriptions.py b/tests/unit/proxy/_experimental/mcp_server/test_subscriptions.py index a97ab04c506..fadb709e00f 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_subscriptions.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_subscriptions.py @@ -10,7 +10,7 @@ from mcp.server.context import ServerRequestContext from mcp.server.session import ServerSession from mcp.server.subscriptions import ListenHandler from mcp.shared.subscriptions import PromptsListChanged, ToolsListChanged -from mcp.types import SubscriptionFilter, SubscriptionsListenRequestParams +from mcp.types import ListToolsRequest, ListToolsResult, SubscriptionFilter, SubscriptionsListenRequestParams from litellm.experimental_mcp_client.client import MCPClient from litellm.proxy._experimental.mcp_server import subscriptions @@ -45,6 +45,18 @@ def context(session: ServerSession, revision: str = "2026-07-28") -> ServerReque ) +class _ListingGateway: + def __init__(self, *, result: ListToolsResult | None = None, error: MCPError | None = None) -> None: + self._result: ListToolsResult | None = result + self._error: MCPError | None = error + + async def execute(self, _operation: ListToolsRequest, _context: OperationContext) -> ListToolsResult: + if self._error is not None: + raise self._error + assert self._result is not None + return self._result + + @pytest.mark.asyncio async def test_queued_notification_rechecks_authority_before_delivery() -> None: bus: Final = subscriptions._ListenerBus() @@ -486,3 +498,39 @@ async def test_legacy_listener_catalog_refreshes_do_not_consume_server_rate_limi finally: listener.cancel() await asyncio.gather(listener, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_all_failed_tools_listing_keeps_empty_catalog_digest(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy._experimental.mcp_server.catalog import ToolListingFailedOutright + + listener_context: Final = OperationContext(_caller=UserAPIKeyAuth(user_id="listener")) + empty_gateway: Final = _ListingGateway(result=ListToolsResult(tools=[])) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.operations.GatewayOperations", lambda: empty_gateway + ) + empty_digest: Final = await subscriptions._listing(ToolsListChanged(), listener_context) + + failed_gateway: Final = _ListingGateway( + error=ToolListingFailedOutright(code=-32603, message="Every MCP server in scope failed") + ) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.operations.GatewayOperations", lambda: failed_gateway + ) + failed_digest: Final = await subscriptions._listing(ToolsListChanged(), listener_context) + + assert failed_digest == empty_digest + + +@pytest.mark.asyncio +async def test_tools_listing_propagates_other_mcp_errors(monkeypatch: pytest.MonkeyPatch) -> None: + failure: Final = MCPError(code=-32602, message="Bad cursor") + gateway: Final = _ListingGateway(error=failure) + monkeypatch.setattr("litellm.proxy._experimental.mcp_server.operations.GatewayOperations", lambda: gateway) + listener_context: Final = OperationContext(_caller=UserAPIKeyAuth(user_id="listener")) + + expected_error: Final = pytest.raises(MCPError) + with expected_error as exc_info: + await subscriptions._listing(ToolsListChanged(), listener_context) + + assert exc_info.value is failure