From 73d034b12e1e6820a2eb46ba559c025cfd6ef353 Mon Sep 17 00:00:00 2001 From: Yucheng He Date: Sun, 4 Oct 2026 01:26:17 -0700 Subject: [PATCH] fix(guardrails): keep no-op PUT and logging_only scan-failure semantics stable Three regressions from the logging_only_scope feature, fixed while keeping the input/output/both selection working: 1. PUT with byte-identical litellm_params no longer forces a teardown + re-init of the live callback. reject_invalid_logging_only_scope now only validates; re-init still happens exactly when params/name change. Previously a description-only PUT re-appended the callback at the END of litellm.callbacks, reordering guardrails: with a BLOCK guardrail created before a MASK one, the mask started winning and blocked requests started succeeding (400 -> PUT 200 -> 200). Invalid unchanged scopes are still rejected with 422 without touching the live instance. 2. CustomGuardrail._scan_logged_call no longer swallows input-scan exceptions per branch: a raising or BLOCKED input scan aborts the logging_only hook again (one policy call, one verdict) for guardrails that never selected a logging_only_scope. Explicit input/output/both selection keeps selecting which scans run. 3. A failed PATCH now rolls the DB row back to exactly the stored raw litellm_params (a legacy 4-key row stays 4 keys) instead of expanding it to a full LitellmParams dump; pinned with a test. This matches the PUT rollback shape and is the faithful rollback. --- litellm/integrations/custom_guardrail.py | 14 +-- .../proxy/guardrails/guardrail_registry.py | 31 ++++- .../integrations/test_custom_guardrail.py | 29 ++++- .../guardrails/test_guardrail_registry.py | 116 ++++++++++++++++++ 4 files changed, 173 insertions(+), 17 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 8b09c91f6e7..b34823b31cb 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -1033,18 +1033,12 @@ class CustomGuardrail(CustomLogger): "metadata": scratch_metadata, } if self.logging_only_scope != "output": - try: - await translation.process_input_messages(data=scratch_request, guardrail_to_apply=self) - except Exception as e: - verbose_logger.warning("Guardrail %s: logging_only scan raised: %s", self.guardrail_name, e) + await translation.process_input_messages(data=scratch_request, guardrail_to_apply=self) if response is None or self.logging_only_scope == "input": return - try: - await output_translation.process_output_response( - response=copy.deepcopy(response), guardrail_to_apply=self, request_data=scratch_request - ) - except Exception as e: - verbose_logger.warning("Guardrail %s: logging_only scan raised: %s", self.guardrail_name, e) + await output_translation.process_output_response( + response=copy.deepcopy(response), guardrail_to_apply=self, request_data=scratch_request + ) def supports_scan_only_tool_results(self) -> bool: """Whether this guardrail can scan tool-result content. diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index b4e0b564514..06ede9c97fb 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -669,6 +669,28 @@ class InMemoryGuardrailHandler: siblings: Final = self.guardrail_id_to_sibling_callbacks.get(guardrail_id, ()) return (() if primary is None else (primary,)) + siblings + def _reject_invalid_logging_only_scope(self, guardrail_id: str, guardrail: Guardrail) -> None: + """ + Strictly validate logging_only_scope on a row whose params are otherwise + unchanged, without rebuilding the live callback. + + API write paths send the whole object, so an invalid scope must still be + rejected even when the write changed nothing else. But an unchanged row + must not force a teardown + re-append: initialize_guardrail appends the + rebuilt callback at the END of litellm.callbacks, so a no-op PUT would + reorder guardrails and change which one wins between a BLOCK and a MASK + guardrail over the same content. + """ + params: Final = guardrail.get("litellm_params") + if not isinstance(params, (dict, LitellmParams)): + return + litellm_params: Final = LitellmParams(**params) if isinstance(params, dict) else params + guardrail_name: Final = guardrail.get("guardrail_name", "Unknown") + for custom_guardrail_callback in self._tracked_callbacks(guardrail_id): + error: Final = _logging_only_scope_error(custom_guardrail_callback, guardrail_name, litellm_params) + if error is not None: + raise ValueError(error) + def initialize_custom_guardrail( self, guardrail: Guardrail, @@ -737,13 +759,15 @@ class InMemoryGuardrailHandler: previous instance and raises), anything else only refreshes the stored row """ updated_guardrail: Final = cast(Guardrail, {**guardrail, "guardrail_id": guardrail_id}) - if reject_invalid_logging_only_scope or self._has_guardrail_params_changed(guardrail_id, updated_guardrail): + if self._has_guardrail_params_changed(guardrail_id, updated_guardrail): self.reinitialize_guardrail( guardrail=updated_guardrail, source=source, reject_invalid_logging_only_scope=reject_invalid_logging_only_scope, ) return + if reject_invalid_logging_only_scope: + self._reject_invalid_logging_only_scope(guardrail_id, updated_guardrail) self.IN_MEMORY_GUARDRAILS[guardrail_id] = updated_guardrail self._sources[guardrail_id] = source @@ -967,7 +991,7 @@ class InMemoryGuardrailHandler: verbose_proxy_logger.error("Cannot sync guardrail without guardrail_id") return None - if reject_invalid_logging_only_scope or self._has_guardrail_params_changed(guardrail_id, guardrail): + if self._has_guardrail_params_changed(guardrail_id, guardrail): guardrail_name: Final = guardrail.get("guardrail_name", "Unknown") verbose_proxy_logger.info( "Guardrail '%s' (ID: %s) params changed, re-initializing...", guardrail_name, guardrail_id @@ -979,6 +1003,9 @@ class InMemoryGuardrailHandler: reject_invalid_logging_only_scope=reject_invalid_logging_only_scope, ) + if reject_invalid_logging_only_scope: + self._reject_invalid_logging_only_scope(guardrail_id, guardrail) + # Params unchanged but the entry is still DB-backed; make sure the # source marker reflects that even if it was previously set differently # (e.g. a config entry whose UUID later collided with a DB row). diff --git a/tests/unit/integrations/test_custom_guardrail.py b/tests/unit/integrations/test_custom_guardrail.py index 1dd73feb446..f91709df35f 100644 --- a/tests/unit/integrations/test_custom_guardrail.py +++ b/tests/unit/integrations/test_custom_guardrail.py @@ -2654,12 +2654,31 @@ class TestLoggingOnlyApplyGuardrail: out_kwargs, _ = await guardrail.async_logging_hook(kwargs, response, CallTypes.acompletion.value) - assert guardrail.calls == [ - ("request", ["flagged content"]), - ("response", ["general kenobi"]), - ] + # A raising input scan aborts the hook: the response scan must not run, + # and exactly one intervened verdict is recorded (base semantics for a + # guardrail that never selected a logging_only_scope). + assert guardrail.calls == [("request", ["flagged content"])] entries = out_kwargs["standard_logging_object"]["guardrail_information"] - assert [entry["guardrail_status"] for entry in entries] == ["guardrail_intervened", "guardrail_intervened"] + assert [e["guardrail_status"] for e in entries] == ["guardrail_intervened"] + + @pytest.mark.asyncio + async def test_input_scan_error_aborts_the_response_scan(self): + class _FailingObserver(_ApplyOnlyObserver): + @log_guardrail_information + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + self.calls.append((input_type, list(inputs.get("texts") or []))) + if input_type == "request": + raise RuntimeError("guardrail service unavailable") + return GenericGuardrailAPIInputs(texts=[]) + + guardrail = _FailingObserver() + kwargs, response = _logged_call([{"role": "user", "content": "hello there"}]) + + out_kwargs, _ = await guardrail.async_logging_hook(kwargs, response, CallTypes.acompletion.value) + + assert guardrail.calls == [("request", ["hello there"])] + entries = out_kwargs["standard_logging_object"]["guardrail_information"] + assert [e["guardrail_status"] for e in entries] == ["guardrail_failed_to_respond"] @pytest.mark.asyncio async def test_call_type_without_translation_is_skipped(self): diff --git a/tests/unit/proxy/guardrails/test_guardrail_registry.py b/tests/unit/proxy/guardrails/test_guardrail_registry.py index 1d7116ce64a..e2230076097 100644 --- a/tests/unit/proxy/guardrails/test_guardrail_registry.py +++ b/tests/unit/proxy/guardrails/test_guardrail_registry.py @@ -1,4 +1,5 @@ from collections.abc import Iterable +import json from typing import ClassVar, Final from unittest.mock import AsyncMock, MagicMock @@ -404,6 +405,86 @@ def test_sync_guardrail_from_db_marks_source_db_when_unchanged(): assert handler.get_source("collide") == "db" +def test_sync_guardrail_from_db_reject_flag_keeps_callback_order_on_noop_update(): + """ + The PUT endpoint syncs the whole object with reject_invalid_logging_only_scope=True. + That strictness must not force a teardown + re-append of an unchanged guardrail: + initialize_guardrail appends the rebuilt callback at the END of litellm.callbacks, + so a description-only PUT would reorder guardrails and change which one wins + between a BLOCK and a MASK guardrail over the same content. + """ + import litellm + + registry_module = _register_mode_following_initializer("mode_following_test") + lists = _all_callback_lists() + snapshots = [list(cb_list) for cb_list in lists] + sentinel: Final = CustomGuardrail( + guardrail_name="order-sentinel", + supported_event_hooks=[GuardrailEventHooks.pre_call], + event_hook=GuardrailEventHooks.pre_call, + ) + try: + handler = InMemoryGuardrailHandler() + handler.initialize_guardrail(guardrail=_mode_following_db_row("123", "pre_call"), source="db") + original = handler.guardrail_id_to_custom_guardrail["123"] + assert original is not None + # Park another callback after the guardrail so a re-append would be visible. + litellm.callbacks.append(sentinel) + index_before = litellm.callbacks.index(original) + + handler.sync_guardrail_from_db( + guardrail=_mode_following_db_row("123", "pre_call", "description-only edit"), + reject_invalid_logging_only_scope=True, + ) + + assert handler.guardrail_id_to_custom_guardrail["123"] is original + assert litellm.callbacks.index(original) == index_before + assert _live_instances_named("mode-following") == 1 + finally: + registry_module.guardrail_initializer_registry.pop("mode_following_test", None) + for cb_list, snapshot in zip(lists, snapshots): + cb_list[:] = snapshot + + +def test_sync_guardrail_from_db_reject_flag_still_rejects_invalid_unchanged_scope(): + """ + A PUT sends the whole object, so an unchanged row that already carries an + invalid logging_only_scope (tolerated at load) must still be rejected on the + strict sync path, without rebuilding the live callback. + """ + registry_module = _register_mode_following_initializer("mode_following_test") + lists = _all_callback_lists() + snapshots = [list(cb_list) for cb_list in lists] + row: Final = Guardrail( + guardrail_id="123", + guardrail_name="mode-following", + litellm_params={ + "guardrail": "mode_following_test", + "mode": "pre_call", + "default_on": True, + "logging_only_scope": "input", + }, + guardrail_info={}, + ) + try: + handler = InMemoryGuardrailHandler() + handler.initialize_guardrail(guardrail=row, source="db") + original = handler.guardrail_id_to_custom_guardrail["123"] + assert original is not None + assert original.logging_only_scope is None # tolerated at load + + with pytest.raises(ValueError, match="logging_only_scope is set"): + handler.sync_guardrail_from_db(guardrail=row, reject_invalid_logging_only_scope=True) + + # Rejected without touching the live instance. + assert handler.guardrail_id_to_custom_guardrail["123"] is original + assert _live_instances_named("mode-following") == 1 + finally: + registry_module.guardrail_initializer_registry.pop("mode_following_test", None) + for cb_list, snapshot in zip(lists, snapshots): + cb_list[:] = snapshot + + def _db_litellm_params() -> dict: """ Shape produced by GuardrailRegistry.get_all_guardrails_from_db: litellm_params @@ -1260,6 +1341,41 @@ async def test_update_guardrail_in_db_raises_when_row_missing(): ) +@pytest.mark.asyncio +async def test_update_guardrail_in_db_persists_raw_sparse_params_verbatim(): + """ + After a rejected PATCH, the endpoint rolls back by writing the stored row's + raw litellm_params through update_guardrail_in_db. A raw dict must be + persisted exactly as stored — a legacy 4-key row stays a 4-key row — instead + of being round-tripped through LitellmParams.model_dump(), which materializes + every field default and rewrites a row the admin never wrote. + """ + prisma_client = MagicMock() + prisma_client.db.litellm_guardrailstable.update = AsyncMock( + return_value={"guardrail_id": "legacy-row", "guardrail_name": "legacy-one"} + ) + legacy_params: Final = { + "guardrail": "litellm_content_filter", + "mode": "pre_call", + "guardrail_name": "legacy-one", + "blocked_words": [{"keyword": "x", "action": "BLOCK"}], + } + + await GuardrailRegistry().update_guardrail_in_db( + guardrail_id="legacy-row", + guardrail=Guardrail( + guardrail_id="legacy-row", + guardrail_name="legacy-one", + litellm_params=legacy_params, + guardrail_info={}, + ), + prisma_client=prisma_client, + ) + + persisted: Final = prisma_client.db.litellm_guardrailstable.update.call_args.kwargs["data"] + assert json.loads(persisted["litellm_params"]) == legacy_params + + def test_reinitialize_guardrail_restores_previous_on_failure(): """A reinitialization whose new params make the guardrail constructor raise must restore the previous instance instead of leaving the guardrail silently removed: