diff --git a/litellm/mcp_client/client.py b/litellm/mcp_client/client.py index 0461f1f06ff..01c365d808c 100644 --- a/litellm/mcp_client/client.py +++ b/litellm/mcp_client/client.py @@ -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, diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index c484e2c1bb8..7765a628ad5 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -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", diff --git a/litellm/proxy/mcp_server/mcp_server_manager.py b/litellm/proxy/mcp_server/mcp_server_manager.py index 013336e1d5d..b890e5fde1c 100644 --- a/litellm/proxy/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/mcp_server/mcp_server_manager.py @@ -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 diff --git a/litellm/proxy/mcp_server/outbound_credentials/adapter.py b/litellm/proxy/mcp_server/outbound_credentials/adapter.py index 1aac70cd847..eca55919270 100644 --- a/litellm/proxy/mcp_server/outbound_credentials/adapter.py +++ b/litellm/proxy/mcp_server/outbound_credentials/adapter.py @@ -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: diff --git a/litellm/proxy/mcp_server/rest_endpoints.py b/litellm/proxy/mcp_server/rest_endpoints.py index 2737b567985..ca805ab7982 100644 --- a/litellm/proxy/mcp_server/rest_endpoints.py +++ b/litellm/proxy/mcp_server/rest_endpoints.py @@ -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): diff --git a/litellm/proxy/mcp_server/upstream.py b/litellm/proxy/mcp_server/upstream.py index 3dd245947db..c3850dee126 100644 --- a/litellm/proxy/mcp_server/upstream.py +++ b/litellm/proxy/mcp_server/upstream.py @@ -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, diff --git a/litellm/types/mcp.py b/litellm/types/mcp.py index e1cd37790b3..d01d630ff20 100644 --- a/litellm/types/mcp.py +++ b/litellm/types/mcp.py @@ -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 diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 1d3d5911eb1..2d0e552633f 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -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 diff --git a/tests/unit/proxy/mcp_server/test_mcp_secret_headers.py b/tests/unit/proxy/mcp_server/test_mcp_secret_headers.py new file mode 100644 index 00000000000..3eec457b1b9 --- /dev/null +++ b/tests/unit/proxy/mcp_server/test_mcp_secret_headers.py @@ -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) == () diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index a7749800820..aa96a07d209 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -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; /**