scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-25 05:54:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 126

fix(mcp): retry upstream 429 on tools/list and error when every server fails (#46649)

Browse source 
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Yassin Kortam 2026-10-19 16:01:23 -07:00 • committed by GitHub
parent 24d92f77d5
commit ad1f3c5e91
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 903 additions and 25 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
litellm/constants.py
View file

@ -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)

28
litellm/proxy/_experimental/mcp_server/catalog.py
View file

@ -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,
)

121
litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py
View file

@ -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."""

13
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -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:

22
litellm/proxy/_experimental/mcp_server/operations.py
View file

@ -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:

12
litellm/proxy/_experimental/mcp_server/subscriptions.py
View file

@ -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)

285
tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py
View file

@ -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)})

150
tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py
View file

@ -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

201
tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py
View file

@ -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,
*,

39
tests/unit/proxy/_experimental/mcp_server/test_operations.py
View file

@ -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

50
tests/unit/proxy/_experimental/mcp_server/test_subscriptions.py
View file

@ -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