mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
fix(mcp): skip health check for servers with per-user env-var headers
This commit is contained in:
parent
2e948527ce
commit
3f62063314
2 changed files with 104 additions and 1 deletions
|
|
@ -432,7 +432,8 @@ class MCPServerManager:
|
|||
- server is OpenAPI (spec_path),
|
||||
- non-empty upstream instructions are already cached,
|
||||
- auth preconditions match health_check_server's skip rules
|
||||
(per-user auth / missing static auth token),
|
||||
(per-user auth / missing static auth token / static headers that
|
||||
reference a per-user env var),
|
||||
- a prior probe attempt for this server is within
|
||||
MCP_HEALTH_CHECK_TIMEOUT seconds (the probe is a health-check-shaped
|
||||
op and already uses this knob for its inner call timeout; reusing it
|
||||
|
|
@ -447,6 +448,8 @@ class MCPServerManager:
|
|||
return
|
||||
if server.requires_per_user_auth:
|
||||
return
|
||||
if self._references_per_user_env_var(server):
|
||||
return
|
||||
if (
|
||||
server.auth_type
|
||||
and server.auth_type != MCPAuth.none
|
||||
|
|
@ -1587,6 +1590,25 @@ class MCPServerManager:
|
|||
|
||||
return resolved_env
|
||||
|
||||
def _references_per_user_env_var(self, server: MCPServer) -> bool:
|
||||
"""True when ``server.static_headers`` reference a per-user ``${NAME}`` env var.
|
||||
|
||||
Such placeholders can only be filled from a calling user's stored values,
|
||||
so a userless probe (health check / instructions prefetch) would forward
|
||||
the literal ``${NAME}`` upstream and get rejected. Callers skip the probe
|
||||
and report ``unknown`` instead of a misleading ``unhealthy``.
|
||||
"""
|
||||
static_headers = server.static_headers
|
||||
env_vars = getattr(server, "env_vars", None)
|
||||
if not static_headers or not env_vars:
|
||||
return False
|
||||
_global_values, user_specs = parse_admin_env_vars(env_vars)
|
||||
user_var_names = {spec["name"] for spec in user_specs}
|
||||
if not user_var_names:
|
||||
return False
|
||||
referenced = collect_env_var_references(strings=static_headers.values())
|
||||
return bool(referenced & user_var_names)
|
||||
|
||||
async def _resolve_static_headers_with_env_vars(
|
||||
self,
|
||||
server: MCPServer,
|
||||
|
|
@ -4043,6 +4065,11 @@ class MCPServerManager:
|
|||
and not server.authentication_token
|
||||
):
|
||||
should_skip_health_check = True
|
||||
# Skip if static_headers reference a per-user env var: a userless probe
|
||||
# can't fill ${NAME} and would forward the literal placeholder upstream,
|
||||
# flipping the server to unhealthy even though real user calls succeed.
|
||||
elif self._references_per_user_env_var(server):
|
||||
should_skip_health_check = True
|
||||
|
||||
if not should_skip_health_check:
|
||||
resolved_static_headers = await self._resolve_static_headers_with_env_vars(
|
||||
|
|
|
|||
|
|
@ -446,6 +446,82 @@ async def test_resolve_static_headers_stale_user_value_cannot_override_global(
|
|||
assert headers == {"X-DB-URL": "admin-db/alice"}
|
||||
|
||||
|
||||
# ── health-check skip for per-user-env-var-backed headers ──────────────────
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"static_headers, env_vars, expected",
|
||||
[
|
||||
(
|
||||
{"Authorization": "Bearer ${GITHUB_TOKEN}"},
|
||||
[{"name": "GITHUB_TOKEN", "value": "", "scope": "user"}],
|
||||
True,
|
||||
),
|
||||
(
|
||||
{"Authorization": "Bearer ${SHARED_TOKEN}"},
|
||||
[{"name": "SHARED_TOKEN", "value": "abc", "scope": "global"}],
|
||||
False,
|
||||
),
|
||||
(
|
||||
{"X-Static": "literal"},
|
||||
[{"name": "GITHUB_TOKEN", "value": "", "scope": "user"}],
|
||||
False,
|
||||
),
|
||||
(None, [{"name": "GITHUB_TOKEN", "value": "", "scope": "user"}], False),
|
||||
({"Authorization": "Bearer ${GITHUB_TOKEN}"}, None, False),
|
||||
],
|
||||
)
|
||||
def test_references_per_user_env_var(static_headers, env_vars, expected):
|
||||
"""Only headers that actually reference a *per-user* var count: globals and
|
||||
declared-but-unreferenced user vars do not, since the userless probe can
|
||||
still resolve (or simply not need) them."""
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
MCPServerManager,
|
||||
)
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
||||
manager = MCPServerManager()
|
||||
server = MCPServer(
|
||||
server_id="srv-x",
|
||||
name="srv",
|
||||
transport="http",
|
||||
url="https://example.com",
|
||||
static_headers=static_headers,
|
||||
env_vars=env_vars,
|
||||
)
|
||||
assert manager._references_per_user_env_var(server) is expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_health_check_skips_servers_referencing_per_user_env_var(
|
||||
mock_server, monkeypatch
|
||||
):
|
||||
"""A userless health probe cannot fill per-user ${NAME} placeholders, so a
|
||||
server whose static_headers reference one must report 'unknown' without
|
||||
connecting. Otherwise it forwards the literal placeholder upstream, gets a
|
||||
401, and flips to 'unhealthy' even though real user calls succeed."""
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
MCPServerManager,
|
||||
)
|
||||
|
||||
manager = MCPServerManager()
|
||||
manager.registry[mock_server.server_id] = mock_server
|
||||
|
||||
created = []
|
||||
|
||||
async def fake_create_client(*args, **kwargs):
|
||||
created.append((args, kwargs))
|
||||
raise RuntimeError("upstream rejected literal ${NAME}")
|
||||
|
||||
monkeypatch.setattr(manager, "_create_mcp_client", fake_create_client)
|
||||
|
||||
result = await manager.health_check_server(mock_server.server_id)
|
||||
|
||||
assert created == []
|
||||
assert result.status == "unknown"
|
||||
assert result.health_check_error is None
|
||||
|
||||
|
||||
# ── _load_user_env_vars guard paths ────────────────────────────────────────
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue