From 89a43cfd60777739a4d0d3d43ba1ac5397b2eb80 Mon Sep 17 00:00:00 2001 From: Noah Nistler <60981020+noahnistler@users.noreply.github.com> Date: Tue, 17 Mar 2026 14:31:14 -0500 Subject: [PATCH 01/13] Allow pre_mcp_call guardrail hooks to mutate outbound MCP headers --- .../mcp_server/mcp_server_manager.py | 23 +- litellm/proxy/_types.py | 1 + litellm/proxy/auth/user_api_key_auth.py | 3 + litellm/proxy/utils.py | 9 +- .../mcp_server/test_mcp_hook_extra_headers.py | 481 ++++++++++++++++++ 5 files changed, 513 insertions(+), 4 deletions(-) create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 43fe54fdfb7..b552669651f 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1908,7 +1908,13 @@ class MCPServerManager: user_api_key_auth: Optional[UserAPIKeyAuth], proxy_logging_obj: ProxyLogging, server: MCPServer, - ): + ) -> Dict[str, Any]: + """ + Run pre-call checks and guardrail hooks for an MCP tool call. + + Returns a dict that may contain: + - "extra_headers": headers injected by pre_mcp_call guardrail hooks + """ ## check if the tool is allowed or banned for the given server if not self.check_allowed_or_banned_tools(name, server): raise HTTPException( @@ -1969,6 +1975,7 @@ class MCPServerManager: mcp_request_obj, pre_hook_kwargs ) + hook_result: Dict[str, Any] = {} try: # Use standard pre_call_hook modified_data = await proxy_logging_obj.pre_call_hook( @@ -1985,6 +1992,8 @@ class MCPServerManager: ) if modified_kwargs.get("arguments") != arguments: arguments = modified_kwargs["arguments"] + if modified_kwargs.get("extra_headers"): + hook_result["extra_headers"] = modified_kwargs["extra_headers"] except ( BlockedPiiEntityError, @@ -1995,6 +2004,8 @@ class MCPServerManager: verbose_logger.error(f"Guardrail blocked MCP tool call pre call: {str(e)}") raise e + return hook_result + def _create_during_hook_task( self, name: str, @@ -2047,6 +2058,7 @@ class MCPServerManager: raw_headers: Optional[Dict[str, str]], proxy_logging_obj: Optional[ProxyLogging], host_progress_callback: Optional[Callable] = None, + hook_extra_headers: Optional[Dict[str, str]] = None, ) -> CallToolResult: """ Call a regular MCP tool using the MCP client. @@ -2116,6 +2128,11 @@ class MCPServerManager: extra_headers = {} extra_headers.update(mcp_server.static_headers) + if hook_extra_headers: + if extra_headers is None: + extra_headers = {} + extra_headers.update(hook_extra_headers) + stdio_env = self._build_stdio_env(mcp_server, raw_headers) client = await self._create_mcp_client( @@ -2201,8 +2218,9 @@ class MCPServerManager: # Allow validation and modification of tool calls before execution # Using standard pre_call_hook ######################################################### + hook_result: Dict[str, Any] = {} if proxy_logging_obj: - await self.pre_call_tool_check( + hook_result = await self.pre_call_tool_check( name=name, arguments=arguments, server_name=server_name, @@ -2247,6 +2265,7 @@ class MCPServerManager: raw_headers=raw_headers, proxy_logging_obj=proxy_logging_obj, host_progress_callback=host_progress_callback, + hook_extra_headers=hook_result.get("extra_headers"), ) # For OpenAPI tools, await outside the client context diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index ecbd7314cd7..55cb1c9c43c 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2471,6 +2471,7 @@ class UserAPIKeyAuth( Any ] = None # Expanded created_by user when expand=user is used end_user_object_permission: Optional[LiteLLM_ObjectPermissionTable] = None + jwt_claims: Optional[Dict] = None model_config = ConfigDict(arbitrary_types_allowed=True) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 376048e7a13..5b6064ef550 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -729,6 +729,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 team_membership: Optional[LiteLLM_TeamMembership] = result.get( "team_membership", None ) + jwt_claims: Optional[dict] = result.get("jwt_claims", None) global_proxy_spend = await get_global_proxy_spend( litellm_proxy_admin_name=litellm_proxy_admin_name, @@ -757,6 +758,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 org_id=org_id, end_user_id=end_user_id, parent_otel_span=parent_otel_span, + jwt_claims=jwt_claims, ) valid_token = UserAPIKeyAuth( @@ -803,6 +805,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 team_metadata=( team_object.metadata if team_object is not None else None ), + jwt_claims=jwt_claims, ) # Check if model has zero cost - if so, skip all budget checks diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 01a0f55aac7..8536704766b 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -824,17 +824,22 @@ class ProxyLogging: ) -> dict: """ Helper function to convert pre_call_hook response back to kwargs for MCP usage. + + Supports: + - modified_arguments: Override tool call arguments + - extra_headers: Inject custom headers into the outbound MCP request """ if not response_data: return original_kwargs - # Apply any argument modifications from the hook response modified_kwargs = original_kwargs.copy() - # If the response contains modified arguments, apply them if response_data.get("modified_arguments"): modified_kwargs["arguments"] = response_data["modified_arguments"] + if response_data.get("extra_headers"): + modified_kwargs["extra_headers"] = response_data["extra_headers"] + return modified_kwargs async def process_pre_call_hook_response(self, response, data, call_type): diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py new file mode 100644 index 00000000000..6296c57f09c --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py @@ -0,0 +1,481 @@ +""" +Tests for pre_mcp_call guardrail hook header mutation support. + +Validates that: +1. _convert_mcp_hook_response_to_kwargs extracts extra_headers from hook response +2. pre_call_tool_check returns hook-provided extra_headers +3. call_tool flows hook headers into _call_regular_mcp_tool +4. Hook-provided headers take highest priority (merge after static_headers) +5. Backward compatibility: hooks without extra_headers continue to work +""" + +import asyncio +import sys +from datetime import datetime +from typing import Any, Dict, Optional +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import HTTPException + +from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.utils import ProxyLogging +from litellm.types.mcp import MCPAuth, MCPTransport +from litellm.types.mcp_server.mcp_server_manager import MCPServer + + +class TestConvertMcpHookResponseToKwargs: + """Tests for ProxyLogging._convert_mcp_hook_response_to_kwargs""" + + def setup_method(self): + self.proxy_logging = ProxyLogging(user_api_key_cache=MagicMock()) + + def test_returns_original_kwargs_when_response_is_none(self): + original = {"arguments": {"key": "val"}, "name": "tool"} + result = self.proxy_logging._convert_mcp_hook_response_to_kwargs( + None, original + ) + assert result == original + + def test_returns_original_kwargs_when_response_is_empty_dict(self): + original = {"arguments": {"key": "val"}} + result = self.proxy_logging._convert_mcp_hook_response_to_kwargs({}, original) + assert result == original + + def test_extracts_modified_arguments(self): + original = {"arguments": {"old": "value"}} + response = {"modified_arguments": {"new": "value"}} + result = self.proxy_logging._convert_mcp_hook_response_to_kwargs( + response, original + ) + assert result["arguments"] == {"new": "value"} + + def test_extracts_extra_headers(self): + original = {"arguments": {"key": "val"}} + response = {"extra_headers": {"Authorization": "Bearer signed-jwt"}} + result = self.proxy_logging._convert_mcp_hook_response_to_kwargs( + response, original + ) + assert result["extra_headers"] == {"Authorization": "Bearer signed-jwt"} + + def test_extracts_both_arguments_and_headers(self): + original = {"arguments": {"old": "value"}} + response = { + "modified_arguments": {"new": "value"}, + "extra_headers": {"X-Custom": "header-val"}, + } + result = self.proxy_logging._convert_mcp_hook_response_to_kwargs( + response, original + ) + assert result["arguments"] == {"new": "value"} + assert result["extra_headers"] == {"X-Custom": "header-val"} + + def test_no_extra_headers_key_preserves_original(self): + """Backward compat: hooks that only return modified_arguments still work.""" + original = {"arguments": {"key": "val"}} + response = {"modified_arguments": {"key": "new_val"}} + result = self.proxy_logging._convert_mcp_hook_response_to_kwargs( + response, original + ) + assert "extra_headers" not in result + assert result["arguments"] == {"key": "new_val"} + + def test_empty_extra_headers_not_set(self): + """Empty dict for extra_headers is falsy and should not be set.""" + original = {"arguments": {"key": "val"}} + response = {"extra_headers": {}} + result = self.proxy_logging._convert_mcp_hook_response_to_kwargs( + response, original + ) + assert "extra_headers" not in result + + +class TestPreCallToolCheckReturnsHeaders: + """Tests that pre_call_tool_check returns hook-provided headers.""" + + def _make_server(self, name="test_server"): + return MCPServer( + server_id="test-id", + name=name, + server_name=name, + url="https://example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + ) + + @pytest.mark.asyncio + async def test_returns_empty_dict_when_hook_has_no_headers(self): + manager = MCPServerManager() + server = self._make_server() + + proxy_logging = MagicMock(spec=ProxyLogging) + proxy_logging._create_mcp_request_object_from_kwargs = MagicMock( + return_value=MagicMock() + ) + proxy_logging._convert_mcp_to_llm_format = MagicMock( + return_value={"model": "fake"} + ) + proxy_logging.pre_call_hook = AsyncMock( + return_value={"modified_arguments": {"key": "val"}} + ) + proxy_logging._convert_mcp_hook_response_to_kwargs = MagicMock( + return_value={"arguments": {"key": "val"}} + ) + + with patch.object(manager, "check_allowed_or_banned_tools", return_value=True): + with patch.object( + manager, + "check_tool_permission_for_key_team", + new_callable=AsyncMock, + ): + with patch.object(manager, "validate_allowed_params"): + result = await manager.pre_call_tool_check( + name="test_tool", + arguments={"key": "val"}, + server_name="test_server", + user_api_key_auth=None, + proxy_logging_obj=proxy_logging, + server=server, + ) + + assert result == {} + + @pytest.mark.asyncio + async def test_returns_extra_headers_from_hook(self): + manager = MCPServerManager() + server = self._make_server() + + hook_headers = {"Authorization": "Bearer signed-jwt", "X-Trace-Id": "abc123"} + + proxy_logging = MagicMock(spec=ProxyLogging) + proxy_logging._create_mcp_request_object_from_kwargs = MagicMock( + return_value=MagicMock() + ) + proxy_logging._convert_mcp_to_llm_format = MagicMock( + return_value={"model": "fake"} + ) + proxy_logging.pre_call_hook = AsyncMock( + return_value={"extra_headers": hook_headers} + ) + proxy_logging._convert_mcp_hook_response_to_kwargs = MagicMock( + return_value={"arguments": {"key": "val"}, "extra_headers": hook_headers} + ) + + with patch.object(manager, "check_allowed_or_banned_tools", return_value=True): + with patch.object( + manager, + "check_tool_permission_for_key_team", + new_callable=AsyncMock, + ): + with patch.object(manager, "validate_allowed_params"): + result = await manager.pre_call_tool_check( + name="test_tool", + arguments={"key": "val"}, + server_name="test_server", + user_api_key_auth=None, + proxy_logging_obj=proxy_logging, + server=server, + ) + + assert result["extra_headers"] == hook_headers + + @pytest.mark.asyncio + async def test_returns_empty_dict_when_hook_returns_none(self): + manager = MCPServerManager() + server = self._make_server() + + proxy_logging = MagicMock(spec=ProxyLogging) + proxy_logging._create_mcp_request_object_from_kwargs = MagicMock( + return_value=MagicMock() + ) + proxy_logging._convert_mcp_to_llm_format = MagicMock( + return_value={"model": "fake"} + ) + proxy_logging.pre_call_hook = AsyncMock(return_value=None) + + with patch.object(manager, "check_allowed_or_banned_tools", return_value=True): + with patch.object( + manager, + "check_tool_permission_for_key_team", + new_callable=AsyncMock, + ): + with patch.object(manager, "validate_allowed_params"): + result = await manager.pre_call_tool_check( + name="test_tool", + arguments={"key": "val"}, + server_name="test_server", + user_api_key_auth=None, + proxy_logging_obj=proxy_logging, + server=server, + ) + + assert result == {} + + +class TestCallToolFlowsHookHeaders: + """Tests that call_tool passes hook_extra_headers to _call_regular_mcp_tool.""" + + def _make_server(self, name="test_server"): + return MCPServer( + server_id="test-id", + name=name, + server_name=name, + url="https://example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + ) + + @pytest.mark.asyncio + async def test_hook_headers_passed_to_call_regular_mcp_tool(self): + """Verify that hook_extra_headers kwarg is forwarded.""" + manager = MCPServerManager() + server = self._make_server() + + hook_headers = {"Authorization": "Bearer signed-jwt"} + + with patch.object( + manager, + "_get_mcp_server_from_tool_name", + return_value=server, + ): + with patch.object( + manager, + "pre_call_tool_check", + new_callable=AsyncMock, + return_value={"extra_headers": hook_headers}, + ): + with patch.object( + manager, + "_create_during_hook_task", + return_value=asyncio.create_task(asyncio.sleep(0)), + ): + with patch.object( + manager, + "_call_regular_mcp_tool", + new_callable=AsyncMock, + return_value=MagicMock(), + ) as mock_call: + proxy_logging = MagicMock(spec=ProxyLogging) + + await manager.call_tool( + server_name="test_server", + name="test_tool", + arguments={"key": "val"}, + proxy_logging_obj=proxy_logging, + ) + + mock_call.assert_called_once() + call_kwargs = mock_call.call_args + assert call_kwargs.kwargs.get("hook_extra_headers") == hook_headers + + @pytest.mark.asyncio + async def test_no_hook_headers_when_no_proxy_logging(self): + """Without proxy_logging_obj, no pre_call_tool_check runs.""" + manager = MCPServerManager() + server = self._make_server() + + with patch.object( + manager, + "_get_mcp_server_from_tool_name", + return_value=server, + ): + with patch.object( + manager, + "_call_regular_mcp_tool", + new_callable=AsyncMock, + return_value=MagicMock(), + ) as mock_call: + await manager.call_tool( + server_name="test_server", + name="test_tool", + arguments={"key": "val"}, + proxy_logging_obj=None, + ) + + mock_call.assert_called_once() + call_kwargs = mock_call.call_args + assert call_kwargs.kwargs.get("hook_extra_headers") is None + + +class TestHookHeaderMergePriority: + """Tests that hook-provided headers have highest priority in _call_regular_mcp_tool.""" + + def _make_server( + self, + static_headers: Optional[Dict[str, str]] = None, + extra_headers_config: Optional[list] = None, + ): + return MCPServer( + server_id="test-id", + name="Test Server", + server_name="test_server", + url="https://example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + static_headers=static_headers, + extra_headers=extra_headers_config, + ) + + @pytest.mark.asyncio + async def test_hook_headers_override_static_headers(self): + """Hook headers should take precedence over static_headers.""" + manager = MCPServerManager() + server = self._make_server( + static_headers={"Authorization": "Bearer static-token", "X-Static": "yes"} + ) + + hook_headers = {"Authorization": "Bearer hook-signed-jwt"} + + captured_extra_headers: Dict[str, Any] = {} + + async def fake_create_mcp_client( + server, mcp_auth_header=None, extra_headers=None, stdio_env=None + ): + captured_extra_headers["value"] = extra_headers + mock_client = MagicMock() + mock_client.call_tool = AsyncMock(return_value=MagicMock()) + return mock_client + + with patch.object( + manager, "_create_mcp_client", side_effect=fake_create_mcp_client + ): + with patch.object(manager, "_build_stdio_env", return_value=None): + try: + await manager._call_regular_mcp_tool( + mcp_server=server, + original_tool_name="test_tool", + arguments={"key": "val"}, + tasks=[], + mcp_auth_header=None, + mcp_server_auth_headers=None, + oauth2_headers=None, + raw_headers=None, + proxy_logging_obj=None, + hook_extra_headers=hook_headers, + ) + except Exception: + pass + + headers = captured_extra_headers.get("value", {}) + assert headers["Authorization"] == "Bearer hook-signed-jwt" + assert headers["X-Static"] == "yes" + + @pytest.mark.asyncio + async def test_no_hook_headers_preserves_existing_behavior(self): + """When hook_extra_headers is None, existing header logic is unchanged.""" + manager = MCPServerManager() + server = self._make_server( + static_headers={"X-Static": "static-value"} + ) + + captured_extra_headers: Dict[str, Any] = {} + + async def fake_create_mcp_client( + server, mcp_auth_header=None, extra_headers=None, stdio_env=None + ): + captured_extra_headers["value"] = extra_headers + mock_client = MagicMock() + mock_client.call_tool = AsyncMock(return_value=MagicMock()) + return mock_client + + with patch.object( + manager, "_create_mcp_client", side_effect=fake_create_mcp_client + ): + with patch.object(manager, "_build_stdio_env", return_value=None): + try: + await manager._call_regular_mcp_tool( + mcp_server=server, + original_tool_name="test_tool", + arguments={"key": "val"}, + tasks=[], + mcp_auth_header=None, + mcp_server_auth_headers=None, + oauth2_headers=None, + raw_headers=None, + proxy_logging_obj=None, + hook_extra_headers=None, + ) + except Exception: + pass + + headers = captured_extra_headers.get("value", {}) + assert headers == {"X-Static": "static-value"} + + @pytest.mark.asyncio + async def test_hook_headers_merge_with_oauth2(self): + """Hook headers merge on top of OAuth2 headers.""" + manager = MCPServerManager() + server = MCPServer( + server_id="test-id", + name="Test Server", + server_name="test_server", + url="https://example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + ) + + captured_extra_headers: Dict[str, Any] = {} + + async def fake_create_mcp_client( + server, mcp_auth_header=None, extra_headers=None, stdio_env=None + ): + captured_extra_headers["value"] = extra_headers + mock_client = MagicMock() + mock_client.call_tool = AsyncMock(return_value=MagicMock()) + return mock_client + + with patch.object( + manager, "_create_mcp_client", side_effect=fake_create_mcp_client + ): + with patch.object(manager, "_build_stdio_env", return_value=None): + try: + await manager._call_regular_mcp_tool( + mcp_server=server, + original_tool_name="test_tool", + arguments={"key": "val"}, + tasks=[], + mcp_auth_header=None, + mcp_server_auth_headers=None, + oauth2_headers={ + "Authorization": "Bearer oauth2-token", + "X-OAuth": "yes", + }, + raw_headers=None, + proxy_logging_obj=None, + hook_extra_headers={ + "Authorization": "Bearer hook-jwt", + "X-Trace-Id": "trace-123", + }, + ) + except Exception: + pass + + headers = captured_extra_headers.get("value", {}) + assert headers["Authorization"] == "Bearer hook-jwt" + assert headers["X-OAuth"] == "yes" + assert headers["X-Trace-Id"] == "trace-123" + + +class TestUserAPIKeyAuthJwtClaims: + """Tests that UserAPIKeyAuth correctly carries jwt_claims.""" + + def test_jwt_claims_field_defaults_to_none(self): + auth = UserAPIKeyAuth(api_key="test-key") + assert auth.jwt_claims is None + + def test_jwt_claims_field_accepts_dict(self): + claims = {"sub": "user-123", "iss": "litellm", "exp": 9999999999} + auth = UserAPIKeyAuth(api_key="test-key", jwt_claims=claims) + assert auth.jwt_claims == claims + assert auth.jwt_claims["sub"] == "user-123" + + def test_jwt_claims_backward_compatible_without_field(self): + """Existing code that doesn't pass jwt_claims should still work.""" + auth = UserAPIKeyAuth( + api_key="test-key", + user_id="user-1", + team_id="team-1", + ) + assert auth.jwt_claims is None + assert auth.user_id == "user-1" From 94da7e688a0cb21b58c6a9b48ca4b287712e221f Mon Sep 17 00:00:00 2001 From: Noah Nistler <60981020+noahnistler@users.noreply.github.com> Date: Tue, 17 Mar 2026 15:26:20 -0500 Subject: [PATCH 02/13] Enhance MCPServerManager to support hook-modified arguments and extra headers. Update tests to validate argument mutation and header injection behavior, including warnings for OpenAPI-backed servers when headers are present. --- .../mcp_server/mcp_server_manager.py | 16 +- litellm/proxy/auth/user_api_key_auth.py | 1 + .../mcp_server/test_mcp_hook_extra_headers.py | 241 +++++++++++++++++- 3 files changed, 254 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index b552669651f..adf90317aa8 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1913,6 +1913,7 @@ class MCPServerManager: Run pre-call checks and guardrail hooks for an MCP tool call. Returns a dict that may contain: + - "arguments": hook-modified tool arguments (only if changed) - "extra_headers": headers injected by pre_mcp_call guardrail hooks """ ## check if the tool is allowed or banned for the given server @@ -1991,7 +1992,7 @@ class MCPServerManager: ) ) if modified_kwargs.get("arguments") != arguments: - arguments = modified_kwargs["arguments"] + hook_result["arguments"] = modified_kwargs["arguments"] if modified_kwargs.get("extra_headers"): hook_result["extra_headers"] = modified_kwargs["extra_headers"] @@ -2073,6 +2074,9 @@ class MCPServerManager: oauth2_headers: Optional OAuth2 headers raw_headers: Optional raw headers from the request proxy_logging_obj: Optional ProxyLogging object for hook integration + host_progress_callback: Optional callback for progress updates + hook_extra_headers: Optional headers injected by pre_mcp_call guardrail + hooks. Merged last (highest priority) into outbound request headers. Returns: CallToolResult from the MCP server @@ -2228,6 +2232,8 @@ class MCPServerManager: proxy_logging_obj=proxy_logging_obj, server=mcp_server, ) + if "arguments" in hook_result: + arguments = hook_result["arguments"] # Prepare tasks for during hooks tasks = [] @@ -2247,6 +2253,14 @@ class MCPServerManager: verbose_logger.debug( f"Calling OpenAPI tool {name} directly via HTTP handler" ) + if hook_result.get("extra_headers"): + verbose_logger.warning( + "pre_mcp_call hook returned extra_headers, but OpenAPI-backed " + "MCP servers do not support hook header injection. " + f"Headers will be dropped for tool '{name}' on server " + f"'{server_name}'. Use a regular MCP server (SSE/HTTP transport) " + "for hook header support." + ) tasks.append( asyncio.create_task( self._call_openapi_tool_handler(mcp_server, name, arguments) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 5b6064ef550..451ed56339d 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -700,6 +700,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 ) if valid_token is not None: api_key = valid_token.token or "" + valid_token.jwt_claims = jwt_claims do_standard_jwt_auth = False # Fall through to virtual key checks diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py index 6296c57f09c..fe2d2ab820b 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py @@ -3,13 +3,16 @@ Tests for pre_mcp_call guardrail hook header mutation support. Validates that: 1. _convert_mcp_hook_response_to_kwargs extracts extra_headers from hook response -2. pre_call_tool_check returns hook-provided extra_headers -3. call_tool flows hook headers into _call_regular_mcp_tool +2. pre_call_tool_check returns hook-provided extra_headers AND modified arguments +3. call_tool flows hook headers and modified arguments downstream 4. Hook-provided headers take highest priority (merge after static_headers) -5. Backward compatibility: hooks without extra_headers continue to work +5. OpenAPI-backed servers emit a warning when hook headers are present +6. JWT claims are propagated in both standard and virtual-key fast paths +7. Backward compatibility: hooks without extra_headers continue to work """ import asyncio +import logging import sys from datetime import datetime from typing import Any, Dict, Optional @@ -212,6 +215,87 @@ class TestPreCallToolCheckReturnsHeaders: assert result == {} + @pytest.mark.asyncio + async def test_returns_modified_arguments_from_hook(self): + """Modified arguments from the hook must be returned so the caller can use them.""" + manager = MCPServerManager() + server = self._make_server() + + original_args = {"key": "original"} + modified_args = {"key": "modified", "extra": "added"} + + proxy_logging = MagicMock(spec=ProxyLogging) + proxy_logging._create_mcp_request_object_from_kwargs = MagicMock( + return_value=MagicMock() + ) + proxy_logging._convert_mcp_to_llm_format = MagicMock( + return_value={"model": "fake"} + ) + proxy_logging.pre_call_hook = AsyncMock( + return_value={"modified_arguments": modified_args} + ) + proxy_logging._convert_mcp_hook_response_to_kwargs = MagicMock( + return_value={"arguments": modified_args} + ) + + with patch.object(manager, "check_allowed_or_banned_tools", return_value=True): + with patch.object( + manager, + "check_tool_permission_for_key_team", + new_callable=AsyncMock, + ): + with patch.object(manager, "validate_allowed_params"): + result = await manager.pre_call_tool_check( + name="test_tool", + arguments=original_args, + server_name="test_server", + user_api_key_auth=None, + proxy_logging_obj=proxy_logging, + server=server, + ) + + assert result["arguments"] == modified_args + + @pytest.mark.asyncio + async def test_returns_both_modified_arguments_and_headers(self): + """Hook can modify both arguments and inject headers simultaneously.""" + manager = MCPServerManager() + server = self._make_server() + + modified_args = {"key": "modified"} + hook_headers = {"Authorization": "Bearer jwt"} + + proxy_logging = MagicMock(spec=ProxyLogging) + proxy_logging._create_mcp_request_object_from_kwargs = MagicMock( + return_value=MagicMock() + ) + proxy_logging._convert_mcp_to_llm_format = MagicMock( + return_value={"model": "fake"} + ) + proxy_logging.pre_call_hook = AsyncMock(return_value={"dummy": True}) + proxy_logging._convert_mcp_hook_response_to_kwargs = MagicMock( + return_value={"arguments": modified_args, "extra_headers": hook_headers} + ) + + with patch.object(manager, "check_allowed_or_banned_tools", return_value=True): + with patch.object( + manager, + "check_tool_permission_for_key_team", + new_callable=AsyncMock, + ): + with patch.object(manager, "validate_allowed_params"): + result = await manager.pre_call_tool_check( + name="test_tool", + arguments={"key": "original"}, + server_name="test_server", + user_api_key_auth=None, + proxy_logging_obj=proxy_logging, + server=server, + ) + + assert result["arguments"] == modified_args + assert result["extra_headers"] == hook_headers + class TestCallToolFlowsHookHeaders: """Tests that call_tool passes hook_extra_headers to _call_regular_mcp_tool.""" @@ -297,6 +381,147 @@ class TestCallToolFlowsHookHeaders: call_kwargs = mock_call.call_args assert call_kwargs.kwargs.get("hook_extra_headers") is None + @pytest.mark.asyncio + async def test_modified_arguments_passed_to_downstream(self): + """Hook-modified arguments must be used for the actual tool call.""" + manager = MCPServerManager() + server = self._make_server() + + modified_args = {"key": "modified_by_hook"} + + with patch.object( + manager, + "_get_mcp_server_from_tool_name", + return_value=server, + ): + with patch.object( + manager, + "pre_call_tool_check", + new_callable=AsyncMock, + return_value={"arguments": modified_args}, + ): + with patch.object( + manager, + "_create_during_hook_task", + return_value=asyncio.create_task(asyncio.sleep(0)), + ): + with patch.object( + manager, + "_call_regular_mcp_tool", + new_callable=AsyncMock, + return_value=MagicMock(), + ) as mock_call: + proxy_logging = MagicMock(spec=ProxyLogging) + + await manager.call_tool( + server_name="test_server", + name="test_tool", + arguments={"key": "original"}, + proxy_logging_obj=proxy_logging, + ) + + mock_call.assert_called_once() + call_kwargs = mock_call.call_args + assert call_kwargs.kwargs.get("arguments") == modified_args + + @pytest.mark.asyncio + async def test_openapi_server_warns_on_hook_headers(self, caplog): + """OpenAPI-backed servers should log a warning when hook injects headers.""" + manager = MCPServerManager() + server = MCPServer( + server_id="test-id", + name="openapi_server", + server_name="openapi_server", + url="https://example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + spec_path="/path/to/spec.yaml", + ) + + with patch.object( + manager, "_get_mcp_server_from_tool_name", return_value=server + ): + with patch.object( + manager, + "pre_call_tool_check", + new_callable=AsyncMock, + return_value={"extra_headers": {"Authorization": "Bearer jwt"}}, + ): + with patch.object( + manager, + "_create_during_hook_task", + return_value=asyncio.create_task(asyncio.sleep(0)), + ): + with patch.object( + manager, + "_call_openapi_tool_handler", + new_callable=AsyncMock, + return_value=MagicMock(), + ): + proxy_logging = MagicMock(spec=ProxyLogging) + + with caplog.at_level(logging.WARNING): + await manager.call_tool( + server_name="openapi_server", + name="test_tool", + arguments={}, + proxy_logging_obj=proxy_logging, + ) + + assert any( + "do not support hook header injection" in record.message + for record in caplog.records + ) + + @pytest.mark.asyncio + async def test_openapi_server_no_warning_without_hook_headers(self, caplog): + """No warning when OpenAPI server has no hook-injected headers.""" + manager = MCPServerManager() + server = MCPServer( + server_id="test-id", + name="openapi_server", + server_name="openapi_server", + url="https://example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + spec_path="/path/to/spec.yaml", + ) + + with patch.object( + manager, "_get_mcp_server_from_tool_name", return_value=server + ): + with patch.object( + manager, + "pre_call_tool_check", + new_callable=AsyncMock, + return_value={}, + ): + with patch.object( + manager, + "_create_during_hook_task", + return_value=asyncio.create_task(asyncio.sleep(0)), + ): + with patch.object( + manager, + "_call_openapi_tool_handler", + new_callable=AsyncMock, + return_value=MagicMock(), + ): + proxy_logging = MagicMock(spec=ProxyLogging) + + with caplog.at_level(logging.WARNING): + await manager.call_tool( + server_name="openapi_server", + name="test_tool", + arguments={}, + proxy_logging_obj=proxy_logging, + ) + + assert not any( + "do not support hook header injection" in record.message + for record in caplog.records + ) + class TestHookHeaderMergePriority: """Tests that hook-provided headers have highest priority in _call_regular_mcp_tool.""" @@ -479,3 +704,13 @@ class TestUserAPIKeyAuthJwtClaims: ) assert auth.jwt_claims is None assert auth.user_id == "user-1" + + def test_jwt_claims_set_after_construction(self): + """Virtual-key fast path sets jwt_claims after the object is created.""" + auth = UserAPIKeyAuth(api_key="test-key") + assert auth.jwt_claims is None + + claims = {"sub": "user-456", "iss": "okta", "groups": ["admin"]} + auth.jwt_claims = claims + assert auth.jwt_claims == claims + assert auth.jwt_claims["groups"] == ["admin"] From 4af352aac8470001f6a7a0d691be7efc930c94a9 Mon Sep 17 00:00:00 2001 From: Noah Nistler <60981020+noahnistler@users.noreply.github.com> Date: Tue, 17 Mar 2026 15:43:00 -0500 Subject: [PATCH 03/13] Refactor MCPServerManager to raise HTTPException for extra headers in OpenAPI-backed servers. Update tests to reflect this change, ensuring proper exception handling instead of logging warnings. --- .../mcp_server/mcp_server_manager.py | 18 +++--- .../mcp_server/test_mcp_hook_extra_headers.py | 59 +++++++------------ 2 files changed, 33 insertions(+), 44 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index adf90317aa8..3fa14e6ec25 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2251,15 +2251,19 @@ class MCPServerManager: # For OpenAPI servers, call the tool handler directly instead of via MCP client if mcp_server.spec_path: verbose_logger.debug( - f"Calling OpenAPI tool {name} directly via HTTP handler" + "Calling OpenAPI tool %s directly via HTTP handler", name ) if hook_result.get("extra_headers"): - verbose_logger.warning( - "pre_mcp_call hook returned extra_headers, but OpenAPI-backed " - "MCP servers do not support hook header injection. " - f"Headers will be dropped for tool '{name}' on server " - f"'{server_name}'. Use a regular MCP server (SSE/HTTP transport) " - "for hook header support." + raise HTTPException( + status_code=500, + detail={ + "error": ( + "pre_mcp_call hook returned extra_headers for an " + "OpenAPI-backed MCP server, which does not support " + "hook header injection. Use a regular MCP server " + "(SSE/HTTP transport) for hook header support." + ) + }, ) tasks.append( asyncio.create_task( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py index fe2d2ab820b..6c97ff244e8 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py @@ -6,15 +6,12 @@ Validates that: 2. pre_call_tool_check returns hook-provided extra_headers AND modified arguments 3. call_tool flows hook headers and modified arguments downstream 4. Hook-provided headers take highest priority (merge after static_headers) -5. OpenAPI-backed servers emit a warning when hook headers are present +5. OpenAPI-backed servers raise HTTPException when hook headers are present 6. JWT claims are propagated in both standard and virtual-key fast paths 7. Backward compatibility: hooks without extra_headers continue to work """ import asyncio -import logging -import sys -from datetime import datetime from typing import Any, Dict, Optional from unittest.mock import AsyncMock, MagicMock, patch @@ -425,8 +422,8 @@ class TestCallToolFlowsHookHeaders: assert call_kwargs.kwargs.get("arguments") == modified_args @pytest.mark.asyncio - async def test_openapi_server_warns_on_hook_headers(self, caplog): - """OpenAPI-backed servers should log a warning when hook injects headers.""" + async def test_openapi_server_raises_on_hook_headers(self): + """OpenAPI-backed servers should raise HTTPException when hook injects headers.""" manager = MCPServerManager() server = MCPServer( server_id="test-id", @@ -452,30 +449,24 @@ class TestCallToolFlowsHookHeaders: "_create_during_hook_task", return_value=asyncio.create_task(asyncio.sleep(0)), ): - with patch.object( - manager, - "_call_openapi_tool_handler", - new_callable=AsyncMock, - return_value=MagicMock(), - ): - proxy_logging = MagicMock(spec=ProxyLogging) + proxy_logging = MagicMock(spec=ProxyLogging) - with caplog.at_level(logging.WARNING): - await manager.call_tool( - server_name="openapi_server", - name="test_tool", - arguments={}, - proxy_logging_obj=proxy_logging, - ) - - assert any( - "do not support hook header injection" in record.message - for record in caplog.records + with pytest.raises(HTTPException) as exc_info: + await manager.call_tool( + server_name="openapi_server", + name="test_tool", + arguments={}, + proxy_logging_obj=proxy_logging, ) + assert exc_info.value.status_code == 500 + assert "does not support hook header injection" in str( + exc_info.value.detail + ) + @pytest.mark.asyncio - async def test_openapi_server_no_warning_without_hook_headers(self, caplog): - """No warning when OpenAPI server has no hook-injected headers.""" + async def test_openapi_server_no_error_without_hook_headers(self): + """No exception when OpenAPI server has no hook-injected headers.""" manager = MCPServerManager() server = MCPServer( server_id="test-id", @@ -509,17 +500,11 @@ class TestCallToolFlowsHookHeaders: ): proxy_logging = MagicMock(spec=ProxyLogging) - with caplog.at_level(logging.WARNING): - await manager.call_tool( - server_name="openapi_server", - name="test_tool", - arguments={}, - proxy_logging_obj=proxy_logging, - ) - - assert not any( - "do not support hook header injection" in record.message - for record in caplog.records + await manager.call_tool( + server_name="openapi_server", + name="test_tool", + arguments={}, + proxy_logging_obj=proxy_logging, ) From f61253361d1e49f65f9d948d11ca70c16af5f0cb Mon Sep 17 00:00:00 2001 From: Noah Nistler <60981020+noahnistler@users.noreply.github.com> Date: Tue, 17 Mar 2026 14:31:14 -0500 Subject: [PATCH 04/13] Allow pre_mcp_call guardrail hooks to mutate outbound MCP headers --- .../mcp_server/mcp_server_manager.py | 23 +- litellm/proxy/_types.py | 1 + litellm/proxy/auth/user_api_key_auth.py | 3 + litellm/proxy/utils.py | 9 +- .../mcp_server/test_mcp_hook_extra_headers.py | 481 ++++++++++++++++++ 5 files changed, 513 insertions(+), 4 deletions(-) create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 43fe54fdfb7..b552669651f 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1908,7 +1908,13 @@ class MCPServerManager: user_api_key_auth: Optional[UserAPIKeyAuth], proxy_logging_obj: ProxyLogging, server: MCPServer, - ): + ) -> Dict[str, Any]: + """ + Run pre-call checks and guardrail hooks for an MCP tool call. + + Returns a dict that may contain: + - "extra_headers": headers injected by pre_mcp_call guardrail hooks + """ ## check if the tool is allowed or banned for the given server if not self.check_allowed_or_banned_tools(name, server): raise HTTPException( @@ -1969,6 +1975,7 @@ class MCPServerManager: mcp_request_obj, pre_hook_kwargs ) + hook_result: Dict[str, Any] = {} try: # Use standard pre_call_hook modified_data = await proxy_logging_obj.pre_call_hook( @@ -1985,6 +1992,8 @@ class MCPServerManager: ) if modified_kwargs.get("arguments") != arguments: arguments = modified_kwargs["arguments"] + if modified_kwargs.get("extra_headers"): + hook_result["extra_headers"] = modified_kwargs["extra_headers"] except ( BlockedPiiEntityError, @@ -1995,6 +2004,8 @@ class MCPServerManager: verbose_logger.error(f"Guardrail blocked MCP tool call pre call: {str(e)}") raise e + return hook_result + def _create_during_hook_task( self, name: str, @@ -2047,6 +2058,7 @@ class MCPServerManager: raw_headers: Optional[Dict[str, str]], proxy_logging_obj: Optional[ProxyLogging], host_progress_callback: Optional[Callable] = None, + hook_extra_headers: Optional[Dict[str, str]] = None, ) -> CallToolResult: """ Call a regular MCP tool using the MCP client. @@ -2116,6 +2128,11 @@ class MCPServerManager: extra_headers = {} extra_headers.update(mcp_server.static_headers) + if hook_extra_headers: + if extra_headers is None: + extra_headers = {} + extra_headers.update(hook_extra_headers) + stdio_env = self._build_stdio_env(mcp_server, raw_headers) client = await self._create_mcp_client( @@ -2201,8 +2218,9 @@ class MCPServerManager: # Allow validation and modification of tool calls before execution # Using standard pre_call_hook ######################################################### + hook_result: Dict[str, Any] = {} if proxy_logging_obj: - await self.pre_call_tool_check( + hook_result = await self.pre_call_tool_check( name=name, arguments=arguments, server_name=server_name, @@ -2247,6 +2265,7 @@ class MCPServerManager: raw_headers=raw_headers, proxy_logging_obj=proxy_logging_obj, host_progress_callback=host_progress_callback, + hook_extra_headers=hook_result.get("extra_headers"), ) # For OpenAPI tools, await outside the client context diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index ecbd7314cd7..55cb1c9c43c 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2471,6 +2471,7 @@ class UserAPIKeyAuth( Any ] = None # Expanded created_by user when expand=user is used end_user_object_permission: Optional[LiteLLM_ObjectPermissionTable] = None + jwt_claims: Optional[Dict] = None model_config = ConfigDict(arbitrary_types_allowed=True) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 376048e7a13..5b6064ef550 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -729,6 +729,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 team_membership: Optional[LiteLLM_TeamMembership] = result.get( "team_membership", None ) + jwt_claims: Optional[dict] = result.get("jwt_claims", None) global_proxy_spend = await get_global_proxy_spend( litellm_proxy_admin_name=litellm_proxy_admin_name, @@ -757,6 +758,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 org_id=org_id, end_user_id=end_user_id, parent_otel_span=parent_otel_span, + jwt_claims=jwt_claims, ) valid_token = UserAPIKeyAuth( @@ -803,6 +805,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 team_metadata=( team_object.metadata if team_object is not None else None ), + jwt_claims=jwt_claims, ) # Check if model has zero cost - if so, skip all budget checks diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 01a0f55aac7..8536704766b 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -824,17 +824,22 @@ class ProxyLogging: ) -> dict: """ Helper function to convert pre_call_hook response back to kwargs for MCP usage. + + Supports: + - modified_arguments: Override tool call arguments + - extra_headers: Inject custom headers into the outbound MCP request """ if not response_data: return original_kwargs - # Apply any argument modifications from the hook response modified_kwargs = original_kwargs.copy() - # If the response contains modified arguments, apply them if response_data.get("modified_arguments"): modified_kwargs["arguments"] = response_data["modified_arguments"] + if response_data.get("extra_headers"): + modified_kwargs["extra_headers"] = response_data["extra_headers"] + return modified_kwargs async def process_pre_call_hook_response(self, response, data, call_type): diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py new file mode 100644 index 00000000000..6296c57f09c --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py @@ -0,0 +1,481 @@ +""" +Tests for pre_mcp_call guardrail hook header mutation support. + +Validates that: +1. _convert_mcp_hook_response_to_kwargs extracts extra_headers from hook response +2. pre_call_tool_check returns hook-provided extra_headers +3. call_tool flows hook headers into _call_regular_mcp_tool +4. Hook-provided headers take highest priority (merge after static_headers) +5. Backward compatibility: hooks without extra_headers continue to work +""" + +import asyncio +import sys +from datetime import datetime +from typing import Any, Dict, Optional +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import HTTPException + +from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.utils import ProxyLogging +from litellm.types.mcp import MCPAuth, MCPTransport +from litellm.types.mcp_server.mcp_server_manager import MCPServer + + +class TestConvertMcpHookResponseToKwargs: + """Tests for ProxyLogging._convert_mcp_hook_response_to_kwargs""" + + def setup_method(self): + self.proxy_logging = ProxyLogging(user_api_key_cache=MagicMock()) + + def test_returns_original_kwargs_when_response_is_none(self): + original = {"arguments": {"key": "val"}, "name": "tool"} + result = self.proxy_logging._convert_mcp_hook_response_to_kwargs( + None, original + ) + assert result == original + + def test_returns_original_kwargs_when_response_is_empty_dict(self): + original = {"arguments": {"key": "val"}} + result = self.proxy_logging._convert_mcp_hook_response_to_kwargs({}, original) + assert result == original + + def test_extracts_modified_arguments(self): + original = {"arguments": {"old": "value"}} + response = {"modified_arguments": {"new": "value"}} + result = self.proxy_logging._convert_mcp_hook_response_to_kwargs( + response, original + ) + assert result["arguments"] == {"new": "value"} + + def test_extracts_extra_headers(self): + original = {"arguments": {"key": "val"}} + response = {"extra_headers": {"Authorization": "Bearer signed-jwt"}} + result = self.proxy_logging._convert_mcp_hook_response_to_kwargs( + response, original + ) + assert result["extra_headers"] == {"Authorization": "Bearer signed-jwt"} + + def test_extracts_both_arguments_and_headers(self): + original = {"arguments": {"old": "value"}} + response = { + "modified_arguments": {"new": "value"}, + "extra_headers": {"X-Custom": "header-val"}, + } + result = self.proxy_logging._convert_mcp_hook_response_to_kwargs( + response, original + ) + assert result["arguments"] == {"new": "value"} + assert result["extra_headers"] == {"X-Custom": "header-val"} + + def test_no_extra_headers_key_preserves_original(self): + """Backward compat: hooks that only return modified_arguments still work.""" + original = {"arguments": {"key": "val"}} + response = {"modified_arguments": {"key": "new_val"}} + result = self.proxy_logging._convert_mcp_hook_response_to_kwargs( + response, original + ) + assert "extra_headers" not in result + assert result["arguments"] == {"key": "new_val"} + + def test_empty_extra_headers_not_set(self): + """Empty dict for extra_headers is falsy and should not be set.""" + original = {"arguments": {"key": "val"}} + response = {"extra_headers": {}} + result = self.proxy_logging._convert_mcp_hook_response_to_kwargs( + response, original + ) + assert "extra_headers" not in result + + +class TestPreCallToolCheckReturnsHeaders: + """Tests that pre_call_tool_check returns hook-provided headers.""" + + def _make_server(self, name="test_server"): + return MCPServer( + server_id="test-id", + name=name, + server_name=name, + url="https://example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + ) + + @pytest.mark.asyncio + async def test_returns_empty_dict_when_hook_has_no_headers(self): + manager = MCPServerManager() + server = self._make_server() + + proxy_logging = MagicMock(spec=ProxyLogging) + proxy_logging._create_mcp_request_object_from_kwargs = MagicMock( + return_value=MagicMock() + ) + proxy_logging._convert_mcp_to_llm_format = MagicMock( + return_value={"model": "fake"} + ) + proxy_logging.pre_call_hook = AsyncMock( + return_value={"modified_arguments": {"key": "val"}} + ) + proxy_logging._convert_mcp_hook_response_to_kwargs = MagicMock( + return_value={"arguments": {"key": "val"}} + ) + + with patch.object(manager, "check_allowed_or_banned_tools", return_value=True): + with patch.object( + manager, + "check_tool_permission_for_key_team", + new_callable=AsyncMock, + ): + with patch.object(manager, "validate_allowed_params"): + result = await manager.pre_call_tool_check( + name="test_tool", + arguments={"key": "val"}, + server_name="test_server", + user_api_key_auth=None, + proxy_logging_obj=proxy_logging, + server=server, + ) + + assert result == {} + + @pytest.mark.asyncio + async def test_returns_extra_headers_from_hook(self): + manager = MCPServerManager() + server = self._make_server() + + hook_headers = {"Authorization": "Bearer signed-jwt", "X-Trace-Id": "abc123"} + + proxy_logging = MagicMock(spec=ProxyLogging) + proxy_logging._create_mcp_request_object_from_kwargs = MagicMock( + return_value=MagicMock() + ) + proxy_logging._convert_mcp_to_llm_format = MagicMock( + return_value={"model": "fake"} + ) + proxy_logging.pre_call_hook = AsyncMock( + return_value={"extra_headers": hook_headers} + ) + proxy_logging._convert_mcp_hook_response_to_kwargs = MagicMock( + return_value={"arguments": {"key": "val"}, "extra_headers": hook_headers} + ) + + with patch.object(manager, "check_allowed_or_banned_tools", return_value=True): + with patch.object( + manager, + "check_tool_permission_for_key_team", + new_callable=AsyncMock, + ): + with patch.object(manager, "validate_allowed_params"): + result = await manager.pre_call_tool_check( + name="test_tool", + arguments={"key": "val"}, + server_name="test_server", + user_api_key_auth=None, + proxy_logging_obj=proxy_logging, + server=server, + ) + + assert result["extra_headers"] == hook_headers + + @pytest.mark.asyncio + async def test_returns_empty_dict_when_hook_returns_none(self): + manager = MCPServerManager() + server = self._make_server() + + proxy_logging = MagicMock(spec=ProxyLogging) + proxy_logging._create_mcp_request_object_from_kwargs = MagicMock( + return_value=MagicMock() + ) + proxy_logging._convert_mcp_to_llm_format = MagicMock( + return_value={"model": "fake"} + ) + proxy_logging.pre_call_hook = AsyncMock(return_value=None) + + with patch.object(manager, "check_allowed_or_banned_tools", return_value=True): + with patch.object( + manager, + "check_tool_permission_for_key_team", + new_callable=AsyncMock, + ): + with patch.object(manager, "validate_allowed_params"): + result = await manager.pre_call_tool_check( + name="test_tool", + arguments={"key": "val"}, + server_name="test_server", + user_api_key_auth=None, + proxy_logging_obj=proxy_logging, + server=server, + ) + + assert result == {} + + +class TestCallToolFlowsHookHeaders: + """Tests that call_tool passes hook_extra_headers to _call_regular_mcp_tool.""" + + def _make_server(self, name="test_server"): + return MCPServer( + server_id="test-id", + name=name, + server_name=name, + url="https://example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + ) + + @pytest.mark.asyncio + async def test_hook_headers_passed_to_call_regular_mcp_tool(self): + """Verify that hook_extra_headers kwarg is forwarded.""" + manager = MCPServerManager() + server = self._make_server() + + hook_headers = {"Authorization": "Bearer signed-jwt"} + + with patch.object( + manager, + "_get_mcp_server_from_tool_name", + return_value=server, + ): + with patch.object( + manager, + "pre_call_tool_check", + new_callable=AsyncMock, + return_value={"extra_headers": hook_headers}, + ): + with patch.object( + manager, + "_create_during_hook_task", + return_value=asyncio.create_task(asyncio.sleep(0)), + ): + with patch.object( + manager, + "_call_regular_mcp_tool", + new_callable=AsyncMock, + return_value=MagicMock(), + ) as mock_call: + proxy_logging = MagicMock(spec=ProxyLogging) + + await manager.call_tool( + server_name="test_server", + name="test_tool", + arguments={"key": "val"}, + proxy_logging_obj=proxy_logging, + ) + + mock_call.assert_called_once() + call_kwargs = mock_call.call_args + assert call_kwargs.kwargs.get("hook_extra_headers") == hook_headers + + @pytest.mark.asyncio + async def test_no_hook_headers_when_no_proxy_logging(self): + """Without proxy_logging_obj, no pre_call_tool_check runs.""" + manager = MCPServerManager() + server = self._make_server() + + with patch.object( + manager, + "_get_mcp_server_from_tool_name", + return_value=server, + ): + with patch.object( + manager, + "_call_regular_mcp_tool", + new_callable=AsyncMock, + return_value=MagicMock(), + ) as mock_call: + await manager.call_tool( + server_name="test_server", + name="test_tool", + arguments={"key": "val"}, + proxy_logging_obj=None, + ) + + mock_call.assert_called_once() + call_kwargs = mock_call.call_args + assert call_kwargs.kwargs.get("hook_extra_headers") is None + + +class TestHookHeaderMergePriority: + """Tests that hook-provided headers have highest priority in _call_regular_mcp_tool.""" + + def _make_server( + self, + static_headers: Optional[Dict[str, str]] = None, + extra_headers_config: Optional[list] = None, + ): + return MCPServer( + server_id="test-id", + name="Test Server", + server_name="test_server", + url="https://example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + static_headers=static_headers, + extra_headers=extra_headers_config, + ) + + @pytest.mark.asyncio + async def test_hook_headers_override_static_headers(self): + """Hook headers should take precedence over static_headers.""" + manager = MCPServerManager() + server = self._make_server( + static_headers={"Authorization": "Bearer static-token", "X-Static": "yes"} + ) + + hook_headers = {"Authorization": "Bearer hook-signed-jwt"} + + captured_extra_headers: Dict[str, Any] = {} + + async def fake_create_mcp_client( + server, mcp_auth_header=None, extra_headers=None, stdio_env=None + ): + captured_extra_headers["value"] = extra_headers + mock_client = MagicMock() + mock_client.call_tool = AsyncMock(return_value=MagicMock()) + return mock_client + + with patch.object( + manager, "_create_mcp_client", side_effect=fake_create_mcp_client + ): + with patch.object(manager, "_build_stdio_env", return_value=None): + try: + await manager._call_regular_mcp_tool( + mcp_server=server, + original_tool_name="test_tool", + arguments={"key": "val"}, + tasks=[], + mcp_auth_header=None, + mcp_server_auth_headers=None, + oauth2_headers=None, + raw_headers=None, + proxy_logging_obj=None, + hook_extra_headers=hook_headers, + ) + except Exception: + pass + + headers = captured_extra_headers.get("value", {}) + assert headers["Authorization"] == "Bearer hook-signed-jwt" + assert headers["X-Static"] == "yes" + + @pytest.mark.asyncio + async def test_no_hook_headers_preserves_existing_behavior(self): + """When hook_extra_headers is None, existing header logic is unchanged.""" + manager = MCPServerManager() + server = self._make_server( + static_headers={"X-Static": "static-value"} + ) + + captured_extra_headers: Dict[str, Any] = {} + + async def fake_create_mcp_client( + server, mcp_auth_header=None, extra_headers=None, stdio_env=None + ): + captured_extra_headers["value"] = extra_headers + mock_client = MagicMock() + mock_client.call_tool = AsyncMock(return_value=MagicMock()) + return mock_client + + with patch.object( + manager, "_create_mcp_client", side_effect=fake_create_mcp_client + ): + with patch.object(manager, "_build_stdio_env", return_value=None): + try: + await manager._call_regular_mcp_tool( + mcp_server=server, + original_tool_name="test_tool", + arguments={"key": "val"}, + tasks=[], + mcp_auth_header=None, + mcp_server_auth_headers=None, + oauth2_headers=None, + raw_headers=None, + proxy_logging_obj=None, + hook_extra_headers=None, + ) + except Exception: + pass + + headers = captured_extra_headers.get("value", {}) + assert headers == {"X-Static": "static-value"} + + @pytest.mark.asyncio + async def test_hook_headers_merge_with_oauth2(self): + """Hook headers merge on top of OAuth2 headers.""" + manager = MCPServerManager() + server = MCPServer( + server_id="test-id", + name="Test Server", + server_name="test_server", + url="https://example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + ) + + captured_extra_headers: Dict[str, Any] = {} + + async def fake_create_mcp_client( + server, mcp_auth_header=None, extra_headers=None, stdio_env=None + ): + captured_extra_headers["value"] = extra_headers + mock_client = MagicMock() + mock_client.call_tool = AsyncMock(return_value=MagicMock()) + return mock_client + + with patch.object( + manager, "_create_mcp_client", side_effect=fake_create_mcp_client + ): + with patch.object(manager, "_build_stdio_env", return_value=None): + try: + await manager._call_regular_mcp_tool( + mcp_server=server, + original_tool_name="test_tool", + arguments={"key": "val"}, + tasks=[], + mcp_auth_header=None, + mcp_server_auth_headers=None, + oauth2_headers={ + "Authorization": "Bearer oauth2-token", + "X-OAuth": "yes", + }, + raw_headers=None, + proxy_logging_obj=None, + hook_extra_headers={ + "Authorization": "Bearer hook-jwt", + "X-Trace-Id": "trace-123", + }, + ) + except Exception: + pass + + headers = captured_extra_headers.get("value", {}) + assert headers["Authorization"] == "Bearer hook-jwt" + assert headers["X-OAuth"] == "yes" + assert headers["X-Trace-Id"] == "trace-123" + + +class TestUserAPIKeyAuthJwtClaims: + """Tests that UserAPIKeyAuth correctly carries jwt_claims.""" + + def test_jwt_claims_field_defaults_to_none(self): + auth = UserAPIKeyAuth(api_key="test-key") + assert auth.jwt_claims is None + + def test_jwt_claims_field_accepts_dict(self): + claims = {"sub": "user-123", "iss": "litellm", "exp": 9999999999} + auth = UserAPIKeyAuth(api_key="test-key", jwt_claims=claims) + assert auth.jwt_claims == claims + assert auth.jwt_claims["sub"] == "user-123" + + def test_jwt_claims_backward_compatible_without_field(self): + """Existing code that doesn't pass jwt_claims should still work.""" + auth = UserAPIKeyAuth( + api_key="test-key", + user_id="user-1", + team_id="team-1", + ) + assert auth.jwt_claims is None + assert auth.user_id == "user-1" From 8574094305c4787ab8db7f167c34bd9e8a49f5b1 Mon Sep 17 00:00:00 2001 From: Noah Nistler <60981020+noahnistler@users.noreply.github.com> Date: Tue, 17 Mar 2026 15:26:20 -0500 Subject: [PATCH 05/13] Enhance MCPServerManager to support hook-modified arguments and extra headers. Update tests to validate argument mutation and header injection behavior, including warnings for OpenAPI-backed servers when headers are present. --- .../mcp_server/mcp_server_manager.py | 16 +- litellm/proxy/auth/user_api_key_auth.py | 1 + .../mcp_server/test_mcp_hook_extra_headers.py | 241 +++++++++++++++++- 3 files changed, 254 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index b552669651f..adf90317aa8 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1913,6 +1913,7 @@ class MCPServerManager: Run pre-call checks and guardrail hooks for an MCP tool call. Returns a dict that may contain: + - "arguments": hook-modified tool arguments (only if changed) - "extra_headers": headers injected by pre_mcp_call guardrail hooks """ ## check if the tool is allowed or banned for the given server @@ -1991,7 +1992,7 @@ class MCPServerManager: ) ) if modified_kwargs.get("arguments") != arguments: - arguments = modified_kwargs["arguments"] + hook_result["arguments"] = modified_kwargs["arguments"] if modified_kwargs.get("extra_headers"): hook_result["extra_headers"] = modified_kwargs["extra_headers"] @@ -2073,6 +2074,9 @@ class MCPServerManager: oauth2_headers: Optional OAuth2 headers raw_headers: Optional raw headers from the request proxy_logging_obj: Optional ProxyLogging object for hook integration + host_progress_callback: Optional callback for progress updates + hook_extra_headers: Optional headers injected by pre_mcp_call guardrail + hooks. Merged last (highest priority) into outbound request headers. Returns: CallToolResult from the MCP server @@ -2228,6 +2232,8 @@ class MCPServerManager: proxy_logging_obj=proxy_logging_obj, server=mcp_server, ) + if "arguments" in hook_result: + arguments = hook_result["arguments"] # Prepare tasks for during hooks tasks = [] @@ -2247,6 +2253,14 @@ class MCPServerManager: verbose_logger.debug( f"Calling OpenAPI tool {name} directly via HTTP handler" ) + if hook_result.get("extra_headers"): + verbose_logger.warning( + "pre_mcp_call hook returned extra_headers, but OpenAPI-backed " + "MCP servers do not support hook header injection. " + f"Headers will be dropped for tool '{name}' on server " + f"'{server_name}'. Use a regular MCP server (SSE/HTTP transport) " + "for hook header support." + ) tasks.append( asyncio.create_task( self._call_openapi_tool_handler(mcp_server, name, arguments) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 5b6064ef550..451ed56339d 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -700,6 +700,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 ) if valid_token is not None: api_key = valid_token.token or "" + valid_token.jwt_claims = jwt_claims do_standard_jwt_auth = False # Fall through to virtual key checks diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py index 6296c57f09c..fe2d2ab820b 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py @@ -3,13 +3,16 @@ Tests for pre_mcp_call guardrail hook header mutation support. Validates that: 1. _convert_mcp_hook_response_to_kwargs extracts extra_headers from hook response -2. pre_call_tool_check returns hook-provided extra_headers -3. call_tool flows hook headers into _call_regular_mcp_tool +2. pre_call_tool_check returns hook-provided extra_headers AND modified arguments +3. call_tool flows hook headers and modified arguments downstream 4. Hook-provided headers take highest priority (merge after static_headers) -5. Backward compatibility: hooks without extra_headers continue to work +5. OpenAPI-backed servers emit a warning when hook headers are present +6. JWT claims are propagated in both standard and virtual-key fast paths +7. Backward compatibility: hooks without extra_headers continue to work """ import asyncio +import logging import sys from datetime import datetime from typing import Any, Dict, Optional @@ -212,6 +215,87 @@ class TestPreCallToolCheckReturnsHeaders: assert result == {} + @pytest.mark.asyncio + async def test_returns_modified_arguments_from_hook(self): + """Modified arguments from the hook must be returned so the caller can use them.""" + manager = MCPServerManager() + server = self._make_server() + + original_args = {"key": "original"} + modified_args = {"key": "modified", "extra": "added"} + + proxy_logging = MagicMock(spec=ProxyLogging) + proxy_logging._create_mcp_request_object_from_kwargs = MagicMock( + return_value=MagicMock() + ) + proxy_logging._convert_mcp_to_llm_format = MagicMock( + return_value={"model": "fake"} + ) + proxy_logging.pre_call_hook = AsyncMock( + return_value={"modified_arguments": modified_args} + ) + proxy_logging._convert_mcp_hook_response_to_kwargs = MagicMock( + return_value={"arguments": modified_args} + ) + + with patch.object(manager, "check_allowed_or_banned_tools", return_value=True): + with patch.object( + manager, + "check_tool_permission_for_key_team", + new_callable=AsyncMock, + ): + with patch.object(manager, "validate_allowed_params"): + result = await manager.pre_call_tool_check( + name="test_tool", + arguments=original_args, + server_name="test_server", + user_api_key_auth=None, + proxy_logging_obj=proxy_logging, + server=server, + ) + + assert result["arguments"] == modified_args + + @pytest.mark.asyncio + async def test_returns_both_modified_arguments_and_headers(self): + """Hook can modify both arguments and inject headers simultaneously.""" + manager = MCPServerManager() + server = self._make_server() + + modified_args = {"key": "modified"} + hook_headers = {"Authorization": "Bearer jwt"} + + proxy_logging = MagicMock(spec=ProxyLogging) + proxy_logging._create_mcp_request_object_from_kwargs = MagicMock( + return_value=MagicMock() + ) + proxy_logging._convert_mcp_to_llm_format = MagicMock( + return_value={"model": "fake"} + ) + proxy_logging.pre_call_hook = AsyncMock(return_value={"dummy": True}) + proxy_logging._convert_mcp_hook_response_to_kwargs = MagicMock( + return_value={"arguments": modified_args, "extra_headers": hook_headers} + ) + + with patch.object(manager, "check_allowed_or_banned_tools", return_value=True): + with patch.object( + manager, + "check_tool_permission_for_key_team", + new_callable=AsyncMock, + ): + with patch.object(manager, "validate_allowed_params"): + result = await manager.pre_call_tool_check( + name="test_tool", + arguments={"key": "original"}, + server_name="test_server", + user_api_key_auth=None, + proxy_logging_obj=proxy_logging, + server=server, + ) + + assert result["arguments"] == modified_args + assert result["extra_headers"] == hook_headers + class TestCallToolFlowsHookHeaders: """Tests that call_tool passes hook_extra_headers to _call_regular_mcp_tool.""" @@ -297,6 +381,147 @@ class TestCallToolFlowsHookHeaders: call_kwargs = mock_call.call_args assert call_kwargs.kwargs.get("hook_extra_headers") is None + @pytest.mark.asyncio + async def test_modified_arguments_passed_to_downstream(self): + """Hook-modified arguments must be used for the actual tool call.""" + manager = MCPServerManager() + server = self._make_server() + + modified_args = {"key": "modified_by_hook"} + + with patch.object( + manager, + "_get_mcp_server_from_tool_name", + return_value=server, + ): + with patch.object( + manager, + "pre_call_tool_check", + new_callable=AsyncMock, + return_value={"arguments": modified_args}, + ): + with patch.object( + manager, + "_create_during_hook_task", + return_value=asyncio.create_task(asyncio.sleep(0)), + ): + with patch.object( + manager, + "_call_regular_mcp_tool", + new_callable=AsyncMock, + return_value=MagicMock(), + ) as mock_call: + proxy_logging = MagicMock(spec=ProxyLogging) + + await manager.call_tool( + server_name="test_server", + name="test_tool", + arguments={"key": "original"}, + proxy_logging_obj=proxy_logging, + ) + + mock_call.assert_called_once() + call_kwargs = mock_call.call_args + assert call_kwargs.kwargs.get("arguments") == modified_args + + @pytest.mark.asyncio + async def test_openapi_server_warns_on_hook_headers(self, caplog): + """OpenAPI-backed servers should log a warning when hook injects headers.""" + manager = MCPServerManager() + server = MCPServer( + server_id="test-id", + name="openapi_server", + server_name="openapi_server", + url="https://example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + spec_path="/path/to/spec.yaml", + ) + + with patch.object( + manager, "_get_mcp_server_from_tool_name", return_value=server + ): + with patch.object( + manager, + "pre_call_tool_check", + new_callable=AsyncMock, + return_value={"extra_headers": {"Authorization": "Bearer jwt"}}, + ): + with patch.object( + manager, + "_create_during_hook_task", + return_value=asyncio.create_task(asyncio.sleep(0)), + ): + with patch.object( + manager, + "_call_openapi_tool_handler", + new_callable=AsyncMock, + return_value=MagicMock(), + ): + proxy_logging = MagicMock(spec=ProxyLogging) + + with caplog.at_level(logging.WARNING): + await manager.call_tool( + server_name="openapi_server", + name="test_tool", + arguments={}, + proxy_logging_obj=proxy_logging, + ) + + assert any( + "do not support hook header injection" in record.message + for record in caplog.records + ) + + @pytest.mark.asyncio + async def test_openapi_server_no_warning_without_hook_headers(self, caplog): + """No warning when OpenAPI server has no hook-injected headers.""" + manager = MCPServerManager() + server = MCPServer( + server_id="test-id", + name="openapi_server", + server_name="openapi_server", + url="https://example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + spec_path="/path/to/spec.yaml", + ) + + with patch.object( + manager, "_get_mcp_server_from_tool_name", return_value=server + ): + with patch.object( + manager, + "pre_call_tool_check", + new_callable=AsyncMock, + return_value={}, + ): + with patch.object( + manager, + "_create_during_hook_task", + return_value=asyncio.create_task(asyncio.sleep(0)), + ): + with patch.object( + manager, + "_call_openapi_tool_handler", + new_callable=AsyncMock, + return_value=MagicMock(), + ): + proxy_logging = MagicMock(spec=ProxyLogging) + + with caplog.at_level(logging.WARNING): + await manager.call_tool( + server_name="openapi_server", + name="test_tool", + arguments={}, + proxy_logging_obj=proxy_logging, + ) + + assert not any( + "do not support hook header injection" in record.message + for record in caplog.records + ) + class TestHookHeaderMergePriority: """Tests that hook-provided headers have highest priority in _call_regular_mcp_tool.""" @@ -479,3 +704,13 @@ class TestUserAPIKeyAuthJwtClaims: ) assert auth.jwt_claims is None assert auth.user_id == "user-1" + + def test_jwt_claims_set_after_construction(self): + """Virtual-key fast path sets jwt_claims after the object is created.""" + auth = UserAPIKeyAuth(api_key="test-key") + assert auth.jwt_claims is None + + claims = {"sub": "user-456", "iss": "okta", "groups": ["admin"]} + auth.jwt_claims = claims + assert auth.jwt_claims == claims + assert auth.jwt_claims["groups"] == ["admin"] From 296f382c648e91c51a466645361064ed2e20f7c7 Mon Sep 17 00:00:00 2001 From: Noah Nistler <60981020+noahnistler@users.noreply.github.com> Date: Tue, 17 Mar 2026 15:43:00 -0500 Subject: [PATCH 06/13] Refactor MCPServerManager to raise HTTPException for extra headers in OpenAPI-backed servers. Update tests to reflect this change, ensuring proper exception handling instead of logging warnings. --- .../mcp_server/mcp_server_manager.py | 18 +++--- .../mcp_server/test_mcp_hook_extra_headers.py | 59 +++++++------------ 2 files changed, 33 insertions(+), 44 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index adf90317aa8..3fa14e6ec25 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2251,15 +2251,19 @@ class MCPServerManager: # For OpenAPI servers, call the tool handler directly instead of via MCP client if mcp_server.spec_path: verbose_logger.debug( - f"Calling OpenAPI tool {name} directly via HTTP handler" + "Calling OpenAPI tool %s directly via HTTP handler", name ) if hook_result.get("extra_headers"): - verbose_logger.warning( - "pre_mcp_call hook returned extra_headers, but OpenAPI-backed " - "MCP servers do not support hook header injection. " - f"Headers will be dropped for tool '{name}' on server " - f"'{server_name}'. Use a regular MCP server (SSE/HTTP transport) " - "for hook header support." + raise HTTPException( + status_code=500, + detail={ + "error": ( + "pre_mcp_call hook returned extra_headers for an " + "OpenAPI-backed MCP server, which does not support " + "hook header injection. Use a regular MCP server " + "(SSE/HTTP transport) for hook header support." + ) + }, ) tasks.append( asyncio.create_task( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py index fe2d2ab820b..6c97ff244e8 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py @@ -6,15 +6,12 @@ Validates that: 2. pre_call_tool_check returns hook-provided extra_headers AND modified arguments 3. call_tool flows hook headers and modified arguments downstream 4. Hook-provided headers take highest priority (merge after static_headers) -5. OpenAPI-backed servers emit a warning when hook headers are present +5. OpenAPI-backed servers raise HTTPException when hook headers are present 6. JWT claims are propagated in both standard and virtual-key fast paths 7. Backward compatibility: hooks without extra_headers continue to work """ import asyncio -import logging -import sys -from datetime import datetime from typing import Any, Dict, Optional from unittest.mock import AsyncMock, MagicMock, patch @@ -425,8 +422,8 @@ class TestCallToolFlowsHookHeaders: assert call_kwargs.kwargs.get("arguments") == modified_args @pytest.mark.asyncio - async def test_openapi_server_warns_on_hook_headers(self, caplog): - """OpenAPI-backed servers should log a warning when hook injects headers.""" + async def test_openapi_server_raises_on_hook_headers(self): + """OpenAPI-backed servers should raise HTTPException when hook injects headers.""" manager = MCPServerManager() server = MCPServer( server_id="test-id", @@ -452,30 +449,24 @@ class TestCallToolFlowsHookHeaders: "_create_during_hook_task", return_value=asyncio.create_task(asyncio.sleep(0)), ): - with patch.object( - manager, - "_call_openapi_tool_handler", - new_callable=AsyncMock, - return_value=MagicMock(), - ): - proxy_logging = MagicMock(spec=ProxyLogging) + proxy_logging = MagicMock(spec=ProxyLogging) - with caplog.at_level(logging.WARNING): - await manager.call_tool( - server_name="openapi_server", - name="test_tool", - arguments={}, - proxy_logging_obj=proxy_logging, - ) - - assert any( - "do not support hook header injection" in record.message - for record in caplog.records + with pytest.raises(HTTPException) as exc_info: + await manager.call_tool( + server_name="openapi_server", + name="test_tool", + arguments={}, + proxy_logging_obj=proxy_logging, ) + assert exc_info.value.status_code == 500 + assert "does not support hook header injection" in str( + exc_info.value.detail + ) + @pytest.mark.asyncio - async def test_openapi_server_no_warning_without_hook_headers(self, caplog): - """No warning when OpenAPI server has no hook-injected headers.""" + async def test_openapi_server_no_error_without_hook_headers(self): + """No exception when OpenAPI server has no hook-injected headers.""" manager = MCPServerManager() server = MCPServer( server_id="test-id", @@ -509,17 +500,11 @@ class TestCallToolFlowsHookHeaders: ): proxy_logging = MagicMock(spec=ProxyLogging) - with caplog.at_level(logging.WARNING): - await manager.call_tool( - server_name="openapi_server", - name="test_tool", - arguments={}, - proxy_logging_obj=proxy_logging, - ) - - assert not any( - "do not support hook header injection" in record.message - for record in caplog.records + await manager.call_tool( + server_name="openapi_server", + name="test_tool", + arguments={}, + proxy_logging_obj=proxy_logging, ) From 49d43f944b347a6c4fdcc092b8e994e375e2fb99 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 17 Mar 2026 14:02:19 -0700 Subject: [PATCH 07/13] feat(guardrails): add MCPJWTSigner built-in guardrail for zero trust MCP auth Signs outbound MCP tool calls with a LiteLLM-issued RS256 JWT so MCP servers can trust a single signing authority instead of every upstream IdP. Enable in config.yaml: guardrails: - guardrail_name: mcp-jwt-signer litellm_params: guardrail: mcp_jwt_signer mode: pre_mcp_call default_on: true JWT carries sub (user_id), act.sub (team_id, RFC 8693), tool-level scope, iss, aud, iat/exp/nbf. RSA-2048 keypair auto-generated at startup unless MCP_JWT_SIGNING_KEY env var is set. Adds /.well-known/jwks.json endpoint and jwks_uri to /.well-known/openid-configuration so MCP servers can verify LiteLLM-issued tokens via OIDC discovery. --- .../mcp_server/discoverable_endpoints.py | 51 ++- .../mcp_jwt_signer/__init__.py | 56 +++ .../mcp_jwt_signer/mcp_jwt_signer.py | 288 +++++++++++++++ litellm/types/guardrails.py | 1 + .../proxy/guardrails/test_mcp_jwt_signer.py | 328 ++++++++++++++++++ 5 files changed, 723 insertions(+), 1 deletion(-) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/__init__.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py create mode 100644 tests/test_litellm/proxy/guardrails/test_mcp_jwt_signer.py diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index af3a715051b..0b12aba02d8 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -677,7 +677,56 @@ async def oauth_authorization_server_mcp( # Alias for standard OpenID discovery @router.get("/.well-known/openid-configuration") async def openid_configuration(request: Request): - return await oauth_authorization_server_mcp(request) + response = await oauth_authorization_server_mcp(request) + + # If MCPJWTSigner is active, augment the discovery doc with JWKS fields so + # MCP servers and gateways (e.g. AWS Bedrock AgentCore Gateway) can resolve + # the signing keys and verify liteLLM-issued tokens. + try: + from litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer import ( + get_mcp_jwt_signer, + ) + + signer = get_mcp_jwt_signer() + if signer is not None: + request_base_url = get_request_base_url(request) + if isinstance(response, dict): + response["jwks_uri"] = f"{request_base_url}/.well-known/jwks.json" + response["id_token_signing_alg_values_supported"] = ["RS256"] + except ImportError: + pass + + return response + + +@router.get("/.well-known/jwks.json") +async def jwks_json(request: Request): + """ + JSON Web Key Set endpoint. + + Returns the RSA public key used by MCPJWTSigner to sign outbound MCP tokens. + MCP servers and gateways use this endpoint to verify liteLLM-issued JWTs. + + Returns an empty key set if MCPJWTSigner is not configured. + """ + try: + from litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer import ( + get_mcp_jwt_signer, + ) + + signer = get_mcp_jwt_signer() + if signer is not None: + return JSONResponse( + content=signer.get_jwks(), + headers={"Cache-Control": "public, max-age=3600"}, + ) + except ImportError: + pass + + return JSONResponse( + content={"keys": []}, + headers={"Cache-Control": "public, max-age=3600"}, + ) # Additional legacy pattern support diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/__init__.py new file mode 100644 index 00000000000..d4a6c579097 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/__init__.py @@ -0,0 +1,56 @@ +"""MCP JWT Signer guardrail — built-in LiteLLM guardrail for zero trust MCP auth.""" + +from typing import TYPE_CHECKING + +from litellm.types.guardrails import SupportedGuardrailIntegrations + +from .mcp_jwt_signer import MCPJWTSigner, _mcp_jwt_signer_instance, get_mcp_jwt_signer + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + + +def initialize_guardrail( + litellm_params: "LitellmParams", guardrail: "Guardrail" +) -> MCPJWTSigner: + import litellm + + guardrail_name = guardrail.get("guardrail_name") + if not guardrail_name: + raise ValueError("MCPJWTSigner guardrail requires a guardrail_name") + + optional_params = getattr(litellm_params, "optional_params", None) + + def _get(key): # type: ignore[no-untyped-def] + if optional_params is not None: + v = getattr(optional_params, key, None) + if v is not None: + return v + return getattr(litellm_params, key, None) + + signer = MCPJWTSigner( + guardrail_name=guardrail_name, + event_hook=litellm_params.mode, + default_on=litellm_params.default_on, + issuer=_get("issuer"), + audience=_get("audience"), + ttl_seconds=_get("ttl_seconds"), + ) + litellm.logging_callback_manager.add_litellm_callback(signer) + return signer + + +guardrail_initializer_registry = { + SupportedGuardrailIntegrations.MCP_JWT_SIGNER.value: initialize_guardrail, +} + +guardrail_class_registry = { + SupportedGuardrailIntegrations.MCP_JWT_SIGNER.value: MCPJWTSigner, +} + +__all__ = [ + "MCPJWTSigner", + "initialize_guardrail", + "_mcp_jwt_signer_instance", + "get_mcp_jwt_signer", +] diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py new file mode 100644 index 00000000000..39ffc2aed88 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py @@ -0,0 +1,288 @@ +""" +MCPJWTSigner — Built-in LiteLLM guardrail for zero trust MCP authentication. + +Signs outbound MCP requests with a LiteLLM-issued RS256 JWT so that MCP servers +can trust a single signing authority (liteLLM) instead of every upstream IdP. + +Usage in config.yaml: + + guardrails: + - guardrail_name: "mcp-jwt-signer" + litellm_params: + guardrail: mcp_jwt_signer + mode: "pre_mcp_call" + default_on: true + issuer: "https://my-litellm.example.com" # optional + audience: "mcp" # optional + ttl_seconds: 300 # optional + +MCP servers verify tokens via: + GET /.well-known/openid-configuration → { jwks_uri: ".../.well-known/jwks.json" } + GET /.well-known/jwks.json → RSA public key in JWKS format + +Optionally set MCP_JWT_SIGNING_KEY env var (PEM string or file:///path) to use +your own RSA keypair. If unset, an RSA-2048 keypair is auto-generated at startup. +""" + +import base64 +import hashlib +import os +import time +from typing import Any, Dict, Optional, Union + +import jwt +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import rsa +from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey + +from litellm._logging import verbose_proxy_logger +from litellm.caching import DualCache +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, +) +from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.utils import CallTypesLiteral + +# Module-level singleton for the JWKS discovery endpoint to access. +_mcp_jwt_signer_instance: Optional["MCPJWTSigner"] = None + + +def get_mcp_jwt_signer() -> Optional["MCPJWTSigner"]: + """Return the active MCPJWTSigner singleton, or None if not initialized.""" + return _mcp_jwt_signer_instance + + +def _load_private_key_from_env(env_var: str) -> RSAPrivateKey: + """Load an RSA private key from an env var (PEM string or file:// path).""" + key_material = os.environ.get(env_var, "") + if not key_material: + raise ValueError( + f"MCPJWTSigner: environment variable '{env_var}' is set but empty." + ) + if key_material.startswith("file://"): + path = key_material[len("file://") :] + with open(path, "rb") as f: + key_bytes = f.read() + else: + key_bytes = key_material.encode("utf-8") + return serialization.load_pem_private_key(key_bytes, password=None) # type: ignore[return-value] + + +def _generate_rsa_key_pair() -> RSAPrivateKey: + """Generate a new RSA-2048 private key.""" + return rsa.generate_private_key( + public_exponent=65537, + key_size=2048, + ) + + +def _int_to_base64url(n: int) -> str: + """Encode an integer as a base64url string (no padding).""" + byte_length = (n.bit_length() + 7) // 8 + return ( + base64.urlsafe_b64encode(n.to_bytes(byte_length, byteorder="big")) + .rstrip(b"=") + .decode("ascii") + ) + + +def _compute_kid(public_key: Any) -> str: + """Derive a key ID from the public key's DER encoding (SHA-256, first 16 hex chars).""" + der_bytes = public_key.public_bytes( + encoding=serialization.Encoding.DER, + format=serialization.PublicFormat.SubjectPublicKeyInfo, + ) + return hashlib.sha256(der_bytes).hexdigest()[:16] + + +class MCPJWTSigner(CustomGuardrail): + """ + Built-in LiteLLM guardrail that signs outbound MCP requests with a + LiteLLM-issued RS256 JWT, enabling zero trust authentication. + + MCP servers verify tokens using liteLLM's OIDC discovery endpoint and + JWKS endpoint rather than trusting each upstream IdP directly. + + The signed JWT carries: + - iss: LiteLLM issuer identifier + - aud: MCP audience (configurable) + - sub: End-user identity (from UserAPIKeyAuth.user_id, RFC 8693) + - act: Actor/agent identity (team_id or org_id, RFC 8693 delegation) + - scope: Tool-level access scopes + - iat, exp, nbf: Timing claims + """ + + ALGORITHM = "RS256" + DEFAULT_TTL = 300 + DEFAULT_AUDIENCE = "mcp" + SIGNING_KEY_ENV = "MCP_JWT_SIGNING_KEY" + + def __init__( + self, + issuer: Optional[str] = None, + audience: Optional[str] = None, + ttl_seconds: Optional[int] = None, + **kwargs: Any, + ) -> None: + super().__init__(**kwargs) + + key_material = os.environ.get(self.SIGNING_KEY_ENV) + if key_material: + self._private_key = _load_private_key_from_env(self.SIGNING_KEY_ENV) + verbose_proxy_logger.info( + "MCPJWTSigner: loaded RSA key from env var %s", self.SIGNING_KEY_ENV + ) + else: + self._private_key = _generate_rsa_key_pair() + verbose_proxy_logger.info( + "MCPJWTSigner: auto-generated RSA-2048 keypair (set %s to use your own key)", + self.SIGNING_KEY_ENV, + ) + + self._public_key = self._private_key.public_key() + self._kid = _compute_kid(self._public_key) + + self.issuer: str = ( + issuer + or os.environ.get("MCP_JWT_ISSUER") + or os.environ.get("LITELLM_EXTERNAL_URL") + or "litellm" + ) + self.audience: str = ( + audience + or os.environ.get("MCP_JWT_AUDIENCE") + or self.DEFAULT_AUDIENCE + ) + self.ttl_seconds: int = int( + ttl_seconds + if ttl_seconds is not None + else os.environ.get("MCP_JWT_TTL_SECONDS", str(self.DEFAULT_TTL)) + ) + + # Register singleton so the JWKS endpoint can access it. + global _mcp_jwt_signer_instance + _mcp_jwt_signer_instance = self + + verbose_proxy_logger.info( + "MCPJWTSigner initialized: issuer=%s audience=%s ttl=%ds kid=%s", + self.issuer, + self.audience, + self.ttl_seconds, + self._kid, + ) + + # ------------------------------------------------------------------ + # Public helpers (used by /.well-known/jwks.json endpoint) + # ------------------------------------------------------------------ + + def get_jwks(self) -> Dict[str, Any]: + """ + Return the JWKS (JSON Web Key Set) for the RSA public key. + Used by GET /.well-known/jwks.json so MCP servers can verify tokens. + """ + public_numbers = self._public_key.public_numbers() + return { + "keys": [ + { + "kty": "RSA", + "alg": self.ALGORITHM, + "use": "sig", + "kid": self._kid, + "n": _int_to_base64url(public_numbers.n), + "e": _int_to_base64url(public_numbers.e), + } + ] + } + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + def _build_claims( + self, + user_api_key_dict: UserAPIKeyAuth, + data: dict, + ) -> Dict[str, Any]: + """ + Build JWT claims from the authenticated user context and MCP request data. + Follows RFC 8693 (OAuth 2.0 Token Exchange) for sub/act semantics. + """ + now = int(time.time()) + claims: Dict[str, Any] = { + "iss": self.issuer, + "aud": self.audience, + "iat": now, + "exp": now + self.ttl_seconds, + "nbf": now, + } + + # sub: End-user identity (RFC 8693) + user_id = getattr(user_api_key_dict, "user_id", None) + if user_id: + claims["sub"] = user_id + + user_email = getattr(user_api_key_dict, "user_email", None) + if user_email: + claims["email"] = user_email + + # act: Requester/agent identity (RFC 8693 delegation) + team_id = getattr(user_api_key_dict, "team_id", None) + org_id = getattr(user_api_key_dict, "org_id", None) + act_sub = team_id or org_id or "litellm-proxy" + claims["act"] = {"sub": act_sub} + + # end_user_id (if set separately from user_id) + end_user_id = getattr(user_api_key_dict, "end_user_id", None) + if end_user_id: + claims["end_user_id"] = end_user_id + + # scope: tool-level access + tool_name: str = data.get("mcp_tool_name", "") + scopes = ["mcp:tools/call", "mcp:tools/list"] + if tool_name: + scopes.append(f"mcp:tools/{tool_name}:call") + claims["scope"] = " ".join(scopes) + + return claims + + # ------------------------------------------------------------------ + # Guardrail hook + # ------------------------------------------------------------------ + + @log_guardrail_information + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict, + call_type: CallTypesLiteral, + ) -> Optional[Union[Exception, str, dict]]: + """ + Signs a JWT and injects it as the outbound Authorization header for + MCP tool calls. All other call types pass through unchanged. + """ + if call_type != "call_mcp_tool": + return data + + claims = self._build_claims(user_api_key_dict, data) + + signed_token = jwt.encode( + claims, + self._private_key, + algorithm=self.ALGORITHM, + ) + + data["extra_headers"] = { + "Authorization": f"Bearer {signed_token}", + } + + verbose_proxy_logger.debug( + "MCPJWTSigner: signed JWT sub=%s act=%s tool=%s exp=%d", + claims.get("sub"), + claims.get("act", {}).get("sub"), + data.get("mcp_tool_name"), + claims["exp"], + ) + + return data diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 27fa27e6da3..f798f05380d 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -79,6 +79,7 @@ class SupportedGuardrailIntegrations(Enum): SEMANTIC_GUARD = "semantic_guard" MCP_END_USER_PERMISSION = "mcp_end_user_permission" BLOCK_CODE_EXECUTION = "block_code_execution" + MCP_JWT_SIGNER = "mcp_jwt_signer" class Role(Enum): diff --git a/tests/test_litellm/proxy/guardrails/test_mcp_jwt_signer.py b/tests/test_litellm/proxy/guardrails/test_mcp_jwt_signer.py new file mode 100644 index 00000000000..eca7299b7bc --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/test_mcp_jwt_signer.py @@ -0,0 +1,328 @@ +""" +Tests for the MCPJWTSigner built-in guardrail. + +Tests cover: + - RSA key generation and loading + - JWT signing and JWKS format + - Claim building (sub, act, scope) + - Hook fires for call_mcp_tool, skips other call types + - get_mcp_jwt_signer() singleton pattern +""" + +import base64 +import time +from typing import Any, Dict, Optional +from unittest.mock import MagicMock, patch + +import jwt +import pytest +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import rsa + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_user_api_key_dict( + user_id: str = "user-123", + team_id: str = "team-abc", + user_email: str = "user@example.com", + end_user_id: Optional[str] = None, +) -> MagicMock: + mock = MagicMock() + mock.user_id = user_id + mock.team_id = team_id + mock.user_email = user_email + mock.end_user_id = end_user_id + mock.org_id = None + return mock + + +def _decode_unverified(token: str) -> Dict[str, Any]: + return jwt.decode(token, options={"verify_signature": False}) + + +# --------------------------------------------------------------------------- +# Import target (inline so we can reset the singleton between tests) +# --------------------------------------------------------------------------- + + +def _make_signer(**kwargs: Any): + # Reset singleton before each signer creation to avoid cross-test pollution + import litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer as mod + + mod._mcp_jwt_signer_instance = None + + from litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer import ( + MCPJWTSigner, + ) + + return MCPJWTSigner( + guardrail_name="test-jwt-signer", + event_hook="pre_mcp_call", + default_on=True, + **kwargs, + ) + + +# --------------------------------------------------------------------------- +# Key generation tests +# --------------------------------------------------------------------------- + + +def test_auto_generates_rsa_keypair(): + """MCPJWTSigner auto-generates an RSA-2048 keypair when env var is unset.""" + signer = _make_signer() + assert signer._private_key is not None + assert signer._public_key is not None + assert signer._kid is not None and len(signer._kid) == 16 + + +def test_kid_is_deterministic(): + """Two signers built from the same key have the same kid.""" + signer1 = _make_signer() + private_pem = signer1._private_key.private_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PrivateFormat.TraditionalOpenSSL, + encryption_algorithm=serialization.NoEncryption(), + ).decode("utf-8") + + with patch.dict("os.environ", {"MCP_JWT_SIGNING_KEY": private_pem}): + signer2 = _make_signer() + + assert signer1._kid == signer2._kid + + +def test_load_key_from_env_var(): + """MCPJWTSigner loads a user-provided RSA key from the env var.""" + private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + pem = private_key.private_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PrivateFormat.TraditionalOpenSSL, + encryption_algorithm=serialization.NoEncryption(), + ).decode("utf-8") + + with patch.dict("os.environ", {"MCP_JWT_SIGNING_KEY": pem}): + signer = _make_signer() + + assert signer._kid is not None + + +# --------------------------------------------------------------------------- +# JWKS tests +# --------------------------------------------------------------------------- + + +def test_get_jwks_format(): + """get_jwks() returns a valid JWKS dict with RSA fields.""" + signer = _make_signer() + jwks = signer.get_jwks() + + assert "keys" in jwks + assert len(jwks["keys"]) == 1 + key = jwks["keys"][0] + + assert key["kty"] == "RSA" + assert key["alg"] == "RS256" + assert key["use"] == "sig" + assert key["kid"] == signer._kid + assert "n" in key and len(key["n"]) > 0 + assert "e" in key and key["e"] == "AQAB" # 65537 in base64url + + +def test_jwks_public_key_can_verify_signed_jwt(): + """A JWT signed by MCPJWTSigner can be verified using the JWKS public key.""" + signer = _make_signer(issuer="https://litellm.example.com", audience="mcp") + now = int(time.time()) + claims = {"iss": "https://litellm.example.com", "aud": "mcp", "iat": now, "exp": now + 300} + + token = jwt.encode(claims, signer._private_key, algorithm="RS256") + + # Reconstruct public key from JWKS + jwks = signer.get_jwks() + key_data = jwks["keys"][0] + n = int.from_bytes(base64.urlsafe_b64decode(key_data["n"] + "=="), byteorder="big") + e = int.from_bytes(base64.urlsafe_b64decode(key_data["e"] + "=="), byteorder="big") + from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicNumbers + pub_key = RSAPublicNumbers(e=e, n=n).public_key() + + decoded = jwt.decode( + token, + pub_key, + algorithms=["RS256"], + audience="mcp", + issuer="https://litellm.example.com", + ) + assert decoded["iss"] == "https://litellm.example.com" + + +# --------------------------------------------------------------------------- +# Claim building tests +# --------------------------------------------------------------------------- + + +def test_build_claims_standard_fields(): + """_build_claims() populates iss, aud, iat, exp, nbf.""" + signer = _make_signer(issuer="https://litellm.example.com", audience="mcp", ttl_seconds=300) + user_dict = _make_user_api_key_dict() + data = {"mcp_tool_name": "get_weather"} + + claims = signer._build_claims(user_dict, data) + + assert claims["iss"] == "https://litellm.example.com" + assert claims["aud"] == "mcp" + assert "iat" in claims + assert "exp" in claims + assert claims["exp"] - claims["iat"] == 300 + assert "nbf" in claims + + +def test_build_claims_identity(): + """_build_claims() sets sub from user_id and act from team_id (RFC 8693).""" + signer = _make_signer() + user_dict = _make_user_api_key_dict(user_id="user-xyz", team_id="team-eng") + data: Dict[str, Any] = {} + + claims = signer._build_claims(user_dict, data) + + assert claims["sub"] == "user-xyz" + assert claims["act"]["sub"] == "team-eng" + assert claims["email"] == "user@example.com" + + +def test_build_claims_scope_with_tool(): + """_build_claims() encodes tool-specific scope when mcp_tool_name is set.""" + signer = _make_signer() + user_dict = _make_user_api_key_dict() + data = {"mcp_tool_name": "search_web"} + + claims = signer._build_claims(user_dict, data) + + scopes = set(claims["scope"].split()) + assert "mcp:tools/call" in scopes + assert "mcp:tools/list" in scopes + assert "mcp:tools/search_web:call" in scopes + + +def test_build_claims_scope_without_tool(): + """_build_claims() omits per-tool scope when mcp_tool_name is not set.""" + signer = _make_signer() + user_dict = _make_user_api_key_dict() + data: Dict[str, Any] = {} + + claims = signer._build_claims(user_dict, data) + + scopes = set(claims["scope"].split()) + assert "mcp:tools/call" in scopes + assert "mcp:tools/list" in scopes + # No per-tool scope + for scope in scopes: + assert ":" not in scope.replace("mcp:", "") or scope.endswith(":call") is False or scope == "mcp:tools/call" + + +def test_build_claims_act_fallback_to_litellm_proxy(): + """_build_claims() falls back to 'litellm-proxy' when team_id and org_id are absent.""" + signer = _make_signer() + user_dict = _make_user_api_key_dict() + user_dict.team_id = None + user_dict.org_id = None + + claims = signer._build_claims(user_dict, {}) + + assert claims["act"]["sub"] == "litellm-proxy" + + +# --------------------------------------------------------------------------- +# Hook dispatch tests +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_hook_fires_for_call_mcp_tool(): + """async_pre_call_hook() injects Authorization header for call_mcp_tool.""" + signer = _make_signer(issuer="https://litellm.example.com", audience="mcp") + user_dict = _make_user_api_key_dict() + data = {"mcp_tool_name": "do_thing"} + + result = await signer.async_pre_call_hook( + user_api_key_dict=user_dict, + cache=MagicMock(), + data=data, + call_type="call_mcp_tool", + ) + + assert isinstance(result, dict) + assert "extra_headers" in result + assert result["extra_headers"]["Authorization"].startswith("Bearer ") + + +@pytest.mark.asyncio +async def test_hook_skips_non_mcp_call_types(): + """async_pre_call_hook() leaves data unchanged for non-MCP call types.""" + signer = _make_signer() + user_dict = _make_user_api_key_dict() + data = {"messages": [{"role": "user", "content": "hello"}]} + + for call_type in ("completion", "acompletion", "embedding", "list_mcp_tools"): + original_data = {**data} + result = await signer.async_pre_call_hook( + user_api_key_dict=user_dict, + cache=MagicMock(), + data=original_data, + call_type=call_type, # type: ignore[arg-type] + ) + assert "extra_headers" not in (result or {}), f"extra_headers should not be set for {call_type}" + + +@pytest.mark.asyncio +async def test_signed_token_is_verifiable(): + """The JWT injected by the hook can be verified against the JWKS public key.""" + signer = _make_signer(issuer="https://litellm.example.com", audience="mcp", ttl_seconds=300) + user_dict = _make_user_api_key_dict(user_id="alice", team_id="backend") + data = {"mcp_tool_name": "search"} + + result = await signer.async_pre_call_hook( + user_api_key_dict=user_dict, + cache=MagicMock(), + data=data, + call_type="call_mcp_tool", + ) + + assert isinstance(result, dict) + token = result["extra_headers"]["Authorization"].removeprefix("Bearer ") + + decoded = _decode_unverified(token) + assert decoded["sub"] == "alice" + assert decoded["act"]["sub"] == "backend" + assert "mcp:tools/search:call" in decoded["scope"] + assert decoded["iss"] == "https://litellm.example.com" + assert decoded["aud"] == "mcp" + + +# --------------------------------------------------------------------------- +# Singleton tests +# --------------------------------------------------------------------------- + + +def test_get_mcp_jwt_signer_returns_none_before_init(): + """get_mcp_jwt_signer() returns None before any MCPJWTSigner is created.""" + import litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer as mod + + mod._mcp_jwt_signer_instance = None + + from litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer import ( + get_mcp_jwt_signer, + ) + + assert get_mcp_jwt_signer() is None + + +def test_get_mcp_jwt_signer_returns_instance_after_init(): + """get_mcp_jwt_signer() returns the initialized signer instance.""" + from litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer import ( + get_mcp_jwt_signer, + ) + + signer = _make_signer() + assert get_mcp_jwt_signer() is signer From 9f443a31aa2558d4048fdce62150d11705194a1b Mon Sep 17 00:00:00 2001 From: Noah Nistler <60981020+noahnistler@users.noreply.github.com> Date: Tue, 17 Mar 2026 16:07:06 -0500 Subject: [PATCH 08/13] Update MCPServerManager to raise HTTPException with status code 400 for extra headers in OpenAPI-backed servers. Adjust tests to verify the correct status code and exception message. --- .../mcp_server/mcp_server_manager.py | 29 ++++++++++--------- .../mcp_server/test_mcp_hook_extra_headers.py | 2 +- 2 files changed, 17 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 3fa14e6ec25..75f0642b950 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2235,6 +2235,21 @@ class MCPServerManager: if "arguments" in hook_result: arguments = hook_result["arguments"] + # OpenAPI-backed servers cannot forward hook-injected headers — reject early + # before scheduling any background tasks to avoid orphaned asyncio.Tasks. + if mcp_server.spec_path and hook_result.get("extra_headers"): + raise HTTPException( + status_code=400, + detail={ + "error": ( + "pre_mcp_call hook returned extra_headers for an " + "OpenAPI-backed MCP server, which does not support " + "hook header injection. Use a regular MCP server " + "(SSE/HTTP transport) for hook header support." + ) + }, + ) + # Prepare tasks for during hooks tasks = [] if proxy_logging_obj: @@ -2251,20 +2266,8 @@ class MCPServerManager: # For OpenAPI servers, call the tool handler directly instead of via MCP client if mcp_server.spec_path: verbose_logger.debug( - "Calling OpenAPI tool %s directly via HTTP handler", name + f"Calling OpenAPI tool {name} directly via HTTP handler" ) - if hook_result.get("extra_headers"): - raise HTTPException( - status_code=500, - detail={ - "error": ( - "pre_mcp_call hook returned extra_headers for an " - "OpenAPI-backed MCP server, which does not support " - "hook header injection. Use a regular MCP server " - "(SSE/HTTP transport) for hook header support." - ) - }, - ) tasks.append( asyncio.create_task( self._call_openapi_tool_handler(mcp_server, name, arguments) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py index 6c97ff244e8..8362193bb0b 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py @@ -459,7 +459,7 @@ class TestCallToolFlowsHookHeaders: proxy_logging_obj=proxy_logging, ) - assert exc_info.value.status_code == 500 + assert exc_info.value.status_code == 400 assert "does not support hook header injection" in str( exc_info.value.detail ) From 8574543f81c631f9a3b9ad64aab65c521cf33654 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 17 Mar 2026 14:15:07 -0700 Subject: [PATCH 09/13] fix: address P1 issues in MCPJWTSigner - OpenAPI servers: warn + skip header injection instead of 500 - JWKS Cache-Control: 5min for auto-generated keys, 1h for persistent - sub claim: fallback to apikey:{token_hash} for anonymous callers - ttl_seconds: validate > 0 at init time --- .../mcp_server/discoverable_endpoints.py | 5 +- .../mcp_server/mcp_server_manager.py | 16 ++--- .../mcp_jwt_signer/mcp_jwt_signer.py | 32 +++++++++- .../mcp_server/test_mcp_hook_extra_headers.py | 34 +++++----- .../proxy/guardrails/test_mcp_jwt_signer.py | 62 +++++++++++++++++++ 5 files changed, 121 insertions(+), 28 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 0b12aba02d8..0e2a4c257e7 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -718,14 +718,15 @@ async def jwks_json(request: Request): if signer is not None: return JSONResponse( content=signer.get_jwks(), - headers={"Cache-Control": "public, max-age=3600"}, + headers={"Cache-Control": f"public, max-age={signer.jwks_max_age}"}, ) except ImportError: pass + # No signer active — return empty key set; short cache so activation is picked up quickly. return JSONResponse( content={"keys": []}, - headers={"Cache-Control": "public, max-age=3600"}, + headers={"Cache-Control": "public, max-age=60"}, ) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 3fa14e6ec25..0208f52968d 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2254,16 +2254,12 @@ class MCPServerManager: "Calling OpenAPI tool %s directly via HTTP handler", name ) if hook_result.get("extra_headers"): - raise HTTPException( - status_code=500, - detail={ - "error": ( - "pre_mcp_call hook returned extra_headers for an " - "OpenAPI-backed MCP server, which does not support " - "hook header injection. Use a regular MCP server " - "(SSE/HTTP transport) for hook header support." - ) - }, + verbose_logger.warning( + "pre_mcp_call hook returned extra_headers for OpenAPI-backed " + "MCP server '%s' — header injection is not supported for " + "OpenAPI servers; headers will be ignored. Use SSE/HTTP " + "transport to enable hook header injection.", + server_name, ) tasks.append( asyncio.create_task( diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py index 39ffc2aed88..12a06da2897 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py @@ -130,11 +130,13 @@ class MCPJWTSigner(CustomGuardrail): key_material = os.environ.get(self.SIGNING_KEY_ENV) if key_material: self._private_key = _load_private_key_from_env(self.SIGNING_KEY_ENV) + self._persistent_key: bool = True verbose_proxy_logger.info( "MCPJWTSigner: loaded RSA key from env var %s", self.SIGNING_KEY_ENV ) else: self._private_key = _generate_rsa_key_pair() + self._persistent_key = False verbose_proxy_logger.info( "MCPJWTSigner: auto-generated RSA-2048 keypair (set %s to use your own key)", self.SIGNING_KEY_ENV, @@ -154,11 +156,16 @@ class MCPJWTSigner(CustomGuardrail): or os.environ.get("MCP_JWT_AUDIENCE") or self.DEFAULT_AUDIENCE ) - self.ttl_seconds: int = int( + resolved_ttl = int( ttl_seconds if ttl_seconds is not None else os.environ.get("MCP_JWT_TTL_SECONDS", str(self.DEFAULT_TTL)) ) + if resolved_ttl <= 0: + raise ValueError( + f"MCPJWTSigner: ttl_seconds must be > 0, got {resolved_ttl}" + ) + self.ttl_seconds: int = resolved_ttl # Register singleton so the JWKS endpoint can access it. global _mcp_jwt_signer_instance @@ -176,6 +183,17 @@ class MCPJWTSigner(CustomGuardrail): # Public helpers (used by /.well-known/jwks.json endpoint) # ------------------------------------------------------------------ + @property + def jwks_max_age(self) -> int: + """ + Recommended Cache-Control max-age for the JWKS response (seconds). + + Use 1 hour for persistent keys (loaded from env var) — safe to cache long. + Use 5 minutes for auto-generated keys — key rotates on every restart, so + MCP servers must re-fetch quickly to avoid verifying with a stale key. + """ + return 3600 if self._persistent_key else 300 + def get_jwks(self) -> Dict[str, Any]: """ Return the JWKS (JSON Web Key Set) for the RSA public key. @@ -217,10 +235,20 @@ class MCPJWTSigner(CustomGuardrail): "nbf": now, } - # sub: End-user identity (RFC 8693) + # sub: End-user identity (RFC 8693). + # Falls back to a stable hash of the API token for service-account / anonymous + # callers so strict JWT consumers (which require sub) always get a value. user_id = getattr(user_api_key_dict, "user_id", None) if user_id: claims["sub"] = user_id + else: + token = getattr(user_api_key_dict, "token", None) or getattr( + user_api_key_dict, "api_key", None + ) + if token: + claims["sub"] = "apikey:" + hashlib.sha256(token.encode()).hexdigest()[:16] + else: + claims["sub"] = "litellm-proxy" user_email = getattr(user_api_key_dict, "user_email", None) if user_email: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py index 6c97ff244e8..40242f70e38 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py @@ -422,8 +422,8 @@ class TestCallToolFlowsHookHeaders: assert call_kwargs.kwargs.get("arguments") == modified_args @pytest.mark.asyncio - async def test_openapi_server_raises_on_hook_headers(self): - """OpenAPI-backed servers should raise HTTPException when hook injects headers.""" + async def test_openapi_server_warns_and_continues_on_hook_headers(self): + """OpenAPI-backed servers log a warning and continue when hook injects headers.""" manager = MCPServerManager() server = MCPServer( server_id="test-id", @@ -449,20 +449,26 @@ class TestCallToolFlowsHookHeaders: "_create_during_hook_task", return_value=asyncio.create_task(asyncio.sleep(0)), ): - proxy_logging = MagicMock(spec=ProxyLogging) + with patch.object( + manager, + "_call_openapi_tool_handler", + new_callable=AsyncMock, + return_value=MagicMock(), + ): + import litellm.proxy._experimental.mcp_server.mcp_server_manager as mgr_mod - with pytest.raises(HTTPException) as exc_info: - await manager.call_tool( - server_name="openapi_server", - name="test_tool", - arguments={}, - proxy_logging_obj=proxy_logging, - ) + proxy_logging = MagicMock(spec=ProxyLogging) - assert exc_info.value.status_code == 500 - assert "does not support hook header injection" in str( - exc_info.value.detail - ) + with patch.object(mgr_mod, "verbose_logger") as mock_logger: + # Should NOT raise — just warn and proceed + await manager.call_tool( + server_name="openapi_server", + name="test_tool", + arguments={}, + proxy_logging_obj=proxy_logging, + ) + mock_logger.warning.assert_called_once() + assert "header injection is not supported" in mock_logger.warning.call_args[0][0] @pytest.mark.asyncio async def test_openapi_server_no_error_without_hook_headers(self): diff --git a/tests/test_litellm/proxy/guardrails/test_mcp_jwt_signer.py b/tests/test_litellm/proxy/guardrails/test_mcp_jwt_signer.py index eca7299b7bc..bc0b398cd84 100644 --- a/tests/test_litellm/proxy/guardrails/test_mcp_jwt_signer.py +++ b/tests/test_litellm/proxy/guardrails/test_mcp_jwt_signer.py @@ -233,6 +233,68 @@ def test_build_claims_act_fallback_to_litellm_proxy(): assert claims["act"]["sub"] == "litellm-proxy" +def test_build_claims_sub_fallback_to_token_hash(): + """_build_claims() sets sub to an apikey: hash when user_id is absent.""" + signer = _make_signer() + user_dict = _make_user_api_key_dict(user_id="") + user_dict.user_id = None + user_dict.token = "sk-test-api-key-abc123" + + claims = signer._build_claims(user_dict, {}) + + assert claims["sub"].startswith("apikey:") + assert len(claims["sub"]) == len("apikey:") + 16 # sha256 hex[:16] + + +def test_build_claims_sub_fallback_to_litellm_proxy_when_no_token(): + """_build_claims() falls back to 'litellm-proxy' when user_id and token are both absent.""" + signer = _make_signer() + user_dict = _make_user_api_key_dict(user_id="") + user_dict.user_id = None + user_dict.token = None + user_dict.api_key = None + + claims = signer._build_claims(user_dict, {}) + + assert claims["sub"] == "litellm-proxy" + + +def test_init_raises_on_zero_ttl(): + """MCPJWTSigner raises ValueError when ttl_seconds is 0.""" + with pytest.raises(ValueError, match="ttl_seconds must be > 0"): + _make_signer(ttl_seconds=0) + + +def test_init_raises_on_negative_ttl(): + """MCPJWTSigner raises ValueError when ttl_seconds is negative.""" + with pytest.raises(ValueError, match="ttl_seconds must be > 0"): + _make_signer(ttl_seconds=-60) + + +def test_jwks_max_age_persistent_key(): + """jwks_max_age is 3600 when key loaded from env var.""" + from cryptography.hazmat.primitives import serialization + from cryptography.hazmat.primitives.asymmetric import rsa as crsa + + private_key = crsa.generate_private_key(public_exponent=65537, key_size=2048) + pem = private_key.private_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PrivateFormat.TraditionalOpenSSL, + encryption_algorithm=serialization.NoEncryption(), + ).decode("utf-8") + + with patch.dict("os.environ", {"MCP_JWT_SIGNING_KEY": pem}): + signer = _make_signer() + + assert signer.jwks_max_age == 3600 + + +def test_jwks_max_age_auto_generated_key(): + """jwks_max_age is 300 for auto-generated (ephemeral) keys.""" + signer = _make_signer() + assert signer.jwks_max_age == 300 + + # --------------------------------------------------------------------------- # Hook dispatch tests # --------------------------------------------------------------------------- From bb6a9aa9e46e0927697104bcd4f71e4ccdde8867 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 17 Mar 2026 14:16:07 -0700 Subject: [PATCH 10/13] docs: add MCP zero trust auth guide with architecture diagram --- docs/my-website/docs/mcp_zero_trust.md | 99 ++++++++++++++++++++++++++ docs/my-website/sidebars.js | 1 + 2 files changed, 100 insertions(+) create mode 100644 docs/my-website/docs/mcp_zero_trust.md diff --git a/docs/my-website/docs/mcp_zero_trust.md b/docs/my-website/docs/mcp_zero_trust.md new file mode 100644 index 00000000000..aefd14f6eae --- /dev/null +++ b/docs/my-website/docs/mcp_zero_trust.md @@ -0,0 +1,99 @@ +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +# MCP Zero Trust Auth (JWT Signer) + +The `MCPJWTSigner` guardrail signs every outbound MCP tool call with a LiteLLM-issued RS256 JWT. MCP servers validate tokens against LiteLLM's JWKS endpoint instead of trusting each upstream IdP directly. + +## Architecture + +```mermaid +sequenceDiagram + participant Client + participant LiteLLM + participant JWKS as LiteLLM JWKS
/.well-known/jwks.json + participant MCP as MCP Server + + Client->>LiteLLM: tool call (Bearer API key / JWT) + Note over LiteLLM: MCPJWTSigner.async_pre_call_hook()
builds RS256 JWT:
sub=user_id, act=team_id,
scope=mcp:tools/{name}:call + + LiteLLM->>MCP: call_tool(args)
Authorization: Bearer + MCP->>JWKS: GET /.well-known/jwks.json + JWKS-->>MCP: RSA public key (JWKS) + MCP->>MCP: verify JWT signature + claims + MCP-->>LiteLLM: tool result + LiteLLM-->>Client: response +``` + +### OIDC Discovery + +LiteLLM publishes standard OIDC discovery so MCP servers can find the signing key automatically: + +``` +GET /.well-known/openid-configuration +→ { "jwks_uri": "https:///.well-known/jwks.json", ... } + +GET /.well-known/jwks.json +→ { "keys": [{ "kty": "RSA", "alg": "RS256", "kid": "...", "n": "...", "e": "..." }] } +``` + +## Setup + +### 1. Enable in `config.yaml` + +```yaml title="config.yaml" +guardrails: + - guardrail_name: "mcp-jwt-signer" + litellm_params: + guardrail: mcp_jwt_signer + mode: pre_mcp_call + default_on: true + issuer: "https://my-litellm.example.com" # optional — defaults to request base URL + audience: "mcp" # optional — default: "mcp" + ttl_seconds: 300 # optional — default: 300 +``` + +### 2. (Optional) Bring your own RSA key + +If unset, LiteLLM auto-generates an RSA-2048 keypair at startup (lost on restart). + +```bash +# PEM string +export MCP_JWT_SIGNING_KEY="-----BEGIN RSA PRIVATE KEY-----\n..." + +# Or point to a file +export MCP_JWT_SIGNING_KEY="file:///secrets/mcp-signing-key.pem" +``` + +### 3. Configure your MCP server to verify tokens + +Point your MCP server at LiteLLM's OIDC discovery endpoint: + +``` +https:///.well-known/openid-configuration +``` + +Most JWT middleware (e.g. `python-jose`, `jsonwebtoken`, AWS Lambda authorizers) supports OIDC auto-discovery. + +## JWT Claims + +| Claim | Value | RFC | +|-------|-------|-----| +| `iss` | LiteLLM issuer URL | RFC 7519 | +| `aud` | configured `audience` | RFC 7519 | +| `sub` | `user_api_key_dict.user_id` | RFC 8693 | +| `act.sub` | `team_id` → `org_id` → `"litellm-proxy"` | RFC 8693 delegation | +| `email` | `user_api_key_dict.user_email` (if set) | — | +| `scope` | `mcp:tools/call mcp:tools/list mcp:tools/{name}:call` | — | +| `iat`, `exp`, `nbf` | standard timing | RFC 7519 | + +## Limitations + +- **OpenAPI-backed MCP servers** (`spec_path` set) do not support hook header injection. Calls to these servers will fail with a 500 if `MCPJWTSigner` is active with `default_on: true`. Use SSE/HTTP transport MCP servers instead. +- The keypair is **in-memory by default** — rotated on every restart unless `MCP_JWT_SIGNING_KEY` is set. MCP servers should re-fetch JWKS on verification failure (short JWKS cache TTL recommended). + +## Related + +- [MCP Guardrails](./mcp_guardrail) — PII masking and blocking for MCP calls +- [MCP OAuth](./mcp_oauth) — upstream OAuth2 for MCP server access +- [MCP AWS SigV4](./mcp_aws_sigv4) — AWS-signed requests to MCP servers diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index 1362745a91f..46e1bd041e6 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -636,6 +636,7 @@ const sidebars = { "mcp_control", "mcp_cost", "mcp_guardrail", + "mcp_zero_trust", "mcp_troubleshoot", ] }, From 18fc3066eed8559607afa8eb097e8697a6209551 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 17 Mar 2026 14:37:01 -0700 Subject: [PATCH 11/13] docs: add FastMCP JWT verification guide to zero trust doc --- docs/my-website/docs/mcp_zero_trust.md | 80 +++++++++++++++++++++++--- 1 file changed, 73 insertions(+), 7 deletions(-) diff --git a/docs/my-website/docs/mcp_zero_trust.md b/docs/my-website/docs/mcp_zero_trust.md index aefd14f6eae..282d8b29e9e 100644 --- a/docs/my-website/docs/mcp_zero_trust.md +++ b/docs/my-website/docs/mcp_zero_trust.md @@ -65,15 +65,81 @@ export MCP_JWT_SIGNING_KEY="-----BEGIN RSA PRIVATE KEY-----\n..." export MCP_JWT_SIGNING_KEY="file:///secrets/mcp-signing-key.pem" ``` -### 3. Configure your MCP server to verify tokens +### 3. Build a verified MCP server with FastMCP -Point your MCP server at LiteLLM's OIDC discovery endpoint: +[FastMCP](https://gofastmcp.com) has a built-in `JWTVerifier` that fetches LiteLLM's JWKS automatically, handles key rotation, and enforces `iss`/`aud`/`exp` — zero boilerplate. -``` -https:///.well-known/openid-configuration +**Install:** +```bash +pip install fastmcp PyJWT cryptography ``` -Most JWT middleware (e.g. `python-jose`, `jsonwebtoken`, AWS Lambda authorizers) supports OIDC auto-discovery. +**`weather_server.py`:** +```python +from fastmcp import FastMCP, Context +from fastmcp.server.auth.providers.jwt import JWTVerifier + +LITELLM_BASE_URL = "https://my-litellm.example.com" + +# Point JWTVerifier at LiteLLM's JWKS endpoint. +# It auto-fetches and caches the RSA public key — no key material to manage. +auth = JWTVerifier( + jwks_uri=f"{LITELLM_BASE_URL}/.well-known/jwks.json", + issuer=LITELLM_BASE_URL, # must match MCPJWTSigner `issuer:` in config.yaml + audience="mcp", # must match MCPJWTSigner `audience:` + algorithm="RS256", +) + +mcp = FastMCP("weather-server", auth=auth) + + +@mcp.tool() +async def get_weather(city: str, ctx: Context) -> str: + """Return weather for a city. Caller identity comes from the verified JWT.""" + caller = ctx.client_id # = JWT `sub` claim (user_id or apikey hash) + await ctx.info(f"Request from {caller}") + return f"Weather in {city}: sunny, 72°F" + + +if __name__ == "__main__": + mcp.run(transport="http", host="0.0.0.0", port=8000) +``` + +`ctx.client_id` is populated from the JWT `sub` claim after verification — you get the caller's identity for free with no extra code. + +**Wire it into LiteLLM `config.yaml`:** +```yaml title="config.yaml" +mcp_servers: + - server_name: weather + url: http://localhost:8000/mcp + transport: http + +guardrails: + - guardrail_name: mcp-jwt-signer + litellm_params: + guardrail: mcp_jwt_signer + mode: pre_mcp_call + default_on: true + issuer: "https://my-litellm.example.com" + audience: "mcp" +``` + +**Run and test:** +```bash +# Terminal 1 — start the MCP server +python weather_server.py + +# Terminal 2 — start LiteLLM +litellm --config config.yaml + +# Terminal 3 — call through LiteLLM (JWT is injected automatically) +curl -X POST http://localhost:4000/mcp/weather/call_tool \ + -H "Authorization: Bearer $LITELLM_API_KEY" \ + -H "Content-Type: application/json" \ + -d '{"name": "get_weather", "arguments": {"city": "San Francisco"}}' +``` + +LiteLLM signs the JWT, sends it to the weather server, and FastMCP verifies it in one round-trip. A request without a valid token gets a `401` back from FastMCP before any tool code runs. ## JWT Claims @@ -89,8 +155,8 @@ Most JWT middleware (e.g. `python-jose`, `jsonwebtoken`, AWS Lambda authorizers) ## Limitations -- **OpenAPI-backed MCP servers** (`spec_path` set) do not support hook header injection. Calls to these servers will fail with a 500 if `MCPJWTSigner` is active with `default_on: true`. Use SSE/HTTP transport MCP servers instead. -- The keypair is **in-memory by default** — rotated on every restart unless `MCP_JWT_SIGNING_KEY` is set. MCP servers should re-fetch JWKS on verification failure (short JWKS cache TTL recommended). +- **OpenAPI-backed MCP servers** (`spec_path` set) do not support hook header injection. When `MCPJWTSigner` is active, calls to these servers log a warning and the JWT header is skipped. Use SSE/HTTP transport MCP servers to get full JWT injection. +- The keypair is **in-memory by default** — rotated on every restart unless `MCP_JWT_SIGNING_KEY` is set. FastMCP's `JWTVerifier` automatically re-fetches JWKS on key ID miss, so rotation is handled transparently. ## Related From 9cceff757d178fe7cf2572e33cac98b0e50bc3f6 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 17 Mar 2026 14:44:27 -0700 Subject: [PATCH 12/13] fix: address remaining Greptile review issues (round 2) - mcp_server_manager: warn when hook Authorization overwrites existing header - __init__: remove _mcp_jwt_signer_instance from __all__ (private internal) - discoverable_endpoints: copy dict instead of mutating in-place on OIDC augmentation - test docstring: reflect warn-and-continue behavior for OpenAPI servers - test: update scope assertions for least-privilege (no mcp:tools/list on tool-call JWTs) --- .../mcp_server/discoverable_endpoints.py | 7 +++-- .../mcp_server/mcp_server_manager.py | 6 ++++ .../mcp_jwt_signer/__init__.py | 3 +- .../mcp_jwt_signer/mcp_jwt_signer.py | 31 ++++++++++++++----- .../mcp_server/test_mcp_hook_extra_headers.py | 2 +- .../proxy/guardrails/test_mcp_jwt_signer.py | 10 +++--- 6 files changed, 42 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 0e2a4c257e7..3385e7feef6 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -691,8 +691,11 @@ async def openid_configuration(request: Request): if signer is not None: request_base_url = get_request_base_url(request) if isinstance(response, dict): - response["jwks_uri"] = f"{request_base_url}/.well-known/jwks.json" - response["id_token_signing_alg_values_supported"] = ["RS256"] + response = { + **response, + "jwks_uri": f"{request_base_url}/.well-known/jwks.json", + "id_token_signing_alg_values_supported": ["RS256"], + } except ImportError: pass diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 0208f52968d..c66dc307089 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2135,6 +2135,12 @@ class MCPServerManager: if hook_extra_headers: if extra_headers is None: extra_headers = {} + if "Authorization" in extra_headers and "Authorization" in hook_extra_headers: + verbose_logger.warning( + "MCPServerManager: hook_extra_headers contains 'Authorization' which will " + "overwrite the existing Authorization header set by static_headers or server " + "authentication. The hook JWT will take precedence." + ) extra_headers.update(hook_extra_headers) stdio_env = self._build_stdio_env(mcp_server, raw_headers) diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/__init__.py index d4a6c579097..230edaec855 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/__init__.py @@ -4,7 +4,7 @@ from typing import TYPE_CHECKING from litellm.types.guardrails import SupportedGuardrailIntegrations -from .mcp_jwt_signer import MCPJWTSigner, _mcp_jwt_signer_instance, get_mcp_jwt_signer +from .mcp_jwt_signer import MCPJWTSigner, get_mcp_jwt_signer if TYPE_CHECKING: from litellm.types.guardrails import Guardrail, LitellmParams @@ -51,6 +51,5 @@ guardrail_class_registry = { __all__ = [ "MCPJWTSigner", "initialize_guardrail", - "_mcp_jwt_signer_instance", "get_mcp_jwt_signer", ] diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py index 12a06da2897..4fe51ba2fbc 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py @@ -27,6 +27,7 @@ your own RSA keypair. If unset, an RSA-2048 keypair is auto-generated at startup import base64 import hashlib import os +import re import time from typing import Any, Dict, Optional, Union @@ -169,6 +170,12 @@ class MCPJWTSigner(CustomGuardrail): # Register singleton so the JWKS endpoint can access it. global _mcp_jwt_signer_instance + if _mcp_jwt_signer_instance is not None: + verbose_proxy_logger.warning( + "MCPJWTSigner: replacing existing singleton — previously issued tokens " + "signed with the old key will fail JWKS verification. " + "Avoid configuring multiple mcp_jwt_signer guardrails." + ) _mcp_jwt_signer_instance = self verbose_proxy_logger.info( @@ -265,11 +272,19 @@ class MCPJWTSigner(CustomGuardrail): if end_user_id: claims["end_user_id"] = end_user_id - # scope: tool-level access - tool_name: str = data.get("mcp_tool_name", "") - scopes = ["mcp:tools/call", "mcp:tools/list"] + # scope: minimal tool-level access. + # Only grant mcp:tools/list when no specific tool is being called — + # tool call JWTs should not carry enumeration permissions. + # Tool names are sanitized (alphanumeric + _ and -) before embedding + # so path-traversal or malformed scope values cannot be injected. + import re + + raw_tool_name: str = data.get("mcp_tool_name", "") + tool_name = re.sub(r"[^a-zA-Z0-9_\-]", "_", raw_tool_name) if raw_tool_name else "" if tool_name: - scopes.append(f"mcp:tools/{tool_name}:call") + scopes = ["mcp:tools/call", f"mcp:tools/{tool_name}:call"] + else: + scopes = ["mcp:tools/call", "mcp:tools/list"] claims["scope"] = " ".join(scopes) return claims @@ -301,9 +316,11 @@ class MCPJWTSigner(CustomGuardrail): algorithm=self.ALGORITHM, ) - data["extra_headers"] = { - "Authorization": f"Bearer {signed_token}", - } + # Merge into existing extra_headers rather than replacing — a prior guardrail + # in the chain may have already injected headers (e.g. tracing, correlation IDs). + # MCPJWTSigner sets Authorization last so its JWT takes precedence. + existing_headers: Dict[str, str] = data.get("extra_headers") or {} + data["extra_headers"] = {**existing_headers, "Authorization": f"Bearer {signed_token}"} verbose_proxy_logger.debug( "MCPJWTSigner: signed JWT sub=%s act=%s tool=%s exp=%d", diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py index 40242f70e38..32f3a340855 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py @@ -6,7 +6,7 @@ Validates that: 2. pre_call_tool_check returns hook-provided extra_headers AND modified arguments 3. call_tool flows hook headers and modified arguments downstream 4. Hook-provided headers take highest priority (merge after static_headers) -5. OpenAPI-backed servers raise HTTPException when hook headers are present +5. OpenAPI-backed servers log a warning and continue (skip injection) when hook headers are present 6. JWT claims are propagated in both standard and virtual-key fast paths 7. Backward compatibility: hooks without extra_headers continue to work """ diff --git a/tests/test_litellm/proxy/guardrails/test_mcp_jwt_signer.py b/tests/test_litellm/proxy/guardrails/test_mcp_jwt_signer.py index bc0b398cd84..991f813c675 100644 --- a/tests/test_litellm/proxy/guardrails/test_mcp_jwt_signer.py +++ b/tests/test_litellm/proxy/guardrails/test_mcp_jwt_signer.py @@ -201,12 +201,13 @@ def test_build_claims_scope_with_tool(): scopes = set(claims["scope"].split()) assert "mcp:tools/call" in scopes - assert "mcp:tools/list" in scopes assert "mcp:tools/search_web:call" in scopes + # Tool-call JWTs must NOT carry mcp:tools/list — least-privilege + assert "mcp:tools/list" not in scopes def test_build_claims_scope_without_tool(): - """_build_claims() omits per-tool scope when mcp_tool_name is not set.""" + """_build_claims() includes mcp:tools/list when no specific tool is called.""" signer = _make_signer() user_dict = _make_user_api_key_dict() data: Dict[str, Any] = {} @@ -216,9 +217,8 @@ def test_build_claims_scope_without_tool(): scopes = set(claims["scope"].split()) assert "mcp:tools/call" in scopes assert "mcp:tools/list" in scopes - # No per-tool scope - for scope in scopes: - assert ":" not in scope.replace("mcp:", "") or scope.endswith(":call") is False or scope == "mcp:tools/call" + # No per-tool call scope when no tool name was given + assert not any(s.endswith(":call") and s != "mcp:tools/call" for s in scopes) def test_build_claims_act_fallback_to_litellm_proxy(): From 8acb06dc68eb840aa2f8d8f966801be0bef20032 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 17 Mar 2026 14:56:45 -0700 Subject: [PATCH 13/13] fix: address Greptile round 3 feedback MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - initialize_guardrail: validate mode='pre_mcp_call' at init time — misconfigured mode silently bypasses JWT injection, which is a zero-trust bypass - _build_claims: remove duplicate inline 'import re' (module-level import already present) - _types.py: add TODO comment explaining jwt_claims is forward-compat plumbing for a follow-up PR that will forward upstream IdP claims into outbound MCP JWTs --- litellm/proxy/_types.py | 3 +++ .../guardrails/guardrail_hooks/mcp_jwt_signer/__init__.py | 7 +++++++ .../guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py | 2 -- 3 files changed, 10 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 55cb1c9c43c..240123ba6df 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2471,6 +2471,9 @@ class UserAPIKeyAuth( Any ] = None # Expanded created_by user when expand=user is used end_user_object_permission: Optional[LiteLLM_ObjectPermissionTable] = None + # TODO: jwt_claims carries decoded upstream IdP claims (groups, roles, etc.) so + # guardrails can forward them into outbound tokens (e.g. MCPJWTSigner). Currently + # populated but not yet consumed — forward-compat hook for a follow-up PR. jwt_claims: Optional[Dict] = None model_config = ConfigDict(arbitrary_types_allowed=True) diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/__init__.py index 230edaec855..81364448997 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/__init__.py @@ -19,6 +19,13 @@ def initialize_guardrail( if not guardrail_name: raise ValueError("MCPJWTSigner guardrail requires a guardrail_name") + mode = litellm_params.mode + if mode != "pre_mcp_call": + raise ValueError( + f"MCPJWTSigner guardrail '{guardrail_name}' has mode='{mode}' but must use " + "mode='pre_mcp_call'. JWT injection only fires for MCP tool calls." + ) + optional_params = getattr(litellm_params, "optional_params", None) def _get(key): # type: ignore[no-untyped-def] diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py index 4fe51ba2fbc..3321515d0ff 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py @@ -277,8 +277,6 @@ class MCPJWTSigner(CustomGuardrail): # tool call JWTs should not carry enumeration permissions. # Tool names are sanitized (alphanumeric + _ and -) before embedding # so path-traversal or malformed scope values cannot be injected. - import re - raw_tool_name: str = data.get("mcp_tool_name", "") tool_name = re.sub(r"[^a-zA-Z0-9_\-]", "_", raw_tool_name) if raw_tool_name else "" if tool_name: