mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-28 01:32:17 +00:00
fix(mcp): probe openapi-backed servers by their spec in health checks
An OpenAPI-backed server stores the REST base url in `url` and only looks like http transport, so the health check opened an MCP session against a plain REST API. The upstream answers 404, which the MCP SDK reports as "Session terminated", and the server showed as unhealthy even though every tool had loaded from its spec These servers now probe the spec: reachable and parseable means healthy, and a spec that stopped loading reports the real upstream error. Per-user auth still skips to unknown, since fetching the spec userless could 401 and report a false failure _ensure_upstream_initialize_instructions_cached already skips OpenAPI servers for the same reason, and its docstring says it mirrors health_check_server's skip rules Fixes #40079
This commit is contained in:
parent
168a0055a2
commit
6d1c9b1f7b
2 changed files with 107 additions and 17 deletions
|
|
@ -6665,25 +6665,33 @@ class MCPServerManager:
|
|||
|
||||
if not should_skip_health_check:
|
||||
try:
|
||||
resolved_static_headers: Final = await self._resolve_static_headers_with_env_vars(
|
||||
server=server,
|
||||
user_api_key_auth=None,
|
||||
raise_on_missing=False,
|
||||
)
|
||||
extra_headers: Final = dict(resolved_static_headers) if resolved_static_headers else {}
|
||||
client: Final = await self._create_mcp_client(
|
||||
server=server,
|
||||
mcp_auth_header=None,
|
||||
extra_headers=extra_headers,
|
||||
stdio_env=None,
|
||||
)
|
||||
if server.spec_path:
|
||||
# url is a REST base, not an MCP endpoint, so the spec is the only thing to probe
|
||||
from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import (
|
||||
load_openapi_spec_async,
|
||||
)
|
||||
|
||||
async def _noop(session):
|
||||
return "ok"
|
||||
await asyncio.wait_for(load_openapi_spec_async(server.spec_path), timeout=MCP_HEALTH_CHECK_TIMEOUT)
|
||||
else:
|
||||
resolved_static_headers: Final = await self._resolve_static_headers_with_env_vars(
|
||||
server=server,
|
||||
user_api_key_auth=None,
|
||||
raise_on_missing=False,
|
||||
)
|
||||
extra_headers: Final = dict(resolved_static_headers) if resolved_static_headers else {}
|
||||
client: Final = await self._create_mcp_client(
|
||||
server=server,
|
||||
mcp_auth_header=None,
|
||||
extra_headers=extra_headers,
|
||||
stdio_env=None,
|
||||
)
|
||||
|
||||
# Add timeout wrapper to prevent hanging
|
||||
await asyncio.wait_for(client.run_with_session(_noop), timeout=MCP_HEALTH_CHECK_TIMEOUT)
|
||||
self._remember_upstream_initialize_instructions(server, client)
|
||||
async def _noop(session):
|
||||
return "ok"
|
||||
|
||||
# Add timeout wrapper to prevent hanging
|
||||
await asyncio.wait_for(client.run_with_session(_noop), timeout=MCP_HEALTH_CHECK_TIMEOUT)
|
||||
self._remember_upstream_initialize_instructions(server, client)
|
||||
status = "healthy"
|
||||
except asyncio.TimeoutError:
|
||||
health_check_error = f"Health check timed out after {MCP_HEALTH_CHECK_TIMEOUT} seconds"
|
||||
|
|
|
|||
|
|
@ -4493,6 +4493,88 @@ class TestMCPServerManager:
|
|||
assert result.status == "unhealthy"
|
||||
assert "OAuth discovery unavailable" in (result.health_check_error or "")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_health_check_server_openapi_spec_is_probed_instead_of_mcp_session(self):
|
||||
"""An OpenAPI server is healthy when its spec loads, and is never asked to speak MCP.
|
||||
|
||||
Regression for #40079: these servers store the REST base url in ``url`` and only look like
|
||||
http transport, so opening an MCP session against them always failed with "Session
|
||||
terminated" even though every tool had loaded from the spec.
|
||||
"""
|
||||
manager = MCPServerManager()
|
||||
server = MCPServer(
|
||||
server_id="openapi-server",
|
||||
name="openapi-server",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.none,
|
||||
url="http://rest.example.com",
|
||||
spec_path="http://rest.example.com/openapi.json",
|
||||
)
|
||||
manager.get_mcp_server_by_id = MagicMock(return_value=server)
|
||||
manager._create_mcp_client = AsyncMock()
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator.load_openapi_spec_async",
|
||||
new=AsyncMock(return_value={"openapi": "3.0.0", "paths": {}}),
|
||||
) as mock_load:
|
||||
result = await manager.health_check_server("openapi-server")
|
||||
|
||||
assert result.status == "healthy"
|
||||
assert result.health_check_error is None
|
||||
mock_load.assert_awaited_once_with("http://rest.example.com/openapi.json")
|
||||
manager._create_mcp_client.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_health_check_server_openapi_spec_failure_is_unhealthy(self):
|
||||
"""A spec that stopped loading is unhealthy, and the upstream error is what gets reported."""
|
||||
manager = MCPServerManager()
|
||||
server = MCPServer(
|
||||
server_id="openapi-server",
|
||||
name="openapi-server",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.none,
|
||||
url="http://rest.example.com",
|
||||
spec_path="http://rest.example.com/openapi.json",
|
||||
)
|
||||
manager.get_mcp_server_by_id = MagicMock(return_value=server)
|
||||
manager._create_mcp_client = AsyncMock()
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator.load_openapi_spec_async",
|
||||
new=AsyncMock(side_effect=httpx.ConnectError("Cannot connect to host rest.example.com")),
|
||||
):
|
||||
result = await manager.health_check_server("openapi-server")
|
||||
|
||||
assert result.status == "unhealthy"
|
||||
assert result.health_check_error == "Cannot connect to host rest.example.com"
|
||||
manager._create_mcp_client.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_health_check_server_openapi_spec_keeps_per_user_auth_skip(self):
|
||||
"""Per-user auth still wins: fetching the spec userless could 401 and report a false failure."""
|
||||
manager = MCPServerManager()
|
||||
server = MCPServer(
|
||||
server_id="openapi-server",
|
||||
name="openapi-server",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.bearer_token,
|
||||
url="http://rest.example.com",
|
||||
spec_path="http://rest.example.com/openapi.json",
|
||||
)
|
||||
manager.get_mcp_server_by_id = MagicMock(return_value=server)
|
||||
manager._create_mcp_client = AsyncMock()
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator.load_openapi_spec_async",
|
||||
new=AsyncMock(return_value={"openapi": "3.0.0", "paths": {}}),
|
||||
) as mock_load:
|
||||
result = await manager.health_check_server("openapi-server")
|
||||
|
||||
assert result.status == "unknown"
|
||||
assert result.health_check_error is None
|
||||
mock_load.assert_not_called()
|
||||
manager._create_mcp_client.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_health_check_server_not_found(self):
|
||||
"""Test health check for a server that doesn't exist"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue