From ade5a425e8bb3ab60858d35f2c473a0b1d830b93 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Mon, 10 Aug 2026 14:37:09 -0700 Subject: [PATCH] fix(proxy): isolate guardrail load failures per row (#36432) * fix(proxy): isolate guardrail load failures per row One DB guardrail row that fails to initialize aborted the whole _init_guardrails_in_db loop, so a single typo'd guardrail type or a missing required param left the proxy running with zero DB guardrails registered and requests that should have been blocked reaching the provider. Catch per row around sync_guardrail_from_db, log the guardrail name, id and error, and continue with the remaining rows. The failing row's id is still added to db_guardrail_ids before the attempt so reconcile_db_guardrails cannot mistake a live row for a deleted one. * test(proxy): drop inline note and record reconcile via a handler double Replaces the patched bound method with an InMemoryGuardrailHandler subclass that records what reconcile_db_guardrails received, so the test injects a double instead of swapping a method on a live object. --- litellm/proxy/proxy_server.py | 16 ++++- .../proxy/proxy_server/test_proxy_config.py | 64 +++++++++++++++++++ 2 files changed, 77 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3a4896dca9e..bc980934f9f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6860,9 +6860,19 @@ class ProxyConfig: guardrail_id = guardrail.get("guardrail_id") if guardrail_id: db_guardrail_ids.add(guardrail_id) - IN_MEMORY_GUARDRAIL_HANDLER.sync_guardrail_from_db( - guardrail=cast(Guardrail, guardrail), - ) + try: + IN_MEMORY_GUARDRAIL_HANDLER.sync_guardrail_from_db( + guardrail=cast(Guardrail, guardrail), + ) + except Exception as e: # noqa: BLE001 # one unloadable row must not stop the remaining guardrails + verbose_proxy_logger.error( + "litellm.proxy.proxy_server.py::ProxyConfig:_init_guardrails_in_db - " + "skipping guardrail '%s' (ID: %s): %s: %s", + guardrail.get("guardrail_name"), + guardrail_id, + type(e).__name__, + e, + ) # Drop in-memory DB-backed entries whose row was deleted on another # pod. Config-loaded entries are never touched. diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 91a7e1bc2c2..f70be17eb95 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -2639,3 +2639,67 @@ async def test_ProxyConfig__init_non_llm_configs_empty_agents_key_clears_remembe assert clean_agent_registry.config_agents == () clean_agent_registry.load_agents_from_db_and_config(db_agents=None) assert clean_agent_registry.get_agent_list() == () + + +# --------------------------------------------------------------------------- +# _init_guardrails_in_db +# --------------------------------------------------------------------------- + + +def _db_guardrail_row(guardrail_id: str, guardrail_type: str) -> dict[str, object]: + return { + "guardrail_id": guardrail_id, + "guardrail_name": f"name-{guardrail_id}", + "litellm_params": {"guardrail": guardrail_type, "mode": "pre_call"}, + "guardrail_info": None, + "team_id": None, + } + + +@pytest.mark.asyncio +async def test_ProxyConfig__init_guardrails_in_db_skips_only_the_unloadable_row(monkeypatch): + """ + A single DB row that fails to initialize used to abort the whole loop, so one + typo'd guardrail type left the proxy running with zero guardrails loaded. + + The failing row's id must still reach reconcile_db_guardrails so that eviction + pass cannot treat a row that is alive in the DB as one that was deleted. + """ + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.proxy.guardrails import guardrail_registry as registry_module + from litellm.types.guardrails import Guardrail, GuardrailEventHooks, LitellmParams + + class _RecordingHandler(registry_module.InMemoryGuardrailHandler): + def __init__(self) -> None: + super().__init__() + self.reconciled_with: list[set[str]] = [] + + def reconcile_db_guardrails(self, db_guardrail_ids: set[str]) -> list[str]: + self.reconciled_with.append(set(db_guardrail_ids)) + return super().reconcile_db_guardrails(db_guardrail_ids) + + handler = _RecordingHandler() + monkeypatch.setattr(registry_module, "IN_MEMORY_GUARDRAIL_HANDLER", handler) + + def _initializer(litellm_params: LitellmParams, guardrail: Guardrail) -> CustomGuardrail: + return CustomGuardrail( + guardrail_name=guardrail["guardrail_name"], + event_hook=GuardrailEventHooks.pre_call, + default_on=False, + ) + + monkeypatch.setitem(registry_module.guardrail_initializer_registry, "lit5367_ok", _initializer) + + prisma_client = MagicMock() + prisma_client.db.litellm_guardrailstable.find_many = AsyncMock( + return_value=[ + _db_guardrail_row("first", "lit5367_ok"), + _db_guardrail_row("broken", "litellm_tool_permission"), + _db_guardrail_row("last", "lit5367_ok"), + ] + ) + + await ProxyConfig()._init_guardrails_in_db(prisma_client=prisma_client) + + assert sorted(handler.IN_MEMORY_GUARDRAILS) == ["first", "last"] + assert handler.reconciled_with == [{"first", "broken", "last"}]