Skip to content

fix(guardrails): run pre_call hook once for model-level guardrails (#30543) - #31144

Closed
mubashir1osmani wants to merge 1 commit into
BerriAI:litellm_cherry-pick-pr-29946-v1.89.0-rc.2from
mubashir1osmani:litellm_cherry-pick-pr-30543-v1.89.0-rc.2
Closed

fix(guardrails): run pre_call hook once for model-level guardrails (#30543)#31144
mubashir1osmani wants to merge 1 commit into
BerriAI:litellm_cherry-pick-pr-29946-v1.89.0-rc.2from
mubashir1osmani:litellm_cherry-pick-pr-30543-v1.89.0-rc.2

Conversation

@mubashir1osmani

Copy link
Copy Markdown
Collaborator

Relevant issues

Backports #30543 onto the litellm_cherry-pick-pr-29946-v1.89.0-rc.2 realtime line to close the second half of the Bedrock-guardrail high-CPU regression seen on the v1.89.0-dev.x builds. The branch tip already carries #30542 (stop re-initializing DB guardrails on every poll); this adds the companion fix it was missing

Linear ticket

Pre-Submission checklist

  • I have added meaningful tests
  • My PR passes all unit tests
  • My PR's scope is as isolated as possible; it only solves 1 specific problem
  • I have requested a Greptile review

Screenshots / Proof of Fix

This is a faithful cherry-pick of the already-merged #30543 (3ab500c79e, the same content shipped in v1.89.3); the diff is byte-identical to upstream, so the behavioral verification lives in the original PR. The two regression tests that ship with the fix pass unchanged on this branch:

pytest tests/test_litellm/integrations/test_custom_guardrail.py \
       tests/test_litellm/proxy/test_model_level_guardrails.py
59 passed in 3.41s

If you want a live-proxy reproduction before merge (model-level Bedrock guardrail, count ApplyGuardrail calls per request before vs after), say the word and I will stand one up against real Bedrock

Type

Bug Fix

Changes

A CustomGuardrail attached to a deployment via litellm_params.guardrails had 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. For a Bedrock guardrail that means two full ApplyGuardrail round trips per request, each with its own SigV4 signing and response redaction, which is pure wasted CPU on the hot path and doubles the cost of every guarded call

The fix records in request metadata that the proxy pre-call loop already ran a given guardrail, and has 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, and it is recorded on the post-hook request data so a guardrail that returns a brand-new request dict does not drop the marker and re-trigger the second run

This is the model-level / chat-completions counterpart to the callback-accumulation fix already on this branch; together they remove both sources of the duplicated Bedrock guardrail work behind the high CPU

…erriAI#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 4faeabc)
@CLAassistant

Copy link
Copy Markdown

CLA assistant check
Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
You have signed the CLA already but the status is still pending? Let us recheck it.

@codecov

codecov Bot commented Jun 23, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 97.67442% with 1 line in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
litellm/integrations/custom_guardrail.py 97.14% 1 Missing ⚠️

📢 Thoughts on this report? Let us know!

@greptile-apps

greptile-apps Bot commented Jun 23, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR cherry-picks the already-merged #30543 onto the litellm_cherry-pick-pr-29946-v1.89.0-rc.2 branch to fix a double-invocation bug where a CustomGuardrail attached via litellm_params.guardrails had its async_pre_call_hook fired twice per request — once by the proxy pre-call loop and again by async_pre_call_deployment_hook after the router spread model-level guardrails into the top-level kwargs.

  • Introduces a per-process secret token and a metadata marker (PRE_CALL_EXECUTED_GUARDRAILS_KEY) written by the proxy pre-call loop; the deployment hook reads the marker and skips re-invocation when present. The marker is included in both _UNTRUSTED_METADATA_CONTROL_FIELDS (strips caller-injected forgeries) and LITELLM_PROXY_INTERNAL_METADATA_KEYS (hides it from downstream/logs).
  • Handles the fresh-dict return case by stamping the marker on both the original data dict and the returned response dict, ensuring the marker survives even when a guardrail replaces the entire request object.
  • Direct-SDK usage (no proxy loop) is unaffected: the deployment hook remains the sole invocation path and still fires exactly once.

Confidence Score: 5/5

Safe to merge — the change is a faithful cherry-pick of an already-merged and tested fix with no new logic surface.

The deduplication logic is correct across all three execution paths (proxy loop, pipeline executor, direct SDK). The per-process token design prevents forge attacks from caller-supplied metadata. Both the _UNTRUSTED_METADATA_CONTROL_FIELDS strip and LITELLM_PROXY_INTERNAL_METADATA_KEYS exclusion are properly updated. The fresh-dict case is explicitly covered. The new tests are thorough, use in-memory fakes, and leave all existing tests untouched.

No files require special attention.

Important Files Changed

Filename Overview
litellm/constants.py Adds PRE_CALL_EXECUTED_GUARDRAILS_KEY to constants.py, correctly following the project convention for sentinel-like keys.
litellm/integrations/custom_guardrail.py Core of the fix: adds per-process _PRE_CALL_EXECUTED_TOKEN, mark_pre_call_hook_ran, and _pre_call_hook_already_ran methods to CustomGuardrail, plus an early-exit check in async_pre_call_deployment_hook. No FastAPI imports introduced.
litellm/proxy/utils.py Calls callback.mark_pre_call_hook_ran(data) after the guardrail runs and after process_pre_call_hook_response updates data, ensuring the marker lands on the correct dict including the fresh-dict case.
litellm/proxy/policy_engine/pipeline_executor.py Marks on both data (original dict) and response (fresh dict if returned) in the pre_call pipeline path, covering both the mutate-in-place and return-new-dict cases.
litellm/proxy/litellm_pre_call_utils.py Adds PRE_CALL_EXECUTED_GUARDRAILS_KEY to _UNTRUSTED_METADATA_CONTROL_FIELDS, preventing a caller from injecting a pre-crafted marker to suppress a guardrail.
litellm/proxy/common_utils/callback_utils.py Adds PRE_CALL_EXECUTED_GUARDRAILS_KEY to LITELLM_PROXY_INTERNAL_METADATA_KEYS, keeping internal metadata hidden from downstream and log outputs.
tests/test_litellm/integrations/test_custom_guardrail.py Adds four new unit tests: skip-when-marked, run-when-unmarked, litellm_metadata bucket, and forge-rejection. All use in-memory fakes; no real network calls. Existing tests untouched.
tests/test_litellm/proxy/test_model_level_guardrails.py Adds three integration-style regression tests exercising the full proxy pre-call loop to deployment hook path, the fresh-dict return case, and direct-SDK-only path. Uses patch and DualCache — no real network calls.

Reviews (1): Last reviewed commit: "fix(guardrails): run pre_call hook once ..." | Re-trigger Greptile

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants