mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-05 02:41:56 +00:00
Allow pre_mcp_call guardrail hooks to mutate outbound MCP headers
This commit is contained in:
Noah Nistler
parent
ef9cc33ee3
commit
89a43cfd60
5 changed files with 513 additions and 4 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
Loading…
Add table
Reference in a new issue