scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-20 00:11:50 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 134

fix(mcp): preserve request-selected guardrails during tool execution

Browse source
This commit is contained in:
Joshua Valluru 2026-09-17 09:57:08 -07:00
parent 4b368bf066
commit 743684bdbe
19 changed files with 363 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py
View file

@ -8,6 +8,7 @@ tool through a ``tool_use`` content block, and results are fed back as
"""
from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mapping, Sequence
from types import MappingProxyType
from typing import Any, Final, NamedTuple
from litellm._logging import verbose_logger
@ -94,7 +95,7 @@ async def anthropic_messages_with_mcp(
**kwargs,
)
context: Final = MCPRequestContext.resolve(kwargs=dict(kwargs), tools=tools)
context: Final = MCPRequestContext.resolve(kwargs=MappingProxyType({**kwargs, "model": model}), tools=tools)
(
deduplicated_mcp_tools,
@ -155,6 +156,7 @@ async def anthropic_messages_with_mcp(
litellm_call_id=context.litellm_call_id,
litellm_trace_id=context.litellm_trace_id,
request_tags=list(context.request_tags) if context.request_tags else None,
guardrail_context=context.guardrail_context,
)
# Every tool call was skipped, so there is nothing to feed back; a

7
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -5592,6 +5592,7 @@ class MCPServerManager:
server: MCPServer,
raw_headers: dict[str, str] | None = None,
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
guardrail_context: Mapping[str, object] | None = None,
) -> dict[str, Any]:
"""
Run pre-call checks and guardrail hooks for an MCP tool call.
@ -5645,6 +5646,7 @@ class MCPServerManager:
incoming_bearer_token = auth_hdr[len("bearer ") :]
pre_hook_kwargs: Final = {
"guardrail_context": guardrail_context,
"name": name,
"arguments": arguments,
"server_name": server_name,
@ -5712,6 +5714,7 @@ class MCPServerManager:
proxy_logging_obj: ProxyLogging,
start_time: datetime.datetime,
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
guardrail_context: Mapping[str, object] | None = None,
):
"""Create and return a during hook task for MCP tool calls.
@ -5731,6 +5734,7 @@ class MCPServerManager:
)
during_hook_kwargs: Final = {
"guardrail_context": guardrail_context,
"name": name,
"arguments": arguments,
"server_name": server_name_from_prefix,
@ -6276,6 +6280,7 @@ class MCPServerManager:
raw_headers: dict[str, str] | None = None,
host_progress_callback: Callable | None = None,
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
guardrail_context: Mapping[str, object] | None = None,
) -> CallToolResult:
"""
Call a tool with the given name and arguments
@ -6322,6 +6327,7 @@ class MCPServerManager:
server=mcp_server,
raw_headers=raw_headers,
litellm_logging_obj=litellm_logging_obj,
guardrail_context=guardrail_context,
)
if "arguments" in hook_result:
arguments = hook_result["arguments"]
@ -6337,6 +6343,7 @@ class MCPServerManager:
proxy_logging_obj=proxy_logging_obj,
start_time=start_time,
litellm_logging_obj=litellm_logging_obj,
guardrail_context=guardrail_context,
)
tasks.append(during_hook_task)

6
litellm/proxy/_experimental/mcp_server/rest_endpoints.py
View file

@ -51,6 +51,7 @@ from litellm.proxy._types import (
)
from litellm.proxy.auth.ip_address_utils import IPAddressUtils
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.responses.mcp.request_context import MCPRequestContext
if TYPE_CHECKING:
from mcp.types import CallToolResult
@ -1168,6 +1169,7 @@ if MCP_AVAILABLE:
oauth2_headers=user_oauth_extra_headers or data.get("oauth2_headers"),
raw_headers=data.get("raw_headers"),
litellm_logging_obj=data.get("litellm_logging_obj"),
guardrail_context=MCPRequestContext.resolve_guardrail_context(data),
requested_server_id=canonical_server_id,
)
except Exception as e:
@ -1212,8 +1214,8 @@ if MCP_AVAILABLE:
"guardrail_name": getattr(e, "guardrail_name", None),
},
)
except GuardrailRaisedException as e:
verbose_logger.error("GuardrailRaisedException in MCP tool call: %s", e)
except (GuardrailRaisedException, ModifyResponseException) as e:
verbose_logger.error("Guardrail violation in MCP tool call: %s", e)
raise HTTPException(
status_code=400,
detail={

6
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -2927,6 +2927,7 @@ if MCP_AVAILABLE:
oauth2_headers: dict[str, str] | None = None,
raw_headers: dict[str, str] | None = None,
host_progress_callback: Callable | None = None,
guardrail_context: Mapping[str, object] | None = None,
**kwargs: Any,
) -> CallToolResult:
"""
@ -3115,6 +3116,7 @@ if MCP_AVAILABLE:
server=mcp_server,
raw_headers=raw_headers,
litellm_logging_obj=litellm_logging_obj,
guardrail_context=guardrail_context,
)
# `pre_call_tool_check` may return guardrail-modified
# arguments; honor them on the local path too.
@ -3168,6 +3170,7 @@ if MCP_AVAILABLE:
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
litellm_logging_obj=litellm_logging_obj,
guardrail_context=guardrail_context,
host_progress_callback=host_progress_callback,
)
@ -3221,6 +3224,7 @@ if MCP_AVAILABLE:
server=prefix_server,
raw_headers=raw_headers,
litellm_logging_obj=litellm_logging_obj,
guardrail_context=guardrail_context,
)
if "arguments" in hook_result:
arguments = hook_result["arguments"] # pyright: ignore[reportAny] # hook returns untyped args
@ -3598,6 +3602,7 @@ if MCP_AVAILABLE:
raw_headers: dict[str, str] | None = None,
litellm_logging_obj: LiteLLMLoggingObj | None = None,
host_progress_callback: Callable | None = None,
guardrail_context: Mapping[str, object] | None = None,
) -> CallToolResult:
"""Handle tool execution for managed server tools"""
# Import here to avoid circular import
@ -3615,6 +3620,7 @@ if MCP_AVAILABLE:
proxy_logging_obj=proxy_logging_obj,
host_progress_callback=host_progress_callback,
litellm_logging_obj=litellm_logging_obj,
guardrail_context=guardrail_context,
)
verbose_logger.debug("CALL TOOL RESULT: %s", call_tool_result)
return call_tool_result

34
litellm/proxy/utils.py
View file

@ -1246,15 +1246,31 @@ class ProxyLogging:
"""
from litellm.types.llms.openai import ChatCompletionUserMessage
guardrail_context: Final = TypeAdapter(Mapping[str, object]).validate_python(
kwargs.get("guardrail_context") or MappingProxyType({})
)
parent_metadata: Final = copy.deepcopy(
TypeAdapter(dict[str, object]).validate_python(guardrail_context.get("metadata") or MappingProxyType({}))
)
# Create a synthetic message that represents the tool call
tool_call_content: Final = f"Tool: {request_obj.tool_name}\nArguments: {request_obj.arguments}"
synthetic_message: Final = ChatCompletionUserMessage(role="user", content=tool_call_content)
synthetic_metadata: Final[dict[str, object]] = { # mutable-ok: existing guardrail hooks mutate request metadata
**MappingProxyType({key: value for key, value in parent_metadata.items() if key != "guardrails"}),
"headers": kwargs.get("headers") or {},
"user_api_key_user_id": kwargs.get("user_api_key_user_id"),
"user_api_key_team_id": kwargs.get("user_api_key_team_id"),
"user_api_key_end_user_id": kwargs.get("user_api_key_end_user_id"),
}
# Create synthetic LLM data that guardrails can process
synthetic_data: Final = {
"messages": [synthetic_message],
"model": kwargs.get("model", "mcp-tool-call"),
"model": guardrail_context.get("model", kwargs.get("model", "mcp-tool-call")),
"user_api_key_user_id": kwargs.get("user_api_key_user_id"),
"user_api_key_team_id": kwargs.get("user_api_key_team_id"),
"user_api_key_end_user_id": kwargs.get("user_api_key_end_user_id"),
@ -1271,12 +1287,7 @@ class ProxyLogging:
# (e.g. MCPJWTSigner) to independently verify the caller's identity
# before re-signing an outbound token (FR-5 verify+re-sign).
"incoming_bearer_token": kwargs.get("incoming_bearer_token"),
"metadata": {
"headers": kwargs.get("headers") or {},
"user_api_key_user_id": kwargs.get("user_api_key_user_id"),
"user_api_key_team_id": kwargs.get("user_api_key_team_id"),
"user_api_key_end_user_id": kwargs.get("user_api_key_end_user_id"),
},
"metadata": synthetic_metadata,
}
user_api_key_auth: Final = kwargs.get("user_api_key_auth")
if isinstance(user_api_key_auth, UserAPIKeyAuth):
@ -1285,6 +1296,15 @@ class ProxyLogging:
data=synthetic_data,
metadata_variable_name="metadata",
)
synthetic_metadata["user_api_key_metadata"] = copy.deepcopy(user_api_key_auth.metadata)
synthetic_metadata["user_api_key_team_metadata"] = copy.deepcopy(user_api_key_auth.team_metadata)
merged_guardrails: Final = (
*TypeAdapter(tuple[object, ...]).validate_python(synthetic_metadata.get("guardrails") or ()),
*TypeAdapter(tuple[object, ...]).validate_python(parent_metadata.get("guardrails") or ()),
)
synthetic_metadata["guardrails"] = [ # mutable-ok: existing guardrail selection and policy hooks require a list
selection for index, selection in enumerate(merged_guardrails) if selection not in merged_guardrails[:index]
]
return synthetic_data
def _convert_llm_result_to_mcp_response(self, llm_result, request_obj) -> MCPPreCallResponseObject | None:

4
litellm/responses/main.py
View file

@ -31,6 +31,7 @@ from litellm.llms.openai_like.responses.transformation import OpenAILikeResponse
from litellm.responses.litellm_completion_transformation.handler import (
LiteLLMCompletionTransformationHandler,
)
from litellm.responses.mcp.request_context import MCPRequestContext
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.llms.openai import (
PromptObject,
@ -331,6 +332,9 @@ async def aresponses_api_with_mcp(
litellm_call_id=kwargs.get("litellm_call_id"),
litellm_trace_id=kwargs.get("litellm_trace_id"),
request_tags=LiteLLM_Proxy_MCP_Handler._get_parent_request_tags(kwargs),
guardrail_context=MCPRequestContext.resolve_guardrail_context(
MappingProxyType({**kwargs, "metadata": metadata, "model": model})
),
)
if tool_results:

5
litellm/responses/mcp/chat_completions_handler.py
View file

@ -1,6 +1,7 @@
"""Helpers for handling MCP-aware `/chat/completions` requests."""
import logging
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, cast
from typing_extensions import TypedDict, Unpack
@ -118,7 +119,7 @@ async def acompletion_with_mcp(
**kwargs,
)
context: Final = MCPRequestContext.resolve(kwargs=kwargs, tools=tools)
context: Final = MCPRequestContext.resolve(kwargs=MappingProxyType({**kwargs, "model": model}), tools=tools)
user_api_key_auth: Final[UserAPIKeyAuth | None] = context.user_api_key_auth
request_tags: Final = list(context.request_tags) if context.request_tags else None
mcp_auth_header: Final = context.mcp_auth_header
@ -442,6 +443,7 @@ async def acompletion_with_mcp(
litellm_call_id=self.litellm_call_id,
litellm_trace_id=self.litellm_trace_id,
request_tags=self.request_tags,
guardrail_context=context.guardrail_context,
)
async def _prepare_follow_up_call(self):
@ -614,6 +616,7 @@ async def acompletion_with_mcp(
litellm_call_id=context.litellm_call_id,
litellm_trace_id=context.litellm_trace_id,
request_tags=request_tags,
guardrail_context=context.guardrail_context,
)
if not tool_results:

2
litellm/responses/mcp/litellm_proxy_mcp_handler.py
View file

@ -691,6 +691,7 @@ class LiteLLM_Proxy_MCP_Handler:
litellm_call_id: str | None = None,
litellm_trace_id: str | None = None,
request_tags: list[str] | None = None,
guardrail_context: Mapping[str, object] | None = None,
) -> list[MCPToolResult]:
"""Execute tool calls and return results."""
from fastapi import HTTPException
@ -854,6 +855,7 @@ class LiteLLM_Proxy_MCP_Handler:
raw_headers=raw_headers,
proxy_logging_obj=proxy_logging_obj,
litellm_logging_obj=litellm_logging_obj,
guardrail_context=guardrail_context,
)
if proxy_logging_obj:

2
litellm/responses/mcp/mcp_streaming_iterator.py
View file

@ -3,6 +3,7 @@ from typing import TYPE_CHECKING, Any, Final, cast
from litellm._logging import verbose_logger
from litellm._uuid import uuid
from litellm.responses.mcp.request_context import MCPRequestContext
from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator
from litellm.types.llms.openai import (
BaseLiteLLMOpenAIResponseObject,
@ -698,6 +699,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
litellm_call_id=self.litellm_call_id,
litellm_trace_id=self.litellm_trace_id,
request_tags=LiteLLM_Proxy_MCP_Handler._get_parent_request_tags(self.original_request_params),
guardrail_context=MCPRequestContext.resolve_guardrail_context(self.original_request_params),
)
# Create completion events and output_item.done events for tool execution

57
litellm/responses/mcp/request_context.py
View file

@ -9,9 +9,12 @@ still executes the tool, just with no credentials.
"""
from collections.abc import Iterable, Mapping, Sequence
from copy import deepcopy
from dataclasses import dataclass
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final
from pydantic import TypeAdapter
from typing_extensions import NotRequired, ReadOnly, TypedDict
if TYPE_CHECKING:
@ -36,6 +39,7 @@ class MCPRequestContext:
request_tags: Sequence[str] | None = None
litellm_trace_id: str | None = None
litellm_call_id: str | None = None
guardrail_context: Mapping[str, object] | None = None
@classmethod
def resolve(
@ -82,4 +86,57 @@ class MCPRequestContext:
request_tags=LiteLLM_Proxy_MCP_Handler._get_parent_request_tags(dict(kwargs)),
litellm_trace_id=kwargs.get("litellm_trace_id"),
litellm_call_id=kwargs.get("litellm_call_id"),
guardrail_context=cls.resolve_guardrail_context(kwargs),
)
@staticmethod
def resolve_guardrail_context(kwargs: Mapping[str, object]) -> Mapping[str, object]:
metadata_keys: Final = (
"guardrails",
"guardrail_config",
"_guardrail_pipelines",
"_pipeline_managed_guardrails",
"applied_policies",
"policy_sources",
"tags",
)
buckets: Final = tuple(
TypeAdapter(dict[str, object]).validate_python(kwargs[key])
for key in ("litellm_metadata", "metadata")
if isinstance(kwargs.get(key), Mapping)
)
sources: Final = (*buckets, kwargs)
metadata: Final = MappingProxyType(
{
**MappingProxyType(
{
key: deepcopy(value)
for bucket in buckets
for key, value in bucket.items()
if key in metadata_keys
}
),
"guardrails": deepcopy(
tuple(
selection
for source in sources
for selection in TypeAdapter(list[object]).validate_python(source.get("guardrails") or ())
)
),
"guardrail_config": deepcopy(
{ # mutable-ok: per-request guardrail configuration is a mutable JSON object in existing callbacks
key: value
for source in sources
for key, value in TypeAdapter(dict[str, object])
.validate_python(source.get("guardrail_config") or MappingProxyType({}))
.items()
}
),
}
)
return MappingProxyType(
{
**MappingProxyType({key: kwargs[key] for key in ("model",) if key in kwargs}),
"metadata": metadata,
}
)

3
tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_mcp_handler.py
View file

@ -147,6 +147,7 @@ async def test_anthropic_messages_with_mcp_forwards_the_callers_mcp_credentials(
request_tags=["team-a"],
litellm_trace_id="trace-123",
litellm_call_id="call-456",
guardrail_context={"metadata": {"guardrails": ("block-all",)}},
)
process = AsyncMock(return_value=([], {}))
@ -193,6 +194,8 @@ async def test_anthropic_messages_with_mcp_forwards_the_callers_mcp_credentials(
assert execution["litellm_trace_id"] == "trace-123"
assert execution["request_tags"] == ["team-a"]
assert execution["guardrail_context"] == {"metadata": {"guardrails": ("block-all",)}}
@pytest.mark.asyncio
async def test_anthropic_messages_with_mcp_stops_when_every_tool_call_is_skipped():

2
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
View file

@ -6675,8 +6675,10 @@ async def test_execute_mcp_tool_rest_server_id_authoritative_for_unprefixed_tool
allowed_mcp_servers=[api_key_server, oauth_server],
start_time=datetime.now(),
requested_server_id=api_key_server.server_id,
guardrail_context={"metadata": {"guardrails": ("block-all",)}},
)
assert captured["guardrail_context"] == {"metadata": {"guardrails": ("block-all",)}}
assert captured["server_name"] == "echo_api_key"
assert captured["name"] == "echo"

48
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
View file

@ -13891,3 +13891,51 @@ class TestProtectedCredentialPreparation:
client: Final = await MCPServerManager()._create_mcp_client(server)
request: Final = await client.prepare_request_auth()
assert request.headers["Authorization"] == expected
@pytest.mark.asyncio
@pytest.mark.parametrize("selected", [False, True])
async def test_request_selected_during_guardrail_runs_concurrently_with_tool(monkeypatch, selected):
from litellm.responses.mcp.request_context import MCPRequestContext
from litellm.proxy._experimental.mcp_server import tool_registry
tool_started = asyncio.Event()
guardrail_started = asyncio.Event()
class ObserveDuring(CustomGuardrail):
async def async_moderation_hook(self, data, user_api_key_dict, call_type):
if not self.should_run_guardrail(data, GuardrailEventHooks.during_mcp_call):
return data
assert data["mcp_tool_name"] == "execute"
assert data["mcp_arguments"] == {"text": "hello"}
guardrail_started.set()
await tool_started.wait()
return data
async def upstream(text):
assert text == "hello"
tool_started.set()
if selected:
await guardrail_started.wait()
return "executed"
guardrail = ObserveDuring(guardrail_name="observe", event_hook="during_mcp_call", default_on=False)
monkeypatch.setattr(litellm, "callbacks", [guardrail])
registry = tool_registry.MCPToolRegistry()
registry.register_tool("observer-execute", "Execute", {"type": "object"}, upstream)
monkeypatch.setattr(tool_registry, "global_mcp_tool_registry", registry)
manager = MCPServerManager()
manager.registry = {"observer": MCPServer(
server_id="observer", name="observer", server_name="observer", transport="http",
url="https://observer.example/mcp", spec_path="observer.json", auth_type="none",
)}
manager.tool_name_to_mcp_server_name_mapping = {"observer-execute": "observer"}
result = await asyncio.wait_for(manager.call_tool(
server_name="observer", name="execute", arguments={"text": "hello"},
user_api_key_auth=UserAPIKeyAuth(), proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()),
guardrail_context=MCPRequestContext.resolve_guardrail_context({"metadata": {"guardrails": ["observe"] if selected else []}}),
), timeout=5)
assert tool_started.is_set()
assert guardrail_started.is_set() is selected
assert result.isError is False
assert result.content[0].text == "executed"

2
tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py
View file

@ -78,6 +78,7 @@ async def test_openapi_local_tool_runs_pre_call_tool_check():
allowed_mcp_servers=[fake_server],
start_time=datetime.now(timezone.utc),
user_api_key_auth=user,
guardrail_context={"metadata": {"guardrails": ("block-all",)}},
)
pre_call.assert_awaited_once()
@ -88,6 +89,7 @@ async def test_openapi_local_tool_runs_pre_call_tool_check():
# records call order indirectly — we already asserted both were
# called; the relative ordering is enforced by the source change.
pre_call_kwargs = pre_call.await_args.kwargs
assert pre_call_kwargs["guardrail_context"] == {"metadata": {"guardrails": ("block-all",)}}
assert pre_call_kwargs["name"] == "list_pets"
assert pre_call_kwargs["server"] is fake_server
assert pre_call_kwargs["user_api_key_auth"] is user

16
tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py
View file

@ -2839,7 +2839,8 @@ class TestCallToolRestAPI:
assert not any("relaying upstream" in m for m in info_messages)
@pytest.mark.parametrize("raise_site", ["pre_call_hook", "execute_mcp_tool"])
async def test_guardrail_block_runs_failure_logging_before_http_translation(self, monkeypatch, raise_site):
@pytest.mark.parametrize("custom_code", [False, True])
async def test_guardrail_block_runs_failure_logging_before_http_translation(self, monkeypatch, raise_site, custom_code):
"""A pre_mcp_call guardrail block, whether raised by the pre-call hook or from inside
execute_mcp_tool, must reach proxy_logging_obj.post_call_failure_hook (the only path that
writes the failure spend-log row) with the logging object's failure payload already built,
@ -2870,6 +2871,11 @@ class TestCallToolRestAPI:
detail={"error": "Content blocked: keyword 'confidential' detected", "keyword": "confidential"},
)
if custom_code:
guardrail_error = rest_endpoints.ModifyResponseException(
message="Content blocked", model="mcp-tool-call", request_data={}, guardrail_name="block-all"
)
async def passthrough_pre_call_hook(user_api_key_dict, data, call_type):
return data
@ -2924,7 +2930,13 @@ class TestCallToolRestAPI:
with pytest.raises(HTTPException) as exc_info:
await rest_endpoints.call_tool_rest_api(request, user_api_key_dict=user_api_key_dict)
assert exc_info.value is guardrail_error
assert exc_info.value.status_code == 400
if custom_code:
assert exc_info.value.detail == {
"error": "guardrail_violation", "message": "Content blocked", "guardrail_name": "block-all"
}
else:
assert exc_info.value is guardrail_error
post_call_failure_hook.assert_awaited_once()
hook_kwargs = post_call_failure_hook.await_args.kwargs

82
tests/test_litellm/proxy/test_proxy_utils.py
View file

@ -2277,12 +2277,15 @@ def test_create_model_info_response_resolves_mode_through_deployment_model():
],
)
def test_convert_mcp_to_llm_format_carries_key_and_team_guardrails(key_metadata, team_metadata, expected_to_run):
from litellm.responses.mcp.request_context import MCPRequestContext
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
guardrail = CustomGuardrail(guardrail_name="key-scoped-guardrail", event_hook="pre_mcp_call", default_on=False)
kwargs = {
"name": "ask_question",
"arguments": {"question": "hello"},
"server_name": "deepwiki",
"guardrail_context": MCPRequestContext.resolve_guardrail_context({"guardrails": ["parent-rule"]}),
"user_api_key_auth": UserAPIKeyAuth(metadata=key_metadata, team_metadata=team_metadata),
}
request_obj = proxy_logging._create_mcp_request_object_from_kwargs(kwargs)
@ -2294,6 +2297,8 @@ def test_convert_mcp_to_llm_format_carries_key_and_team_guardrails(key_metadata,
assert guardrail.should_run_guardrail(synthetic, GuardrailEventHooks.pre_mcp_call) is expected_to_run
assert "parent-rule" in synthetic["metadata"]["guardrails"]
class _TracebackRecordingLogger(CustomLogger):
def __init__(self) -> None:
@ -2391,3 +2396,80 @@ class TestPrismaClientTokenAuthBehindThePool:
assert isinstance(client.db, RoutingPrismaWrapper)
assert client.db.writer.iam_token_db_auth is True
assert client.db.reader.iam_token_db_auth is True
@pytest.mark.parametrize("bucket", ["metadata", "litellm_metadata"])
def test_mcp_conversion_preserves_request_policy_and_isolates_guardrail_data(bucket):
from copy import deepcopy
from litellm.responses.mcp.request_context import MCPRequestContext
parent = {
"model": "parent-model",
bucket: {
"guardrails": ["policy-rule"], "guardrail_config": {"language": "en"},
"applied_policies": ["parent-policy"], "policy_sources": {"parent-policy": "model"},
"_guardrail_pipelines": [], "_pipeline_managed_guardrails": ["pipeline-rule"], "tags": ["review"],
},
"guardrails": [{"request-rule": {"extra_body": {"threshold": 0.9}}}],
"guardrail_config": {"entities": ["EMAIL_ADDRESS"]},
}
original = deepcopy(parent)
context = MCPRequestContext.resolve(kwargs=parent, tools=None)
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
kwargs = {"name": "execute", "arguments": {"text": "hello"}, "guardrail_context": context.guardrail_context}
request_obj = proxy_logging._create_mcp_request_object_from_kwargs(kwargs)
first = proxy_logging._convert_mcp_to_llm_format(request_obj, kwargs)
assert first["model"] == "parent-model"
assert first["metadata"]["guardrails"] == ["policy-rule", {"request-rule": {"extra_body": {"threshold": 0.9}}}]
assert first["metadata"]["guardrail_config"] == {"language": "en", "entities": ["EMAIL_ADDRESS"]}
assert first["metadata"]["applied_policies"] == ["parent-policy"]
assert first["metadata"]["policy_sources"] == {"parent-policy": "model"}
assert first["metadata"]["_pipeline_managed_guardrails"] == ["pipeline-rule"]
first["metadata"]["guardrails"].clear()
first["metadata"]["guardrail_config"]["entities"].clear()
assert parent == original
second = proxy_logging._convert_mcp_to_llm_format(request_obj, kwargs)
assert second["metadata"]["guardrails"] == ["policy-rule", {"request-rule": {"extra_body": {"threshold": 0.9}}}]
assert second["metadata"]["guardrail_config"]["entities"] == ["EMAIL_ADDRESS"]
@pytest.mark.parametrize("opt_out", [False, True])
def test_mcp_conversion_honors_only_authenticated_global_guardrail_opt_outs(opt_out):
from litellm.responses.mcp.request_context import MCPRequestContext
auth = UserAPIKeyAuth(metadata={"opted_out_global_guardrails": ["global-rule"] if opt_out else []})
context = MCPRequestContext.resolve(kwargs={"metadata": {
"user_api_key_auth": auth, "disable_global_guardrails": True,
"user_api_key_metadata": {"disable_global_guardrails": True},
}}, tools=None)
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
kwargs = {"name": "execute", "arguments": {}, "user_api_key_auth": auth, "guardrail_context": context.guardrail_context}
synthetic = proxy_logging._convert_mcp_to_llm_format(proxy_logging._create_mcp_request_object_from_kwargs(kwargs), kwargs)
guardrail = CustomGuardrail(guardrail_name="global-rule", event_hook="pre_mcp_call", default_on=True)
assert guardrail.should_run_guardrail(synthetic, GuardrailEventHooks.pre_mcp_call) is (not opt_out)
synthetic["metadata"]["user_api_key_metadata"]["opted_out_global_guardrails"].append("unrelated")
assert auth.metadata == {"opted_out_global_guardrails": ["global-rule"] if opt_out else []}
@pytest.mark.parametrize("model, expected", [("parent-model", True), ("unmatched-model", False)])
def test_mcp_auth_policy_uses_original_request_model(monkeypatch, model, expected):
from litellm.responses.mcp.request_context import MCPRequestContext
from litellm.proxy.policy_engine import policy_registry
from litellm.types.proxy.policy_engine import Policy, PolicyCondition, PolicyGuardrails
registry = policy_registry.PolicyRegistry()
registry._policies = {"model-policy": Policy(
condition=PolicyCondition(model="parent-model"), guardrails=PolicyGuardrails(add=["model-rule"])
)}
registry._initialized = True
monkeypatch.setattr(policy_registry, "_policy_registry", registry)
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
kwargs = {
"name": "execute", "arguments": {},
"user_api_key_auth": UserAPIKeyAuth(metadata={"policies": ["model-policy"]}),
"guardrail_context": MCPRequestContext.resolve_guardrail_context({"model": model, "guardrails": ["request-rule"]}),
}
synthetic = proxy_logging._convert_mcp_to_llm_format(proxy_logging._create_mcp_request_object_from_kwargs(kwargs), kwargs)
assert ("model-rule" in synthetic["metadata"]["guardrails"]) is expected
assert "request-rule" in synthetic["metadata"]["guardrails"]

91
tests/test_litellm/responses/mcp/test_chat_completions_handler.py
View file

@ -1387,3 +1387,94 @@ async def test_acompletion_with_mcp_forwards_unserved_external_mcp_tool_to_the_p
assert isinstance(result, ModelResponse)
assert result.id == "chatcmpl-zapier"
assert json.loads(provider.calls.last.request.content)["tools"] == [zapier_tool]
@pytest.mark.asyncio
@pytest.mark.parametrize("selected", [False, True])
@pytest.mark.parametrize("stream", [False, True])
@pytest.mark.parametrize("selection_source", ["metadata", "litellm_metadata", "body"])
@pytest.mark.parametrize("logging_failure", [False, True])
async def test_request_selected_mcp_guardrail_blocks_before_upstream(monkeypatch, selected, stream, selection_source, logging_failure):
from fastapi import HTTPException
from mcp.types import Tool
from litellm.caching.caching import DualCache
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.proxy import proxy_server
from litellm.proxy._types import UserAPIKeyAuth, LiteLLM_ObjectPermissionTable
from litellm.proxy._experimental.mcp_server import mcp_server_manager, server, tool_registry
from litellm.proxy._experimental.mcp_server.faults.list_outcomes import AggregateToolListing
from litellm.proxy.utils import ProxyLogging
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.mcp_server.mcp_server_manager import MCPServer
class BlockSelected(CustomGuardrail):
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
if self.should_run_guardrail(data, GuardrailEventHooks.pre_mcp_call):
raise HTTPException(status_code=400, detail="request-selected MCP block")
return data
guardrail = BlockSelected(guardrail_name="block-all", event_hook="pre_mcp_call", default_on=False)
monkeypatch.setattr(litellm, "callbacks", [guardrail])
manager = mcp_server_manager.MCPServerManager()
manager.registry = {"observer": MCPServer(
server_id="observer", name="observer", server_name="observer", transport="http",
url="https://observer.example/mcp", spec_path="observer.json", auth_type="none",
)}
manager.tool_name_to_mcp_server_name_mapping = {"observer-execute": "observer"}
upstream = AsyncMock(return_value={"executed": True})
registry = tool_registry.MCPToolRegistry()
registry.register_tool("observer-execute", "Execute", {"type": "object"}, upstream)
monkeypatch.setattr(tool_registry, "global_mcp_tool_registry", registry)
monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager)
monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(user_api_key_cache=DualCache()))
monkeypatch.setattr(server, "_get_tools_from_mcp_servers", AsyncMock(return_value=AggregateToolListing(
tools=[Tool(name="observer-execute", inputSchema={"type": "object"})], outcomes={}
)))
responses = [
ModelResponse(choices=[{"message": {"role": "assistant", "content": None, "tool_calls": [
{"id": "call-1", "type": "function", "function": {"name": "observer-execute", "arguments": "{}"}}
]}, "finish_reason": "tool_calls"}]),
ModelResponse(choices=[{"message": {"role": "assistant", "content": "done"}}]),
]
if stream:
from litellm.types.utils import ModelResponseStream
responses = [
await litellm.acompletion(
model="openai/gpt-5", messages=[{"role": "user", "content": "execute"}], stream=True,
mock_response=ModelResponseStream(choices=[{"index": 0, "delta": {
"role": "assistant", "content": None, "tool_calls": [{
"index": 0, "id": "call-1", "type": "function",
"function": {"name": "observer-execute", "arguments": "{}"},
}],
}, "finish_reason": "tool_calls"}]),
),
await litellm.acompletion(
model="openai/gpt-5", messages=[{"role": "user", "content": "done"}],
stream=True, mock_response="done",
),
]
if logging_failure:
from litellm.responses.mcp import litellm_proxy_mcp_handler
def fail_logging(*args, **kwargs):
raise RuntimeError("logging initialization failed")
monkeypatch.setattr(litellm_proxy_mcp_handler, "function_setup", fail_logging)
model_call = AsyncMock(side_effect=responses)
monkeypatch.setattr(litellm, "acompletion", model_call)
result = await acompletion_with_mcp(
model="test-model", messages=[{"role": "user", "content": "execute"}],
tools=[{"type": "mcp", "server_url": "litellm_proxy/observer", "require_approval": "never"}],
stream=stream,
user_api_key_auth=UserAPIKeyAuth(
object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="test", mcp_servers=["observer"])
),
**({"guardrails": ["block-all"] if selected else []} if selection_source == "body" else {
selection_source: {"guardrails": ["block-all"] if selected else []}
}),
)
if stream:
chunks = [chunk async for chunk in result]
assert chunks
assert model_call.await_count == 2
assert upstream.await_count == (0 if selected else 1)
tool_message = model_call.await_args.kwargs["messages"][-1]
assert ("request-selected MCP block" in tool_message["content"]) is selected

3
tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py
View file

@ -1077,6 +1077,8 @@ async def test_mcp_follow_up_call_is_stateless_when_store_is_false(
return ([], {"foo": "litellm_proxy"})
async def fake_execute(**kwargs: Any) -> list[dict[str, Any]]:
assert kwargs["guardrail_context"]["metadata"]["guardrails"] == ("block-all",)
assert kwargs["guardrail_context"]["model"] == "gpt-5"
return [{"tool_call_id": "call-1", "name": "foo", "result": "done"}]
monkeypatch.setattr(responses_main, "aresponses", fake_aresponses)
@ -1090,6 +1092,7 @@ async def test_mcp_follow_up_call_is_stateless_when_store_is_false(
input="hi",
model="gpt-5",
tools=[{"type": "mcp", "server_url": "litellm_proxy", "require_approval": "never"}],
litellm_metadata={"guardrails": ["block-all"]},
store=store,
previous_response_id=caller_previous_response_id,
)

2
tests/test_litellm/responses/mcp/test_mcp_streaming_iterator.py
View file

@ -127,10 +127,12 @@ async def test_second_round_tool_call_is_executed_and_reaches_final_text(monkeyp
]
)
iterator.original_request_params["litellm_metadata"] = {"guardrails": ["block-all"]}
chunks = [chunk async for chunk in iterator]
# Both rounds' tool calls were actually executed, not just streamed unexecuted.
assert call_tool.call_count == 2
assert all(call.kwargs["guardrail_context"]["metadata"]["guardrails"] == ("block-all",) for call in call_tool.call_args_list)
assert iterator.tool_call_round == 2
# The stream reached round 3 and produced the final text response instead