From 05efbc64542dbba2a85be4e4475dd208281684fd Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Tue, 16 Jun 2026 11:17:03 -0700 Subject: [PATCH 1/4] fix(guardrails): run pre_call hook once for model-level guardrails (#30543) * fix(guardrails): run pre_call hook once for model-level guardrails A CustomGuardrail attached to a deployment via litellm_params.guardrails gets its async_pre_call_hook invoked twice per request: once by the proxy pre-call loop and again by async_pre_call_deployment_hook after the router spreads the model-level guardrails into the top-level request kwargs. Record in request metadata that the proxy pre-call loop already ran a given guardrail, and have the deployment hook skip it when the marker is present. Direct-SDK usage never runs the proxy loop, so the deployment hook stays the sole invocation there and still fires exactly once. The marker key is stripped from untrusted caller metadata so a request body cannot suppress a model-only guardrail by pre-seeding it. * fix(guardrails): mark pre_call dedup on the post-hook request data Record the exactly-once marker after async_pre_call_hook runs, on the data object that flows downstream, rather than before it. A guardrail whose hook returns a brand-new request dict (instead of mutating or spreading the one it received) would otherwise discard the marker, letting the deployment hook re-run the guardrail a second time. (cherry picked from commit 4faeabc2541912bfec38c2d755cbad0ae3394671) --- litellm/constants.py | 4 + litellm/integrations/custom_guardrail.py | 54 +++++++ litellm/proxy/common_utils/callback_utils.py | 2 + litellm/proxy/litellm_pre_call_utils.py | 2 + .../proxy/policy_engine/pipeline_executor.py | 4 + litellm/proxy/utils.py | 2 + .../integrations/test_custom_guardrail.py | 106 ++++++++++++ .../proxy/test_model_level_guardrails.py | 152 +++++++++++++++++- 8 files changed, 325 insertions(+), 1 deletion(-) diff --git a/litellm/constants.py b/litellm/constants.py index c08cb3b60d6..33adc7a67e4 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -190,6 +190,10 @@ DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH_LITE = int( # Override with LITELLM_MAX_CALLBACKS env var for large deployments (e.g., many teams with guardrails) MAX_CALLBACKS = get_env_int("LITELLM_MAX_CALLBACKS", 100) +# Metadata key recording which pre_call guardrails the proxy loop already ran, +# so the deployment-level hook does not re-run them for the same request +PRE_CALL_EXECUTED_GUARDRAILS_KEY = "_pre_call_executed_guardrails" + # Generic fallback for unknown models DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET = int( os.getenv("DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET", 128) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 82a35f2eedd..586831a078f 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -1,3 +1,4 @@ +import secrets from datetime import datetime from typing import ( TYPE_CHECKING, @@ -43,12 +44,19 @@ if TYPE_CHECKING: dc = DualCache() +from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY from litellm.exceptions import ( BlockedPiiEntityError, GuardrailRaisedException, ModifyResponseException, ) +# Per-process secret tagging each recorded marker. The deployment hook only +# honors markers carrying this token, so a caller cannot forge the metadata +# field to suppress a guardrail on the direct-SDK path that never reaches the +# proxy's metadata sanitizer. +_PRE_CALL_EXECUTED_TOKEN = secrets.token_hex(16) + class CustomGuardrail(CustomLogger): # If True, during_call runs async_moderation_hook instead of the unified apply_guardrail path. @@ -325,6 +333,49 @@ class CustomGuardrail(CustomLogger): return False + def _pre_call_marker(self) -> Optional[str]: + name = self.guardrail_name + if not name: + return None + return f"{_PRE_CALL_EXECUTED_TOKEN}:{name}" + + def mark_pre_call_hook_ran(self, data: Dict[str, Any]) -> None: + """ + Record that this guardrail's ``async_pre_call_hook`` already ran for this + request, so the deployment-level hook does not run it a second time. + + The proxy runs pre-call guardrails in ``ProxyLogging.pre_call_hook``. The + router later spreads a deployment's model-level ``guardrails`` into the + top-level request kwargs, which would otherwise re-trigger the same hook + from ``async_pre_call_deployment_hook``. + """ + marker = self._pre_call_marker() + if marker is None: + return + for meta_key in ("metadata", "litellm_metadata"): + meta = data.get(meta_key) + if isinstance(meta, dict): + executed = meta.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY) + if isinstance(executed, list): + if marker not in executed: + executed.append(marker) + else: + meta[PRE_CALL_EXECUTED_GUARDRAILS_KEY] = [marker] + return + data["metadata"] = {PRE_CALL_EXECUTED_GUARDRAILS_KEY: [marker]} + + def _pre_call_hook_already_ran(self, data: Dict[str, Any]) -> bool: + marker = self._pre_call_marker() + if marker is None: + return False + for meta_key in ("metadata", "litellm_metadata"): + meta = data.get(meta_key) + if isinstance(meta, dict): + executed = meta.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY) + if isinstance(executed, list) and marker in executed: + return True + return False + async def async_pre_call_deployment_hook( self, kwargs: Dict[str, Any], call_type: Optional[CallTypes] ) -> Optional[dict]: @@ -335,6 +386,9 @@ class CustomGuardrail(CustomLogger): if litellm_guardrails is None or not isinstance(litellm_guardrails, list): return kwargs + if self._pre_call_hook_already_ran(kwargs): + return kwargs + if ( self.should_run_guardrail( data=kwargs, event_type=GuardrailEventHooks.pre_call diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index a65e737f248..ffd2ac68c54 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Literal, import litellm from litellm import get_secret from litellm._logging import verbose_proxy_logger +from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker from litellm.proxy._types import CommonProxyErrors, LiteLLMPromptInjectionParams @@ -490,6 +491,7 @@ LITELLM_PROXY_INTERNAL_METADATA_KEYS = frozenset( "guardrail_config", "_guardrail_pipelines", "_pipeline_managed_guardrails", + PRE_CALL_EXECUTED_GUARDRAILS_KEY, "disable_global_guardrails", "disable_global_guardrail", "opted_out_global_guardrails", diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 7666b23f2af..22657f6250c 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -13,6 +13,7 @@ from starlette.datastructures import Headers import litellm from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._service_logger import ServiceLogging +from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.litellm_core_utils.url_utils import is_url_destination_allowed_by_host @@ -161,6 +162,7 @@ _UNTRUSTED_METADATA_CONTROL_FIELDS = ( "secret_fields", "_guardrail_pipelines", "_pipeline_managed_guardrails", + PRE_CALL_EXECUTED_GUARDRAILS_KEY, ) _UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS = frozenset( diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 3c5a1d67be4..e46e3e1dc9f 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -171,6 +171,10 @@ class PipelineExecutor: data=data, call_type=call_type, # type: ignore ) + if isinstance(callback, CustomGuardrail): + callback.mark_pre_call_hook_ran(data) + if isinstance(response, dict): + callback.mark_pre_call_hook_ran(response) elif mode == "post_call": response = await target.async_post_call_success_hook( user_api_key_dict=user_api_key_dict, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 0720389e1f2..e7ad5a9155f 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1150,6 +1150,8 @@ class ProxyLogging: response=response, data=data, call_type=call_type ) + callback.mark_pre_call_hook_ran(data) + except Exception as e: status = "error" error_type = type(e).__name__ diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index a881044dc18..09469cd33e5 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -84,6 +84,112 @@ class TestCustomGuardrailDeploymentHook: assert result["messages"] == mock_result["messages"] assert result["messages"] != original_messages + @pytest.mark.asyncio + async def test_deployment_hook_skips_when_pre_call_already_ran(self): + """The deployment hook must not re-run async_pre_call_hook once the proxy + pre-call loop has already run it for this request.""" + + class CountingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__(guardrail_name="g1", default_on=True) + self.pre_call_count = 0 + + async def async_pre_call_hook( + self, user_api_key_dict, cache, data, call_type + ): + self.pre_call_count += 1 + return data + + guardrail = CountingGuardrail() + kwargs = { + "messages": [{"role": "user", "content": "hi"}], + "model": "gpt-3.5-turbo", + "guardrails": ["g1"], + "metadata": {}, + } + + guardrail.mark_pre_call_hook_ran(kwargs) + await guardrail.async_pre_call_deployment_hook( + kwargs=kwargs, call_type=CallTypes.completion + ) + + assert guardrail.pre_call_count == 0 + + @pytest.mark.asyncio + async def test_deployment_hook_runs_when_not_marked(self): + """Without the proxy marker (direct-SDK usage) the deployment hook is the + only execution path and must still run the guardrail exactly once.""" + + class CountingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__(guardrail_name="g1", default_on=True) + self.pre_call_count = 0 + + async def async_pre_call_hook( + self, user_api_key_dict, cache, data, call_type + ): + self.pre_call_count += 1 + return data + + guardrail = CountingGuardrail() + kwargs = { + "messages": [{"role": "user", "content": "hi"}], + "model": "gpt-3.5-turbo", + "guardrails": ["g1"], + "metadata": {}, + } + + await guardrail.async_pre_call_deployment_hook( + kwargs=kwargs, call_type=CallTypes.completion + ) + + assert guardrail.pre_call_count == 1 + + def test_mark_pre_call_hook_ran_uses_litellm_metadata(self): + """The marker is recorded in litellm_metadata when that is the metadata + bucket in use, and is then visible to the skip check.""" + from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY + + guardrail = CustomGuardrail(guardrail_name="g1") + kwargs = {"litellm_metadata": {}} + + guardrail.mark_pre_call_hook_ran(kwargs) + + assert kwargs["litellm_metadata"][PRE_CALL_EXECUTED_GUARDRAILS_KEY] + assert guardrail._pre_call_hook_already_ran(kwargs) is True + + @pytest.mark.asyncio + async def test_deployment_hook_ignores_forged_caller_marker(self): + """A direct-SDK caller controls request metadata but cannot know the + per-process token, so a hand-crafted marker must not suppress a + requested guardrail in async_pre_call_deployment_hook.""" + from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY + + class CountingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__(guardrail_name="g1", default_on=True) + self.pre_call_count = 0 + + async def async_pre_call_hook( + self, user_api_key_dict, cache, data, call_type + ): + self.pre_call_count += 1 + return data + + guardrail = CountingGuardrail() + kwargs = { + "messages": [{"role": "user", "content": "hi"}], + "model": "gpt-3.5-turbo", + "guardrails": ["g1"], + "metadata": {PRE_CALL_EXECUTED_GUARDRAILS_KEY: ["g1"]}, + } + + await guardrail.async_pre_call_deployment_hook( + kwargs=kwargs, call_type=CallTypes.completion + ) + + assert guardrail.pre_call_count == 1 + class TestCustomGuardrailShouldRunGuardrail: diff --git a/tests/test_litellm/proxy/test_model_level_guardrails.py b/tests/test_litellm/proxy/test_model_level_guardrails.py index 3d74edd772b..9a79fa7f496 100644 --- a/tests/test_litellm/proxy/test_model_level_guardrails.py +++ b/tests/test_litellm/proxy/test_model_level_guardrails.py @@ -19,7 +19,6 @@ from litellm.proxy.utils import ( _merge_guardrails_with_existing, ) - # --------------------------------------------------------------------------- # Unit tests for _check_and_merge_model_level_guardrails # --------------------------------------------------------------------------- @@ -159,6 +158,157 @@ class TestCheckAndMergeModelLevelGuardrails: assert "existing" in result["metadata"]["guardrails"] +# --------------------------------------------------------------------------- +# Regression test: pre_call hook must run exactly once with model-level guardrails +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_pre_call_hook_runs_once_with_model_level_guardrails(): + """ + A guardrail attached at the model level (litellm_params.guardrails) is + spread into the top-level request kwargs by the router. The proxy pre-call + loop (async_pre_call_hook) and the deployment-level hook + (async_pre_call_deployment_hook) must together invoke async_pre_call_hook + exactly once, not twice. + """ + from litellm.caching.caching import DualCache + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.proxy._types import CallTypes, UserAPIKeyAuth + from litellm.proxy.utils import ProxyLogging + from litellm.types.guardrails import GuardrailEventHooks + + class CountingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="counting-guardrail", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + ) + self.pre_call_count = 0 + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + self.pre_call_count += 1 + return data + + guardrail = CountingGuardrail() + + with patch("litellm.callbacks", [guardrail]): + ProxyLogging._callback_capabilities_cache.clear() + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + user_api_key_dict = UserAPIKeyAuth(api_key="test-key") + + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "hello"}], + "metadata": {}, + } + + # Path A: proxy pre-call loop runs the guardrail and records that it ran + data = await proxy_logging.pre_call_hook( + user_api_key_dict=user_api_key_dict, + data=data, + call_type="acompletion", + ) + + # Path B: the router spreads the deployment's model-level guardrails into + # the top-level kwargs, then litellm.acompletion fires the deployment hook + data["guardrails"] = ["counting-guardrail"] + await guardrail.async_pre_call_deployment_hook(data, CallTypes.acompletion) + + assert guardrail.pre_call_count == 1 + + +@pytest.mark.asyncio +async def test_pre_call_hook_runs_once_when_hook_returns_fresh_dict(): + """ + async_pre_call_hook may return a brand-new request dict instead of mutating + or spreading the one it received. The exactly-once marker must live on the + data that flows downstream, so the deployment hook still skips the guardrail + even when the proxy loop swapped in a fresh dict that never carried it. + """ + from litellm.caching.caching import DualCache + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.proxy._types import CallTypes, UserAPIKeyAuth + from litellm.proxy.utils import ProxyLogging + from litellm.types.guardrails import GuardrailEventHooks + + class FreshDictGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="counting-guardrail", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + ) + self.pre_call_count = 0 + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + self.pre_call_count += 1 + return {"model": data["model"], "messages": data["messages"]} + + guardrail = FreshDictGuardrail() + + with patch("litellm.callbacks", [guardrail]): + ProxyLogging._callback_capabilities_cache.clear() + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + user_api_key_dict = UserAPIKeyAuth(api_key="test-key") + + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "hello"}], + "metadata": {}, + } + + data = await proxy_logging.pre_call_hook( + user_api_key_dict=user_api_key_dict, + data=data, + call_type="acompletion", + ) + + data["guardrails"] = ["counting-guardrail"] + await guardrail.async_pre_call_deployment_hook(data, CallTypes.acompletion) + + assert guardrail.pre_call_count == 1 + + +@pytest.mark.asyncio +async def test_deployment_hook_runs_pre_call_without_proxy_loop(): + """ + Direct-SDK usage (litellm.acompletion(..., guardrails=[...]) without the + proxy) never runs the proxy pre-call loop, so the deployment hook is the + only place the guardrail executes and it must still run exactly once. + """ + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.proxy._types import CallTypes + from litellm.types.guardrails import GuardrailEventHooks + + class CountingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="counting-guardrail", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + ) + self.pre_call_count = 0 + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + self.pre_call_count += 1 + return data + + guardrail = CountingGuardrail() + + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "hello"}], + "guardrails": ["counting-guardrail"], + "metadata": {}, + } + + await guardrail.async_pre_call_deployment_hook(data, CallTypes.acompletion) + + assert guardrail.pre_call_count == 1 + + # --------------------------------------------------------------------------- # Integration test: post_call_success_hook with model-level guardrails # --------------------------------------------------------------------------- From a68db8b5bbfb26b7b3156ea6f6bd26f8aff2afed Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Tue, 16 Jun 2026 11:17:49 -0700 Subject: [PATCH 2/4] fix(guardrails): stop re-initializing DB guardrails on every poll (#30542) * fix(guardrails): stop re-initializing DB guardrails on every poll InMemoryGuardrailHandler._has_guardrail_params_changed compared the in-memory LitellmParams against the raw dict loaded from the DB. The in-memory side carries every field default and coerces enums via model_dump(), while the DB side only holds the keys originally stored, so the two shapes never compared equal and the guardrail was rebuilt on every poll cycle. Each rebuild created a fresh instance, but delete_in_memory_guardrail only removed the old callback from litellm.callbacks. Request handling promotes guardrail callbacks into the success/failure/async lists, so the previous instance stayed referenced there and instances accumulated. Normalize both sides through LitellmParams(...).model_dump() before diffing, and purge the callback from every callback list on delete. * refactor(guardrails): narrow params-normalization fallback to ValidationError The comparison normalizer caught a bare Exception and silently fell back to the raw dict, which hid the cause and quietly degraded the affected guardrail back to re-initializing on every poll. Catch only the ValidationError that LitellmParams construction can raise, log a warning so the offending row is diagnosable, and let any other error surface instead of being swallowed. * refactor(callbacks): add remove_callback_from_all_lists helper to manager Move the knowledge of which callback lists a callback can be promoted into out of the guardrail registry and into LoggingCallbackManager, where the rest of the callback-list bookkeeping already lives. delete_in_memory_guardrail now delegates to the new helper instead of iterating the lists itself. (cherry picked from commit 9fa74ad8b4a3cf206c847d92d20c1bd20daa2b69) --- .../logging_callback_manager.py | 16 ++ .../proxy/guardrails/guardrail_registry.py | 64 +++++-- .../test_logging_callback_manager.py | 23 +++ .../guardrails/test_guardrail_registry.py | 169 ++++++++++++++++++ 4 files changed, 253 insertions(+), 19 deletions(-) diff --git a/litellm/litellm_core_utils/logging_callback_manager.py b/litellm/litellm_core_utils/logging_callback_manager.py index 6c749118dec..b7adda3a9a4 100644 --- a/litellm/litellm_core_utils/logging_callback_manager.py +++ b/litellm/litellm_core_utils/logging_callback_manager.py @@ -394,6 +394,22 @@ class LoggingCallbackManager: + litellm._async_failure_callback ) + def remove_callback_from_all_lists(self, obj, require_self=False) -> None: + """ + Remove a callback object from every callback list it may have been + promoted into, so a re-initialized callback leaves no stale instance behind. + """ + for callback_list in ( + litellm.callbacks, + litellm.success_callback, + litellm.failure_callback, + litellm._async_success_callback, + litellm._async_failure_callback, + ): + self.remove_callback_from_list_by_object( + callback_list, obj, require_self=require_self + ) + def get_active_additional_logging_utils_from_custom_logger( self, ) -> Set[AdditionalLoggingUtils]: diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index aafcc5f1819..d589217a7ea 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -5,6 +5,8 @@ import os from datetime import datetime, timezone from typing import Any, Dict, List, Literal, Optional, Set, Type, cast +from pydantic import ValidationError + import litellm from litellm import Router from litellm._logging import verbose_proxy_logger @@ -598,21 +600,25 @@ class InMemoryGuardrailHandler: def delete_in_memory_guardrail(self, guardrail_id: str) -> None: """ Delete a guardrail in memory and remove from litellm callbacks. + + The callback is purged from every callback list, not just + litellm.callbacks: request handling promotes guardrail callbacks into the + success/failure/async lists, so removing it from only litellm.callbacks + leaves the old instance stranded in those lists on every re-initialization. """ # Remove from in-memory storage self.IN_MEMORY_GUARDRAILS.pop(guardrail_id, None) self._sources.pop(guardrail_id, None) - # Remove the callback from litellm.callbacks custom_guardrail_callback = self.guardrail_id_to_custom_guardrail.pop( guardrail_id, None ) - if custom_guardrail_callback: - litellm.logging_callback_manager.remove_callback_from_list_by_object( - callback_list=litellm.callbacks, - obj=custom_guardrail_callback, - require_self=False, - ) + if custom_guardrail_callback is None: + return + + litellm.logging_callback_manager.remove_callback_from_all_lists( + custom_guardrail_callback + ) def list_in_memory_guardrails(self) -> List[Guardrail]: """ @@ -654,6 +660,34 @@ class InMemoryGuardrailHandler: self.delete_in_memory_guardrail(guardrail_id) return stale_ids + @staticmethod + def _normalize_litellm_params_for_comparison( + params: Optional[Any], + ) -> Optional[Dict[str, Any]]: + """ + Render litellm_params to a canonical dict so an in-memory LitellmParams and + the raw dict loaded from the DB compare equal when they describe the same + config. The in-memory side is a LitellmParams whose model_dump() carries + every field default and coerces enums, while the DB side is the raw stored + dict holding only the keys originally provided. Comparing those two shapes + directly never matches, so each DB poll would re-initialize the guardrail + forever; normalizing both through LitellmParams keeps the diff meaningful. + """ + if params is None: + return None + if isinstance(params, LitellmParams): + return params.model_dump() + if isinstance(params, dict): + try: + return LitellmParams(**params).model_dump() + except ValidationError as e: + verbose_proxy_logger.warning( + f"Could not normalize guardrail litellm_params for comparison; " + f"treating the guardrail as changed. Error: {e}" + ) + return params + return params + def _has_guardrail_params_changed( self, guardrail_id: str, new_guardrail: Guardrail ) -> bool: @@ -670,19 +704,11 @@ class InMemoryGuardrailHandler: return True # Compare litellm_params - existing_params = existing.get("litellm_params") - new_params = new_guardrail.get("litellm_params") - - # Convert to dicts for comparison - existing_dict = ( - existing_params.model_dump() - if isinstance(existing_params, LitellmParams) - else existing_params + existing_dict = self._normalize_litellm_params_for_comparison( + existing.get("litellm_params") ) - new_dict = ( - new_params.model_dump() - if isinstance(new_params, LitellmParams) - else new_params + new_dict = self._normalize_litellm_params_for_comparison( + new_guardrail.get("litellm_params") ) # Compare and identify specific differences diff --git a/tests/litellm_utils_tests/test_logging_callback_manager.py b/tests/litellm_utils_tests/test_logging_callback_manager.py index d9540f8f850..d9bfca425e4 100644 --- a/tests/litellm_utils_tests/test_logging_callback_manager.py +++ b/tests/litellm_utils_tests/test_logging_callback_manager.py @@ -192,6 +192,29 @@ def test_remove_callback_from_list_by_object(): assert len(litellm._async_failure_callback) == 0 +def test_remove_callback_from_all_lists(): + manager = LoggingCallbackManager() + manager._reset_all_callbacks() + + class TestLogger(CustomLogger): + pass + + obj = TestLogger() + manager.add_litellm_callback(obj) + manager.add_litellm_success_callback(obj) + manager.add_litellm_failure_callback(obj) + manager.add_litellm_async_success_callback(obj) + manager.add_litellm_async_failure_callback(obj) + + manager.remove_callback_from_all_lists(obj) + + assert obj not in litellm.callbacks + assert obj not in litellm.success_callback + assert obj not in litellm.failure_callback + assert obj not in litellm._async_success_callback + assert obj not in litellm._async_failure_callback + + def test_reset_callbacks(callback_manager): # Add various callbacks callback_manager.add_litellm_callback("test") diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py index 9f7173383b0..0ef9ad857f9 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py @@ -180,3 +180,172 @@ def test_sync_guardrail_from_db_marks_source_db_when_unchanged(): handler.sync_guardrail_from_db(g) assert handler.get_source("collide") == "db" + + +def _db_litellm_params() -> dict: + """ + Shape produced by GuardrailRegistry.get_all_guardrails_from_db: litellm_params + is a raw dict (not a LitellmParams), holding only the keys originally stored, + a non-schema extra key, and plain-string enum values. + """ + return { + "guardrail": "litellm_content_filter", + "mode": "pre_call", + "default_on": True, + "version": 2, + "blocked_words": [{"keyword": "secret", "action": "BLOCK"}], + } + + +def test_unchanged_db_params_do_not_register_as_changed(): + """ + A DB poll returns litellm_params as a raw dict while the in-memory copy is a + LitellmParams whose model_dump() fills every field default and coerces enums. + The two shapes must compare equal when the config is identical; otherwise + every poll cycle re-initializes the guardrail indefinitely. + """ + handler = InMemoryGuardrailHandler() + raw = _db_litellm_params() + gid = "11111111-1111-1111-1111-111111111111" + handler.IN_MEMORY_GUARDRAILS[gid] = Guardrail( + guardrail_id=gid, + guardrail_name="cf", + litellm_params=LitellmParams(**raw), + ) + + new = Guardrail(guardrail_id=gid, guardrail_name="cf", litellm_params=dict(raw)) + assert handler._has_guardrail_params_changed(gid, new) is False + + +def test_changed_db_params_register_as_changed(): + """Normalizing both sides must still surface a genuine config change.""" + handler = InMemoryGuardrailHandler() + raw = _db_litellm_params() + gid = "22222222-2222-2222-2222-222222222222" + handler.IN_MEMORY_GUARDRAILS[gid] = Guardrail( + guardrail_id=gid, + guardrail_name="cf", + litellm_params=LitellmParams(**raw), + ) + + changed = {**raw, "blocked_words": [{"keyword": "different", "action": "BLOCK"}]} + new = Guardrail(guardrail_id=gid, guardrail_name="cf", litellm_params=changed) + assert handler._has_guardrail_params_changed(gid, new) is True + + +def test_unnormalizable_db_params_register_as_changed_without_raising(): + """ + A DB row whose litellm_params fail LitellmParams validation must not crash the + poll loop. The comparison falls back to treating the guardrail as changed so it + re-initializes (and surfaces the bad row in logs) rather than propagating the + validation error up through the polling cycle. + """ + handler = InMemoryGuardrailHandler() + raw = _db_litellm_params() + gid = "55555555-5555-5555-5555-555555555555" + handler.IN_MEMORY_GUARDRAILS[gid] = Guardrail( + guardrail_id=gid, + guardrail_name="cf", + litellm_params=LitellmParams(**raw), + ) + + malformed = {**raw, "default_on": "not-a-bool-xyz"} + new = Guardrail(guardrail_id=gid, guardrail_name="cf", litellm_params=malformed) + assert handler._has_guardrail_params_changed(gid, new) is True + + +def _all_callback_lists(): + import litellm + + return [ + litellm.callbacks, + litellm.success_callback, + litellm.failure_callback, + litellm._async_success_callback, + litellm._async_failure_callback, + ] + + +def test_delete_in_memory_guardrail_removes_callback_from_all_lists(): + """ + Request handling promotes guardrail callbacks from litellm.callbacks into the + success/failure/async lists. delete_in_memory_guardrail must purge the callback + from every list, otherwise a re-initialized guardrail leaves its old instance + stranded in those lists and instances accumulate. + """ + handler = InMemoryGuardrailHandler() + callback = CustomGuardrail( + guardrail_name="cf-delete", + default_on=True, + event_hook=GuardrailEventHooks.pre_call, + ) + gid = "33333333-3333-3333-3333-333333333333" + handler.IN_MEMORY_GUARDRAILS[gid] = _make_guardrail(gid, "cf-delete") + handler._sources[gid] = "db" + handler.guardrail_id_to_custom_guardrail[gid] = callback + + lists = _all_callback_lists() + snapshots = [list(cb_list) for cb_list in lists] + try: + for cb_list in lists: + cb_list.append(callback) + + handler.delete_in_memory_guardrail(gid) + + for cb_list in lists: + assert callback not in cb_list + finally: + for cb_list, snapshot in zip(lists, snapshots): + cb_list[:] = snapshot + + +def test_repeated_db_sync_does_not_accumulate_runner_instances(): + """ + End-to-end regression for the OOM: across repeated DB polls (with the config + genuinely changing each cycle to force re-initialization), exactly one live + guardrail instance must exist across all callback lists. On the unfixed code + the stale instance lingers in the success/failure lists and the distinct count + climbs above one. + """ + import litellm + + handler = InMemoryGuardrailHandler() + gid = "44444444-4444-4444-4444-444444444444" + name = "cf-accum" + + def db_guardrail(word: str) -> Guardrail: + params = { + **_db_litellm_params(), + "blocked_words": [{"keyword": word, "action": "BLOCK"}], + } + return Guardrail(guardrail_id=gid, guardrail_name=name, litellm_params=params) + + def promote_into_request_lists() -> None: + manager = litellm.logging_callback_manager + for callback in list(litellm.callbacks): + manager.add_litellm_success_callback(callback) + manager.add_litellm_failure_callback(callback) + manager.add_litellm_async_success_callback(callback) + manager.add_litellm_async_failure_callback(callback) + + def distinct_runner_instances() -> int: + seen = set() + for callback in litellm.logging_callback_manager._get_all_callbacks(): + if ( + isinstance(callback, CustomGuardrail) + and getattr(callback, "guardrail_name", None) == name + ): + seen.add(id(callback)) + return len(seen) + + lists = _all_callback_lists() + snapshots = [list(cb_list) for cb_list in lists] + try: + for cycle in range(5): + handler.sync_guardrail_from_db(db_guardrail(f"word-{cycle}")) + promote_into_request_lists() + + assert distinct_runner_instances() == 1 + finally: + for cb_list, snapshot in zip(lists, snapshots): + cb_list[:] = snapshot From 072d757b543af7605559dccfbb7b5fbc1b38d37c Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 17 Jun 2026 12:21:47 -0700 Subject: [PATCH 3/4] =?UTF-8?q?bump:=20version=201.88.2=20=E2=86=92=201.88?= =?UTF-8?q?.3?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- pyproject.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index bd7e002b717..164327e0152 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm" -version = "1.88.2" +version = "1.88.3" description = "Library to easily interface with LLM API providers" readme = "README.md" requires-python = ">=3.10, <3.14" @@ -264,7 +264,7 @@ source-exclude = [ profile = "black" [tool.commitizen] -version = "1.88.2" +version = "1.88.3" version_files = [ "pyproject.toml:^version", ] From 7678078be27733ed5043f75e29b57411e0f448ea Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 17 Jun 2026 12:22:23 -0700 Subject: [PATCH 4/4] chore: refresh uv.lock for 1.88.3 --- uv.lock | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/uv.lock b/uv.lock index e254ead54e1..0e6579a7e54 100644 --- a/uv.lock +++ b/uv.lock @@ -9,7 +9,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-06-11T00:01:45.852753Z" +exclude-newer = "2026-06-14T19:22:16.739045Z" exclude-newer-span = "P3D" [manifest] @@ -3280,7 +3280,7 @@ wheels = [ [[package]] name = "litellm" -version = "1.88.2" +version = "1.88.3" source = { editable = "." } dependencies = [ { name = "aiohttp" },