mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
feat(mcp): add secret_headers auth type that sends credential headers upstream
This commit is contained in:
parent
a92fd56a67
commit
796c57e335
10 changed files with 402 additions and 21 deletions
|
|
@ -514,6 +514,7 @@ class MCPClient:
|
|||
protocol_version: MCPUpstreamProtocol = "auto",
|
||||
skip_tool_output_validation: bool = False,
|
||||
upstream_transport: httpx2.AsyncBaseTransport | None = None,
|
||||
credential_header_names: tuple[str, ...] = (),
|
||||
):
|
||||
self.protocol_version: MCPUpstreamProtocol = TypeAdapter[MCPUpstreamProtocol](
|
||||
MCPUpstreamProtocol
|
||||
|
|
@ -531,6 +532,7 @@ class MCPClient:
|
|||
# Every consumer reads this rather than re-deriving it, since each re-derivation so far
|
||||
# picked up a different bug.
|
||||
self._credential_slot: str | None = auth_header_name or getattr(resolved_auth, "header_name", None)
|
||||
self._credential_header_names: Final = credential_header_names
|
||||
self.stdio_config: MCPStdioConfig | None = stdio_config
|
||||
self.extra_headers: dict[str, str] | None = extra_headers
|
||||
self.ssl_verify: VerifyTypes | None = ssl_verify
|
||||
|
|
@ -651,7 +653,12 @@ class MCPClient:
|
|||
if self._upstream_transport is None
|
||||
else self._create_httpx_client_factory(transport=self._upstream_transport)
|
||||
)
|
||||
verbose_logger.debug("litellm headers for streamable_http_client: %s", headers)
|
||||
outbound_headers: Final = cast(dict[str, str], headers) # cast-ok: _get_auth_headers returns an untyped dict
|
||||
secret_names: Final = {name.lower() for name in self._credential_header_names}
|
||||
verbose_logger.debug(
|
||||
"litellm headers for streamable_http_client: %s",
|
||||
{name: value for name, value in outbound_headers.items() if name.lower() not in secret_names},
|
||||
)
|
||||
http_client = httpx_client_factory(
|
||||
headers=headers,
|
||||
timeout=httpx2.Timeout(self.timeout),
|
||||
|
|
@ -1003,7 +1010,9 @@ class MCPClient:
|
|||
# SigV4 aws_auth. Both are None for the common case — no behavior change.
|
||||
fallback_auth: Final = self._resolved_auth if self._resolved_auth is not None else self._aws_auth
|
||||
effective_auth: Final = auth if auth is not None else fallback_auth
|
||||
guard: Final = credential_redirect_hook(self.server_url, self._credential_slot)
|
||||
guard: Final = credential_redirect_hook(
|
||||
self.server_url, self._credential_slot, self._credential_header_names
|
||||
)
|
||||
return _MCPHTTPClient(
|
||||
transport=transport,
|
||||
headers=headers,
|
||||
|
|
|
|||
|
|
@ -37635,6 +37635,7 @@
|
|||
"authorization",
|
||||
"oauth2",
|
||||
"aws_sigv4",
|
||||
"secret_headers",
|
||||
"token",
|
||||
"oauth2_token_exchange",
|
||||
"oauth2_id_jag",
|
||||
|
|
@ -39429,6 +39430,7 @@
|
|||
"authorization",
|
||||
"oauth2",
|
||||
"aws_sigv4",
|
||||
"secret_headers",
|
||||
"token",
|
||||
"oauth2_token_exchange",
|
||||
"oauth2_id_jag",
|
||||
|
|
@ -40477,6 +40479,7 @@
|
|||
"authorization",
|
||||
"oauth2",
|
||||
"aws_sigv4",
|
||||
"secret_headers",
|
||||
"token",
|
||||
"oauth2_token_exchange",
|
||||
"oauth2_id_jag",
|
||||
|
|
@ -42642,6 +42645,7 @@
|
|||
"authorization",
|
||||
"oauth2",
|
||||
"aws_sigv4",
|
||||
"secret_headers",
|
||||
"token",
|
||||
"oauth2_token_exchange",
|
||||
"oauth2_id_jag",
|
||||
|
|
@ -44534,6 +44538,7 @@
|
|||
"authorization",
|
||||
"oauth2",
|
||||
"aws_sigv4",
|
||||
"secret_headers",
|
||||
"token",
|
||||
"oauth2_token_exchange",
|
||||
"oauth2_id_jag",
|
||||
|
|
@ -45089,6 +45094,7 @@
|
|||
"authorization",
|
||||
"oauth2",
|
||||
"aws_sigv4",
|
||||
"secret_headers",
|
||||
"token",
|
||||
"oauth2_token_exchange",
|
||||
"oauth2_id_jag",
|
||||
|
|
@ -48012,6 +48018,7 @@
|
|||
"authorization",
|
||||
"oauth2",
|
||||
"aws_sigv4",
|
||||
"secret_headers",
|
||||
"token",
|
||||
"oauth2_token_exchange",
|
||||
"oauth2_id_jag",
|
||||
|
|
|
|||
|
|
@ -2493,6 +2493,7 @@ class MCPServerManager:
|
|||
server.auth_type
|
||||
and server.auth_type != MCPAuth.none
|
||||
and server.auth_type != MCPAuth.aws_sigv4
|
||||
and server.auth_type != MCPAuth.secret_headers
|
||||
and not server.authentication_token
|
||||
):
|
||||
return
|
||||
|
|
@ -3276,6 +3277,11 @@ class MCPServerManager:
|
|||
|
||||
# AWS SigV4 credential fields
|
||||
aws_creds: Final = self._extract_aws_credentials(credentials_dict, credentials_are_encrypted)
|
||||
credential_headers: Final = (
|
||||
self._extract_credential_headers(credentials_dict, credentials_are_encrypted)
|
||||
if mcp_server.auth_type == MCPAuth.secret_headers
|
||||
else {}
|
||||
)
|
||||
|
||||
scopes: list[str] | None = None
|
||||
if credentials_dict:
|
||||
|
|
@ -3366,7 +3372,12 @@ class MCPServerManager:
|
|||
authentication_token=auth_value,
|
||||
mcp_info=mcp_info,
|
||||
extra_headers=getattr(mcp_server, "extra_headers", None),
|
||||
static_headers=static_headers_dict,
|
||||
static_headers=(
|
||||
self._with_credential_headers(static_headers_dict, credential_headers)
|
||||
if credential_headers
|
||||
else static_headers_dict
|
||||
),
|
||||
credential_header_names=tuple(credential_headers),
|
||||
env_vars=env_vars_list,
|
||||
dcr_issuer=credentials_dict.get("dcr_issuer") if credentials_dict else None,
|
||||
dcr_server_url=credentials_dict.get("dcr_server_url") if credentials_dict else None,
|
||||
|
|
@ -4165,7 +4176,7 @@ class MCPServerManager:
|
|||
the literal ``${NAME}`` upstream and get rejected. Callers skip the probe
|
||||
and report ``unknown`` instead of a misleading ``unhealthy``.
|
||||
"""
|
||||
static_headers: Final = server.static_headers
|
||||
static_headers: Final = server.listed_static_headers
|
||||
env_vars: Final[_EnvVarList | None] = getattr(server, "env_vars", None)
|
||||
if not static_headers or not env_vars:
|
||||
return False
|
||||
|
|
@ -4221,7 +4232,7 @@ class MCPServerManager:
|
|||
# also carries a global value is always covered by that global (globals
|
||||
# win in the merge below), so it can never be genuinely "missing" even if
|
||||
# the user hasn't filled it in -- only vars without a global fallback do.
|
||||
referenced: Final = collect_env_var_references(strings=(static_headers or {}).values())
|
||||
referenced: Final = collect_env_var_references(strings=(server.listed_static_headers or {}).values())
|
||||
referenced_user_vars: Final = referenced & user_var_names
|
||||
required_user_vars: Final = {name for name in referenced_user_vars if name not in global_values}
|
||||
|
||||
|
|
@ -4267,6 +4278,11 @@ class MCPServerManager:
|
|||
merged_vars: Final[dict[str, str]] = {**scoped_user_values, **global_values}
|
||||
if not static_headers:
|
||||
return static_headers
|
||||
if server.credential_header_names:
|
||||
return {
|
||||
**interpolate_headers(server.listed_static_headers or {}, merged_vars),
|
||||
**server.credential_static_headers,
|
||||
}
|
||||
return interpolate_headers(static_headers, merged_vars)
|
||||
|
||||
async def _load_user_env_vars(
|
||||
|
|
@ -5627,6 +5643,30 @@ class MCPServerManager:
|
|||
"aws_session_name": credentials_dict.get("aws_session_name"),
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _with_credential_headers(
|
||||
static_headers: Mapping[str, str] | None, credential_headers: Mapping[str, str]
|
||||
) -> dict[str, str]:
|
||||
secret: Final = {name.lower() for name in credential_headers}
|
||||
kept: Final = {name: value for name, value in (static_headers or {}).items() if name.lower() not in secret}
|
||||
return {**kept, **credential_headers}
|
||||
|
||||
def _extract_credential_headers(
|
||||
self,
|
||||
credentials_dict: dict[str, str] | None,
|
||||
credentials_are_encrypted: bool,
|
||||
) -> dict[str, str]:
|
||||
raw_headers: Final[object] = credentials_dict.get("headers") if credentials_dict else None
|
||||
if not isinstance(raw_headers, dict):
|
||||
return {}
|
||||
headers: Final = cast(Mapping[str, object], raw_headers) # cast-ok: dict shape validated above
|
||||
return {
|
||||
name: value
|
||||
for name, stored in headers.items()
|
||||
if isinstance(stored, str)
|
||||
and (value := self._decrypt_credential_field(stored, "headers", credentials_are_encrypted))
|
||||
}
|
||||
|
||||
def _extract_scopes(self, scopes_value: str | Sequence[object] | None) -> list[str] | None:
|
||||
if isinstance(scopes_value, str):
|
||||
scopes = [s.strip() for s in scopes_value.split() if s.strip()]
|
||||
|
|
@ -7641,6 +7681,7 @@ class MCPServerManager:
|
|||
server.auth_type
|
||||
and server.auth_type != MCPAuth.none
|
||||
and server.auth_type != MCPAuth.aws_sigv4
|
||||
and server.auth_type != MCPAuth.secret_headers
|
||||
and not server.authentication_token
|
||||
)
|
||||
or self._references_per_user_env_var(server)
|
||||
|
|
@ -7702,7 +7743,7 @@ class MCPServerManager:
|
|||
allowed_tools=server.allowed_tools or [],
|
||||
extra_headers=server.extra_headers or [],
|
||||
mcp_info=server.mcp_info,
|
||||
static_headers=server.static_headers,
|
||||
static_headers=server.listed_static_headers,
|
||||
env_vars=self._env_vars_to_models(server.env_vars),
|
||||
status=status,
|
||||
last_health_check=datetime.now(),
|
||||
|
|
@ -7819,7 +7860,7 @@ class MCPServerManager:
|
|||
pinned_tools=server.pinned_tools,
|
||||
extra_headers=server.extra_headers or [],
|
||||
mcp_info=server.mcp_info,
|
||||
static_headers=server.static_headers,
|
||||
static_headers=server.listed_static_headers,
|
||||
env_vars=self._env_vars_to_models(server.env_vars),
|
||||
status=None, # No health check performed
|
||||
last_health_check=None, # No health check performed
|
||||
|
|
|
|||
|
|
@ -103,6 +103,8 @@ def to_server_spec(server: MCPServer) -> ServerSpec | None:
|
|||
if server.is_oauth_passthrough:
|
||||
return None # passthrough is not migrated yet -> defer to v1
|
||||
return ServerSpec(server_id=server.server_id, resource=resource, config=NoneConfig())
|
||||
case MCPAuth.secret_headers:
|
||||
return ServerSpec(server_id=server.server_id, resource=resource, config=NoneConfig())
|
||||
case MCPAuth.api_key:
|
||||
return _shared_key_spec(server, resource, "X-API-Key", "")
|
||||
case MCPAuth.bearer_token:
|
||||
|
|
|
|||
|
|
@ -1575,6 +1575,20 @@ if MCP_AVAILABLE:
|
|||
mcp_auth_header: str | None
|
||||
oauth2_headers: dict[str, str] | None
|
||||
|
||||
def _preview_credential_headers(request: NewMCPServerRequest) -> Mapping[str, str]:
|
||||
return (request.credentials or {}).get("headers") or {} if request.auth_type == MCPAuth.secret_headers else {}
|
||||
|
||||
def _preview_static_headers(
|
||||
request: NewMCPServerRequest,
|
||||
) -> dict[str, str] | None: # mutable-ok: feeds MCPServer.static_headers
|
||||
credential_headers: Final = _preview_credential_headers(request)
|
||||
return (
|
||||
{**(request.static_headers or {}), **credential_headers} if credential_headers else request.static_headers
|
||||
)
|
||||
|
||||
def _preview_credential_header_names(request: NewMCPServerRequest) -> tuple[str, ...]:
|
||||
return tuple(_preview_credential_headers(request))
|
||||
|
||||
def _preview_origin(url: str | None) -> tuple[str, str, int | None] | None:
|
||||
if not url:
|
||||
return None
|
||||
|
|
@ -1608,9 +1622,26 @@ if MCP_AVAILABLE:
|
|||
may_inherit: Final = new_mcp_server_request.auth_type not in _STAGED_AUTH_VALUE_AUTH_TYPES or (
|
||||
saved_origin is not None and saved_origin == preview_origin
|
||||
)
|
||||
request: Final = (
|
||||
inherited: Final = (
|
||||
inherit_credentials_from_existing_server(new_mcp_server_request) if may_inherit else new_mcp_server_request
|
||||
)
|
||||
keeps_saved_headers: Final = (
|
||||
inherited.auth_type == MCPAuth.secret_headers
|
||||
and not (inherited.credentials or {}).get("headers")
|
||||
and saved_server is not None
|
||||
and saved_server.auth_type == MCPAuth.secret_headers
|
||||
and saved_origin is not None
|
||||
and saved_origin == preview_origin
|
||||
)
|
||||
request: Final = (
|
||||
inherited.model_copy(
|
||||
update={
|
||||
"credentials": {**(inherited.credentials or {}), "headers": saved_server.credential_static_headers}
|
||||
}
|
||||
)
|
||||
if keeps_saved_headers and saved_server is not None
|
||||
else inherited
|
||||
)
|
||||
mcp_auth_header: Final = (
|
||||
request.credentials.get("auth_value")
|
||||
if request.auth_type in _STAGED_AUTH_VALUE_AUTH_TYPES and isinstance(request.credentials, dict)
|
||||
|
|
@ -1686,7 +1717,8 @@ if MCP_AVAILABLE:
|
|||
# the saved record during client creation and discard the edited connection settings.
|
||||
preview_server_id: Final = (
|
||||
""
|
||||
if request.auth_type in _STAGED_AUTH_VALUE_AUTH_TYPES or request.auth_type in (None, MCPAuth.none)
|
||||
if request.auth_type in _STAGED_AUTH_VALUE_AUTH_TYPES
|
||||
or request.auth_type in (None, MCPAuth.none, MCPAuth.secret_headers)
|
||||
else request.server_id or ""
|
||||
)
|
||||
server_model: Final = MCPServer(
|
||||
|
|
@ -1700,7 +1732,8 @@ if MCP_AVAILABLE:
|
|||
command=request.command,
|
||||
args=request.args,
|
||||
env=request.env,
|
||||
static_headers=request.static_headers,
|
||||
static_headers=_preview_static_headers(request),
|
||||
credential_header_names=_preview_credential_header_names(request),
|
||||
client_id=client_id,
|
||||
client_secret=client_secret,
|
||||
issuer=request.issuer,
|
||||
|
|
@ -1761,7 +1794,7 @@ if MCP_AVAILABLE:
|
|||
|
||||
merged_headers: Final = merge_mcp_headers(
|
||||
extra_headers=(None if preview_cred_provider else effective_oauth2_headers),
|
||||
static_headers=request.static_headers,
|
||||
static_headers=_preview_static_headers(request),
|
||||
)
|
||||
|
||||
with anyio.fail_after(timeout_seconds):
|
||||
|
|
|
|||
|
|
@ -269,6 +269,7 @@ async def prepare_upstream_client(
|
|||
skip_tool_output_validation=server.skip_tool_output_validation,
|
||||
upstream_transport=upstream_transport,
|
||||
auth_type=server.auth_type,
|
||||
credential_header_names=server.credential_header_names,
|
||||
timeout=timeout,
|
||||
extra_headers=resolved_headers,
|
||||
resolved_auth=resolved_auth,
|
||||
|
|
@ -313,6 +314,7 @@ async def prepare_upstream_client(
|
|||
skip_tool_output_validation=server.skip_tool_output_validation,
|
||||
upstream_transport=upstream_transport,
|
||||
auth_type=server.auth_type,
|
||||
credential_header_names=server.credential_header_names,
|
||||
auth_value=auth_value,
|
||||
auth_header_name=auth_header_name,
|
||||
timeout=timeout,
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ from __future__ import annotations
|
|||
|
||||
import enum
|
||||
import re
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from collections.abc import Awaitable, Callable, Iterable, Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, TypeAlias
|
||||
from urllib.parse import urlsplit
|
||||
|
|
@ -46,6 +46,7 @@ class MCPAuth(str, enum.Enum):
|
|||
authorization = "authorization"
|
||||
oauth2 = "oauth2"
|
||||
aws_sigv4 = "aws_sigv4"
|
||||
secret_headers = "secret_headers"
|
||||
token = "token"
|
||||
oauth2_token_exchange = "oauth2_token_exchange"
|
||||
oauth2_id_jag = "oauth2_id_jag"
|
||||
|
|
@ -91,6 +92,7 @@ MCPAuthType = (
|
|||
MCPAuth.authorization,
|
||||
MCPAuth.oauth2,
|
||||
MCPAuth.aws_sigv4,
|
||||
MCPAuth.secret_headers,
|
||||
MCPAuth.token,
|
||||
MCPAuth.oauth2_token_exchange,
|
||||
MCPAuth.oauth2_id_jag,
|
||||
|
|
@ -395,20 +397,26 @@ def custom_credential_slot(headers: Mapping[str, str] | None) -> str | None:
|
|||
|
||||
|
||||
def credential_redirect_hook(
|
||||
configured_url: str, slot: str | None
|
||||
configured_url: str, slot: str | None, secret_names: Iterable[str] = ()
|
||||
) -> Callable[[httpx.Request | httpx2.Request], Awaitable[None]] | None:
|
||||
"""An httpx request hook dropping ``slot`` once a redirect leaves ``configured_url``'s origin.
|
||||
"""An httpx request hook dropping ``slot`` and ``secret_names`` once a redirect leaves ``configured_url``'s origin.
|
||||
|
||||
None when no guard is needed, so callers do not each repeat the exemption: HTTP clients already
|
||||
strip ``Authorization`` across origins, but forward every other header, so only a credential an
|
||||
operator moved to its own slot can be replayed to whatever host the upstream redirects to.
|
||||
"""
|
||||
if not configured_url or not slot or same_header(slot, DEFAULT_CREDENTIAL_HEADER):
|
||||
names: Final = tuple(
|
||||
name for name in (slot, *secret_names) if name and not same_header(name, DEFAULT_CREDENTIAL_HEADER)
|
||||
)
|
||||
if not configured_url or not names:
|
||||
return None
|
||||
|
||||
async def guard(request: httpx.Request | httpx2.Request) -> None:
|
||||
if slot in request.headers and crosses_origin(configured_url, str(request.url)):
|
||||
del request.headers[slot]
|
||||
if not crosses_origin(configured_url, str(request.url)):
|
||||
return
|
||||
for name in names:
|
||||
if name in request.headers:
|
||||
del request.headers[name] # rebind-ok: the guard strips the credential in place
|
||||
|
||||
return guard
|
||||
|
||||
|
|
|
|||
|
|
@ -116,6 +116,7 @@ class MCPServer(LiteLLMBaseModel):
|
|||
pinned_tools: dict[str, PinnedMCPTool] | None = None
|
||||
allowed_params: dict[str, list[str]] | None = None # map of tool names to allowed parameter lists
|
||||
static_headers: dict[str, str] | None = None # static headers to forward to the MCP server
|
||||
credential_header_names: tuple[str, ...] = ()
|
||||
# Admin-configured env vars. Each entry is {name, value, scope, description}.
|
||||
# scope=="global" values are interpolated into static_headers using ${NAME}.
|
||||
# scope=="user" values must be supplied per-user.
|
||||
|
|
@ -356,6 +357,7 @@ class MCPServer(LiteLLMBaseModel):
|
|||
MCPAuth.authorization,
|
||||
MCPAuth.token,
|
||||
MCPAuth.aws_sigv4,
|
||||
MCPAuth.secret_headers,
|
||||
):
|
||||
return False
|
||||
return not any(
|
||||
|
|
@ -363,6 +365,18 @@ class MCPServer(LiteLLMBaseModel):
|
|||
for header in (self.extra_headers or ())
|
||||
)
|
||||
|
||||
@property
|
||||
def listed_static_headers(self) -> dict[str, str] | None: # mutable-ok: feeds the API table model
|
||||
if not self.static_headers:
|
||||
return self.static_headers
|
||||
return {name: value for name, value in self.static_headers.items() if name not in self.credential_header_names}
|
||||
|
||||
@property
|
||||
def credential_static_headers(self) -> dict[str, str]: # mutable-ok: merged back into request headers
|
||||
return {
|
||||
name: value for name, value in (self.static_headers or {}).items() if name in self.credential_header_names
|
||||
}
|
||||
|
||||
@property
|
||||
def is_true_passthrough(self) -> bool:
|
||||
"""True for the transparent-proxy mode: LiteLLM performs no admission auth and forwards the
|
||||
|
|
|
|||
265
tests/unit/proxy/mcp_server/test_mcp_secret_headers.py
Normal file
265
tests/unit/proxy/mcp_server/test_mcp_secret_headers.py
Normal file
|
|
@ -0,0 +1,265 @@
|
|||
import io
|
||||
import logging
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
import httpx2
|
||||
import pytest
|
||||
from starlette.datastructures import Headers
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.mcp_client.client import MCPClient
|
||||
from litellm.proxy._types import LiteLLM_MCPServerTable, NewMCPServerRequest
|
||||
from litellm.proxy.mcp_server import rest_endpoints
|
||||
from litellm.proxy.mcp_server.db import encrypt_credentials
|
||||
from litellm.proxy.mcp_server.mcp_server_manager import MCPServerManager, global_mcp_server_manager
|
||||
from litellm.types.mcp import MCPAuth, MCPTransport
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
||||
_URL: Final = "https://upstream.example/mcp"
|
||||
_CREDENTIAL_HEADERS: Final = {"DD-API-KEY": "dd-api-secret", "DD-APPLICATION-KEY": "dd-app-secret"}
|
||||
_STATIC_HEADERS: Final = {"X-Datadog-Site": "datadoghq.com"}
|
||||
_CREDENTIAL_HEADER_TEXT: Final = (*_CREDENTIAL_HEADERS, *_CREDENTIAL_HEADERS.values())
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _salt_key(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("LITELLM_SALT_KEY", "test-salt-key-for-secret-header-tests-1234")
|
||||
|
||||
|
||||
def _table(*, encrypted: bool) -> LiteLLM_MCPServerTable:
|
||||
credentials: Final = {"headers": dict(_CREDENTIAL_HEADERS)}
|
||||
return LiteLLM_MCPServerTable(
|
||||
server_id="srv-dd",
|
||||
server_name="dd",
|
||||
url=_URL,
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.secret_headers,
|
||||
credentials=encrypt_credentials(credentials, encryption_key=None) if encrypted else credentials,
|
||||
static_headers=dict(_STATIC_HEADERS),
|
||||
)
|
||||
|
||||
|
||||
async def _server(*, encrypted: bool) -> MCPServer:
|
||||
return await MCPServerManager().build_mcp_server_from_table(
|
||||
_table(encrypted=encrypted), credentials_are_encrypted=encrypted
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("encrypted", [True, False])
|
||||
async def test_secret_header_server_sends_credential_and_static_headers_upstream(encrypted: bool) -> None:
|
||||
manager: Final = MCPServerManager()
|
||||
server: Final = await _server(encrypted=encrypted)
|
||||
extra_headers: Final = await manager._resolve_static_headers_with_env_vars(server, None)
|
||||
client: Final = await manager._create_mcp_client(server, extra_headers=extra_headers)
|
||||
sent: Final = (await client.prepare_request_auth()).headers
|
||||
expected: Final = {**_STATIC_HEADERS, **_CREDENTIAL_HEADERS}
|
||||
assert {name: sent[name] for name in expected} == expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_admin_table_of_secret_header_server_lists_only_non_secret_static_headers() -> None:
|
||||
server: Final = await _server(encrypted=True)
|
||||
table: Final = MCPServerManager().build_mcp_server_table(server)
|
||||
assert table.static_headers == _STATIC_HEADERS
|
||||
assert table.credentials is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_instructions_prefetch_authenticates_secret_header_server_with_its_headers() -> None:
|
||||
sent: Final[list[httpx2.Headers]] = [] # mutable-ok: records upstream requests
|
||||
|
||||
def upstream(request: httpx2.Request) -> httpx2.Response:
|
||||
sent.append(request.headers)
|
||||
return httpx2.Response(500)
|
||||
|
||||
manager: Final = MCPServerManager(upstream_transport=httpx2.MockTransport(upstream))
|
||||
await manager.ensure_upstream_initialize_instructions_cached(await _server(encrypted=True))
|
||||
expected: Final = {**_STATIC_HEADERS, **_CREDENTIAL_HEADERS}
|
||||
assert {name: sent[0][name] for name in expected} == expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_health_check_authenticates_secret_header_server_with_its_headers() -> None:
|
||||
sent: Final[list[httpx2.Headers]] = [] # mutable-ok: records upstream requests
|
||||
|
||||
def upstream(request: httpx2.Request) -> httpx2.Response:
|
||||
sent.append(request.headers)
|
||||
return httpx2.Response(500)
|
||||
|
||||
manager: Final = MCPServerManager(upstream_transport=httpx2.MockTransport(upstream))
|
||||
server: Final = await _server(encrypted=True)
|
||||
manager.registry[server.server_id] = server
|
||||
result: Final = await manager.health_check_server(server.server_id)
|
||||
expected: Final = {**_STATIC_HEADERS, **_CREDENTIAL_HEADERS}
|
||||
assert {name: sent[0][name] for name in expected} == expected
|
||||
assert result.static_headers == _STATIC_HEADERS
|
||||
|
||||
|
||||
def test_secret_header_server_advertises_the_gateway_authorization_server_like_an_api_key_server() -> None:
|
||||
def advertises(auth_type: MCPAuth) -> bool:
|
||||
return MCPServer(
|
||||
server_id="s", name="s", url=_URL, transport=MCPTransport.http, auth_type=auth_type
|
||||
).advertises_gateway_authorization_server
|
||||
|
||||
assert advertises(MCPAuth.secret_headers) is advertises(MCPAuth.api_key) is True
|
||||
|
||||
|
||||
def test_preview_sends_the_credential_headers_with_the_static_headers() -> None:
|
||||
request: Final = NewMCPServerRequest(
|
||||
url=_URL,
|
||||
auth_type=MCPAuth.secret_headers,
|
||||
credentials={"headers": dict(_CREDENTIAL_HEADERS)},
|
||||
static_headers=dict(_STATIC_HEADERS),
|
||||
)
|
||||
assert rest_endpoints._preview_static_headers(request) == {**_STATIC_HEADERS, **_CREDENTIAL_HEADERS}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_credential_headers_are_ignored_unless_the_auth_type_is_secret_headers() -> None:
|
||||
table: Final = _table(encrypted=True).model_copy(update={"auth_type": MCPAuth.api_key})
|
||||
server: Final = await MCPServerManager().build_mcp_server_from_table(table, credentials_are_encrypted=True)
|
||||
assert server.static_headers == _STATIC_HEADERS
|
||||
assert server.credential_header_names == ()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_credential_header_replaces_a_static_header_of_the_same_name_in_any_case() -> None:
|
||||
table: Final = _table(encrypted=True).model_copy(
|
||||
update={"static_headers": {**_STATIC_HEADERS, "dd-api-key": "stale-plaintext"}}
|
||||
)
|
||||
manager: Final = MCPServerManager()
|
||||
server: Final = await manager.build_mcp_server_from_table(table, credentials_are_encrypted=True)
|
||||
assert server.static_headers == {**_STATIC_HEADERS, **_CREDENTIAL_HEADERS}
|
||||
assert manager.build_mcp_server_table(server).static_headers == _STATIC_HEADERS
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_env_var_placeholders_are_filled_in_static_headers_but_not_in_secret_values() -> None:
|
||||
server: Final = MCPServer(
|
||||
server_id="srv-env",
|
||||
name="srv-env",
|
||||
url=_URL,
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.secret_headers,
|
||||
static_headers={"X-Site": "${SITE}", "DD-API-KEY": "key-${SITE}"},
|
||||
credential_header_names=("DD-API-KEY",),
|
||||
env_vars=[{"name": "SITE", "value": "datadoghq.com", "scope": "global"}],
|
||||
)
|
||||
resolved: Final = await MCPServerManager()._resolve_static_headers_with_env_vars(server, None)
|
||||
assert resolved == {"X-Site": "datadoghq.com", "DD-API-KEY": "key-${SITE}"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("caller_token", [None, "caller-token"])
|
||||
@pytest.mark.parametrize(
|
||||
"redirect_to,expected",
|
||||
[
|
||||
("https://upstream.example/other", {"DD-API-KEY": "dd-api-secret", "DD-APPLICATION-KEY": "dd-app-secret"}),
|
||||
("https://elsewhere.example/collect", {"DD-API-KEY": None, "DD-APPLICATION-KEY": None}),
|
||||
],
|
||||
)
|
||||
async def test_credential_headers_do_not_follow_a_redirect_to_another_origin(
|
||||
redirect_to: str, expected: Mapping[str, str | None], caller_token: str | None
|
||||
) -> None:
|
||||
seen: Final[list[httpx2.Headers]] = [] # mutable-ok: records upstream requests
|
||||
|
||||
def upstream(request: httpx2.Request) -> httpx2.Response:
|
||||
if request.url.path == "/mcp":
|
||||
return httpx2.Response(307, headers={"location": redirect_to})
|
||||
seen.append(request.headers)
|
||||
return httpx2.Response(200)
|
||||
|
||||
manager: Final = MCPServerManager()
|
||||
server: Final = await _server(encrypted=True)
|
||||
client: Final = await manager._create_mcp_client(
|
||||
server,
|
||||
mcp_auth_header=caller_token,
|
||||
extra_headers=await manager._resolve_static_headers_with_env_vars(server, None),
|
||||
)
|
||||
async with client._create_httpx_client_factory(transport=httpx2.MockTransport(upstream))(
|
||||
headers=client._get_auth_headers()
|
||||
) as http:
|
||||
await http.get(_URL)
|
||||
assert {name: seen[0].get(name) for name in expected} == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"preview_url,headers",
|
||||
[(_URL, _CREDENTIAL_HEADERS), ("https://elsewhere.example/mcp", {})],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_preview_of_a_saved_server_reuses_its_headers_only_for_the_same_origin(
|
||||
monkeypatch: pytest.MonkeyPatch, preview_url: str, headers: Mapping[str, str]
|
||||
) -> None:
|
||||
saved: Final = await _server(encrypted=True)
|
||||
monkeypatch.setitem(global_mcp_server_manager.registry, saved.server_id, saved)
|
||||
request: Final = NewMCPServerRequest(
|
||||
server_id=saved.server_id, url=preview_url, transport=MCPTransport.http, auth_type=MCPAuth.secret_headers
|
||||
)
|
||||
staged: Final = rest_endpoints._stage_server_test(request, Headers({}))
|
||||
assert rest_endpoints._preview_static_headers(staged.request) == (headers or None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_debug_log_of_the_upstream_headers_leaves_out_credential_headers() -> None:
|
||||
manager: Final = MCPServerManager()
|
||||
server: Final = await _server(encrypted=True)
|
||||
client: Final = await manager._create_mcp_client(
|
||||
server, extra_headers=await manager._resolve_static_headers_with_env_vars(server, None)
|
||||
)
|
||||
captured: Final = io.StringIO()
|
||||
handler: Final = logging.StreamHandler(captured)
|
||||
previous_level: Final = verbose_logger.level
|
||||
verbose_logger.addHandler(handler)
|
||||
verbose_logger.setLevel(logging.DEBUG)
|
||||
try:
|
||||
client._create_transport_context()
|
||||
finally:
|
||||
verbose_logger.removeHandler(handler)
|
||||
verbose_logger.setLevel(previous_level)
|
||||
logged: Final = captured.getvalue()
|
||||
assert all(text not in logged for text in _CREDENTIAL_HEADER_TEXT)
|
||||
assert _STATIC_HEADERS["X-Datadog-Site"] in logged
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_preview_client_keeps_credential_headers_out_of_the_debug_log() -> None:
|
||||
request: Final = NewMCPServerRequest(
|
||||
url=_URL,
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.secret_headers,
|
||||
credentials={"headers": dict(_CREDENTIAL_HEADERS)},
|
||||
static_headers=dict(_STATIC_HEADERS),
|
||||
)
|
||||
|
||||
async def debug_log_of_transport(client: MCPClient) -> Mapping[str, object]:
|
||||
captured: Final = io.StringIO()
|
||||
handler: Final = logging.StreamHandler(captured)
|
||||
previous_level: Final = verbose_logger.level
|
||||
verbose_logger.addHandler(handler)
|
||||
verbose_logger.setLevel(logging.DEBUG)
|
||||
try:
|
||||
client._create_transport_context()
|
||||
finally:
|
||||
verbose_logger.removeHandler(handler)
|
||||
verbose_logger.setLevel(previous_level)
|
||||
return {"log": captured.getvalue()}
|
||||
|
||||
result: Final = await rest_endpoints._execute_with_mcp_client(request, debug_log_of_transport)
|
||||
logged: Final = str(result["log"])
|
||||
assert all(text not in logged for text in _CREDENTIAL_HEADER_TEXT)
|
||||
assert _STATIC_HEADERS["X-Datadog-Site"] in logged
|
||||
|
||||
|
||||
@pytest.mark.parametrize("auth_type", [MCPAuth.none, MCPAuth.api_key])
|
||||
def test_preview_ignores_credential_headers_for_any_other_auth_type(auth_type: MCPAuth) -> None:
|
||||
request: Final = NewMCPServerRequest(
|
||||
url=_URL,
|
||||
auth_type=auth_type,
|
||||
credentials={"headers": dict(_CREDENTIAL_HEADERS)},
|
||||
static_headers=dict(_STATIC_HEADERS),
|
||||
)
|
||||
assert rest_endpoints._preview_static_headers(request) == _STATIC_HEADERS
|
||||
assert rest_endpoints._preview_credential_header_names(request) == ()
|
||||
8
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
8
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -35792,7 +35792,7 @@ export interface components {
|
|||
/** Audience */
|
||||
audience?: string | null;
|
||||
/** Auth Type */
|
||||
auth_type?: ("none" | "api_key" | "bearer_token" | "basic" | "authorization" | "oauth2" | "aws_sigv4" | "token" | "oauth2_token_exchange" | "oauth2_id_jag" | "true_passthrough" | "oauth_delegate") | null;
|
||||
auth_type?: ("none" | "api_key" | "bearer_token" | "basic" | "authorization" | "oauth2" | "aws_sigv4" | "secret_headers" | "token" | "oauth2_token_exchange" | "oauth2_id_jag" | "true_passthrough" | "oauth_delegate") | null;
|
||||
/** Authorization Url */
|
||||
authorization_url?: string | null;
|
||||
/**
|
||||
|
|
@ -38974,7 +38974,7 @@ export interface components {
|
|||
/** Alias */
|
||||
alias?: string | null;
|
||||
/** Auth Type */
|
||||
auth_type?: ("none" | "api_key" | "bearer_token" | "basic" | "authorization" | "oauth2" | "aws_sigv4" | "token" | "oauth2_token_exchange" | "oauth2_id_jag" | "true_passthrough" | "oauth_delegate") | null;
|
||||
auth_type?: ("none" | "api_key" | "bearer_token" | "basic" | "authorization" | "oauth2" | "aws_sigv4" | "secret_headers" | "token" | "oauth2_token_exchange" | "oauth2_id_jag" | "true_passthrough" | "oauth_delegate") | null;
|
||||
/** Mcp Info */
|
||||
mcp_info?: {
|
||||
[key: string]: unknown;
|
||||
|
|
@ -40091,7 +40091,7 @@ export interface components {
|
|||
/** Audience */
|
||||
audience?: string | null;
|
||||
/** Auth Type */
|
||||
auth_type?: ("none" | "api_key" | "bearer_token" | "basic" | "authorization" | "oauth2" | "aws_sigv4" | "token" | "oauth2_token_exchange" | "oauth2_id_jag" | "true_passthrough" | "oauth_delegate") | null;
|
||||
auth_type?: ("none" | "api_key" | "bearer_token" | "basic" | "authorization" | "oauth2" | "aws_sigv4" | "secret_headers" | "token" | "oauth2_token_exchange" | "oauth2_id_jag" | "true_passthrough" | "oauth_delegate") | null;
|
||||
/** Authorization Url */
|
||||
authorization_url?: string | null;
|
||||
/**
|
||||
|
|
@ -50141,7 +50141,7 @@ export interface components {
|
|||
/** Audience */
|
||||
audience?: string | null;
|
||||
/** Auth Type */
|
||||
auth_type?: ("none" | "api_key" | "bearer_token" | "basic" | "authorization" | "oauth2" | "aws_sigv4" | "token" | "oauth2_token_exchange" | "oauth2_id_jag" | "true_passthrough" | "oauth_delegate") | null;
|
||||
auth_type?: ("none" | "api_key" | "bearer_token" | "basic" | "authorization" | "oauth2" | "aws_sigv4" | "secret_headers" | "token" | "oauth2_token_exchange" | "oauth2_id_jag" | "true_passthrough" | "oauth_delegate") | null;
|
||||
/** Authorization Url */
|
||||
authorization_url?: string | null;
|
||||
/**
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue