fix(mcp): skip health check for servers with per-user env-var headers

This commit is contained in:
mateo-berri 2026-06-04 16:42:00 +00:00
parent 2e948527ce
commit 3f62063314
No known key found for this signature in database
2 changed files with 104 additions and 1 deletions

View file

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

View file

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