fix(policy_engine): keep a policy-added guardrail's other stages when a pipeline steps it

A policy that both adds a guardrail and steps it in a post_call pipeline
used to drop the guardrail from the request's guardrail list outright, so
its pre_call stage never ran. The per-hook loops already skip guardrails
by pipeline mode, so the mode-agnostic subtraction only lost coverage
This commit is contained in:
mateo-berri 2026-09-07 21:51:34 -07:00
parent 91fc1b2010
commit d08a177bc7
2 changed files with 44 additions and 5 deletions

View file

@ -3079,10 +3079,9 @@ def _apply_resolved_guardrails_to_metadata(
if metadata_variable_name not in data:
data[metadata_variable_name] = {}
# Track pipeline-managed guardrails to exclude from independent execution
pipeline_managed_guardrails: set = set()
# Record the pipelines and the guardrails they step; the hook loops skip those per pipeline mode
if pipelines:
pipeline_managed_guardrails = PolicyResolver.get_pipeline_managed_guardrails(pipelines)
pipeline_managed_guardrails: Final = PolicyResolver.get_pipeline_managed_guardrails(pipelines)
data[metadata_variable_name]["_guardrail_pipelines"] = pipelines
data[metadata_variable_name]["_pipeline_managed_guardrails"] = pipeline_managed_guardrails
verbose_proxy_logger.debug(
@ -3099,10 +3098,8 @@ def _apply_resolved_guardrails_to_metadata(
existing_guardrails = []
# Combine existing guardrails with policy-resolved guardrails (no duplicates)
# Exclude pipeline-managed guardrails from the flat list
combined = set(existing_guardrails)
combined.update(resolved_guardrails)
combined -= pipeline_managed_guardrails
data[metadata_variable_name]["guardrails"] = list(combined)
verbose_proxy_logger.debug("Policy engine: added guardrails to request metadata: %s", list(combined))

View file

@ -4148,6 +4148,48 @@ async def test_add_guardrails_from_policy_engine():
attachment_registry._initialized = False
@pytest.mark.asyncio
async def test_add_guardrails_from_policy_engine_keeps_a_policy_added_guardrail_its_pipeline_also_steps():
from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
from litellm.types.proxy.policy_engine import (
GuardrailPipeline,
PipelineStep,
Policy,
PolicyAttachment,
PolicyGuardrails,
)
data = {"model": "gpt-4", "messages": [{"role": "user", "content": "Hello"}], "metadata": {}}
policy_registry = get_policy_registry()
policy_registry._policies = {
"response-governance": Policy(
guardrails=PolicyGuardrails(add=["pii_blocker"]),
pipeline=GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="pii_blocker")]),
),
}
policy_registry._initialized = True
attachment_registry = get_attachment_registry()
attachment_registry._attachments = [PolicyAttachment(policy="response-governance", scope="*")]
attachment_registry._initialized = True
try:
await add_guardrails_from_policy_engine(
data=data,
metadata_variable_name="metadata",
user_api_key_dict=UserAPIKeyAuth(api_key="test-key"),
)
finally:
policy_registry._policies = {}
policy_registry._initialized = False
attachment_registry._attachments = []
attachment_registry._initialized = False
assert data["metadata"]["guardrails"] == ["pii_blocker"]
assert data["metadata"]["_pipeline_managed_guardrails"] == {"pii_blocker"}
assert [pipeline.mode for _policy_name, pipeline in data["metadata"]["_guardrail_pipelines"]] == ["post_call"]
@pytest.mark.asyncio
async def test_add_guardrails_from_policy_engine_accepts_dynamic_policies_and_pops_from_data():
"""