From e6311b95e081c6283c09c8332b1cf07c0af97ee5 Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Thu, 27 Aug 2026 06:50:12 -0400 Subject: [PATCH] fix(guardrails): address maintainer review findings on Lakera v2 advisory mode - Gate advisory_system_message template validation on on_flagged= 'inject_system_message', since block/monitor mode never reads it. - Allow on_flagged='inject_system_message' with mode='during_call' at construction/hot-reload instead of rejecting it; async_moderation_hook already degrades gracefully (masks if possible, else logs a warning). - reinitialize_guardrail now restores the previous live instance when the new config fails to initialize, instead of leaving the guardrail deleted entirely with nothing enforcing it. - PATCH /guardrails/{id} rolls back the DB write and returns 422 when the in-memory sync rejects the new config, instead of persisting a config that never actually took effect and returning 200. - Qualifire now rejects on_flagged values it doesn't implement (only Lakera should accept 'inject_system_message'; LitellmParams flattens the field across every guardrail config mixin). --- .../proxy/guardrails/guardrail_endpoints.py | 21 ++++ .../guardrail_hooks/lakera_ai_v2.py | 58 ++------- .../guardrail_hooks/qualifire/qualifire.py | 10 ++ .../proxy/guardrails/guardrail_registry.py | 10 +- .../guardrail_hooks/test_lakera_ai_v2.py | 119 ++++++++++++------ .../guardrail_hooks/test_qualifire.py | 34 +++++ .../guardrails/test_guardrail_endpoints.py | 38 +++++- .../guardrails/test_guardrail_registry.py | 61 +++++++++ 8 files changed, 260 insertions(+), 91 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 20efbe06ecc..3d00d9a0025 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -1218,6 +1218,27 @@ async def patch_guardrail( verbose_proxy_logger.info( "Immediate sync: Successfully updated guardrail '%s' (ID: %s)", guardrail_name, guardrail_id ) + except (ValueError, TypeError) as update_error: + # The new config is invalid (e.g. an unsupported on_flagged combination): + # reinitialize_guardrail already restored the previous live instance, but + # update_guardrail_in_db above already persisted the rejected config to + # the DB. Roll that back too, so the DB and the live guardrail never + # disagree about what's actually enforcing, and surface the rejection to + # the caller instead of a misleading 200. + await GUARDRAIL_REGISTRY.update_guardrail_in_db( + guardrail_id=guardrail_id, + guardrail=Guardrail( + guardrail_id=guardrail_id, + guardrail_name=existing_guardrail.get("guardrail_name") or "", + litellm_params=LitellmParams(**existing_litellm_params), + guardrail_info=existing_guardrail.get("guardrail_info", {}), + ), + prisma_client=prisma_client, + ) + raise HTTPException( + status_code=422, + detail=f"Invalid guardrail configuration, update rejected: {update_error}", + ) from update_error except Exception as update_error: verbose_proxy_logger.warning( "Immediate sync: Failed to update '%s' (ID: %s) in memory: %s", diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index 0146283ed1e..ba45dc1c0b4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -31,7 +31,7 @@ from litellm.proxy.guardrails._content_utils import ( has_non_string_content, ) from litellm.secret_managers.main import get_secret_str -from litellm.types.guardrails import GuardrailEventHooks, LitellmParams, Mode +from litellm.types.guardrails import GuardrailEventHooks, LitellmParams from litellm.types.llms.openai import AllMessageValues from litellm.types.proxy.guardrails.guardrail_hooks.lakera_ai_v2 import ( LakeraAIBreakdownItem, @@ -82,28 +82,6 @@ def _template_uses_reason_placeholder(template: str) -> bool: return any(field_name == "reason" for _, field_name, _, _ in Formatter().parse(template)) -def _event_hook_includes_during_call( - event_hook: GuardrailEventHooks | Sequence[GuardrailEventHooks] | Mode | str | Sequence[str] | None, -) -> bool: - """True if ``event_hook`` could ever resolve to during_call, covering a plain - value, a list of values, or a tag-based Mode (checked across every tag value - and the default).""" - candidates: Final = ( - tuple(event_hook.tags.values()) + (event_hook.default,) - if isinstance(event_hook, Mode) - else tuple(event_hook) - if isinstance(event_hook, list) - else (event_hook,) - ) - - flattened: Final = tuple( - value - for candidate in candidates - for value in (tuple(candidate) if isinstance(candidate, list) else (candidate,)) - ) - return any(value == GuardrailEventHooks.during_call for value in flattened if value is not None) - - def _pre_masking_scope_indices( guardrail: "LakeraAIGuardrail", messages: Sequence[object], @@ -303,7 +281,6 @@ class LakeraAIGuardrail(CustomGuardrail): self._validate_advisory_config( on_flagged=self.on_flagged, advisory_system_message=self.advisory_system_message, - event_hook=self.event_hook, payload=self.payload, breakdown=self.breakdown, ) @@ -311,21 +288,18 @@ class LakeraAIGuardrail(CustomGuardrail): def update_in_memory_litellm_params(self, litellm_params: LitellmParams) -> None: """ The base implementation blindly ``setattr``s every field on ``litellm_params`` - (including ``on_flagged``/``advisory_system_message``) onto this live instance - with no revalidation, so an in-place config update (via the DB/UI, without a - restart) could otherwise reintroduce the exact invalid on_flagged/event_hook - combinations __init__ rejects. Validate the prospective post-update state - *before* mutating, so a rejected update leaves the live instance untouched + (including ``on_flagged``/``advisory_system_message``/``payload``/``breakdown``) + onto this live instance with no revalidation, so an in-place config update (via + the DB/UI, without a restart) could otherwise reintroduce the exact invalid + on_flagged combinations __init__ rejects. Validate the prospective post-update + state *before* mutating, so a rejected update leaves the live instance untouched instead of raising after it's already been corrupted. - The base setattr also writes ``litellm_params.mode`` onto a new - ``self.mode`` attribute rather than the ``self.event_hook`` dispatch - actually reads (LitellmParams has no field literally named - ``event_hook``), so without the explicit sync below a hot reload that - moves this guardrail off during_call would pass validation but still - dispatch as during_call afterward -- inject_system_message would then - run against a live instance validation had confirmed was safe, but - whose real dispatch hook never changed. + The base setattr also writes ``litellm_params.mode`` onto a new ``self.mode`` + attribute rather than the ``self.event_hook`` dispatch actually reads + (LitellmParams has no field literally named ``event_hook``), so without the + explicit sync below a hot reload that changes mode would pass validation but + keep dispatching on the stale event_hook. """ new_event_hook: Final = getattr(litellm_params, "mode", None) or self.event_hook prospective_payload: Final = getattr(litellm_params, "payload", None) @@ -333,7 +307,6 @@ class LakeraAIGuardrail(CustomGuardrail): self._validate_advisory_config( on_flagged=getattr(litellm_params, "on_flagged", None) or self.on_flagged, advisory_system_message=getattr(litellm_params, "advisory_system_message", None), - event_hook=new_event_hook, payload=self.payload if prospective_payload is None else prospective_payload, breakdown=self.breakdown if prospective_breakdown is None else prospective_breakdown, ) @@ -344,11 +317,10 @@ class LakeraAIGuardrail(CustomGuardrail): self, on_flagged: str, advisory_system_message: str | None, - event_hook: GuardrailEventHooks | Sequence[GuardrailEventHooks] | Mode | str | Sequence[str] | None, payload: bool | None, breakdown: bool | None, ) -> None: - if advisory_system_message is not None: + if on_flagged == "inject_system_message" and advisory_system_message is not None: if not _template_uses_reason_placeholder(advisory_system_message): raise ValueError( "Invalid advisory_system_message template: must include a real {reason} " @@ -361,12 +333,6 @@ class LakeraAIGuardrail(CustomGuardrail): f"Invalid advisory_system_message template: {e}. The template must be a valid " "str.format() string using only the {reason} placeholder." ) from e - if on_flagged == "inject_system_message" and _event_hook_includes_during_call(event_hook): - raise ValueError( - "on_flagged='inject_system_message' is not supported for mode='during_call': during_call " - "runs concurrently with the LLM dispatch with no pre-call barrier, so the advisory message " - "cannot reliably reach the request. Use mode='pre_call' instead." - ) if on_flagged == "inject_system_message" and not (payload and breakdown): raise ValueError( "on_flagged='inject_system_message' requires payload=True and breakdown=True: advisory " diff --git a/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py b/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py index d6fb1378da0..5a16481894e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py +++ b/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py @@ -87,6 +87,16 @@ class QualifireGuardrail(CustomGuardrail): self.tool_selection_quality_check = tool_selection_quality_check self.assertions = assertions self.on_flagged = on_flagged or "block" + if self.on_flagged not in ("block", "monitor"): + # on_flagged is defined on LakeraV2GuardrailConfigModel but LitellmParams + # flattens every guardrail config mixin together, so a value Lakera + # supports (e.g. "inject_system_message") type-checks for any guardrail, + # including this one, which never implements it. Reject it explicitly + # instead of silently falling through to a block-on-anything-else branch. + raise ValueError( + f"Qualifire guardrail does not support on_flagged={self.on_flagged!r}; " + "only 'block' and 'monitor' are supported." + ) # If no checks are specified and no evaluation_id, default to prompt_injections if not self._has_any_check_enabled() and not self.evaluation_id: diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index fce2b3ec465..8bf85a54300 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -778,15 +778,23 @@ class InMemoryGuardrailHandler: """ Force re-initialization of a guardrail even if it exists in memory. Removes old callback from litellm.callbacks and creates fresh instance. + + If the new config fails to initialize (e.g. an invalid on_flagged + combination), the previous instance is restored rather than left + deleted: initialize_guardrail's own ValueError/TypeError propagate + uncaught, so a caller reaching this point after already deleting the + old instance would otherwise leave the guardrail providing no + protection at all, not merely "still enforcing the old config." """ guardrail_id: Final = guardrail.get("guardrail_id") if not guardrail_id: verbose_proxy_logger.error("Cannot reinitialize guardrail without guardrail_id") return None - # Remove from memory if exists (also removes from callbacks) previous_guardrail: Final = self.IN_MEMORY_GUARDRAILS.get(guardrail_id) previous_source: Final = self._sources.get(guardrail_id, source) + + # Remove from memory if exists (also removes from callbacks) if guardrail_id in self.IN_MEMORY_GUARDRAILS: self.delete_in_memory_guardrail(guardrail_id) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py index 44fc456dd5f..6febd02d895 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py @@ -685,18 +685,29 @@ class TestHumanizeLakeraBlockReasons: class TestAdvisorySystemMessageValidation: """advisory_system_message must be validated eagerly at construction time, - not lazily the first time a real request gets flagged.""" + not lazily the first time a real request gets flagged -- but only when + on_flagged='inject_system_message' actually reads it. Maintainer finding + on BerriAI/litellm#34940: this check previously ran unconditionally, so a + leftover/typo'd advisory_system_message on a guardrail configured + on_flagged='block' (which never calls _build_advisory_message at all) + disabled the entire guardrail for a field it never uses.""" def test_valid_template_constructs_without_error(self): - guardrail = LakeraAIGuardrail(api_key="test_key", advisory_system_message="Flagged for {reason}.") + guardrail = LakeraAIGuardrail( + api_key="test_key", on_flagged="inject_system_message", advisory_system_message="Flagged for {reason}." + ) assert guardrail.advisory_system_message == "Flagged for {reason}." def test_malformed_template_raises_at_construction(self): with pytest.raises(ValueError, match="Invalid advisory_system_message template"): - LakeraAIGuardrail(api_key="test_key", advisory_system_message="Flagged for {typo_field}.") + LakeraAIGuardrail( + api_key="test_key", + on_flagged="inject_system_message", + advisory_system_message="Flagged for {typo_field}.", + ) def test_none_template_is_allowed(self): - guardrail = LakeraAIGuardrail(api_key="test_key", advisory_system_message=None) + guardrail = LakeraAIGuardrail(api_key="test_key", on_flagged="inject_system_message", advisory_system_message=None) assert guardrail.advisory_system_message is None def test_template_missing_reason_placeholder_raises_at_construction(self): @@ -704,41 +715,76 @@ class TestAdvisorySystemMessageValidation: silently never tells the LLM why the request was flagged, defeating the point of advisory mode; this must be rejected too, not just malformed ones.""" with pytest.raises(ValueError, match="must include a real"): - LakeraAIGuardrail(api_key="test_key", advisory_system_message="This request was flagged.") + LakeraAIGuardrail( + api_key="test_key", on_flagged="inject_system_message", advisory_system_message="This request was flagged." + ) def test_escaped_reason_placeholder_raises_at_construction(self): """{{reason}} contains the substring "{reason}" but str.format() treats double braces as an escaped literal, never substituting the real value -- a naive substring check would wrongly accept this.""" with pytest.raises(ValueError, match="must include a real"): - LakeraAIGuardrail(api_key="test_key", advisory_system_message="Flagged for {{reason}}.") - - -class TestAdvisoryModeDuringCallUnsupported: - """inject_system_message cannot deliver its advertised behavior for - mode='during_call' (no pre-call barrier exists to land the mutation before - dispatch), so that combination must be rejected at construction time rather - than silently downgrading to monitor with no clear signal to the operator.""" - - def test_during_call_string_mode_raises_at_construction(self): - with pytest.raises(ValueError, match="not supported for mode='during_call'"): - LakeraAIGuardrail(api_key="test_key", on_flagged="inject_system_message", event_hook="during_call") - - def test_during_call_in_list_mode_raises_at_construction(self): - with pytest.raises(ValueError, match="not supported for mode='during_call'"): LakeraAIGuardrail( - api_key="test_key", - on_flagged="inject_system_message", - event_hook=["pre_call", "during_call"], + api_key="test_key", on_flagged="inject_system_message", advisory_system_message="Flagged for {{reason}}." ) - def test_during_call_in_tag_mode_raises_at_construction(self): - with pytest.raises(ValueError, match="not supported for mode='during_call'"): - LakeraAIGuardrail( - api_key="test_key", - on_flagged="inject_system_message", - event_hook=Mode(tags={"vip": "during_call"}, default="pre_call"), - ) + def test_malformed_template_with_block_mode_constructs_without_error(self): + """Maintainer finding on BerriAI/litellm#34940: on_flagged='block' never + reads advisory_system_message, so a malformed/leftover value there must + not disable the guardrail -- it's dead config, not a real error.""" + guardrail = LakeraAIGuardrail( + api_key="test_key", on_flagged="block", advisory_system_message="This request was flagged." + ) + assert guardrail.on_flagged == "block" + + def test_malformed_template_with_monitor_mode_constructs_without_error(self): + guardrail = LakeraAIGuardrail( + api_key="test_key", on_flagged="monitor", advisory_system_message="Flagged for {typo_field}." + ) + assert guardrail.on_flagged == "monitor" + + def test_in_memory_update_to_block_mode_with_malformed_template_is_allowed(self): + """A hot-reload that turns off advisory mode in the same update that + introduces a malformed advisory_system_message must succeed, not be + rejected for a field the new on_flagged value never reads.""" + guardrail = LakeraAIGuardrail(api_key="test_key", on_flagged="inject_system_message") + updated_params = LitellmParams( + guardrail="lakera_v2", mode="pre_call", on_flagged="block", advisory_system_message="No placeholder here." + ) + guardrail.update_in_memory_litellm_params(litellm_params=updated_params) + assert guardrail.on_flagged == "block" + + +class TestAdvisoryModeDuringCallDegradesGracefully: + """Maintainer finding on BerriAI/litellm#34940: rejecting on_flagged= + 'inject_system_message' + mode='during_call' at construction time disabled + the entire guardrail (via init_guardrails_v2's catch-and-skip) for a + combination async_moderation_hook already handles safely at runtime -- + it masks whatever's maskable and falls back to a log-only warning when + the advisory itself can't be delivered (see TestAdvisoryModeWiring's + during_call coverage). Construction/hot-reload must allow this + combination rather than disabling the guardrail outright.""" + + def test_during_call_string_mode_constructs_without_error(self): + guardrail = LakeraAIGuardrail(api_key="test_key", on_flagged="inject_system_message", event_hook="during_call") + assert guardrail.on_flagged == "inject_system_message" + assert guardrail.event_hook == "during_call" + + def test_during_call_in_list_mode_constructs_without_error(self): + guardrail = LakeraAIGuardrail( + api_key="test_key", + on_flagged="inject_system_message", + event_hook=["pre_call", "during_call"], + ) + assert guardrail.on_flagged == "inject_system_message" + + def test_during_call_in_tag_mode_constructs_without_error(self): + guardrail = LakeraAIGuardrail( + api_key="test_key", + on_flagged="inject_system_message", + event_hook=Mode(tags={"vip": "during_call"}, default="pre_call"), + ) + assert guardrail.on_flagged == "inject_system_message" def test_pre_call_only_mode_constructs_without_error(self): guardrail = LakeraAIGuardrail(api_key="test_key", on_flagged="inject_system_message", event_hook="pre_call") @@ -750,18 +796,11 @@ class TestAdvisoryModeDuringCallUnsupported: assert guardrail.on_flagged == "block" assert guardrail.event_hook == "during_call" - def test_in_memory_update_reintroducing_the_combo_raises(self): - """update_in_memory_litellm_params (the DB/UI hot-reload path) setattrs - every LitellmParams field onto a live instance with no revalidation, so - an update that flips on_flagged to inject_system_message on an instance - already running as during_call must be rejected too, not just the - combination formed at construction time.""" + def test_in_memory_update_reintroducing_the_combo_is_allowed(self): guardrail = LakeraAIGuardrail(api_key="test_key", on_flagged="block", event_hook="during_call") updated_params = LitellmParams(guardrail="lakera_v2", mode="during_call", on_flagged="inject_system_message") - with pytest.raises(ValueError, match="not supported for mode='during_call'"): - guardrail.update_in_memory_litellm_params(litellm_params=updated_params) - - assert guardrail.on_flagged == "block", "a rejected update must leave the live instance untouched" + guardrail.update_in_memory_litellm_params(litellm_params=updated_params) + assert guardrail.on_flagged == "inject_system_message" def test_in_memory_update_moving_off_during_call_in_the_same_update_is_allowed(self): """Bugbot finding on BerriAI/litellm#34940: validation checked the live, diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_qualifire.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_qualifire.py index fd72185d1e7..cee0006db7c 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_qualifire.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_qualifire.py @@ -102,6 +102,40 @@ class TestQualifireGuardrailInit: assert guardrail.qualifire_api_base == "https://custom.qualifire.ai" + def test_on_flagged_defaults_to_block(self): + from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( + QualifireGuardrail, + ) + + guardrail = QualifireGuardrail(api_key="test_key", guardrail_name="test_guardrail") + assert guardrail.on_flagged == "block" + + def test_on_flagged_monitor_is_accepted(self): + from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( + QualifireGuardrail, + ) + + guardrail = QualifireGuardrail(api_key="test_key", guardrail_name="test_guardrail", on_flagged="monitor") + assert guardrail.on_flagged == "monitor" + + def test_on_flagged_inject_system_message_raises_at_construction(self): + """ + Maintainer finding on BerriAI/litellm#34940: on_flagged is defined on + LakeraV2GuardrailConfigModel, but LitellmParams flattens every guardrail + config mixin together, so 'inject_system_message' type-checks for any + guardrail's config, including Qualifire, which never implements it. + Silently accepting it would let an admin believe advisory mode is active + when Qualifire actually just blocks on any unrecognized value. + """ + from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( + QualifireGuardrail, + ) + + with pytest.raises(ValueError, match="does not support on_flagged"): + QualifireGuardrail( + api_key="test_key", guardrail_name="test_guardrail", on_flagged="inject_system_message" + ) + class TestQualifireGuardrailMessageConversion: """Tests for message conversion to API format.""" diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py index 45f5afef1bc..f77f948a595 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -1157,13 +1157,15 @@ async def test_update_guardrail_endpoint( "scenario,expected_result,expected_exception", [ ("success_with_sync", "test-db-guardrail", None), - ("success_sync_fails", "test-db-guardrail", None), + ("success_sync_fails_unexpected_error", "test-db-guardrail", None), + ("sync_fails_invalid_config", None, HTTPException), ("database_failure", None, HTTPException), ("no_prisma_client", None, HTTPException), ], ids=[ "success_with_immediate_sync", - "success_but_sync_fails", + "success_but_sync_fails_with_unexpected_error", + "sync_rejects_invalid_config", "database_error", "missing_prisma_client", ], @@ -1194,7 +1196,10 @@ async def test_patch_guardrail_endpoint( mock_in_memory_handler, ) - elif scenario == "success_sync_fails": + elif scenario == "success_sync_fails_unexpected_error": + # A non-ValueError/TypeError failure (e.g. a transient bug) is not a + # config-rejection signal, so it keeps the pre-existing swallow-and-warn + # behavior rather than rolling back the DB write. mock_prisma_client = mocker.Mock() mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock( side_effect=Exception("Sync failed") @@ -1213,6 +1218,25 @@ async def test_patch_guardrail_endpoint( mock_in_memory_handler, ) + elif scenario == "sync_fails_invalid_config": + # Maintainer finding on BerriAI/litellm#34940: a ValueError from + # sync_guardrail_from_db (e.g. an invalid on_flagged combination) must + # roll back the DB write and surface a 422, not persist the rejected + # config with a 200. + mock_prisma_client = mocker.Mock() + mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock( + side_effect=ValueError("on_flagged='inject_system_message' requires payload=True and breakdown=True") + ) + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mocker.patch( + "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", + mock_guardrail_registry, + ) + mocker.patch( + "litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", + mock_in_memory_handler, + ) + elif scenario == "database_failure": mock_prisma_client = mocker.Mock() mock_guardrail_registry.update_guardrail_in_db.side_effect = Exception( @@ -1241,6 +1265,12 @@ async def test_patch_guardrail_endpoint( assert "Database error" in str(exc_info.value.detail) elif scenario == "no_prisma_client": assert "Prisma client not initialized" in str(exc_info.value.detail) + elif scenario == "sync_fails_invalid_config": + assert exc_info.value.status_code == 422 + assert "update rejected" in str(exc_info.value.detail) + # Rolled back: update_guardrail_in_db is called once for the + # rejected write and once more to restore the previous config. + assert mock_guardrail_registry.update_guardrail_in_db.call_count == 2 else: result = await patch_guardrail( @@ -1256,7 +1286,7 @@ async def test_patch_guardrail_endpoint( guardrail=mocker.ANY ) - if scenario == "success_sync_fails": + if scenario == "success_sync_fails_unexpected_error": assert mock_logger is not None mock_logger.warning.assert_called_once() assert "Failed to update" in str(mock_logger.warning.call_args) diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py index 5ffbcdedf0b..2b1dee32a2c 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py @@ -553,6 +553,67 @@ def test_reinitialized_judge_guardrail_uses_lazy_router_provider(): cb_list[:] = snapshot +def _lakera_guardrail(guardrail_id: str, **litellm_params_overrides) -> Guardrail: + params = {"guardrail": "lakera_v2", "mode": "pre_call", "on_flagged": "block", **litellm_params_overrides} + return Guardrail( + guardrail_id=guardrail_id, + guardrail_name="lakera-test", + litellm_params=LitellmParams(**params), + ) + + +class TestReinitializeGuardrailRestoresOnFailure: + """Maintainer finding on BerriAI/litellm#34940: reinitialize_guardrail deletes + the old in-memory instance and its callback registration before attempting to + construct the new one. initialize_guardrail's own ValueError/TypeError + propagate uncaught, so a rejected hot-reload (e.g. PATCH /guardrails/{id} + with an invalid on_flagged combination) previously left the guardrail + deleted entirely, not merely "still enforcing the old config", while the + DB/API kept reporting the new config as live.""" + + def test_invalid_update_restores_previous_instance(self): + handler = InMemoryGuardrailHandler() + lists = _all_callback_lists() + snapshots = [list(cb_list) for cb_list in lists] + try: + handler.reinitialize_guardrail(_lakera_guardrail("lakera-restore", on_flagged="block"), source="db") + + with pytest.raises(ValueError, match="requires payload=True and breakdown=True"): + handler.reinitialize_guardrail( + _lakera_guardrail("lakera-restore", on_flagged="inject_system_message", payload=False), + source="db", + ) + + assert "lakera-restore" in handler.IN_MEMORY_GUARDRAILS, "a rejected update must not delete the guardrail" + restored_instance = handler.guardrail_id_to_custom_guardrail["lakera-restore"] + assert restored_instance.on_flagged == "block" + finally: + for cb_list, snapshot in zip(lists, snapshots): + cb_list[:] = snapshot + + def test_invalid_update_leaves_dict_metadata_matching_the_restored_instance(self): + """IN_MEMORY_GUARDRAILS's own dict entry (what /guardrails/list-style + reads would see) must reflect the restored config too, not the + rejected one -- otherwise admin-facing reads and the live callback + instance disagree about what's actually configured.""" + handler = InMemoryGuardrailHandler() + lists = _all_callback_lists() + snapshots = [list(cb_list) for cb_list in lists] + try: + handler.reinitialize_guardrail(_lakera_guardrail("lakera-restore-meta", on_flagged="block"), source="db") + + with pytest.raises(ValueError): + handler.reinitialize_guardrail( + _lakera_guardrail("lakera-restore-meta", on_flagged="inject_system_message", breakdown=False), + source="db", + ) + + assert handler.IN_MEMORY_GUARDRAILS["lakera-restore-meta"]["litellm_params"].on_flagged == "block" + finally: + for cb_list, snapshot in zip(lists, snapshots): + cb_list[:] = snapshot + + class TestScanOnlyToolResultsInitRefusal: """A guardrail whose role filtering never scans tool results must be rejected at initialization when configured with scan_only_tool_results, instead of booting a