From 987c89807338701a95b12b7eb18cf67505c5f652 Mon Sep 17 00:00:00 2001 From: Airia Open Source Date: Wed, 21 Oct 2026 04:21:53 +0000 Subject: [PATCH 1/5] feat(guardrails): support during_call and MCP tool call modes in the Airia guardrail --- .../guardrails/guardrail_hooks/airia/airia.py | 39 ++++- .../guardrails/guardrail_hooks/test_airia.py | 149 +++++++++++++++++- 2 files changed, 178 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/airia/airia.py b/litellm/proxy/guardrails/guardrail_hooks/airia/airia.py index 993c1d30949..8270d2a5167 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/airia/airia.py +++ b/litellm/proxy/guardrails/guardrail_hooks/airia/airia.py @@ -1,8 +1,10 @@ import os import uuid from collections.abc import Mapping +from contextvars import ContextVar from typing import ( Any, # noqa: TID251 # the only type CustomGuardrail.__init__ accepts for its open-ended kwargs + ClassVar, Final, Literal, ) @@ -22,9 +24,11 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # the factory's params are untyped upstream httpxSpecialProvider, ) +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import UnifiedLLMGuardrails from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.guardrails.guardrail_hooks.airia import AiriaGuardrailConfigModel -from litellm.types.utils import GenericGuardrailAPIInputs +from litellm.types.utils import CallTypesLiteral, GenericGuardrailAPIInputs GUARDRAIL_PATH: Final = "/v1/guardrails/litellm" @@ -34,7 +38,15 @@ ACTION_INTERVENED: Final = "GUARDRAIL_INTERVENED" DEFAULT_BLOCKED_MESSAGE: Final = "Blocked by your organization's content policy." -SUPPORTED_EVENT_HOOKS: Final = (GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call) +SUPPORTED_EVENT_HOOKS: Final = ( + GuardrailEventHooks.pre_call, + GuardrailEventHooks.during_call, + GuardrailEventHooks.post_call, + GuardrailEventHooks.pre_mcp_call, + GuardrailEventHooks.post_mcp_call, +) + +_EVALUATING_DURING_CALL: Final[ContextVar[bool]] = ContextVar("airia_evaluating_during_call", default=False) def _stream_cannot_carry_a_rewrite(request_data: Mapping[str, object] | None) -> bool: @@ -88,7 +100,9 @@ def _as_object(value: object) -> Mapping[str, object] | None: class AiriaGuardrail(CustomGuardrail): - """Evaluates prompts and responses against your Airia guardrail policy.""" + """Evaluates prompts, responses and MCP tool calls against your Airia guardrail policy.""" + + use_native_during_call_hook: ClassVar[bool] = True def __init__( self, @@ -189,6 +203,10 @@ class AiriaGuardrail(CustomGuardrail): ) if action == ACTION_INTERVENED: + # during_call runs alongside the model call, so the prompt has already been sent and a + # rewrite of it would reach nothing but the logs + if _EVALUATING_DURING_CALL.get(): + raise self._blocked() if input_type == "response" and _stream_cannot_carry_a_rewrite(request_data): raise self._blocked() return self._rewritten(body, inputs) @@ -198,6 +216,21 @@ class AiriaGuardrail(CustomGuardrail): return inputs + async def async_moderation_hook( + self, + data: dict[str, object], # mutable-ok: the unified hook reads guardrail_to_apply out of the request + user_api_key_dict: UserAPIKeyAuth, + call_type: CallTypesLiteral, + ) -> None: + token: Final = _EVALUATING_DURING_CALL.set(True) + try: + data["guardrail_to_apply"] = self # rebind-ok: the unified hook's only way to receive the guardrail + await UnifiedLLMGuardrails().async_moderation_hook( # pyright: ignore[reportUnknownMemberType] # annotates data as a bare dict + data=data, user_api_key_dict=user_api_key_dict, call_type=call_type + ) + finally: + _EVALUATING_DURING_CALL.reset(token) + def _rewritten( self, body: Mapping[str, object], diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_airia.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_airia.py index 61f41b0482a..dce8b861ff7 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_airia.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_airia.py @@ -6,14 +6,19 @@ from unittest.mock import AsyncMock, MagicMock import httpx import pytest +from mcp.types import CallToolResult, TextContent import litellm +from litellm.caching.caching import DualCache from litellm.exceptions import GuardrailRaisedException +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.guardrails.guardrail_hooks.airia.airia import AiriaGuardrail +from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import UnifiedLLMGuardrails from litellm.proxy.guardrails.init_guardrails import ( init_guardrails_v2, # pyright: ignore[reportUnknownVariableType] # untyped upstream ) -from litellm.types.guardrails import GuardrailEventHooks +from litellm.proxy.utils import ProxyLogging from litellm.types.proxy.guardrails.guardrail_hooks.airia import AiriaGuardrailConfigModel from litellm.types.utils import GenericGuardrailAPIInputs @@ -98,13 +103,143 @@ def test_airia_guardrail_config(monkeypatch: pytest.MonkeyPatch) -> None: assert _timeout(registered[0]).read == 45.0 -def test_during_call_is_not_a_supported_event_hook() -> None: - """during_call runs concurrently with the upstream call, so a block can land too late.""" - hooks = _make_guardrail().supported_event_hooks +def test_tool_call_and_during_call_modes_register_through_config(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "guardrail_name_config_map", {}) + monkeypatch.setattr(litellm, "callbacks", []) + modes: Final = ["pre_mcp_call", "during_call", "post_mcp_call"] - assert hooks is not None - assert GuardrailEventHooks.during_call not in hooks - assert list(hooks) == [GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call] + init_guardrails_v2( + all_guardrails=[ + { + "guardrail_name": "airia-guard", + "litellm_params": {"guardrail": "airia", "mode": modes, "api_base": API_BASE, "api_key": API_KEY}, + } + ], + config_file_path="", + ) + + registered: Final = [c for c in litellm.callbacks if isinstance(c, AiriaGuardrail)] + assert len(registered) == 1 + assert registered[0].event_hook == modes + + +@pytest.mark.asyncio +async def test_pre_mcp_call_redaction_rewrites_the_tool_arguments() -> None: + guardrail: Final = _make_guardrail_with_response( + {"action": "GUARDRAIL_INTERVENED", "texts": ["[EmailAddress1]"]}, event_hook="pre_mcp_call", default_on=True + ) + + data: Final = await UnifiedLLMGuardrails().async_pre_call_hook( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # annotates data as a bare dict + user_api_key_dict=UserAPIKeyAuth(api_key="test"), + cache=DualCache(), + data={ # mutable-ok: the hook writes the masked arguments back into the request + "guardrail_to_apply": guardrail, + "mcp_tool_name": "send_email", + "mcp_arguments": {"to": "ada@example.com"}, + }, + call_type="call_mcp_tool", + ) + + assert isinstance(data, dict) + assert data["mcp_arguments"] == {"to": "[EmailAddress1]"} + assert _post_mock(guardrail).call_args.kwargs["json"]["input_type"] == "request" + assert _post_mock(guardrail).call_args.kwargs["json"]["texts"] == ["ada@example.com"] + + +@pytest.mark.asyncio +async def test_pre_mcp_call_block_stops_the_tool_call() -> None: + guardrail: Final = _make_guardrail_with_response( + {"action": "BLOCKED", "blocked_reason": "Tool not allowed"}, event_hook="pre_mcp_call", default_on=True + ) + + with pytest.raises(GuardrailRaisedException) as excinfo: + await UnifiedLLMGuardrails().async_pre_call_hook( # pyright: ignore[reportUnknownMemberType] # annotates data as a bare dict + user_api_key_dict=UserAPIKeyAuth(api_key="test"), + cache=DualCache(), + data={"guardrail_to_apply": guardrail, "mcp_tool_name": "drop_table", "mcp_arguments": {"name": "users"}}, + call_type="call_mcp_tool", + ) + + assert excinfo.value.blocked_content is True + assert "Tool not allowed" in str(excinfo.value) + + +@pytest.mark.asyncio +async def test_post_mcp_call_redaction_masks_the_tool_result(monkeypatch: pytest.MonkeyPatch) -> None: + guardrail: Final = _make_guardrail_with_response( + {"action": "GUARDRAIL_INTERVENED", "texts": ["reach [EmailAddress1]"]}, + event_hook="post_mcp_call", + default_on=True, + ) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + checked: Final = await ProxyLogging(user_api_key_cache=UserApiKeyCache()).post_mcp_call_hook( + response=CallToolResult(content=[TextContent(type="text", text="reach ada@example.com")]), + request_data={}, + user_api_key_dict=UserAPIKeyAuth(api_key="test"), + ) + + assert checked.content == [TextContent(type="text", text="reach [EmailAddress1]")] + assert _post_mock(guardrail).call_args.kwargs["json"]["input_type"] == "response" + + +def _during_call_request() -> dict[str, object]: + return {"model": "gpt-4o", "messages": [{"role": "user", "content": "mail ada@example.com"}]} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "verdict", + [ + {"action": "BLOCKED", "blocked_reason": "Not allowed"}, + {"action": "GUARDRAIL_INTERVENED", "texts": ["mail [EmailAddress1]"]}, + ], +) +async def test_during_call_enforces_every_intervention_as_a_block( + verdict: dict[str, object], monkeypatch: pytest.MonkeyPatch +) -> None: + """The prompt is already on its way to the model, so a redaction can no longer be delivered""" + guardrail: Final = _make_guardrail_with_response(verdict, event_hook="during_call", default_on=True) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + with pytest.raises(GuardrailRaisedException) as excinfo: + await ProxyLogging(user_api_key_cache=UserApiKeyCache()).during_call_hook( # pyright: ignore[reportUnknownMemberType] # annotates data as a bare dict + data=_during_call_request(), user_api_key_dict=UserAPIKeyAuth(api_key="test"), call_type="completion" + ) + + assert excinfo.value.blocked_content is True + assert _post_mock(guardrail).call_args.kwargs["json"]["input_type"] == "request" + + +@pytest.mark.asyncio +async def test_during_call_allow_lets_the_request_through(monkeypatch: pytest.MonkeyPatch) -> None: + guardrail: Final = _make_guardrail_with_response({"action": "NONE"}, event_hook="during_call", default_on=True) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + request: Final = _during_call_request() + + await ProxyLogging(user_api_key_cache=UserApiKeyCache()).during_call_hook( # pyright: ignore[reportUnknownMemberType] # annotates data as a bare dict + data=request, user_api_key_dict=UserAPIKeyAuth(api_key="test"), call_type="completion" + ) + + assert request["messages"] == _during_call_request()["messages"] + _post_mock(guardrail).assert_called_once() + + +@pytest.mark.asyncio +async def test_a_during_call_block_does_not_turn_later_redactions_into_blocks() -> None: + guardrail: Final = _make_guardrail_with_response( + {"action": "GUARDRAIL_INTERVENED", "texts": ["[Redacted]"]}, event_hook="during_call", default_on=True + ) + with pytest.raises(GuardrailRaisedException): + await guardrail.async_moderation_hook( + data=_during_call_request(), user_api_key_dict=UserAPIKeyAuth(api_key="test"), call_type="completion" + ) + + result: Final = await guardrail.apply_guardrail( + inputs=_inputs(texts=["secret"]), request_data={}, input_type="request" + ) + + assert result.get("texts") == ["[Redacted]"] @pytest.mark.parametrize("missing", ["api_base", "api_key"]) From c2a12ef5a74f76f62d2a2afcc4c323e188138eac Mon Sep 17 00:00:00 2001 From: Airia Open Source Date: Wed, 21 Oct 2026 12:56:47 +0000 Subject: [PATCH 2/5] feat(guardrails): tell the Airia gateway which MCP tool an evaluation is for --- .../guardrails/guardrail_hooks/airia/airia.py | 21 +++++++++++-- .../guardrails/guardrail_hooks/test_airia.py | 30 +++++++++++++++++-- 2 files changed, 46 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/airia/airia.py b/litellm/proxy/guardrails/guardrail_hooks/airia/airia.py index 8270d2a5167..d0f6c3a2ccd 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/airia/airia.py +++ b/litellm/proxy/guardrails/guardrail_hooks/airia/airia.py @@ -28,7 +28,7 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import UnifiedLLMGuardrails from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.guardrails.guardrail_hooks.airia import AiriaGuardrailConfigModel -from litellm.types.utils import CallTypesLiteral, GenericGuardrailAPIInputs +from litellm.types.utils import MCP_GUARDRAIL_CALL_TYPES, CallTypesLiteral, GenericGuardrailAPIInputs GUARDRAIL_PATH: Final = "/v1/guardrails/litellm" @@ -93,6 +93,21 @@ def _request_route(request_data: Mapping[str, object]) -> str | None: return None +def _mcp_tool_name(request_data: Mapping[str, object], logging_obj: "LiteLLMLoggingObj | None") -> str | None: + """The MCP tool this evaluation is for, or None for model traffic. + + The logger's call type wins when there is one, because a chat client can put MCP keys in its request body. + """ + if logging_obj is not None and logging_obj.call_type not in MCP_GUARDRAIL_CALL_TYPES: + return None + pre_call_name: Final = request_data.get("mcp_tool_name") + if isinstance(pre_call_name, str) and pre_call_name: + return pre_call_name + call_metadata: Final = _as_object(request_data.get("mcp_tool_call_metadata")) + post_call_name: Final = call_metadata.get("name") if call_metadata is not None else None + return post_call_name if isinstance(post_call_name, str) and post_call_name else None + + def _as_object(value: object) -> Mapping[str, object] | None: """`value` as a JSON object, or None. States the key type once so the checker does not see `dict[Unknown, Unknown]` at every read.""" @@ -171,6 +186,7 @@ class AiriaGuardrail(CustomGuardrail): "tool_calls": inputs.get("tool_calls") or (), "model": inputs.get("model"), "litellm_call_id": call_id, + "mcp_tool_name": _mcp_tool_name(request_data, logging_obj), }, headers={"Authorization": f"Bearer {self.api_key}"}, ) @@ -203,8 +219,7 @@ class AiriaGuardrail(CustomGuardrail): ) if action == ACTION_INTERVENED: - # during_call runs alongside the model call, so the prompt has already been sent and a - # rewrite of it would reach nothing but the logs + # during_call races the model call, so a rewrite of the prompt cannot be relied on to reach it if _EVALUATING_DURING_CALL.get(): raise self._blocked() if input_type == "response" and _stream_cannot_carry_a_rewrite(request_data): diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_airia.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_airia.py index dce8b861ff7..57ba1ed0b33 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_airia.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_airia.py @@ -1,5 +1,6 @@ import json import uuid +from datetime import datetime, timezone from collections.abc import AsyncGenerator from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock @@ -11,6 +12,7 @@ from mcp.types import CallToolResult, TextContent import litellm from litellm.caching.caching import DualCache from litellm.exceptions import GuardrailRaisedException +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.guardrails.guardrail_hooks.airia.airia import AiriaGuardrail @@ -144,6 +146,7 @@ async def test_pre_mcp_call_redaction_rewrites_the_tool_arguments() -> None: assert data["mcp_arguments"] == {"to": "[EmailAddress1]"} assert _post_mock(guardrail).call_args.kwargs["json"]["input_type"] == "request" assert _post_mock(guardrail).call_args.kwargs["json"]["texts"] == ["ada@example.com"] + assert _post_mock(guardrail).call_args.kwargs["json"]["mcp_tool_name"] == "send_email" @pytest.mark.asyncio @@ -175,12 +178,36 @@ async def test_post_mcp_call_redaction_masks_the_tool_result(monkeypatch: pytest checked: Final = await ProxyLogging(user_api_key_cache=UserApiKeyCache()).post_mcp_call_hook( response=CallToolResult(content=[TextContent(type="text", text="reach ada@example.com")]), - request_data={}, + request_data={"mcp_tool_call_metadata": {"name": "lookup_contact", "arguments": {}}}, user_api_key_dict=UserAPIKeyAuth(api_key="test"), ) assert checked.content == [TextContent(type="text", text="reach [EmailAddress1]")] assert _post_mock(guardrail).call_args.kwargs["json"]["input_type"] == "response" + assert _post_mock(guardrail).call_args.kwargs["json"]["mcp_tool_name"] == "lookup_contact" + + +@pytest.mark.asyncio +async def test_mcp_keys_in_a_chat_request_body_do_not_mark_it_as_a_tool_call() -> None: + guardrail: Final = _make_guardrail_with_response({"action": "NONE"}) + logging_obj: Final = LiteLLMLoggingObj( + model="gpt-4o", + messages=[], + stream=False, + call_type="acompletion", + start_time=datetime(2026, 1, 1, tzinfo=timezone.utc), + litellm_call_id="call-1", + function_id="1", + ) + + await guardrail.apply_guardrail( + inputs=_inputs(), + request_data={"mcp_tool_name": "send_email", "mcp_tool_call_metadata": {"name": "send_email"}}, + input_type="request", + logging_obj=logging_obj, + ) + + assert _post_mock(guardrail).call_args.kwargs["json"]["mcp_tool_name"] is None def _during_call_request() -> dict[str, object]: @@ -198,7 +225,6 @@ def _during_call_request() -> dict[str, object]: async def test_during_call_enforces_every_intervention_as_a_block( verdict: dict[str, object], monkeypatch: pytest.MonkeyPatch ) -> None: - """The prompt is already on its way to the model, so a redaction can no longer be delivered""" guardrail: Final = _make_guardrail_with_response(verdict, event_hook="during_call", default_on=True) monkeypatch.setattr(litellm, "callbacks", [guardrail]) From a1361e07c16b99e90d51bcdcfda982b668840b10 Mon Sep 17 00:00:00 2001 From: Airia Open Source Date: Wed, 21 Oct 2026 12:57:11 +0000 Subject: [PATCH 3/5] fix(guardrails): read the logged call type without an untyped attribute access --- litellm/proxy/guardrails/guardrail_hooks/airia/airia.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/airia/airia.py b/litellm/proxy/guardrails/guardrail_hooks/airia/airia.py index d0f6c3a2ccd..d41398e406c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/airia/airia.py +++ b/litellm/proxy/guardrails/guardrail_hooks/airia/airia.py @@ -98,7 +98,8 @@ def _mcp_tool_name(request_data: Mapping[str, object], logging_obj: "LiteLLMLogg The logger's call type wins when there is one, because a chat client can put MCP keys in its request body. """ - if logging_obj is not None and logging_obj.call_type not in MCP_GUARDRAIL_CALL_TYPES: + logged_call_type: Final[object] = getattr(logging_obj, "call_type", None) + if logging_obj is not None and logged_call_type not in MCP_GUARDRAIL_CALL_TYPES: return None pre_call_name: Final = request_data.get("mcp_tool_name") if isinstance(pre_call_name, str) and pre_call_name: From b6d3eb555e1d00fa1712c464587da756563963b7 Mon Sep 17 00:00:00 2001 From: Airia Open Source Date: Wed, 21 Oct 2026 13:09:20 +0000 Subject: [PATCH 4/5] fix(guardrails): mark Airia MCP evaluations by the gated event, not by request data --- .../guardrails/guardrail_hooks/airia/airia.py | 32 ++++++++----- .../guardrails/guardrail_hooks/test_airia.py | 47 +++++++++++++------ 2 files changed, 54 insertions(+), 25 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/airia/airia.py b/litellm/proxy/guardrails/guardrail_hooks/airia/airia.py index d41398e406c..df4aa9ef085 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/airia/airia.py +++ b/litellm/proxy/guardrails/guardrail_hooks/airia/airia.py @@ -28,7 +28,7 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import UnifiedLLMGuardrails from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.guardrails.guardrail_hooks.airia import AiriaGuardrailConfigModel -from litellm.types.utils import MCP_GUARDRAIL_CALL_TYPES, CallTypesLiteral, GenericGuardrailAPIInputs +from litellm.types.utils import CallTypesLiteral, GenericGuardrailAPIInputs GUARDRAIL_PATH: Final = "/v1/guardrails/litellm" @@ -46,7 +46,13 @@ SUPPORTED_EVENT_HOOKS: Final = ( GuardrailEventHooks.post_mcp_call, ) +MCP_EVENT_HOOKS: Final = frozenset((GuardrailEventHooks.pre_mcp_call, GuardrailEventHooks.post_mcp_call)) +UNKNOWN_MCP_TOOL_NAME: Final = "unknown" + _EVALUATING_DURING_CALL: Final[ContextVar[bool]] = ContextVar("airia_evaluating_during_call", default=False) +# Request bodies on several routes are merged into the hook data, so MCP keys in it prove nothing; the event the +# framework gates on is the only signal a client cannot set +_PENDING_EVENT: Final[ContextVar[GuardrailEventHooks | None]] = ContextVar("airia_pending_event", default=None) def _stream_cannot_carry_a_rewrite(request_data: Mapping[str, object] | None) -> bool: @@ -93,20 +99,13 @@ def _request_route(request_data: Mapping[str, object]) -> str | None: return None -def _mcp_tool_name(request_data: Mapping[str, object], logging_obj: "LiteLLMLoggingObj | None") -> str | None: - """The MCP tool this evaluation is for, or None for model traffic. - - The logger's call type wins when there is one, because a chat client can put MCP keys in its request body. - """ - logged_call_type: Final[object] = getattr(logging_obj, "call_type", None) - if logging_obj is not None and logged_call_type not in MCP_GUARDRAIL_CALL_TYPES: - return None +def _mcp_tool_name(request_data: Mapping[str, object]) -> str: pre_call_name: Final = request_data.get("mcp_tool_name") if isinstance(pre_call_name, str) and pre_call_name: return pre_call_name call_metadata: Final = _as_object(request_data.get("mcp_tool_call_metadata")) post_call_name: Final = call_metadata.get("name") if call_metadata is not None else None - return post_call_name if isinstance(post_call_name, str) and post_call_name else None + return post_call_name if isinstance(post_call_name, str) and post_call_name else UNKNOWN_MCP_TOOL_NAME def _as_object(value: object) -> Mapping[str, object] | None: @@ -158,6 +157,15 @@ class AiriaGuardrail(CustomGuardrail): ) verbose_proxy_logger.info("AiriaGuardrail initialized with gateway: %s", self.api_base) + def should_run_guardrail( + self, + data: object, + event_type: GuardrailEventHooks, + ) -> bool: + runs: Final = super().should_run_guardrail(data, event_type) # pyright: ignore[reportUnknownMemberType] # the parent annotates data as untyped + _ = _PENDING_EVENT.set(event_type if runs else None) + return runs + def supports_structured_messages(self) -> bool: return False @@ -169,6 +177,8 @@ class AiriaGuardrail(CustomGuardrail): input_type: Literal["request", "response"], logging_obj: "LiteLLMLoggingObj | None" = None, ) -> GenericGuardrailAPIInputs: + is_mcp_tool_evaluation: Final = _PENDING_EVENT.get() in MCP_EVENT_HOOKS + _ = _PENDING_EVENT.set(None) call_id: Final = ( logging_obj.litellm_call_id if logging_obj @@ -187,7 +197,7 @@ class AiriaGuardrail(CustomGuardrail): "tool_calls": inputs.get("tool_calls") or (), "model": inputs.get("model"), "litellm_call_id": call_id, - "mcp_tool_name": _mcp_tool_name(request_data, logging_obj), + "mcp_tool_name": _mcp_tool_name(request_data or {}) if is_mcp_tool_evaluation else None, }, headers={"Authorization": f"Bearer {self.api_key}"}, ) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_airia.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_airia.py index 57ba1ed0b33..4eb9e2fcaf0 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_airia.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_airia.py @@ -21,6 +21,7 @@ from litellm.proxy.guardrails.init_guardrails import ( init_guardrails_v2, # pyright: ignore[reportUnknownVariableType] # untyped upstream ) from litellm.proxy.utils import ProxyLogging +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.guardrails.guardrail_hooks.airia import AiriaGuardrailConfigModel from litellm.types.utils import GenericGuardrailAPIInputs @@ -188,28 +189,46 @@ async def test_post_mcp_call_redaction_masks_the_tool_result(monkeypatch: pytest @pytest.mark.asyncio -async def test_mcp_keys_in_a_chat_request_body_do_not_mark_it_as_a_tool_call() -> None: - guardrail: Final = _make_guardrail_with_response({"action": "NONE"}) - logging_obj: Final = LiteLLMLoggingObj( - model="gpt-4o", - messages=[], - stream=False, - call_type="acompletion", - start_time=datetime(2026, 1, 1, tzinfo=timezone.utc), - litellm_call_id="call-1", - function_id="1", +@pytest.mark.parametrize("logging_obj", [None, "acompletion"]) +async def test_mcp_keys_in_a_request_body_do_not_mark_model_traffic_as_a_tool_call(logging_obj: str | None) -> None: + guardrail: Final = _make_guardrail_with_response({"action": "NONE"}, event_hook="pre_call", default_on=True) + request_data: Final[dict[str, object]] = {"prompt": "draw a cat", "mcp_tool_name": "send_email"} + logger: Final = ( + LiteLLMLoggingObj( + model="gpt-4o", + messages=[], + stream=False, + call_type=logging_obj, + start_time=datetime(2026, 1, 1, tzinfo=timezone.utc), + litellm_call_id="call-1", + function_id="1", + ) + if logging_obj + else None ) + assert guardrail.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.pre_call) await guardrail.apply_guardrail( - inputs=_inputs(), - request_data={"mcp_tool_name": "send_email", "mcp_tool_call_metadata": {"name": "send_email"}}, - input_type="request", - logging_obj=logging_obj, + inputs=_inputs(), request_data=request_data, input_type="request", logging_obj=logger ) assert _post_mock(guardrail).call_args.kwargs["json"]["mcp_tool_name"] is None +@pytest.mark.asyncio +async def test_an_mcp_evaluation_does_not_mark_the_next_model_evaluation_as_a_tool_call() -> None: + guardrail: Final = _make_guardrail_with_response( + {"action": "NONE"}, event_hook=["pre_mcp_call", "pre_call"], default_on=True + ) + tool_call: Final[dict[str, object]] = {"mcp_tool_name": "send_email"} + assert guardrail.should_run_guardrail(data=tool_call, event_type=GuardrailEventHooks.pre_mcp_call) + await guardrail.apply_guardrail(inputs=_inputs(), request_data=tool_call, input_type="request") + + await guardrail.apply_guardrail(inputs=_inputs(), request_data=tool_call, input_type="request") + + assert _post_mock(guardrail).call_args.kwargs["json"]["mcp_tool_name"] is None + + def _during_call_request() -> dict[str, object]: return {"model": "gpt-4o", "messages": [{"role": "user", "content": "mail ada@example.com"}]} From 93266c971e33834954eee19217028a7cebaae8ed Mon Sep 17 00:00:00 2001 From: Airia Open Source Date: Tue, 3 Nov 2026 02:57:25 -0600 Subject: [PATCH 5/5] test(guardrails): pin that only MCP hooks mark an Airia evaluation as a tool call Cover the cases the marker must not appear in: during_call and post_call evaluations, and client request data shaped like an MCP call (a top-level mcp_tool_name, mcp_tool_call_metadata, or metadata.mcp_tool_name). Also pin the fallback name used when an MCP evaluation carries none, and that a hook that does not run leaves no marker for the next evaluation. --- .../guardrails/guardrail_hooks/test_airia.py | 55 +++++++++++++++++++ 1 file changed, 55 insertions(+) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_airia.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_airia.py index 4eb9e2fcaf0..76e330cca73 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_airia.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_airia.py @@ -229,6 +229,57 @@ async def test_an_mcp_evaluation_does_not_mark_the_next_model_evaluation_as_a_to assert _post_mock(guardrail).call_args.kwargs["json"]["mcp_tool_name"] is None +@pytest.mark.asyncio +@pytest.mark.parametrize( + "marker", + [ + {"mcp_tool_name": "send_email"}, + {"mcp_tool_call_metadata": {"name": "send_email", "arguments": {}}}, + {"metadata": {"mcp_tool_name": "send_email"}}, + ], +) +async def test_client_supplied_mcp_markers_do_not_mark_model_traffic_as_a_tool_call( + marker: dict[str, object], +) -> None: + guardrail: Final = _make_guardrail_with_response({"action": "NONE"}, event_hook="pre_call", default_on=True) + + await UnifiedLLMGuardrails().async_pre_call_hook( # pyright: ignore[reportUnknownMemberType] # annotates data as a bare dict + user_api_key_dict=UserAPIKeyAuth(api_key="test"), + cache=DualCache(), + data={"guardrail_to_apply": guardrail, "messages": [{"role": "user", "content": "hello"}], **marker}, + call_type="completion", + ) + + assert _post_mock(guardrail).call_args.kwargs["json"]["mcp_tool_name"] is None + + +@pytest.mark.asyncio +async def test_a_hook_that_does_not_run_leaves_no_marker_for_the_next_evaluation() -> None: + guardrail: Final = _make_guardrail_with_response({"action": "NONE"}, event_hook="pre_call", default_on=True) + tool_call: Final[dict[str, object]] = {"mcp_tool_name": "send_email"} + + assert not guardrail.should_run_guardrail(data=tool_call, event_type=GuardrailEventHooks.pre_mcp_call) + await guardrail.apply_guardrail(inputs=_inputs(), request_data=tool_call, input_type="request") + + assert _post_mock(guardrail).call_args.kwargs["json"]["mcp_tool_name"] is None + + +@pytest.mark.asyncio +async def test_an_mcp_evaluation_without_a_tool_name_still_names_one_so_tool_filters_apply( + monkeypatch: pytest.MonkeyPatch, +) -> None: + guardrail: Final = _make_guardrail_with_response({"action": "NONE"}, event_hook="post_mcp_call", default_on=True) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + await ProxyLogging(user_api_key_cache=UserApiKeyCache()).post_mcp_call_hook( + response=CallToolResult(content=[TextContent(type="text", text="ok")]), + request_data={"mcp_tool_call_metadata": {"arguments": {}}}, + user_api_key_dict=UserAPIKeyAuth(api_key="test"), + ) + + assert _post_mock(guardrail).call_args.kwargs["json"]["mcp_tool_name"] == "unknown" + + def _during_call_request() -> dict[str, object]: return {"model": "gpt-4o", "messages": [{"role": "user", "content": "mail ada@example.com"}]} @@ -254,6 +305,7 @@ async def test_during_call_enforces_every_intervention_as_a_block( assert excinfo.value.blocked_content is True assert _post_mock(guardrail).call_args.kwargs["json"]["input_type"] == "request" + assert _post_mock(guardrail).call_args.kwargs["json"]["mcp_tool_name"] is None @pytest.mark.asyncio @@ -268,6 +320,7 @@ async def test_during_call_allow_lets_the_request_through(monkeypatch: pytest.Mo assert request["messages"] == _during_call_request()["messages"] _post_mock(guardrail).assert_called_once() + assert _post_mock(guardrail).call_args.kwargs["json"]["mcp_tool_name"] is None @pytest.mark.asyncio @@ -571,6 +624,7 @@ async def test_streaming_redaction_actually_reaches_the_client() -> None: request_data={ # mutable-ok: one-shot request_data the framework requires as a plain dict "messages": [{"role": "user", "content": "share your contact"}], "stream": True, + "mcp_tool_name": "send_email", }, guardrail_to_apply=guardrail, ): @@ -581,6 +635,7 @@ async def test_streaming_redaction_actually_reaches_the_client() -> None: assert "[EmailAddress1]" in assembled assert "ada@example.com" not in assembled _post_mock(guardrail).assert_called_once() + assert _post_mock(guardrail).call_args.kwargs["json"]["mcp_tool_name"] is None def _streamed_request(route: str | None) -> dict[str, object]: