From 205a5e9d6cfe878fb489f360908fc45efcafe4a1 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 4 Sep 2026 12:45:24 -0700 Subject: [PATCH] feat(mcp): use x-mcp--* headers as default upstream credentials for group members (#39717) Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/mcp_server_manager.py | 3 + .../mcp_server/rest_endpoints.py | 7 +- .../proxy/_experimental/mcp_server/server.py | 6 +- .../proxy/_experimental/mcp_server/utils.py | 49 ++++++++++---- .../mcp_server/test_mcp_header_alias_utils.py | 64 +++++++++++++++++++ .../mcp_server/test_mcp_server.py | 37 +++++++++++ .../mcp_server/test_rest_endpoints.py | 37 ++++++++++- 7 files changed, 186 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index ce4928ff83d..bfc5f629faf 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -928,6 +928,7 @@ def _resolve_openapi_tool_auth( mcp_server_auth_headers, alias=mcp_server.alias, server_name=mcp_server.server_name, + access_groups=mcp_server.access_groups, ) if mcp_server_auth_headers else None @@ -3296,6 +3297,7 @@ class MCPServerManager: mcp_server_auth_headers, alias=server.alias, server_name=server.server_name, + access_groups=server.access_groups, ) # Fall back to deprecated mcp_auth_header if no server-specific header found @@ -5373,6 +5375,7 @@ class MCPServerManager: mcp_server_auth_headers, alias=mcp_server.alias, server_name=mcp_server.server_name, + access_groups=mcp_server.access_groups, ) # Fall back to deprecated mcp_auth_header if no server-specific header found diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 37474f85fe7..b3469da9071 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -257,7 +257,7 @@ if MCP_AVAILABLE: ) def _get_server_auth_header( - server, + server: MCPServer, mcp_server_auth_headers: dict[str, dict[str, str]] | None, mcp_auth_header: str | None, ) -> dict[str, str] | str | None: @@ -269,8 +269,9 @@ if MCP_AVAILABLE: if mcp_server_auth_headers: server_auth: Final = lookup_mcp_server_auth_in_headers( mcp_server_auth_headers, - alias=getattr(server, "alias", None), - server_name=getattr(server, "server_name", None), + alias=server.alias, + server_name=server.server_name, + access_groups=server.access_groups, ) if server_auth is not None: return server_auth diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 3d7c947a913..975d9642b36 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -1612,7 +1612,10 @@ if MCP_AVAILABLE: ) server_headers: Final = lookup_mcp_server_auth_in_headers( - mcp_server_auth_headers, alias=server.alias, server_name=server.server_name + mcp_server_auth_headers, + alias=server.alias, + server_name=server.server_name, + access_groups=server.access_groups, ) if isinstance(server_headers, str): return bool(server_headers.strip()) @@ -1712,6 +1715,7 @@ if MCP_AVAILABLE: mcp_server_auth_headers, alias=server.alias, server_name=server.server_name, + access_groups=server.access_groups, ) extra_headers: dict[str, str] | None = None diff --git a/litellm/proxy/_experimental/mcp_server/utils.py b/litellm/proxy/_experimental/mcp_server/utils.py index 83883664df5..252756e0458 100644 --- a/litellm/proxy/_experimental/mcp_server/utils.py +++ b/litellm/proxy/_experimental/mcp_server/utils.py @@ -8,11 +8,12 @@ import json import os import re import typing -from collections.abc import Iterable, Iterator, Mapping, MutableMapping, MutableSequence +from collections.abc import Iterable, Iterator, Mapping, MutableMapping, MutableSequence, Sequence from collections.abc import Set as AbstractSet from typing import Any, Final, Protocol from urllib.parse import quote +from litellm._logging import verbose_logger from litellm.types.mcp_server.mcp_server_manager import MCPServer if typing.TYPE_CHECKING: @@ -169,34 +170,58 @@ def sanitize_mcp_alias_for_header(alias: str) -> str: return sanitized.strip("_") +def _header_keys_for_identifier(identifier: str) -> tuple[str, ...]: + lowered: Final = identifier.lower() + sanitized: Final = sanitize_mcp_alias_for_header(identifier) + return (lowered,) if not sanitized or sanitized == lowered else (lowered, sanitized) + + +def _matching_header_key(normalized_headers: Mapping[str, object], identifier: str) -> str | None: + return next((key for key in _header_keys_for_identifier(identifier) if key in normalized_headers), None) + + def lookup_mcp_server_auth_in_headers( mcp_server_auth_headers: Mapping[str, str | dict[str, str]], *, alias: str | None = None, server_name: str | None = None, + access_groups: Sequence[str] | None = None, ) -> str | dict[str, str] | None: """ Resolve server-specific auth headers with case-insensitive matching. Tries the raw alias/server_name (lowercased) and the header-safe sanitized alias so dashboard clients using sanitize_mcp_alias_for_header() still match. + + When no server-level header matches, an ``x-mcp-{access_group}-*`` header is + used as the default for every server in that group. If the server belongs to + several groups that each carry a different credential, nothing is returned so + a token is never forwarded to a server it may not have been meant for. """ if not mcp_server_auth_headers: return None normalized_headers: Final = {k.lower(): v for k, v in mcp_server_auth_headers.items()} - for identifier in (alias, server_name): - if not identifier: - continue - keys_to_try = [identifier.lower()] - sanitized = sanitize_mcp_alias_for_header(identifier) - if sanitized and sanitized not in keys_to_try: - keys_to_try.append(sanitized) - for key in keys_to_try: - if key in normalized_headers: - return normalized_headers[key] - return None + server_keys: Final = ( + _matching_header_key(normalized_headers, identifier) for identifier in (alias, server_name) if identifier + ) + server_key: Final = next((key for key in server_keys if key is not None), None) + if server_key is not None: + return normalized_headers[server_key] + + group_keys: Final = (_matching_header_key(normalized_headers, group) for group in access_groups or ()) + group_matches: Final = tuple(normalized_headers[key] for key in group_keys if key is not None) + if not group_matches: + return None + if any(match != group_matches[0] for match in group_matches[1:]): + verbose_logger.debug( + "Ambiguous MCP group auth headers for server alias=%s (groups=%s); not forwarding any group credential", + alias, + access_groups, + ) + return None + return group_matches[0] MCP_TOOL_ALLOWLIST_ENFORCED_KEY: Final = "tool_allowlist_enforced" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_header_alias_utils.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_header_alias_utils.py index 2627199570b..6c24205c258 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_header_alias_utils.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_header_alias_utils.py @@ -16,3 +16,67 @@ def test_lookup_mcp_server_auth_in_headers_sanitized_alias(): headers = {"github_mcp": {"Authorization": "Bearer token"}} result = lookup_mcp_server_auth_in_headers(headers, alias="GitHub-MCP") assert result == {"Authorization": "Bearer token"} + + +def test_lookup_mcp_server_auth_in_headers_group_header_is_default_for_members(): + headers = {"shared": {"Authorization": "Bearer group-token"}} + assert lookup_mcp_server_auth_in_headers(headers, alias="alpha", server_name="alpha", access_groups=["shared"]) == { + "Authorization": "Bearer group-token" + } + assert lookup_mcp_server_auth_in_headers(headers, alias="beta", server_name="beta", access_groups=["Shared"]) == { + "Authorization": "Bearer group-token" + } + + +def test_lookup_mcp_server_auth_in_headers_group_header_sanitized_group_name(): + headers = {"dev_group": {"Authorization": "Bearer group-token"}} + assert lookup_mcp_server_auth_in_headers(headers, alias="alpha", access_groups=["Dev Group"]) == { + "Authorization": "Bearer group-token" + } + + +def test_lookup_mcp_server_auth_in_headers_server_header_overrides_group_header(): + headers = { + "shared": {"Authorization": "Bearer group-token"}, + "beta": {"Authorization": "Bearer beta-token"}, + } + assert lookup_mcp_server_auth_in_headers(headers, alias="beta", server_name="beta", access_groups=["shared"]) == { + "Authorization": "Bearer beta-token" + } + + +def test_lookup_mcp_server_auth_in_headers_group_header_not_forwarded_outside_group(): + headers = {"shared": {"Authorization": "Bearer group-token"}} + assert ( + lookup_mcp_server_auth_in_headers(headers, alias="gamma", server_name="gamma", access_groups=["other"]) is None + ) + assert lookup_mcp_server_auth_in_headers(headers, alias="gamma", server_name="gamma", access_groups=None) is None + + +def test_lookup_mcp_server_auth_in_headers_alias_colliding_with_group_name_keeps_server_level_match(): + headers = {"shared": {"Authorization": "Bearer shared-token"}} + assert lookup_mcp_server_auth_in_headers(headers, alias="shared", access_groups=["other"]) == { + "Authorization": "Bearer shared-token" + } + assert lookup_mcp_server_auth_in_headers(headers, alias="alpha", access_groups=["shared"]) == { + "Authorization": "Bearer shared-token" + } + assert lookup_mcp_server_auth_in_headers(headers, alias="gamma", access_groups=["other"]) is None + + +def test_lookup_mcp_server_auth_in_headers_conflicting_group_headers_fail_closed(): + headers = { + "shared": {"Authorization": "Bearer group-token"}, + "other": {"Authorization": "Bearer other-token"}, + } + assert lookup_mcp_server_auth_in_headers(headers, alias="delta", access_groups=["shared", "other"]) is None + + +def test_lookup_mcp_server_auth_in_headers_identical_group_headers_resolve(): + headers = { + "shared": {"Authorization": "Bearer group-token"}, + "other": {"Authorization": "Bearer group-token"}, + } + assert lookup_mcp_server_auth_in_headers(headers, alias="delta", access_groups=["shared", "other"]) == { + "Authorization": "Bearer group-token" + } diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 9a6815a61e5..086ab854e36 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -303,6 +303,43 @@ def test_prepare_mcp_server_headers_case_insensitive_extra_headers(): assert extra_headers == {"Authorization": "Bearer token"} +def test_prepare_mcp_server_headers_group_header_defaults_for_members_only(): + try: + from litellm.proxy._experimental.mcp_server.server import ( + _prepare_mcp_server_headers, + ) + except ImportError: + pytest.skip("MCP server not available") + + def server(alias: str, group: str) -> MCPServer: + return MCPServer( + server_id=f"server-{alias}", + name=alias, + alias=alias, + transport=MCPTransport.http, + access_groups=[group], + ) + + mcp_server_auth_headers = { + "shared": {"Authorization": "Bearer group-token"}, + "beta": {"Authorization": "Bearer beta-token"}, + } + + def resolve(mcp_server: MCPServer): + server_auth_header, _ = _prepare_mcp_server_headers( + server=mcp_server, + mcp_server_auth_headers=mcp_server_auth_headers, + mcp_auth_header=None, + oauth2_headers=None, + raw_headers={"x-litellm-api-key": "Bearer sk-litellm-key"}, + ) + return server_auth_header + + assert resolve(server("alpha", "shared")) == {"Authorization": "Bearer group-token"} + assert resolve(server("beta", "shared")) == {"Authorization": "Bearer beta-token"} + assert resolve(server("gamma", "other")) is None + + def test_prepare_mcp_server_headers_passthrough_strips_authorization_without_admission_header(): try: from litellm.proxy._experimental.mcp_server.server import ( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index d441c05090b..f2c8f8c80c5 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -25,7 +25,8 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.types.mcp import MCPAuth +from litellm.types.mcp import MCPAuth, MCPTransport +from litellm.types.mcp_server.mcp_server_manager import MCPServer def _rendered_log_message(call): @@ -3268,6 +3269,40 @@ class TestConnectionErrorMessage: assert "proxy logs" in message.lower() +class TestGetServerAuthHeaderGroupDefault: + """``x-mcp--authorization`` is the default for group members, the per-server + header still wins, and servers outside the group never see the group credential.""" + + @staticmethod + def _server(alias: str, group: str) -> MCPServer: + return MCPServer( + server_id=f"server-{alias}", + name=alias, + server_name=alias, + alias=alias, + url="https://example.com/mcp", + transport=MCPTransport.http, + access_groups=[group], + ) + + def test_group_header_applies_to_members_and_per_server_header_overrides(self): + headers = { + "shared": {"Authorization": "Bearer group-token"}, + "beta": {"Authorization": "Bearer beta-token"}, + } + assert rest_endpoints._get_server_auth_header(self._server("alpha", "shared"), headers, None) == { + "Authorization": "Bearer group-token" + } + assert rest_endpoints._get_server_auth_header(self._server("beta", "shared"), headers, None) == { + "Authorization": "Bearer beta-token" + } + + def test_group_header_falls_back_to_legacy_header_outside_group(self): + headers = {"shared": {"Authorization": "Bearer group-token"}} + assert rest_endpoints._get_server_auth_header(self._server("gamma", "other"), headers, None) is None + assert rest_endpoints._get_server_auth_header(self._server("gamma", "other"), headers, "legacy") == "legacy" + + class TestToolResponseMcpInfoEnrichment: """The REST tools/list response must expose the user-facing alias and the server_id alongside the internal server_name so clients (agent builder UIs)