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] 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"]