refactor(mcp): extract static_header env resolver and add regression tests

Replace the seed-empty-then-mutate dict in build_mcp_server_from_table with a
typed _resolve_static_header_env_ref helper and a comprehension, clearing the
LIT mutable-collection lint. Add two regression tests covering os.environ/NAME
resolution and the unset-env literal-fallback path
This commit is contained in:
Shriniwas410 2026-06-24 00:58:34 -04:00
parent 152b2290b3
commit 409a84f664
2 changed files with 85 additions and 14 deletions

View file

@ -331,6 +331,23 @@ def _deserialize_json_dict(data: Any) -> Optional[Dict[str, str]]:
return data
def _resolve_static_header_env_ref(key: str, value: str) -> str:
"""Resolve an ``os.environ/NAME`` static header to its environment value.
Non-references pass through untouched; an unset env var keeps the literal so
the misconfiguration stays visible instead of being silently dropped.
"""
if not value.startswith("os.environ/"):
return value
secret = get_secret_str(value)
if secret is not None:
return secret
verbose_logger.warning(
f"MCP static_header {key!r}: env var {value!r} not set; keeping literal value"
)
return value
def _deserialize_json_list(data: Any) -> Optional[List[Dict[str, Any]]]:
"""Deserialize a JSON array stored in the DB (``env_vars`` and friends).
@ -1065,20 +1082,10 @@ class MCPServerManager:
getattr(mcp_server, "static_headers", None)
)
if static_headers_dict:
resolved: dict[str, str] = {}
for k, v in static_headers_dict.items():
if isinstance(v, str) and v.startswith("os.environ/"):
secret = get_secret_str(v)
if secret is None:
verbose_logger.warning(
f"MCP static_header {k!r}: env var {v!r} not set; keeping literal value"
)
resolved[k] = v
else:
resolved[k] = secret
else:
resolved[k] = v
static_headers_dict = resolved
static_headers_dict = {
key: _resolve_static_header_env_ref(key, value)
for key, value in static_headers_dict.items()
}
env_vars_list = self._resolve_env_vars_list(
mcp_server,
env_vars_are_encrypted=(

View file

@ -1709,3 +1709,67 @@ async def test_missing_user_env_vars_error_renders_in_mcp_call_tool():
assert "CorporateDB" in text
assert "CORP_USERNAME" in text
assert "fill_env_vars=srv-99" in text
@pytest.mark.asyncio
async def test_build_mcp_server_from_table_resolves_os_environ_static_headers(
monkeypatch,
):
"""A DB-loaded ``os.environ/NAME`` static header must build into the runtime
server carrying the resolved environment value, not the literal reference, so
upstream requests send the real secret."""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
MCPServerManager,
)
from litellm.proxy._types import LiteLLM_MCPServerTable
monkeypatch.setenv("MCP_STATIC_HEADER_TOKEN", "resolved-secret")
table = LiteLLM_MCPServerTable(
server_id="srv-os-environ",
alias="echo",
url="https://upstream.example.com/mcp",
transport="http",
auth_type="none",
static_headers={
"X-Token": "os.environ/MCP_STATIC_HEADER_TOKEN",
"X-Static": "plain-value",
},
)
server = await MCPServerManager().build_mcp_server_from_table(table)
assert server.static_headers == {
"X-Token": "resolved-secret",
"X-Static": "plain-value",
}
@pytest.mark.asyncio
async def test_build_mcp_server_from_table_keeps_literal_when_env_missing(
monkeypatch, caplog
):
"""When the referenced env var is unset, the literal ``os.environ/NAME`` is
preserved and a warning is logged, so the misconfiguration stays visible
instead of silently collapsing to an empty header."""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
MCPServerManager,
)
from litellm.proxy._types import LiteLLM_MCPServerTable
monkeypatch.delenv("MCP_MISSING_HEADER_TOKEN", raising=False)
table = LiteLLM_MCPServerTable(
server_id="srv-os-environ-missing",
alias="echo",
url="https://upstream.example.com/mcp",
transport="http",
auth_type="none",
static_headers={"X-Token": "os.environ/MCP_MISSING_HEADER_TOKEN"},
)
with caplog.at_level("WARNING"):
server = await MCPServerManager().build_mcp_server_from_table(table)
assert server.static_headers == {"X-Token": "os.environ/MCP_MISSING_HEADER_TOKEN"}
assert "MCP_MISSING_HEADER_TOKEN" in caplog.text