mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
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:
parent
152b2290b3
commit
409a84f664
2 changed files with 85 additions and 14 deletions
|
|
@ -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=(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue