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.
This commit is contained in:
Yassin Kortam 2026-06-16 11:17:03 -07:00 committed by GitHub
parent bed6ce820c
commit 4faeabc254
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 325 additions and 1 deletions

View file

@ -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)

View file

@ -1,3 +1,4 @@
import secrets
from datetime import datetime
from typing import (
TYPE_CHECKING,
@ -43,6 +44,7 @@ if TYPE_CHECKING:
dc = DualCache()
from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY
from litellm.exceptions import (
BlockedPiiEntityError,
GuardrailRaisedException,
@ -50,6 +52,12 @@ from litellm.exceptions import (
SensitiveDataRouteException,
)
# 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)
def get_session_id_from_request_data(request_data: Dict[str, Any]) -> Optional[str]:
"""Extract session_id from request data (litellm_session_id or metadata)."""
@ -458,6 +466,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]:
@ -468,6 +519,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

View file

@ -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
@ -497,6 +498,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",

View file

@ -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(

View file

@ -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,

View file

@ -1171,6 +1171,8 @@ class ProxyLogging:
response=response, data=data, call_type=call_type
)
callback.mark_pre_call_hook_ran(data)
except SensitiveDataRouteException:
status = "intervened"
raise

View file

@ -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:

View file

@ -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
# ---------------------------------------------------------------------------