From cffa202bbf8352ec862e4d58c8ceed8579e4f6fb Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 12:48:47 -0700 Subject: [PATCH] fix(guardrails): keep the resynced event_hook in the plain-string shape the constructor stores update_in_memory_litellm_params validated mode into GuardrailEventHooks members while __init__ stores the plain strings LitellmParams.mode carries, so readers that stringify event_hook (akto, straiker) saw different values on the serving worker than on re-initialized workers. Presidio forced post_call assignments go through the same shape, and Straiker recomputes configured_modes on every update --- litellm/integrations/custom_guardrail.py | 23 +++++++++-- .../guardrails/guardrail_hooks/presidio.py | 11 ++++-- .../guardrail_hooks/straiker/straiker.py | 6 +++ .../integrations/test_custom_guardrail.py | 38 +++++++++++++++++-- .../guardrail_hooks/test_presidio.py | 3 +- .../guardrail_hooks/test_straiker.py | 15 ++++++++ .../guardrails/test_guardrail_registry.py | 4 +- 7 files changed, 86 insertions(+), 14 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index e2754cd7723..3a3783e58a5 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -5,7 +5,7 @@ import os import secrets from collections.abc import Mapping, Sequence from datetime import datetime -from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, get_args +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, cast, get_args from pydantic import TypeAdapter @@ -137,6 +137,23 @@ GUARDRAIL_MODE_ADAPTER: Final[TypeAdapter[GuardrailEventHooks | list[GuardrailEv ) +def event_hook_as_constructed( + validated_mode: GuardrailEventHooks | list[GuardrailEventHooks] | Mode, +) -> GuardrailEventHooks | list[GuardrailEventHooks] | Mode: + """ + Return the shape ``__init__`` stores for the same mode: ``LitellmParams`` + coerces enum members to plain strings, so a resynced ``event_hook`` must + hold plain strings too or workers end up disagreeing on ``str(event_hook)``. + """ + if isinstance(validated_mode, Mode): + return validated_mode + if isinstance(validated_mode, list): + return cast( # cast-ok: __init__ stores the plain strings LitellmParams.mode carries + list[GuardrailEventHooks], [hook.value for hook in validated_mode] + ) + return cast(GuardrailEventHooks, validated_mode.value) # cast-ok: same parity as the list branch + + def get_session_id_from_request_data(request_data: dict[str, Any]) -> str | None: """Extract session_id from request data (litellm_session_id or metadata).""" session_id = request_data.get("litellm_session_id") @@ -1378,7 +1395,7 @@ class CustomGuardrail(CustomLogger): if value is not None: setattr(self, key, value) if new_event_hook is not None: - self.event_hook = new_event_hook + self.event_hook = event_hook_as_constructed(new_event_hook) def get_guardrails_messages_for_call_type( self, call_type: CallTypes, data: dict | None = None @@ -1407,8 +1424,6 @@ class CustomGuardrail(CustomLogger): # User/System messages are stored in the "input" key, use litellm transformation to get the messages ######################################################### if call_type == CallTypes.responses.value or call_type == CallTypes.aresponses.value: - from typing import cast - from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index da2776eda7d..0272247392c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -37,6 +37,7 @@ from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException from litellm.integrations.custom_guardrail import ( GUARDRAIL_MODE_ADAPTER, CustomGuardrail, + event_hook_as_constructed, log_guardrail_information, ) from litellm.proxy._types import UserAPIKeyAuth @@ -1641,12 +1642,16 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): if self.event_hook == GuardrailEventHooks.logging_only: return if self.apply_to_output: - self.event_hook = GuardrailEventHooks.post_call + self.event_hook = event_hook_as_constructed(GuardrailEventHooks.post_call) return if not self.output_parse_pii: return current_hook: Final = self.event_hook if isinstance(current_hook, str) and current_hook != "post_call": - self.event_hook = GUARDRAIL_MODE_ADAPTER.validate_python((current_hook, GuardrailEventHooks.post_call)) + self.event_hook = event_hook_as_constructed( + GUARDRAIL_MODE_ADAPTER.validate_python((current_hook, GuardrailEventHooks.post_call)) + ) elif isinstance(current_hook, list) and "post_call" not in current_hook: - self.event_hook = GUARDRAIL_MODE_ADAPTER.validate_python((*current_hook, GuardrailEventHooks.post_call)) + self.event_hook = event_hook_as_constructed( + GUARDRAIL_MODE_ADAPTER.validate_python((*current_hook, GuardrailEventHooks.post_call)) + ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py b/litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py index 7cca1ae2d63..00ccba11be6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py +++ b/litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio import json import random +from collections.abc import Mapping from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn from urllib.parse import urlsplit @@ -47,6 +48,7 @@ from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.guardrails import LitellmParams from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel GUARDRAIL_NAME: Final = "straiker" @@ -329,6 +331,10 @@ class StraikerGuardrail(CustomGuardrail): self.configured_modes = _configured_modes(self.event_hook) + def update_in_memory_litellm_params(self, litellm_params: LitellmParams | Mapping[str, object]) -> None: + super().update_in_memory_litellm_params(litellm_params) + self.configured_modes = _configured_modes(self.event_hook) + def _webhook_url(self) -> str: return f"{self.api_base}{WEBHOOK_PATH}" diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index d5f94d553a0..4d162f5bc54 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -2261,7 +2261,7 @@ class TestUpdateInMemoryLitellmParams: LitellmParams(guardrail="update-test", mode="post_call", default_on=True) ) - assert guardrail.event_hook is GuardrailEventHooks.post_call + assert guardrail.event_hook == "post_call" assert guardrail.should_run_guardrail(data={}, event_type=GuardrailEventHooks.post_call) is True assert guardrail.should_run_guardrail(data={}, event_type=GuardrailEventHooks.pre_call) is False @@ -2277,10 +2277,40 @@ class TestUpdateInMemoryLitellmParams: } ) - assert guardrail.event_hook is GuardrailEventHooks.post_call + assert guardrail.event_hook == "post_call" assert getattr(guardrail, "api_base", None) == "https://guardrail.example.com" assert guardrail.should_run_guardrail(data={}, event_type=GuardrailEventHooks.post_call) is True + @pytest.mark.parametrize( + "mode", + [ + "post_call", + ["pre_call", "post_call"], + {"default": "post_call", "tags": {"team-a": ["pre_call", "post_call"]}}, + ], + ids=["str", "list", "mode"], + ) + def test_resynced_event_hook_has_the_shape_a_fresh_worker_constructs(self, mode): + """Other workers rebuild the guardrail from the same DB row through + LitellmParams, which coerces enum members to plain strings; the serving + worker's in-place resync must land on that exact shape, or type-sensitive + readers such as str(self.event_hook) disagree across workers.""" + updated = self._guardrail() + updated.update_in_memory_litellm_params({"guardrail": "update-test", "mode": mode, "default_on": True}) + + constructed = CustomGuardrail( + guardrail_name="update-test", + event_hook=LitellmParams(guardrail="update-test", mode=mode).mode, + default_on=True, + supported_event_hooks=[GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call], + ) + + assert updated.event_hook == constructed.event_hook + assert type(updated.event_hook) is type(constructed.event_hook) + assert str(updated.event_hook) == str(constructed.event_hook) + if isinstance(updated.event_hook, list): + assert [type(hook) for hook in updated.event_hook] == [type(hook) for hook in constructed.event_hook] + def test_none_values_do_not_clobber_constructor_state(self): guardrail = self._guardrail() guardrail.additional_provider_specific_params = {"team": "security"} @@ -2297,7 +2327,7 @@ class TestUpdateInMemoryLitellmParams: assert guardrail.additional_provider_specific_params == {"team": "security"} assert guardrail.api_base == "https://guardrail.example.com" - assert guardrail.event_hook is GuardrailEventHooks.post_call + assert guardrail.event_hook == "post_call" def test_strict_mode_rejects_unsupported_mode_without_mutating(self, monkeypatch): monkeypatch.delenv("LITELLM_STRICT_GUARDRAIL_MODES", raising=False) @@ -2317,7 +2347,7 @@ class TestUpdateInMemoryLitellmParams: guardrail.update_in_memory_litellm_params({"mode": "during_call", "api_base": "https://guardrail.example.com"}) - assert guardrail.event_hook is GuardrailEventHooks.during_call + assert guardrail.event_hook == "during_call" assert getattr(guardrail, "api_base", None) == "https://guardrail.example.com" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py index 519df980d22..0be37155e73 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -3179,7 +3179,8 @@ def test_update_in_memory_output_callback_keeps_forced_post_call(): guardrail.update_in_memory_litellm_params({"guardrail": "presidio", "mode": "pre_call", "default_on": True}) - assert guardrail.event_hook is GuardrailEventHooks.post_call + assert guardrail.event_hook == "post_call" + assert type(guardrail.event_hook) is str assert guardrail.should_run_guardrail(data={}, event_type=GuardrailEventHooks.post_call) is True assert guardrail.should_run_guardrail(data={}, event_type=GuardrailEventHooks.pre_call) is False diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_straiker.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_straiker.py index 81604e22c87..679c69ecd91 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_straiker.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_straiker.py @@ -431,6 +431,21 @@ async def test_context_mode_omitted_when_event_hook_absent(): assert "mode" not in _posted_payload(g)["context"] +@pytest.mark.asyncio +async def test_context_mode_follows_an_in_memory_mode_update(): + g = _make_guardrail(event_hook="pre_call") + g.update_in_memory_litellm_params({"guardrail": "straiker", "mode": "post_call", "default_on": True}) + g.async_handler.post.return_value = _mock_response("NONE") + await g.apply_guardrail( + inputs={"texts": ["x"]}, + request_data={"model": "m"}, + input_type="response", + logging_obj=_logging_obj(), + ) + assert g.configured_modes == ["post_call"] + assert _posted_payload(g)["context"]["mode"] == ["post_call"] + + @pytest.mark.asyncio async def test_identity_key_and_team_coalesce_alias_over_id(): g = _make_guardrail() diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py index 015d530257b..bd1e93d8866 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py @@ -176,7 +176,7 @@ def test_update_in_memory_guardrail(): ) is True ) - assert handler.guardrail_id_to_custom_guardrail["123"].event_hook is GuardrailEventHooks.pre_call + assert handler.guardrail_id_to_custom_guardrail["123"].event_hook == "pre_call" def test_update_in_memory_guardrail_raw_db_dict_resyncs_event_hook(): @@ -199,7 +199,7 @@ def test_update_in_memory_guardrail_raw_db_dict_resyncs_event_hook(): handler.update_in_memory_guardrail("123", updated_row) callback = handler.guardrail_id_to_custom_guardrail["123"] - assert callback.event_hook is GuardrailEventHooks.post_call + assert callback.event_hook == "post_call" assert callback.should_run_guardrail(data={}, event_type=GuardrailEventHooks.post_call) is True assert callback.should_run_guardrail(data={}, event_type=GuardrailEventHooks.pre_call) is False assert handler.IN_MEMORY_GUARDRAILS["123"] == updated_row