scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

fix(mcp): cap tools preview and test-connection at the listing timeout and name the unreachable upstream

Browse source
This commit is contained in:
mateo-berri 2026-08-29 12:47:26 -07:00
parent c24f821652
commit 3275459aec
2 changed files with 96 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

34
litellm/proxy/_experimental/mcp_server/rest_endpoints.py
View file

@ -4,10 +4,12 @@ from collections.abc import Awaitable, Callable, Mapping
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final, Literal
import anyio
import httpx
from fastapi import APIRouter, Depends, HTTPException, Query, Request, status
from litellm._logging import verbose_logger
from litellm.constants import MCP_TOOL_LISTING_TIMEOUT
from litellm.exceptions import (
BlockedPiiEntityError,
GuardrailRaisedException,
@ -68,7 +70,13 @@ _MCP_GUARDRAIL_REJECTIONS: Final = (
)
def _connection_error_message(exc: BaseException) -> str:
def _connection_error_message(exc: BaseException, url: str | None, timeout_seconds: float) -> str:
if isinstance(exc, TimeoutError):
return (
f"Failed to connect to MCP server: no response from {url or 'the server'} "
f"within {timeout_seconds:.0f}s. Check that the LiteLLM proxy can reach this URL "
"from its network (DNS, egress rules, firewalls) and that the server answers MCP requests."
)
if isinstance(exc, httpx.LocalProtocolError):
return (
"Failed to connect to MCP server: a request header is malformed. "
@ -1136,6 +1144,7 @@ if MCP_AVAILABLE:
mcp_auth_header: str | dict[str, str] | None = None,
oauth2_headers: dict[str, str] | None = None,
raw_headers: dict[str, str] | None = None,
timeout_seconds: float = MCP_TOOL_LISTING_TIMEOUT,
) -> Mapping[str, object]:
"""
Create a temporary MCP client from *request*, run *operation*, and return the result.
@ -1151,6 +1160,10 @@ if MCP_AVAILABLE:
oauth2_headers: Headers extracted from the incoming request (may contain the
litellm API key — must NOT be forwarded for M2M servers).
raw_headers: Raw request headers forwarded for stdio env construction.
timeout_seconds: Cap on OAuth discovery, connect, handshake, and *operation*
combined. Defaults to ``MCP_TOOL_LISTING_TIMEOUT`` (30s, below common LB
timeouts) so an unreachable upstream yields this endpoint's JSON error
instead of an opaque load-balancer 504 with an empty body.
Returns:
The dict returned by *operation*, or an error dict on failure.
@ -1240,15 +1253,16 @@ if MCP_AVAILABLE:
static_headers=request.static_headers,
)
client: Final = await global_mcp_server_manager._create_mcp_client(
server=server_model,
mcp_auth_header=mcp_auth_header,
extra_headers=merged_headers,
stdio_env=stdio_env,
cred_provider=preview_cred_provider,
)
with anyio.fail_after(timeout_seconds):
client: Final = await global_mcp_server_manager._create_mcp_client(
server=server_model,
mcp_auth_header=mcp_auth_header,
extra_headers=merged_headers,
stdio_env=stdio_env,
cred_provider=preview_cred_provider,
)
return await operation(client)
return await operation(client)
except (KeyboardInterrupt, SystemExit, asyncio.CancelledError):
raise
@ -1257,7 +1271,7 @@ if MCP_AVAILABLE:
return {
"status": "error",
"error": True,
"message": _connection_error_message(e),
"message": _connection_error_message(e, request.url, timeout_seconds),
}
async def _preview_openapi_tools(spec_path: str) -> dict:

77
tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py
View file

@ -1,4 +1,5 @@
import asyncio
import inspect
import json
import sys
from datetime import datetime
@ -13,6 +14,7 @@ import pytest
from fastapi import HTTPException
from starlette.requests import Request
from litellm.constants import MCP_TOOL_LISTING_TIMEOUT
from litellm.proxy._experimental.mcp_server import rest_endpoints
from litellm.proxy._experimental.mcp_server.auth import (
user_api_key_auth_mcp as auth_mcp,
@ -109,6 +111,71 @@ class TestExecuteWithMcpClient:
assert result["status"] == "error"
assert "stack_trace" not in result
@pytest.mark.asyncio
async def test_timeout_caps_hanging_operation_and_names_url(self, monkeypatch):
async def fake_create_client(*args, **kwargs):
return object()
monkeypatch.setattr(
rest_endpoints.global_mcp_server_manager,
"_create_mcp_client",
fake_create_client,
)
async def hanging_operation(client):
await asyncio.Event().wait()
payload = NewMCPServerRequest(
server_name="example",
url="https://mcp.example.com/mcp/",
auth_type=MCPAuth.none,
)
result = await asyncio.wait_for(
rest_endpoints._execute_with_mcp_client(payload, hanging_operation, timeout_seconds=0.05),
timeout=5,
)
assert result["error"] is True
assert "https://mcp.example.com/mcp/" in result["message"]
@pytest.mark.asyncio
async def test_timeout_covers_client_creation(self, monkeypatch):
async def hanging_create_client(*args, **kwargs):
await asyncio.Event().wait()
monkeypatch.setattr(
rest_endpoints.global_mcp_server_manager,
"_create_mcp_client",
hanging_create_client,
)
async def unreached_operation(client):
return {"status": "ok"}
payload = NewMCPServerRequest(
server_name="example",
url="https://mcp.example.com/mcp/",
auth_type=MCPAuth.none,
)
result = await asyncio.wait_for(
rest_endpoints._execute_with_mcp_client(payload, unreached_operation, timeout_seconds=0.05),
timeout=5,
)
assert result["error"] is True
assert "https://mcp.example.com/mcp/" in result["message"]
def test_timeout_defaults_to_tool_listing_timeout(self):
default = inspect.signature(rest_endpoints._execute_with_mcp_client).parameters["timeout_seconds"].default
assert default == MCP_TOOL_LISTING_TIMEOUT
def test_connection_error_message_timeout_names_url_and_budget(self):
message = rest_endpoints._connection_error_message(TimeoutError(), "https://api.example.com/mcp/", 30.0)
assert "https://api.example.com/mcp/" in message
assert "30s" in message
@pytest.mark.asyncio
async def test_forwards_static_headers(self, monkeypatch):
"""Ensure static_headers are forwarded to the MCP client during test calls.
@ -2881,17 +2948,17 @@ class TestConnectionErrorMessage:
secret = "Bearer sk-super-secret-token"
exc = httpx.LocalProtocolError(f"Illegal header value b' {secret}'")
message = rest_endpoints._connection_error_message(exc)
message = rest_endpoints._connection_error_message(exc, "https://example.com", 30.0)
assert "header" in message.lower()
assert secret not in message
def test_connect_error_points_at_reachability(self):
message = rest_endpoints._connection_error_message(httpx.ConnectError("All connection attempts failed"))
message = rest_endpoints._connection_error_message(httpx.ConnectError("All connection attempts failed"), "https://example.com", 30.0)
assert "unreachable" in message.lower()
def test_timeout_error_message(self):
message = rest_endpoints._connection_error_message(httpx.ConnectTimeout("timed out"))
message = rest_endpoints._connection_error_message(httpx.ConnectTimeout("timed out"), "https://example.com", 30.0)
assert "unreachable" in message.lower()
def test_http_status_error_includes_status_code(self):
@ -2901,11 +2968,11 @@ class TestConnectionErrorMessage:
request=httpx.Request("POST", "http://x/"),
response=response,
)
message = rest_endpoints._connection_error_message(exc)
message = rest_endpoints._connection_error_message(exc, "https://example.com", 30.0)
assert "503" in message
def test_unknown_error_falls_back_to_generic(self):
message = rest_endpoints._connection_error_message(RuntimeError("weird"))
message = rest_endpoints._connection_error_message(RuntimeError("weird"), "https://example.com", 30.0)
assert "weird" not in message
assert "proxy logs" in message.lower()