scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

fix(guardrails): apply PUT /guardrails/{id} to the serving worker immediately and reject invalid configs with 422

Browse source
This commit is contained in:
mateo-berri 2026-08-29 22:11:36 -07:00
parent d44d281d1d
commit eac5dc10f3
4 changed files with 158 additions and 60 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

20
litellm/proxy/guardrails/guardrail_endpoints.py
View file

@ -514,12 +514,26 @@ async def update_guardrail(
guardrail_name: Final = result.get("guardrail_name", "Unknown")
try:
IN_MEMORY_GUARDRAIL_HANDLER.update_in_memory_guardrail(
guardrail_id=guardrail_id, guardrail=cast(Guardrail, result)
)
IN_MEMORY_GUARDRAIL_HANDLER.sync_guardrail_from_db(guardrail=cast(Guardrail, result))
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 (a raising guardrail __init__):
# 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=existing_guardrail,
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",

37
litellm/proxy/guardrails/guardrail_registry.py
View file

@ -6,7 +6,7 @@ import os
from collections.abc import Callable, Iterator, Mapping
from datetime import datetime, timezone
from itertools import chain, count
from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, cast
from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol
from pydantic import ValidationError
@ -615,26 +615,6 @@ class InMemoryGuardrailHandler:
return _guardrail_callback
def update_in_memory_guardrail(
self,
guardrail_id: str,
guardrail: Guardrail,
source: Literal["db", "config"] = "db",
) -> None:
"""
Update a guardrail in memory
- updates the guardrail in memory
- updates the guardrail params in litellm.callback_manager
"""
self.IN_MEMORY_GUARDRAILS[guardrail_id] = guardrail
self._sources[guardrail_id] = source
custom_guardrail_callback: Final = self.guardrail_id_to_custom_guardrail.get(guardrail_id)
if custom_guardrail_callback:
updated_litellm_params: Final = cast(LitellmParams, guardrail.get("litellm_params", {}))
custom_guardrail_callback.update_in_memory_litellm_params(litellm_params=updated_litellm_params)
def delete_in_memory_guardrail(self, guardrail_id: str) -> None:
"""
Delete a guardrail in memory and remove from litellm callbacks.
@ -789,11 +769,12 @@ class InMemoryGuardrailHandler:
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."
combination or an invalid regex), the previous instance is restored
rather than left deleted, and the failure is re-raised as ValueError so
every init failure reaches callers as one exception type: 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:
@ -812,7 +793,7 @@ class InMemoryGuardrailHandler:
# that was enforcing must never fail open because an update was bad.
try:
return self.initialize_guardrail(guardrail=guardrail, config_file_path=config_file_path, source=source)
except Exception:
except Exception as init_error:
if previous_guardrail is not None:
verbose_proxy_logger.exception(
"Reinitializing guardrail %s with updated params failed; restoring the previous configuration",
@ -824,7 +805,7 @@ class InMemoryGuardrailHandler:
)
except Exception: # noqa: BLE001 # the original failure must propagate even if the restore breaks
verbose_proxy_logger.exception("Restoring previous guardrail %s also failed", guardrail_id)
raise
raise ValueError(f"Guardrail initialization failed: {init_error}") from init_error
def sync_guardrail_from_db(self, guardrail: Guardrail, config_file_path: str | None = None) -> Guardrail | None:
"""

53
tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py
View file

@ -104,7 +104,7 @@ def mock_in_memory_handler(mocker):
mock_handler.get_guardrail_by_id.return_value = MOCK_CONFIG_GUARDRAIL
mock_handler.get_source.return_value = "config"
mock_handler.initialize_guardrail = mocker.Mock()
mock_handler.update_in_memory_guardrail = mocker.Mock()
mock_handler.sync_guardrail_from_db = mocker.Mock()
mock_handler.delete_in_memory_guardrail = mocker.Mock()
mock_handler.reconcile_db_guardrails = mocker.Mock(return_value=[])
return mock_handler
@ -1045,13 +1045,15 @@ async def test_create_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",
],
@ -1071,6 +1073,7 @@ async def test_update_guardrail_endpoint(
mock_logger = None
if scenario == "success_with_sync":
mock_prisma_client = mocker.Mock()
mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock()
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY",
@ -1081,10 +1084,13 @@ async def test_update_guardrail_endpoint(
mock_in_memory_handler,
)
elif scenario == "success_sync_fails":
elif scenario == "success_sync_fails_unexpected_error":
# A non-ValueError/TypeError failure 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.update_in_memory_guardrail.side_effect = Exception(
"Sync failed"
mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock(
side_effect=Exception("Sync failed")
)
mock_logger = mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger"
@ -1100,6 +1106,25 @@ async def test_update_guardrail_endpoint(
mock_in_memory_handler,
)
elif scenario == "sync_fails_invalid_config":
# Regression for the PUT half of the fix: a TypeError from the sync (the
# deleted update_in_memory_guardrail raised exactly this on every PUT)
# 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=TypeError("vars() argument must have __dict__ attribute")
)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) # test-quality-ok: reused pattern
mocker.patch( # test-quality-ok: reused pattern
"litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY",
mock_guardrail_registry,
)
mocker.patch( # test-quality-ok: reused pattern
"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(
@ -1128,6 +1153,16 @@ async def test_update_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
assert (
mock_guardrail_registry.update_guardrail_in_db.call_args.kwargs["guardrail"]
== MOCK_DB_GUARDRAIL
)
else:
result = await update_guardrail(
@ -1143,11 +1178,11 @@ async def test_update_guardrail_endpoint(
prisma_client=mocker.ANY,
)
mock_in_memory_handler.update_in_memory_guardrail.assert_called_once_with(
guardrail_id="test-guardrail-id", guardrail=mocker.ANY
mock_in_memory_handler.sync_guardrail_from_db.assert_called_once_with(
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)

108
tests/test_litellm/proxy/guardrails/test_guardrail_registry.py
View file

@ -154,29 +154,51 @@ def test_duplicate_config_guardrail_names_get_distinct_stable_ids():
registry_module.guardrail_initializer_registry.pop("dup_name_test", None)
def test_update_in_memory_guardrail():
def test_sync_guardrail_from_db_applies_db_dict_params_to_live_instance():
"""
Regression for PUT /guardrails/{id}: the DB row arrives with litellm_params as
a plain jsonb dict, and the deleted update_in_memory_guardrail cast it to
LitellmParams without constructing one, so vars() raised and the running proxy
kept enforcing the stale config forever. The PUT endpoint now routes through
sync_guardrail_from_db, which must rebuild the live instance from the dict:
new blocked words compiled in, old ones gone, and the event hook re-derived
from mode (the base-class setattr path wrote self.mode while dispatch reads
self.event_hook, so only a full re-init applies a mode change).
"""
from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import (
ContentFilterGuardrail,
)
handler = InMemoryGuardrailHandler()
handler.guardrail_id_to_custom_guardrail["123"] = CustomGuardrail(
guardrail_name="test-guardrail",
default_on=False,
event_hook=GuardrailEventHooks.pre_call,
)
gid = "66666666-6666-6666-6666-666666666666"
handler.update_in_memory_guardrail(
"123",
Guardrail(
guardrail_name="test-guardrail",
litellm_params=LitellmParams(guardrail="test-guardrail", mode="pre_call", default_on=True),
),
)
assert (
handler.guardrail_id_to_custom_guardrail["123"].should_run_guardrail(
data={}, event_type=GuardrailEventHooks.pre_call
def db_guardrail(word: str, mode: str) -> Guardrail:
return Guardrail(
guardrail_id=gid,
guardrail_name="cf-put-sync",
litellm_params={
"guardrail": "litellm_content_filter",
"mode": mode,
"default_on": True,
"blocked_words": [{"keyword": word, "action": "BLOCK"}],
},
)
is True
)
assert handler.guardrail_id_to_custom_guardrail["123"].event_hook is GuardrailEventHooks.pre_call
lists = _all_callback_lists()
snapshots = [list(cb_list) for cb_list in lists]
try:
handler.sync_guardrail_from_db(db_guardrail("foobarblock", "pre_call"))
handler.sync_guardrail_from_db(db_guardrail("quxnewblock", "during_call"))
instance = handler.guardrail_id_to_custom_guardrail[gid]
assert isinstance(instance, ContentFilterGuardrail)
assert instance._check_blocked_words("hello QUXNEWBLOCK") is not None
assert instance._check_blocked_words("hello FOOBARBLOCK") is None
assert instance.event_hook == GuardrailEventHooks.during_call
assert instance.should_run_guardrail(data={}, event_type=GuardrailEventHooks.during_call) is True
finally:
for cb_list, snapshot in zip(lists, snapshots):
cb_list[:] = snapshot
def _make_guardrail(guardrail_id: str, name: str = "g") -> Guardrail:
@ -774,3 +796,49 @@ def test_reinitialize_guardrail_restores_previous_on_failure():
assert restored.guardrail_name == "restore-me"
finally:
registry_module.guardrail_initializer_registry.pop("restore_test", None)
def test_reinitialize_guardrail_raises_value_error_for_non_value_error_init_failures():
"""Regression for the LIT-6479 fix's 422 path: a constructor failure that is not
already a ValueError/TypeError (re.error from an invalid regex has neither in its
MRO) must still surface as ValueError, so the PUT/PATCH endpoints' rollback+422
catch is exhaustive instead of warn-and-200 persisting a broken config."""
import re
from litellm.proxy.guardrails import guardrail_registry as registry_module
def _initializer(litellm_params, guardrail):
if litellm_params.api_key == "bad-regex":
re.compile("([")
return CustomGuardrail(
guardrail_name=guardrail["guardrail_name"],
event_hook=GuardrailEventHooks.pre_call,
default_on=True,
)
registry_module.guardrail_initializer_registry["regex_test"] = _initializer
try:
handler = InMemoryGuardrailHandler()
created = handler.initialize_guardrail(
guardrail={
"guardrail_name": "regex-me",
"litellm_params": {"guardrail": "regex_test", "mode": "pre_call", "api_key": "ok"},
},
)
guardrail_id = created["guardrail_id"]
with pytest.raises(ValueError, match="Guardrail initialization failed") as excinfo:
handler.reinitialize_guardrail(
guardrail={
"guardrail_id": guardrail_id,
"guardrail_name": "regex-me",
"litellm_params": {"guardrail": "regex_test", "mode": "pre_call", "api_key": "bad-regex"},
},
)
assert isinstance(excinfo.value.__cause__, re.error)
assert guardrail_id in handler.IN_MEMORY_GUARDRAILS
restored = handler.guardrail_id_to_custom_guardrail[guardrail_id]
assert restored is not None and restored.guardrail_name == "regex-me"
finally:
registry_module.guardrail_initializer_registry.pop("regex_test", None)