diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 4093f2c5248..c71761a7adb 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -9652,7 +9652,9 @@ { "enum": [ "warn", - "end_session" + "end_session", + "block", + "alert" ], "type": "string" }, @@ -9660,7 +9662,7 @@ "type": "null" } ], - "description": "For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection.", + "description": "For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection. For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning.", "title": "On Violation" }, "only_scan_new_messages": { @@ -11969,7 +11971,9 @@ { "enum": [ "warn", - "end_session" + "end_session", + "block", + "alert" ], "type": "string" }, @@ -11977,7 +11981,7 @@ "type": "null" } ], - "description": "For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection.", + "description": "For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection. For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning.", "title": "On Violation" }, "only_scan_new_messages": { diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 744b2959c73..afb9997f2e6 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -13,7 +13,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, Union, from urllib.parse import urlparse from fastapi import APIRouter, Depends, HTTPException, Request -from pydantic import BaseModel +from pydantic import BaseModel, ValidationError from litellm._logging import verbose_proxy_logger from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH @@ -1202,7 +1202,13 @@ async def patch_guardrail( litellm_params_dict: Final = litellm_params.model_dump(exclude_unset=True) litellm_params_dict.update(requested_litellm_params) merged_litellm_params: Final = _as_str_object_mapping(litellm_params_dict) - litellm_params = LitellmParams(**merged_litellm_params) + try: + litellm_params = LitellmParams(**merged_litellm_params) + except ValidationError as validation_error: + raise HTTPException( + status_code=422, + detail=f"Invalid guardrail configuration, update rejected: {validation_error}", + ) from validation_error # Update guardrail_info if provided guardrail_info: Final = ( diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_security/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_security/__init__.py index d53a4157e0e..1607dfff63e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_security/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_security/__init__.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Final, Literal, Optional, cast +from typing import TYPE_CHECKING, Final, Literal, Optional import litellm from litellm.proxy.guardrails.guardrail_hooks.mcp_security.mcp_security_guardrail import ( @@ -20,10 +20,7 @@ def initialize_guardrail( if not guardrail_name: raise ValueError("MCP Security: guardrail_name is required") - on_violation: Final[Literal["block", "alert"]] = cast( - Literal["block", "alert"], - getattr(litellm_params, "on_violation", "block"), - ) + on_violation: Final[Literal["block", "alert"]] = "block" if litellm_params.on_violation == "block" else "alert" mcp_security_guardrail: Final = MCPSecurityGuardrail( guardrail_name=guardrail_name, diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index ef28181eba5..02dee40f2a3 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -778,6 +778,9 @@ class ContentFilterConfigModel(BaseModel): ) +MCP_SECURITY_ON_VIOLATION: Final = frozenset({"block", "alert"}) + + class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch update guardrails api_key: str | None = Field(default=None, description="API key for the guardrail service") api_base: str | None = Field(default=None, description="Base URL for the guardrail service API") @@ -886,9 +889,13 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up default=None, description="For /v1/realtime sessions: automatically close the session after this many guardrail violations.", ) - on_violation: Literal["warn", "end_session"] | None = Field( + on_violation: Literal["warn", "end_session", "block", "alert"] | None = Field( default=None, - description="For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection.", + description=( + "For /v1/realtime sessions: 'warn' speaks the violation message and continues; " + "'end_session' speaks the message and closes the connection. " + "For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning." + ), ) realtime_violation_message: str | None = Field( default=None, @@ -1093,6 +1100,15 @@ class LitellmParams( # pyright: ignore[reportIncompatibleVariableOverride] # o except (TypeError, ValueError) as e: raise ValueError(f"timeout must be numeric, got {v!r}") from e + @model_validator(mode="after") + def validate_on_violation_for_guardrail(self) -> "LitellmParams": + if ( + self.on_violation in MCP_SECURITY_ON_VIOLATION + and self.guardrail != SupportedGuardrailIntegrations.MCP_SECURITY.value + ): + raise ValueError(f"on_violation={self.on_violation!r} is only supported by guardrail='mcp_security'") + return self + def __init__(self, **kwargs) -> None: default_on: Final = kwargs.pop("default_on", None) if default_on is not None: diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_mcp_security.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_mcp_security.py index 4444cd693ff..d57a91d45bf 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_mcp_security.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_mcp_security.py @@ -6,16 +6,19 @@ and allows requests with only registered servers. Covers both /chat/completions and /responses API paths (same pre_call_hook logic, different call_type). """ +from typing import Literal from unittest.mock import MagicMock, patch import pytest from fastapi import HTTPException +import litellm from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_hooks.mcp_security import initialize_guardrail from litellm.proxy.guardrails.guardrail_hooks.mcp_security.mcp_security_guardrail import ( MCPSecurityGuardrail, ) -from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.guardrails import Guardrail, GuardrailEventHooks, LitellmParams @pytest.fixture @@ -182,3 +185,23 @@ class TestMCPSecurityGuardrailPreCall: call_type="acompletion", ) assert result == data + + +class TestInitializeGuardrail: + @pytest.mark.parametrize( + "configured,expected", + [("block", "block"), ("alert", "alert"), (None, "alert"), ("warn", "alert"), ("end_session", "alert")], + ) + def test_on_violation_from_litellm_params( + self, + configured: Literal["block", "alert", "warn", "end_session"] | None, + expected: Literal["block", "alert"], + ): + litellm_params = LitellmParams(guardrail="mcp_security", mode="pre_call", on_violation=configured) + guardrail = Guardrail(guardrail_name="mcp-security-block", litellm_params=litellm_params) + + result = initialize_guardrail(litellm_params=litellm_params, guardrail=guardrail) + + assert isinstance(result, MCPSecurityGuardrail) + assert result.on_violation == expected + assert result in litellm.callbacks diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py index fe2cd819717..530f8ffd854 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -1361,6 +1361,22 @@ async def test_patch_guardrail_endpoint( assert "Failed to update" in str(mock_logger.warning.call_args) +@pytest.mark.asyncio +async def test_patch_guardrail_rejects_mcp_only_on_violation_with_422(mocker, mock_guardrail_registry): + mocker.patch("litellm.proxy.proxy_server.prisma_client", mocker.Mock()) # test-quality-ok: endpoint has no DI seam + mocker.patch( # test-quality-ok: endpoint has no DI seam + "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_guardrail_registry + ) + request = PatchGuardrailRequest(litellm_params=BaseLitellmParams(on_violation="block")) + + with pytest.raises(HTTPException) as exc_info: + await patch_guardrail("test-guardrail-id", request, user_api_key_dict=MOCK_ADMIN_USER) + + assert exc_info.value.status_code == 422 + assert "only supported by guardrail='mcp_security'" in str(exc_info.value.detail) + mock_guardrail_registry.update_guardrail_in_db.assert_not_called() + + @pytest.mark.parametrize( "scenario,expected_result,expected_exception", [ diff --git a/tests/test_litellm/types/test_guardrails_case_normalization.py b/tests/test_litellm/types/test_guardrails_case_normalization.py index 3e7a573ea8e..26c1d395320 100644 --- a/tests/test_litellm/types/test_guardrails_case_normalization.py +++ b/tests/test_litellm/types/test_guardrails_case_normalization.py @@ -2,6 +2,8 @@ Test case normalization in LitellmParams for all guardrail types """ +from typing import Literal + import pytest from pydantic import ValidationError @@ -93,6 +95,34 @@ class TestLitellmParamsCaseNormalization: assert params.on_disallowed_action.islower() +class TestOnViolationAcceptedValues: + """on_violation is shared by /v1/realtime guardrails and the mcp_security guardrail""" + + @pytest.mark.parametrize("action", ["block", "alert"]) + def test_mcp_security_policy_template_on_violation_is_accepted(self, action: Literal["block", "alert"]): + params = LitellmParams( + guardrail="mcp_security", + mode="pre_call", + default_on=True, + on_violation=action, + ) + assert params.on_violation == action + + @pytest.mark.parametrize("action", ["warn", "end_session"]) + def test_realtime_on_violation_still_accepted(self, action: Literal["warn", "end_session"]): + params = LitellmParams(guardrail="presidio", mode="pre_call", on_violation=action) + assert params.on_violation == action + + @pytest.mark.parametrize("action", ["block", "alert"]) + def test_mcp_only_on_violation_is_rejected_for_other_guardrails(self, action: Literal["block", "alert"]): + with pytest.raises(ValidationError, match="only supported by guardrail='mcp_security'"): + LitellmParams(guardrail="presidio", mode="pre_call", on_violation=action) + + def test_unknown_on_violation_is_rejected(self): + with pytest.raises(ValidationError): + LitellmParams(guardrail="mcp_security", mode="pre_call", on_violation="ignore") + + class TestSensitiveDataRoutingValidation: """on_sensitive_data='route' requires a target model to be set""" diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 9de988fde4c..6fb08445aff 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -23759,9 +23759,9 @@ export interface components { on_sensitive_data?: ("block" | "route") | null; /** * On Violation - * @description For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection. + * @description For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection. For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning. */ - on_violation?: ("warn" | "end_session") | null; + on_violation?: ("warn" | "end_session" | "block" | "alert") | null; /** * Only Scan New Messages * @description When True, the guardrail only scans messages that have not already been scanned earlier in the same session (identified by litellm_session_id / session_id). Message content is hashed per session and cached; only the diff (new or edited messages) is sent to the guardrail provider on follow-up calls. Falls back to a full scan when the request has no session id or the cache is unavailable. Intended for blocking/detection guardrails; not applied when mask_request_content is set. @@ -30801,9 +30801,9 @@ export interface components { on_sensitive_data?: ("block" | "route") | null; /** * On Violation - * @description For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection. + * @description For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection. For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning. */ - on_violation?: ("warn" | "end_session") | null; + on_violation?: ("warn" | "end_session" | "block" | "alert") | null; /** * Only Scan New Messages * @description When True, the guardrail only scans messages that have not already been scanned earlier in the same session (identified by litellm_session_id / session_id). Message content is hashed per session and cached; only the diff (new or edited messages) is sent to the guardrail provider on follow-up calls. Falls back to a full scan when the request has no session id or the cache is unavailable. Intended for blocking/detection guardrails; not applied when mask_request_content is set.