From 92b46538e1f83455a038f5ac047e6ca522db8479 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 16:38:44 -0700 Subject: [PATCH] fix(policy_engine): restore request guardrails list after pipeline allow --- .../proxy/policy_engine/pipeline_executor.py | 52 +++++++-- .../policy_engine/test_pipeline_executor.py | 100 ++++++++++++++++++ 2 files changed, 142 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 50cb813c6fa..4be0f556ed7 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -6,6 +6,7 @@ pass/fail actions (allow, block, next, modify_response) and data forwarding. """ import time +from collections.abc import Sequence from typing import Any, Final, Literal import litellm @@ -114,11 +115,7 @@ class PipelineExecutor: # Handle terminal actions if action == "allow": - return PipelineExecutionResult( - terminal_action="allow", - step_results=step_results, - modified_data=working_data if working_data != data else None, - ) + return _allow_result(step_results=step_results, working_data=working_data, request_data=data) if action == "block": return PipelineExecutionResult( @@ -138,11 +135,7 @@ class PipelineExecutor: # action == "next" → continue to next step # Ran out of steps without a terminal action → default allow - return PipelineExecutionResult( - terminal_action="allow", - step_results=step_results, - modified_data=working_data if working_data != data else None, - ) + return _allow_result(step_results=step_results, working_data=working_data, request_data=data) @staticmethod async def _run_step( @@ -251,6 +244,45 @@ class PipelineExecutor: return None +def _allow_result( + step_results: Sequence[PipelineStepResult], + working_data: dict, # mutable-ok: same request-payload shape as execute_steps' data + request_data: dict, # mutable-ok: same request-payload shape as execute_steps' data +) -> PipelineExecutionResult: + """Build the terminal-allow result, propagating pipeline modifications without the per-step guardrail override.""" + restored: Final = _restore_request_guardrails(working_data, request_data) + return PipelineExecutionResult( + terminal_action="allow", + step_results=list(step_results), # mutable-ok: PipelineExecutionResult field is a list + modified_data=restored if restored != request_data else None, + ) + + +def _restore_request_guardrails( + working_data: dict, # mutable-ok: same request-payload shape as execute_steps' data + request_data: dict, # mutable-ok: same request-payload shape as execute_steps' data +) -> dict: # mutable-ok: merged back into the request dict, which downstream code mutates + """ + Restore the request's own metadata["guardrails"] activation list. + + _run_step overrides it to [step.guardrail] so should_run_guardrail() allows each + step; letting that override escape via modified_data permanently drops every + independently activated guardrail from later lifecycle stages (post_call, etc.). + """ + working_metadata: Final = working_data.get("metadata") + if not isinstance(working_metadata, dict): + return working_data + request_metadata: Final = request_data.get("metadata") + original_guardrails: Final = request_metadata.get("guardrails") if isinstance(request_metadata, dict) else None + stripped: Final = {k: v for k, v in working_metadata.items() if k != "guardrails"} # mutable-ok: request dict + if original_guardrails is not None: + restored: Final = {**stripped, "guardrails": original_guardrails} # mutable-ok: request dict + return {**working_data, "metadata": restored} # mutable-ok: request dict + if not stripped and not isinstance(request_metadata, dict): + return {k: v for k, v in working_data.items() if k != "metadata"} # mutable-ok: request dict + return {**working_data, "metadata": stripped} # mutable-ok: request dict + + def _pipeline_action_for_outcome(step: PipelineStep, outcome: str) -> str: """ Map pipeline step outcome to the configured action. diff --git a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py index 054a5af4148..4fcb7d22588 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -749,6 +749,106 @@ async def test_single_step_pipeline_allow(monkeypatch): assert guard.calls == 1 +@pytest.mark.asyncio +async def test_allow_restores_independent_guardrails_list(monkeypatch): + """ + Request activates an independent guardrail; an unrelated pipeline runs and allows. + Expected: no modified_data escapes, so the request's guardrails list survives + and the independent guardrail still runs at later lifecycle stages (post_call). + Regression: LIT-6587 (pipeline clobbered the list with its last step's guardrail). + """ + pipeline_guard = AlwaysPassGuardrail(guardrail_name="input-scan") + + pipeline = GuardrailPipeline( + mode="pre_call", + steps=[PipelineStep(guardrail="input-scan", on_fail="block", on_pass="allow")], + ) + + monkeypatch.setattr(litellm, "callbacks", [pipeline_guard]) + + data = { + "messages": [{"role": "user", "content": "clean content"}], + "metadata": {"guardrails": ["independent-output-guard"]}, + } + result = await PipelineExecutor.execute_steps( + steps=pipeline.steps, + mode=pipeline.mode, + data=data, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="input-pipeline-policy", + ) + + assert pipeline_guard.calls == 1 + assert result.terminal_action == "allow" + propagated = result.modified_data or data + assert propagated["metadata"]["guardrails"] == ["independent-output-guard"] + assert data["metadata"]["guardrails"] == ["independent-output-guard"] + + +@pytest.mark.asyncio +async def test_allow_does_not_leak_guardrails_into_bare_request(monkeypatch): + """A request without metadata must not gain a metadata.guardrails list from the pipeline.""" + pipeline_guard = AlwaysPassGuardrail(guardrail_name="input-scan") + + pipeline = GuardrailPipeline( + mode="pre_call", + steps=[PipelineStep(guardrail="input-scan", on_fail="block", on_pass="allow")], + ) + + monkeypatch.setattr(litellm, "callbacks", [pipeline_guard]) + + data = {"messages": [{"role": "user", "content": "clean content"}]} + result = await PipelineExecutor.execute_steps( + steps=pipeline.steps, + mode=pipeline.mode, + data=data, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="input-pipeline-policy", + ) + + assert result.terminal_action == "allow" + propagated = result.modified_data or data + assert "guardrails" not in propagated.get("metadata", {}) + assert "metadata" not in data + + +@pytest.mark.asyncio +async def test_data_forwarding_keeps_changes_and_restores_guardrails_list(monkeypatch): + """A pass_data pipeline's modifications propagate while the request's guardrails list is restored.""" + pii_guard = PiiMaskingGuardrail(guardrail_name="pii-masker") + content_guard = ContentCheckGuardrail(guardrail_name="content-check") + + pipeline = GuardrailPipeline( + mode="pre_call", + steps=[ + PipelineStep(guardrail="pii-masker", on_fail="block", on_pass="next", pass_data=True), + PipelineStep(guardrail="content-check", on_fail="block", on_pass="allow"), + ], + ) + + monkeypatch.setattr(litellm, "callbacks", [pii_guard, content_guard]) + + data = { + "messages": [{"role": "user", "content": "Hello John Smith"}], + "metadata": {"guardrails": ["independent-output-guard"]}, + } + result = await PipelineExecutor.execute_steps( + steps=pipeline.steps, + mode=pipeline.mode, + data=data, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="pii-then-safety", + ) + + assert result.terminal_action == "allow" + assert result.modified_data is not None + assert result.modified_data["messages"][0]["content"] == "Hello [REDACTED]" + assert result.modified_data["metadata"]["guardrails"] == ["independent-output-guard"] + + @pytest.mark.asyncio async def test_step_results_include_duration(monkeypatch): """Step results should include timing information."""