diff --git a/litellm/proxy/guardrails/guardrail_hooks/airia/airia.py b/litellm/proxy/guardrails/guardrail_hooks/airia/airia.py index 389706d7417..22bc1b36108 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/airia/airia.py +++ b/litellm/proxy/guardrails/guardrail_hooks/airia/airia.py @@ -1,7 +1,9 @@ import os import uuid from collections.abc import Mapping +from contextvars import ContextVar from typing import ( + ClassVar, Final, Literal, ) @@ -23,9 +25,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" @@ -35,7 +39,21 @@ 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, +) + +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) _DECODED_JSON: Final = TypeAdapter(object) @@ -88,6 +106,15 @@ def _request_route(request_data: Mapping[str, object]) -> str | None: 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 UNKNOWN_MCP_TOOL_NAME + + 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.""" @@ -95,7 +122,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, @@ -135,6 +164,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 @@ -146,6 +184,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 @@ -164,6 +204,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 or {}) if is_mcp_tool_evaluation else None, }, headers={"Authorization": f"Bearer {self.api_key}"}, ) @@ -196,6 +237,9 @@ class AiriaGuardrail(CustomGuardrail): ) if action == ACTION_INTERVENED: + # 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): raise self._blocked() return self._rewritten(body, inputs) @@ -205,6 +249,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..76e330cca73 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_airia.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_airia.py @@ -1,18 +1,26 @@ 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 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.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 +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.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 @@ -98,13 +106,238 @@ 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"] + assert _post_mock(guardrail).call_args.kwargs["json"]["mcp_tool_name"] == "send_email" + + +@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={"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 +@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=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 + + +@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"}]} + + +@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: + 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" + assert _post_mock(guardrail).call_args.kwargs["json"]["mcp_tool_name"] is None + + +@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() + assert _post_mock(guardrail).call_args.kwargs["json"]["mcp_tool_name"] is None + + +@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"]) @@ -391,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, ): @@ -401,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]: