From 492c8753fc04aa503eb11810018cb6673a11b3f8 Mon Sep 17 00:00:00 2001 From: Yana Tokareva <15312436+yanatokareva01@users.noreply.github.com> Date: Tue, 3 Nov 2026 15:30:39 +0100 Subject: [PATCH] feat(mcp): include MCP server, tool and caller in MCP guardrail requests --- .../guardrail_hooks/deepkeep/deepkeep.py | 3 +- .../generic_guardrail_api.py | 4 +- .../guardrail_hooks/straiker/straiker.py | 3 +- .../guardrail_translation/handler.py | 27 ++++++- .../proxy/mcp_server/mcp_server_manager.py | 6 ++ litellm/proxy/utils.py | 11 +++ .../guardrail_hooks/generic_guardrail_api.py | 11 ++- litellm/types/utils.py | 13 ++++ .../test_generic_guardrail_api.py | 75 +++++++++++++++++++ 9 files changed, 144 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/deepkeep/deepkeep.py b/litellm/proxy/guardrails/guardrail_hooks/deepkeep/deepkeep.py index 23803b636f2..37b61cb8695 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/deepkeep/deepkeep.py +++ b/litellm/proxy/guardrails/guardrail_hooks/deepkeep/deepkeep.py @@ -253,8 +253,7 @@ class DeepKeepGuardrail(CustomGuardrail): getattr(logging_obj, "litellm_trace_id", None) if logging_obj else None, exc_info=error, ) - return_inputs: Final[GenericGuardrailAPIInputs] = {} - return_inputs.update(inputs) + return_inputs: Final[GenericGuardrailAPIInputs] = {**inputs} return return_inputs def _handle_guardrail_request_error( diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index 7e052796545..8522f11942e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -413,8 +413,7 @@ class GenericGuardrailAPI(CustomGuardrail): exc_info=error, ) # Keep flow going - treat as action=NONE (no modifications) - return_inputs: Final[GenericGuardrailAPIInputs] = {} - return_inputs.update(inputs) + return_inputs: Final[GenericGuardrailAPIInputs] = {**inputs} return return_inputs def _build_request_headers(self) -> Mapping[str, str]: @@ -592,6 +591,7 @@ class GenericGuardrailAPI(CustomGuardrail): tools=([GuardrailToolParam.model_validate(t) for t in tools] if tools else None), structured_messages=None if self.forward_messages_as_sent else structured_messages, tool_calls=tool_calls, + mcp_tool_call=inputs.get("mcp_tool_call"), additional_provider_specific_params=additional_params, input_type=input_type, model=model, diff --git a/litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py b/litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py index 589b865618b..3fee83e49b3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py +++ b/litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py @@ -1213,8 +1213,7 @@ class StraikerGuardrail(CustomGuardrail): inputs: GenericGuardrailAPIInputs, parsed: StraikerWebhookResponse, ) -> GenericGuardrailAPIInputs: - return_inputs: Final[GenericGuardrailAPIInputs] = {} - return_inputs.update(inputs) + return_inputs: Final[GenericGuardrailAPIInputs] = {**inputs} if parsed.texts is not None: return_inputs["texts"] = parsed.texts return return_inputs diff --git a/litellm/proxy/mcp_server/guardrail_translation/handler.py b/litellm/proxy/mcp_server/guardrail_translation/handler.py index 28c51148da5..a38b56c6eff 100644 --- a/litellm/proxy/mcp_server/guardrail_translation/handler.py +++ b/litellm/proxy/mcp_server/guardrail_translation/handler.py @@ -23,6 +23,7 @@ from typing import TYPE_CHECKING, Any, Final from fastapi import HTTPException from mcp.types import Tool as MCPTool +from pydantic import JsonValue, TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation @@ -43,7 +44,7 @@ from litellm.types.llms.openai import ( ChatCompletionToolParam, ChatCompletionToolParamFunctionChunk, ) -from litellm.types.utils import GenericGuardrailAPIInputs +from litellm.types.utils import GenericGuardrailAPIInputs, GenericGuardrailAPIMCPToolCall if TYPE_CHECKING: from mcp.types import CallToolResult @@ -52,6 +53,29 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +_MCP_ARGUMENTS: Final[TypeAdapter[Mapping[str, JsonValue]]] = TypeAdapter(Mapping[str, JsonValue]) + + +def _optional_text(value: object) -> str | None: + return value if isinstance(value, str) else None + + +def _mcp_tool_call( + tool_name: str, description: object, arguments: object, data: Mapping[str, object] +) -> GenericGuardrailAPIMCPToolCall: + return GenericGuardrailAPIMCPToolCall( + mcp_server_id=_optional_text(data.get("mcp_server_id")), + mcp_server_name=_optional_text(data.get("mcp_server_name")), + mcp_server_url=_optional_text(data.get("mcp_server_url")), + mcp_server_transport=_optional_text(data.get("mcp_server_transport")), + mcp_server_auth_type=_optional_text(data.get("mcp_server_auth_type")), + mcp_server_is_byok=is_byok if isinstance(is_byok := data.get("mcp_server_is_byok"), bool) else None, + tool_name=tool_name, + tool_description=_optional_text(description), + arguments=_MCP_ARGUMENTS.validate_python(arguments) if isinstance(arguments, Mapping) else {}, + ) + + def _blocked(reason: str) -> HTTPException: return HTTPException(status_code=400, detail={"error": f"Content blocked: {reason}"}) @@ -188,6 +212,7 @@ class MCPGuardrailTranslationHandler(BaseTranslation): inputs: Final[GenericGuardrailAPIInputs] = GenericGuardrailAPIInputs( tools=[tool_def], texts=list(scanned_texts), + mcp_tool_call=_mcp_tool_call(str(mcp_tool_name), mcp_tool_description, mcp_arguments, data), ) guarded: Final = await guardrail_to_apply.apply_guardrail( diff --git a/litellm/proxy/mcp_server/mcp_server_manager.py b/litellm/proxy/mcp_server/mcp_server_manager.py index eacba25257e..219c7762dba 100644 --- a/litellm/proxy/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/mcp_server/mcp_server_manager.py @@ -6252,6 +6252,12 @@ class MCPServerManager: "arguments": arguments, "server_name": server_name, "mcp_rate_limit_server_name": server.alias or server.server_name or server.name, + "mcp_server_id": server.server_id, + "mcp_server_url": redact_mcp_resource_url(server.url), + "mcp_server_transport": server.transport, + "mcp_server_auth_type": server.auth_type, + "mcp_server_is_byok": server.is_byok, + "client_ip": client_ip, "user_api_key_auth": checked_auth, "user_api_key_user_id": (getattr(checked_auth, "user_id", None) if checked_auth else None), "user_api_key_team_id": (getattr(checked_auth, "team_id", None) if checked_auth else None), diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 837433fd8cc..e647dc5d537 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -2250,6 +2250,7 @@ class ProxyLogging: "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"), + "requester_ip_address": kwargs.get("client_ip"), } # Create synthetic LLM data that guardrails can process @@ -2270,6 +2271,11 @@ class ProxyLogging: # synthetic call_mcp_tool payload (otherwise a key with # mcp_rpm_limit could exceed it via the MCP path). "mcp_server_name": kwargs.get("mcp_rate_limit_server_name"), + "mcp_server_id": kwargs.get("mcp_server_id"), + "mcp_server_url": kwargs.get("mcp_server_url"), + "mcp_server_transport": kwargs.get("mcp_server_transport"), + "mcp_server_auth_type": kwargs.get("mcp_server_auth_type"), + "mcp_server_is_byok": kwargs.get("mcp_server_is_byok"), # Raw Bearer token from the original HTTP request — allows guardrails # (e.g. MCPJWTSigner) to independently verify the caller's identity # before re-signing an outbound token (FR-5 verify+re-sign). @@ -2278,6 +2284,11 @@ class ProxyLogging: } user_api_key_auth: Final = kwargs.get("user_api_key_auth") if isinstance(user_api_key_auth, UserAPIKeyAuth): + synthetic_metadata.update( + LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_auth) + ) + synthetic_metadata["actor_agent_id"] = user_api_key_auth.agent_id + synthetic_metadata["target_agent_id"] = user_api_key_auth.invoked_agent_id add_guardrails_from_auth_metadata( user_api_key_dict=user_api_key_auth, data=synthetic_data, diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py index 4e4983c53f0..c4ad6a46ff6 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py @@ -2,7 +2,7 @@ from collections.abc import Mapping, Sequence from typing import Any, Final, Literal, cast # noqa: TID251 # JSON chat rows have no typed constructor across roles from pydantic import ConfigDict, Field, SerializerFunctionWrapHandler, model_serializer -from typing_extensions import TypedDict +from typing_extensions import ReadOnly, TypedDict from litellm.types.guardrails import GuardrailNumRetries from litellm.types.llms.base import LiteLLMBaseModel @@ -11,7 +11,7 @@ from litellm.types.llms.openai import ( ChatCompletionToolCallChunk, ) from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel -from litellm.types.utils import ChatCompletionMessageToolCall +from litellm.types.utils import ChatCompletionMessageToolCall, GenericGuardrailAPIMCPToolCall class GuardrailToolParam(LiteLLMBaseModel): @@ -43,6 +43,9 @@ class GenericGuardrailAPIMetadata(TypedDict, total=False): user_api_key_team_alias: str | None user_api_key_end_user_id: str | None user_api_key_org_id: str | None + requester_ip_address: ReadOnly[str | None] + actor_agent_id: ReadOnly[str | None] + target_agent_id: ReadOnly[str | None] class GenericGuardrailAPIOptionalParams(LiteLLMBaseModel): @@ -164,6 +167,10 @@ class GenericGuardrailAPIRequest(LiteLLMBaseModel): additional_provider_specific_params: dict[str, object] | None = None tool_calls: list[ChatCompletionToolCallChunk] | list[ChatCompletionMessageToolCall] | None = None model: str | None = None # the model being used for the LLM call + mcp_tool_call: GenericGuardrailAPIMCPToolCall | None = Field( + default=None, + description="The MCP server, tool, and named arguments of the MCP tool call being checked.", + ) def coerce_stream_holdback_value(value: Any) -> int: diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 57f4d971eb3..a137affd345 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -5329,6 +5329,18 @@ class PriorityReservationSettings(LiteLLMBaseModel): model_config = ConfigDict(protected_namespaces=()) +class GenericGuardrailAPIMCPToolCall(TypedDict, total=False): + mcp_server_id: ReadOnly[str | None] + mcp_server_name: ReadOnly[str | None] + mcp_server_url: ReadOnly[str | None] + mcp_server_transport: ReadOnly[str | None] + mcp_server_auth_type: ReadOnly[str | None] + mcp_server_is_byok: ReadOnly[bool | None] + tool_name: ReadOnly[str] + tool_description: ReadOnly[str | None] + arguments: ReadOnly[Mapping[str, JsonValue]] + + class GenericGuardrailAPIInputs(TypedDict, total=False): texts: list[str] # extracted text from the LLM response - for basic text guardrails images: list[str] # extracted images from the LLM response - for image guardrails @@ -5338,6 +5350,7 @@ class GenericGuardrailAPIInputs(TypedDict, total=False): AllMessageValues ] # structured messages sent to the LLM - indicates if text is from system or user model: str | None # the model being used for the LLM call + mcp_tool_call: ReadOnly[GenericGuardrailAPIMCPToolCall] stream_holdback_chars: list[ int ] # trailing chars to withhold from streaming emission per text (word-boundary safety) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index b5f794cb136..328aea048f1 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -21,6 +21,7 @@ from mcp.types import CallToolResult, TextContent import litellm from litellm import ModelResponse from litellm._version import version as litellm_version +from litellm.caching.caching import DualCache from litellm.exceptions import GuardrailRaisedException, Timeout from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, MaskedHTTPStatusError from litellm.proxy._experimental.mcp_server.guardrail_translation.handler import ( @@ -34,7 +35,9 @@ from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api.generic_guar _HEADER_PRESENT_PLACEHOLDER, ) from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import UnifiedLLMGuardrails +from litellm.proxy.utils import ProxyLogging from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.mcp import MCPAuth, MCPTransport from litellm.types.utils import Choices, GenericGuardrailAPIInputs, Message @@ -2692,3 +2695,75 @@ async def test_during_mcp_call_blocked_raises() -> None: ) assert exc_info.value.message == "credential detected in tool input" + + + +@pytest.mark.asyncio +async def test_pre_mcp_call_sends_the_server_tool_arguments_and_caller() -> None: + sent_bodies: Final[SimpleQueue[bytes]] = SimpleQueue() + + async def answer(request: httpx.Request) -> httpx.Response: + sent_bodies.put(request.content) + return httpx.Response(200, json={"action": "NONE"}, request=request) + + guardrail: Final = _mcp_guardrail("pre_mcp_call") + arguments: Final = {"repo": "octo/hello", "options": {"force": True, "depth": 2}} + caller: Final = UserAPIKeyAuth( + api_key="hashed-key", + key_alias="ci-bot", + user_id="u-1", + user_email="dev@example.com", + team_id="t-1", + team_alias="platform", + org_id="o-1", + end_user_id="e-1", + agent_id="agent-1", + ).model_copy(update={"invoked_agent_id": "agent-2"}) + pre_hook_kwargs: Final = { + "name": "delete_repo", + "arguments": arguments, + "server_name": "github", + "mcp_rate_limit_server_name": "github", + "mcp_server_id": "srv-1", + "mcp_server_url": "https://mcp.example.com", + "mcp_server_transport": MCPTransport.http, + "mcp_server_auth_type": MCPAuth.oauth2, + "mcp_server_is_byok": True, + "client_ip": "203.0.113.7", + "tool_description": "Delete a repository", + "user_api_key_auth": caller, + } + proxy_logging_obj: Final = ProxyLogging(user_api_key_cache=DualCache()) + data: Final = proxy_logging_obj._convert_mcp_to_llm_format( + proxy_logging_obj._create_mcp_request_object_from_kwargs(pre_hook_kwargs), pre_hook_kwargs + ) + + with patch.object(guardrail, "async_handler", AsyncHTTPHandler(transport=httpx.MockTransport(answer))): + await MCPGuardrailTranslationHandler().process_input_messages(data=data, guardrail_to_apply=guardrail) + + body: Final = json.loads(sent_bodies.get_nowait()) + assert body["mcp_tool_call"] == { + "mcp_server_id": "srv-1", + "mcp_server_name": "github", + "mcp_server_url": "https://mcp.example.com", + "mcp_server_transport": "http", + "mcp_server_auth_type": "oauth2", + "mcp_server_is_byok": True, + "tool_name": "delete_repo", + "tool_description": "Delete a repository", + "arguments": arguments, + } + assert body["request_data"] == { + "user_api_key_hash": "hashed-key", + "user_api_key_alias": "ci-bot", + "user_api_key_user_id": "u-1", + "user_api_key_user_email": "dev@example.com", + "user_api_key_team_id": "t-1", + "user_api_key_team_alias": "platform", + "user_api_key_end_user_id": "e-1", + "user_api_key_org_id": "o-1", + "requester_ip_address": "203.0.113.7", + "actor_agent_id": "agent-1", + "target_agent_id": "agent-2", + } + assert body["texts"] == ["Delete a repository", "octo/hello"]