diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index f65310a77e3..0286096e405 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1197,6 +1197,8 @@ def should_strip_caller_authorization( admission credential — a virtual key, an IdP JWT, or an SSO / OIDC / session token whose ``api_key`` is ``None`` — and must never reach the upstream, so it is stripped regardless of the ``api_key`` value. + - **Every other auth-enabled mode**: strip when ``Authorization`` is the + credential admission validated. """ if _authorization_is_admission_jwt(raw_headers, user_api_key_auth): return True @@ -1216,7 +1218,21 @@ def should_strip_caller_authorization( mcp_server.auth_type == MCPAuth.oauth2 and mcp_server.delegate_auth_to_upstream ) if not (mcp_server.is_oauth_passthrough or is_delegated_oauth): - return False + from litellm.proxy.auth.user_api_key_auth import is_no_auth_dev_mode + from litellm.proxy.proxy_server import ( + general_settings, # pyright: ignore[reportUnknownVariableType] # legacy global is declared as a bare dict + master_key, + ) + + proxy_general_settings: Final = _GENERAL_SETTINGS_ADAPTER.validate_python(general_settings) + admission_ran: Final = not is_no_auth_dev_mode(master_key, proxy_general_settings) or ( + user_api_key_auth is not None and user_api_key_auth.authenticated_by_custom_auth + ) + return admission_ran and ( + _authorization_is_litellm_admission_credential_with_settings( + raw_headers, user_api_key_auth, proxy_general_settings + ) + ) has_explicit_litellm_admission_header: Final = _has_explicit_litellm_admission_header(raw_headers) if is_delegated_oauth: @@ -1243,16 +1259,39 @@ def _has_explicit_litellm_admission_header(raw_headers: Mapping[str, str] | None return bool(_raw_header_value(raw_headers, "x-litellm-api-key")) +_GENERAL_SETTINGS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) + + def _authorization_is_litellm_admission_credential( raw_headers: Mapping[str, str] | None, user_api_key_auth: UserAPIKeyAuth | None, +) -> bool: + from litellm.proxy.proxy_server import ( + general_settings, # pyright: ignore[reportUnknownVariableType] # legacy global is declared as a bare dict + ) + + proxy_general_settings: Final = _GENERAL_SETTINGS_ADAPTER.validate_python(general_settings) + return _authorization_is_litellm_admission_credential_with_settings( + raw_headers, user_api_key_auth, proxy_general_settings + ) + + +def _authorization_is_litellm_admission_credential_with_settings( + raw_headers: Mapping[str, str] | None, + user_api_key_auth: UserAPIKeyAuth | None, + proxy_general_settings: Mapping[str, object], ) -> bool: """True when ``Authorization`` carries the LiteLLM key admission validated. - That is the case when no usable ``x-litellm-api-key`` was sent, or when the client repeated the - same key in both headers. + That is the case when no usable admission header was sent, or when the client repeated the same key in both headers. """ - admission_header: Final = _raw_header_value(raw_headers, "x-litellm-api-key") + custom_header_name: Final[object | None] = proxy_general_settings.get("litellm_key_header_name") + custom_admission_header: Final = ( + _raw_header_value(raw_headers, custom_header_name.lower()) + if isinstance(custom_header_name, str) and custom_header_name + else None + ) + admission_header: Final = custom_admission_header or _raw_header_value(raw_headers, "x-litellm-api-key") authorization: Final = _raw_header_value(raw_headers, "authorization") if admission_header and authorization: return strip_auth_scheme(authorization, "Bearer") == strip_auth_scheme(admission_header, "Bearer") diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 4896c2f1706..be70f595fee 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -317,6 +317,204 @@ def test_validated_admission_jwt_is_not_forwarded_as_an_extra_header(auth_type: assert _openapi_forwarded_extra_headers(server, raw, user) == expected +@pytest.mark.parametrize( + "auth_type", + [MCPAuth.none, MCPAuth.api_key, MCPAuth.bearer_token, MCPAuth.basic, MCPAuth.true_passthrough], +) +def test_admission_authorization_is_not_forwarded_for_any_auth_type( + auth_type: MCPAuth, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server.mcp_server_manager import _openapi_forwarded_extra_headers + from litellm.proxy._experimental.mcp_server.operations import _prepare_mcp_server_headers + + monkeypatch.setattr(proxy_server, "master_key", "test-master-key") + server: Final = MCPServer( + server_id="admission-header-boundary", + name="admission-header-boundary", + transport=MCPTransport.http, + auth_type=auth_type, + extra_headers=["Authorization", "x-tenant"], + ) + user: Final = UserAPIKeyAuth(api_key="hashed-key") + admitted_raw: Final = { + "Authorization": "Bearer sk-" + "1234", + "x-tenant": "tenant-control", + } + expected_without_admission_key: Final = {"x-tenant": "tenant-control"} + admission_upstream_headers: Final = _prepare_mcp_server_headers( + server, + mcp_auth_header=None, + mcp_server_auth_headers=None, + oauth2_headers={"Authorization": "Bearer sk-" + "1234"}, + raw_headers=admitted_raw, + user_api_key_auth=user, + )[1] + assert admission_upstream_headers == expected_without_admission_key + assert _openapi_forwarded_extra_headers(server, admitted_raw, user) == expected_without_admission_key + + separate_upstream_raw: Final = { + "x-litellm-api-key": "Bearer another-sk", + "Authorization": "Bearer upstream", + "x-tenant": "tenant-control", + } + expected_with_upstream_authorization: Final = { + "x-tenant": "tenant-control", + "Authorization": "Bearer upstream", + } + distinct_upstream_headers: Final = _prepare_mcp_server_headers( + server, + mcp_auth_header=None, + mcp_server_auth_headers=None, + oauth2_headers={"Authorization": "Bearer upstream"}, + raw_headers=separate_upstream_raw, + user_api_key_auth=user, + )[1] + assert distinct_upstream_headers == expected_with_upstream_authorization + assert _openapi_forwarded_extra_headers(server, separate_upstream_raw, user) == expected_with_upstream_authorization + + +@pytest.mark.parametrize("auth_type", [MCPAuth.none, MCPAuth.bearer_token, MCPAuth.true_passthrough]) +def test_no_master_key_preserves_caller_authorization_unless_auth_is_enabled( + auth_type: MCPAuth, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server.mcp_server_manager import _openapi_forwarded_extra_headers + from litellm.proxy._experimental.mcp_server.operations import _prepare_mcp_server_headers + + monkeypatch.setattr(proxy_server, "master_key", None) + monkeypatch.setattr(proxy_server, "general_settings", {}) + server: Final = MCPServer( + server_id="no-master-key", + name="no-master-key", + transport=MCPTransport.http, + auth_type=auth_type, + extra_headers=["Authorization", "x-tenant"], + ) + raw_headers: Final = {"Authorization": "Bearer upstream", "x-tenant": "t"} + user: Final = UserAPIKeyAuth(api_key="upstream", authenticated_by_custom_auth=False) + expected_forwarded: Final = {"x-tenant": "t", "Authorization": "Bearer upstream"} + forwarded_headers: Final = _prepare_mcp_server_headers( + server, + mcp_auth_header=None, + mcp_server_auth_headers=None, + oauth2_headers={"Authorization": "Bearer upstream"}, + raw_headers=raw_headers, + user_api_key_auth=user, + )[1] + assert forwarded_headers == expected_forwarded + assert _openapi_forwarded_extra_headers(server, raw_headers, user) == expected_forwarded + + monkeypatch.setattr(proxy_server, "general_settings", {"enable_oauth2_auth": True}) + auth_enabled_headers: Final = _prepare_mcp_server_headers( + server, + mcp_auth_header=None, + mcp_server_auth_headers=None, + oauth2_headers={"Authorization": "Bearer upstream"}, + raw_headers=raw_headers, + user_api_key_auth=user, + )[1] + assert auth_enabled_headers == {"x-tenant": "t"} + assert _openapi_forwarded_extra_headers(server, raw_headers, user) == {"x-tenant": "t"} + + +@pytest.mark.parametrize("auth_type", [MCPAuth.none, MCPAuth.bearer_token, MCPAuth.true_passthrough]) +def test_custom_auth_is_stripped_in_no_auth_dev_mode( + auth_type: MCPAuth, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server.mcp_server_manager import _openapi_forwarded_extra_headers + from litellm.proxy._experimental.mcp_server.operations import _prepare_mcp_server_headers + + monkeypatch.setattr(proxy_server, "master_key", None) + monkeypatch.setattr(proxy_server, "general_settings", {}) + server: Final = MCPServer( + server_id="custom-auth-no-master-key", + name="custom-auth-no-master-key", + transport=MCPTransport.http, + auth_type=auth_type, + extra_headers=["Authorization", "x-tenant"], + ) + raw_headers: Final = {"Authorization": "Bearer custom-tok", "x-tenant": "t"} + user: Final = UserAPIKeyAuth(api_key="custom-tok") + user.authenticated_by_custom_auth = True + expected_headers: Final = {"x-tenant": "t"} + forwarded_headers: Final = _prepare_mcp_server_headers( + server, + mcp_auth_header=None, + mcp_server_auth_headers=None, + oauth2_headers={"Authorization": "Bearer custom-tok"}, + raw_headers=raw_headers, + user_api_key_auth=user, + )[1] + + assert forwarded_headers == expected_headers + assert _openapi_forwarded_extra_headers(server, raw_headers, user) == expected_headers + + +@pytest.mark.parametrize("auth_type", [MCPAuth.none, MCPAuth.bearer_token, MCPAuth.true_passthrough]) +@pytest.mark.parametrize( + "raw_headers,oauth2_headers,expected", + [ + ( + { + "X-Company-Key": "sk-" + "1234", + "Authorization": "Bearer upstream", + "x-tenant": "t", + }, + {"Authorization": "Bearer upstream"}, + {"x-tenant": "t", "Authorization": "Bearer upstream"}, + ), + ( + { + "X-Company-Key": "sk-" + "1234", + "Authorization": "Bearer sk-" + "1234", + "x-tenant": "t", + }, + {"Authorization": "Bearer sk-" + "1234"}, + {"x-tenant": "t"}, + ), + ( + {"Authorization": "Bearer sk-" + "1234", "x-tenant": "t"}, + {"Authorization": "Bearer sk-" + "1234"}, + {"x-tenant": "t"}, + ), + ], +) +def test_custom_admission_header_controls_authorization_forwarding( + auth_type: MCPAuth, + raw_headers: dict[str, str], + oauth2_headers: dict[str, str], + expected: dict[str, str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server.mcp_server_manager import _openapi_forwarded_extra_headers + from litellm.proxy._experimental.mcp_server.operations import _prepare_mcp_server_headers + + monkeypatch.setattr(proxy_server, "master_key", "test-master-key") + monkeypatch.setattr(proxy_server, "general_settings", {"litellm_key_header_name": "x-company-key"}) + server: Final = MCPServer( + server_id="custom-admission-header", + name="custom-admission-header", + transport=MCPTransport.http, + auth_type=auth_type, + extra_headers=["Authorization", "x-tenant"], + ) + user: Final = UserAPIKeyAuth(api_key="hashed-key") + + upstream_headers: Final = _prepare_mcp_server_headers( + server, + mcp_auth_header=None, + mcp_server_auth_headers=None, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + user_api_key_auth=user, + )[1] + assert upstream_headers == expected + assert _openapi_forwarded_extra_headers(server, raw_headers, user) == expected + + class TestMCPServerManager: """Test MCP Server Manager stdio functionality""" @@ -3559,14 +3757,14 @@ class TestMCPServerManager: @pytest.mark.asyncio async def test_call_regular_mcp_tool_passthrough_strips_authorization_when_admission_consumed_litellm_key( self, - ): + ) -> None: """OAuth pass-through must not forward the caller's Authorization to upstream when LiteLLM admission consumed the bearer as its API key — otherwise the LiteLLM key the caller used for admission would leak upstream.""" from litellm.proxy._types import UserAPIKeyAuth - manager = MCPServerManager() - server = MCPServer( + manager: Final = MCPServerManager() + server: Final = MCPServer( server_id="server-passthrough-call", name="passthrough-server", url="https://example.com", @@ -3576,20 +3774,19 @@ class TestMCPServerManager: oauth_passthrough=True, ) - mock_client = AsyncMock() + mock_client: Final = AsyncMock() mock_client.call_tool = AsyncMock(return_value=CallToolResult(content=[], isError=False)) - captured_extra_headers = None + capture_extra_headers: Final = AsyncMock() async def capture_create_mcp_client( - server, - mcp_auth_header, - extra_headers, - stdio_env, - subject_token=None, - **kwargs, - ): # pragma: no cover - helper - nonlocal captured_extra_headers - captured_extra_headers = extra_headers + server: MCPServer, + mcp_auth_header: str | dict[str, str] | None, + extra_headers: dict[str, str] | None, + stdio_env: dict[str, str] | None, + subject_token: str | None = None, + **kwargs: object, + ) -> AsyncMock: + await capture_extra_headers(extra_headers) return mock_client manager.create_mcp_client = AsyncMock(side_effect=capture_create_mcp_client) @@ -3610,8 +3807,64 @@ class TestMCPServerManager: user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-key"), ) + assert capture_extra_headers.await_args is not None + captured_extra_headers: Final = capture_extra_headers.await_args.args[0] assert captured_extra_headers == {"x-request-id": "req-123"} + @pytest.mark.asyncio + async def test_call_regular_mcp_tool_never_forwards_admission_key_for_none_auth( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "master_key", "test-master-key") + manager: Final = MCPServerManager() + server: Final = MCPServer( + server_id="server-none-admission-call", + name="none-auth-server", + url="https://example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + extra_headers=["Authorization", "x-tenant"], + ) + + mock_client: Final = AsyncMock() + mock_client.call_tool = AsyncMock(return_value=CallToolResult(content=[], isError=False)) + capture_extra_headers: Final = AsyncMock() + + async def capture_create_mcp_client( + server: MCPServer, + mcp_auth_header: str | dict[str, str] | None, + extra_headers: dict[str, str] | None, + stdio_env: dict[str, str] | None, + subject_token: str | None = None, + **kwargs: object, + ) -> AsyncMock: + await capture_extra_headers(extra_headers) + return mock_client + + manager.create_mcp_client = AsyncMock(side_effect=capture_create_mcp_client) + + await manager._call_regular_mcp_tool( + mcp_server=server, + original_tool_name="tool", + arguments={}, + tasks=[], + mcp_auth_header=None, + mcp_server_auth_headers=None, + oauth2_headers=None, + raw_headers={ + "Authorization": "Bearer sk-" + "1234", + "x-tenant": "tenant-control", + }, + proxy_logging_obj=None, + user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key"), + ) + + assert capture_extra_headers.await_args is not None + captured_extra_headers: Final = capture_extra_headers.await_args.args[0] + assert captured_extra_headers == {"x-tenant": "tenant-control"} + @pytest.mark.asyncio async def test_call_regular_mcp_tool_v2_authz_code_drops_caller_authorization( self, diff --git a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py index b9cf201a47a..60d4d2e89b5 100644 --- a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py +++ b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py @@ -1065,6 +1065,108 @@ async def test_mcp_tool_execution_events_creation(): print("MCP tool execution events creation test passed!") +@pytest.mark.asyncio +async def test_aresponses_mcp_call_does_not_forward_admission_authorization( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import importlib + import sys + import types + from unittest.mock import AsyncMock + + from mcp.types import CallToolResult, TextContent, Tool as MCPTool + + from litellm.proxy._experimental.mcp_server import mcp_server_manager + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.mcp import MCPAuth, MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + responses_main: Final = importlib.import_module("litellm.responses.main") + manager: Final = mcp_server_manager.global_mcp_server_manager + server: Final = MCPServer( + server_id="responses-auth-strip", + name="echo_none", + server_name="echo_none", + url="http://mcp.example", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + extra_headers=["Authorization", "x-tenant"], + allow_all_keys=True, + ) + monkeypatch.setattr(manager, "registry", {server.server_id: server}) + + mcp_client: Final = AsyncMock() + mcp_client.list_tools.return_value = [MCPTool(name="whoami", inputSchema={"type": "object", "properties": {}})] + mcp_client.call_tool.return_value = CallToolResult( + content=[TextContent(type="text", text="ok")], + isError=False, + ) + monkeypatch.setattr(manager, "create_mcp_client", AsyncMock(return_value=mcp_client)) + monkeypatch.setitem( + sys.modules, + "litellm.proxy.proxy_server", + types.SimpleNamespace( + general_settings={}, + master_key="test-master-key", + proxy_logging_obj=None, + prisma_client=None, + user_api_key_cache=None, + ), + ) + + raw_headers: Final = {"authorization": "Bearer sk-" + "1234", "x-tenant": "tenant-a"} + user_api_key_auth: Final = UserAPIKeyAuth(api_key="hashed-key", token="hashed-key") + tools: Final = [ + { + "type": "mcp", + "server_label": "echo_none", + "server_url": "litellm_proxy/mcp/echo_none", + "require_approval": "never", + } + ] + + monkeypatch.setattr( + responses_main, + "aresponses", + AsyncMock( + side_effect=[ + ResponsesAPIResponse( + id="resp_first", + created_at=0, + output=[ + { + "type": "function_call", + "id": "fc_1", + "call_id": "call-1", + "name": "echo_none-whoami", + "arguments": "{}", + "status": "completed", + } + ], + ), + ResponsesAPIResponse(id="resp_final", created_at=0, output=[]), + ] + ), + ) + monkeypatch.setattr( + importlib.import_module("litellm.responses.mcp.litellm_proxy_mcp_handler"), + "aresponses", + AsyncMock(return_value=ResponsesAPIResponse(id="resp_final", created_at=0, output=[])), + ) + + await responses_main.aresponses_api_with_mcp( + input="Call whoami", + model="openai/gpt-5", + tools=tools, + litellm_metadata={"user_api_key_auth": user_api_key_auth}, + secret_fields={"raw_headers": raw_headers}, + ) + + assert mcp_client.call_tool.await_args is not None + assert manager.create_mcp_client.await_args is not None + assert manager.create_mcp_client.await_args.kwargs["extra_headers"] == {"x-tenant": "tenant-a"} + + @pytest.mark.asyncio async def test_no_duplicate_mcp_tools_in_streaming_e2e(): """