scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 124

Merge pull request #38791 from BerriAI/litellm_fix_mcp_oauth_tool_fetch_auth

Browse source 
fix(mcp): cap tools preview and test-connection at the listing timeout and name the unreachable upstream
This commit is contained in:
Mateo Wang 2026-09-02 16:47:10 -07:00 • committed by GitHub
commit c1a26f36ee
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 98 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -74,7 +74,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. "
@ -1154,6 +1160,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.
@ -1169,6 +1176,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.
@ -1259,15 +1270,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
@ -1276,7 +1288,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:

81
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.
@ -3168,17 +3235,21 @@ 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):
@ -3188,11 +3259,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()