Skip to content

Commit f48d219

Browse files
fix(guardrails): run policy pipelines when the caller sends its own metadata (/v1/messages, Claude Code) (#36889)
* fix(guardrails): resolve guardrail pipelines from the canonical metadata bucket Policy-resolved pipelines are stored in litellm_metadata on routes like /v1/messages, but the pre_call reader fell back to the caller-supplied metadata field first, so a request that sends its own top-level metadata (Claude Code sends metadata.user_id) skipped every pipeline-managed guardrail. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(guardrails): drive the pipeline regression through a registered guardrail Exercise the real executor with a guardrail in litellm.callbacks instead of patching PipelineExecutor.execute_steps at class scope. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(guardrails): read pipeline state from the bucket the policy engine wrote Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(guardrails): type the policy pipeline state accessors Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(guardrails): annotate policy pipeline state casts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
1 parent ed02a12 commit f48d219

2 files changed

Lines changed: 87 additions & 5 deletions

File tree

litellm/proxy/utils.py

Lines changed: 44 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111
import threading
1212
import time
1313
import traceback
14-
from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine, Mapping, Sequence
14+
from collections.abc import AsyncGenerator, Awaitable, Callable, Collection, Coroutine, Mapping, Sequence
1515
from dataclasses import dataclass, field
1616
from datetime import date, datetime, timedelta, timezone
1717
from email.mime.multipart import MIMEMultipart
@@ -186,6 +186,7 @@
186186
from litellm.models.team import LiteLLM_TeamTableCachedObj
187187
from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction
188188
from litellm.proxy.db.spend_log_tool_index import ToolUsageTransaction
189+
from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline
189190

190191
Span = _Span | object
191192
else:
@@ -408,6 +409,46 @@ def _exception_changes_request_flow(exc: BaseException) -> bool:
408409
return isinstance(exc, (SensitiveDataRouteException, ModifyResponseException))
409410

410411

412+
def _policy_state_metadata(data: Mapping[str, object]) -> Mapping[str, object]:
413+
"""
414+
Return the metadata bucket the policy engine wrote its pipeline state into.
415+
416+
The route decides the bucket (``litellm_metadata`` for ``/v1/messages``,
417+
responses, batches, files and bedrock, ``metadata`` everywhere else), and both
418+
buckets can be present at once because callers send their own provider-facing
419+
``metadata`` (Claude Code sends ``metadata.user_id``) or their own
420+
``litellm_metadata``. Pipeline slots are stripped from caller input before the
421+
policy engine runs, so whichever bucket carries them is the proxy's own write.
422+
"""
423+
return next(
424+
(
425+
bucket
426+
for bucket in (data.get("metadata"), data.get("litellm_metadata"))
427+
if isinstance(bucket, dict)
428+
and ("_guardrail_pipelines" in bucket or "_pipeline_managed_guardrails" in bucket)
429+
),
430+
{},
431+
)
432+
433+
434+
def _policy_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, "GuardrailPipeline"], ...]:
435+
pipelines: Final = _policy_state_metadata(data).get("_guardrail_pipelines")
436+
return (
437+
tuple(cast("Sequence[tuple[str, GuardrailPipeline]]", pipelines)) # cast-ok: the policy engine wrote the slot
438+
if pipelines
439+
else ()
440+
)
441+
442+
443+
def _pipeline_managed_guardrail_names(data: Mapping[str, object]) -> frozenset[str]:
444+
managed: Final = _policy_state_metadata(data).get("_pipeline_managed_guardrails")
445+
return (
446+
frozenset(cast("Collection[str]", managed)) # cast-ok: the policy engine wrote these guardrail names
447+
if managed
448+
else frozenset()
449+
)
450+
451+
411452
def _prompt_block_text(block: object) -> str:
412453
if isinstance(block, str):
413454
return block
@@ -1446,8 +1487,7 @@ async def _maybe_execute_pipelines(
14461487
14471488
Returns the (possibly modified) data dict.
14481489
"""
1449-
metadata: Final = data.get("metadata", data.get("litellm_metadata", {})) or {}
1450-
pipelines: Final = metadata.get("_guardrail_pipelines")
1490+
pipelines: Final = _policy_pipelines(data)
14511491
if not pipelines:
14521492
return data
14531493

@@ -1631,8 +1671,7 @@ async def pre_call_hook(
16311671
)
16321672

16331673
# Get pipeline-managed guardrails to skip in normal loop
1634-
metadata: Final = data.get("metadata", data.get("litellm_metadata", {})) or {}
1635-
pipeline_managed: Final[set] = metadata.get("_pipeline_managed_guardrails", set())
1674+
pipeline_managed: Final = _pipeline_managed_guardrail_names(data)
16361675

16371676
caps: Final = ProxyLogging._callback_capabilities()
16381677
# Skip the per-request callback walk entirely when nothing in

tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,10 @@
2525
from litellm.integrations.prometheus import PrometheusLogger
2626
from litellm.proxy.utils import ProxyLogging
2727
from litellm.types.guardrails import GuardrailEventHooks
28+
from litellm.types.proxy.policy_engine.pipeline_types import (
29+
GuardrailPipeline,
30+
PipelineStep,
31+
)
2832

2933

3034
@pytest.fixture(autouse=True)
@@ -350,6 +354,45 @@ async def test_maybe_execute_pipelines_skips_pipelines_with_other_mode(proxy_log
350354
assert out is data
351355

352356

357+
@pytest.mark.parametrize(
358+
("policy_state_key", "caller_metadata_key", "call_type"),
359+
[
360+
("litellm_metadata", "metadata", "anthropic_messages"),
361+
("metadata", "litellm_metadata", "completion"),
362+
],
363+
)
364+
@pytest.mark.asyncio
365+
async def test_maybe_execute_pipelines_finds_policy_state_when_caller_sends_own_metadata(
366+
proxy_logging, make_user_api_key_auth, monkeypatch, policy_state_key, caller_metadata_key, call_type
367+
):
368+
"""The route picks the bucket the policy engine writes to (``litellm_metadata`` on
369+
/v1/messages, ``metadata`` on chat completions), and the caller can populate the other
370+
one, e.g. Claude Code sending ``metadata.user_id``. The pipeline must still run and block."""
371+
372+
class BlockingGuardrail(CustomGuardrail):
373+
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
374+
raise HTTPException(status_code=400, detail={"error": "blocked by pipeline"})
375+
376+
monkeypatch.setattr(litellm, "callbacks", [BlockingGuardrail(guardrail_name="gr-1")])
377+
pipeline = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="gr-1", on_fail="block")])
378+
data = {
379+
caller_metadata_key: {"user_id": "user_abc"},
380+
policy_state_key: {"_guardrail_pipelines": [("policy-1", pipeline)]},
381+
"messages": [],
382+
"model": "m",
383+
}
384+
385+
with pytest.raises(HTTPException) as exc_info:
386+
await proxy_logging._maybe_execute_pipelines(
387+
data=data,
388+
user_api_key_dict=make_user_api_key_auth(),
389+
call_type=call_type,
390+
event_hook="pre_call",
391+
)
392+
assert exc_info.value.detail["error"] == "blocked by pipeline"
393+
assert exc_info.value.detail["guardrail_name"] == "gr-1"
394+
395+
353396
@pytest.mark.asyncio
354397
async def test_maybe_execute_pipelines_blocks_on_block_terminal_action_raises(
355398
proxy_logging, make_user_api_key_auth, monkeypatch

0 commit comments

Comments
 (0)