fix(guardrails): match policy-pipeline block response to direct guardrail attachment (#31421)

When a guardrail blocked a request through a flow-builder policy pipeline, the
proxy discarded the guardrail's own exception and synthesized a generic
guardrail_pipeline_error response, so the same guardrail produced a different
HTTP response and trace span depending on whether it was attached directly or
via a policy. The pipeline now carries the guardrail's original exception and
re-raises it verbatim on block, enriching it with the blocking guardrail's name
and mode exactly as the direct path does, so the two attachment methods are
indistinguishable to clients and tracing. The generic pipeline error remains
only as a fallback for blocks with no underlying exception (e.g. a guardrail
that could not be found).

Resolves LIT-4041
This commit is contained in:
Yassin Kortam 2026-06-27 00:25:10 +03:00 • committed by GitHub
parent 5a395ef909
commit c14329128b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 251 additions and 13 deletions

View file

@ -6,7 +6,7 @@ pass/fail actions (allow, block, next, modify_response) and data forwarding.
"""
import time
from typing import Any, List, Optional
from typing import Any, List, Literal, Optional
import litellm
from litellm._logging import verbose_proxy_logger
@ -64,7 +64,12 @@ class PipelineExecutor:
for i, step in enumerate(steps):
start_time = time.perf_counter()
outcome, modified_data, error_detail = await PipelineExecutor._run_step(
(
outcome,
modified_data,
error_detail,
original_exception,
) = await PipelineExecutor._run_step(
step=step,
mode=mode,
data=working_data,
@ -108,6 +113,7 @@ class PipelineExecutor:
terminal_action="block",
step_results=step_results,
error_message=error_detail,
original_exception=original_exception,
)
if action == "modify_response":
@ -134,22 +140,30 @@ class PipelineExecutor:
data: dict,
user_api_key_dict: Any,
call_type: str,
) -> tuple:
) -> tuple[
Literal["pass", "fail", "error"],
Optional[dict],
Optional[str],
Optional[Exception],
]:
"""
Run a single pipeline step's guardrail.
Returns:
Tuple of (outcome, modified_data, error_detail) where:
Tuple of (outcome, modified_data, error_detail, original_exception):
- outcome: "pass", "fail", or "error"
- modified_data: dict if guardrail returned modified data, else None
- error_detail: error message string if fail/error, else None
- original_exception: the exception the guardrail raised, so the
pipeline can re-raise it verbatim and match the direct-attachment
response/trace, else None
"""
callback = PipelineExecutor._find_guardrail_callback(step.guardrail)
callback = PipelineExecutor.find_guardrail_callback(step.guardrail)
if callback is None:
verbose_proxy_logger.warning(
f"Pipeline: guardrail '{step.guardrail}' not found in callbacks"
)
return ("error", None, f"Guardrail '{step.guardrail}' not found")
return ("error", None, f"Guardrail '{step.guardrail}' not found", None)
try:
# Inject guardrail name into metadata so should_run_guardrail() allows it
@ -182,26 +196,26 @@ class PipelineExecutor:
response=data.get("response"), # type: ignore
)
else:
return ("error", None, f"Unsupported pipeline mode: {mode}")
return ("error", None, f"Unsupported pipeline mode: {mode}", None)
# Normal return means pass
modified_data = None
if response is not None and isinstance(response, dict):
modified_data = response
return ("pass", modified_data, None)
return ("pass", modified_data, None, None)
except Exception as e:
if CustomGuardrail._is_guardrail_intervention(e):
error_msg = _extract_error_message(e)
return ("fail", None, error_msg)
return ("fail", None, error_msg, e)
else:
verbose_proxy_logger.error(
f"Pipeline: unexpected error from guardrail '{step.guardrail}': {e}"
)
return ("error", None, str(e))
return ("error", None, str(e), e)
@staticmethod
def _find_guardrail_callback(guardrail_name: str) -> Optional[CustomGuardrail]:
def find_guardrail_callback(guardrail_name: str) -> Optional[CustomGuardrail]:
"""Look up an initialized guardrail callback by name from litellm.callbacks."""
for callback in litellm.callbacks:
if isinstance(callback, CustomGuardrail):

View file

@ -364,6 +364,16 @@ def _enrich_http_exception_with_guardrail_context(
detail.setdefault("guardrail_mode", event_hook)
def _exception_changes_request_flow(exc: BaseException) -> bool:
"""
True for guardrail exceptions the proxy turns into an alternate request flow
(a reroute or a passthrough response) rather than a block. A pipeline step
configured to block must honor that block, so these are surfaced as the
generic pipeline block instead of being re-raised verbatim.
"""
return isinstance(exc, (SensitiveDataRouteException, ModifyResponseException))
@dataclass(frozen=True)
class _CallbackCapabilities:
"""Cached per-hook capability flags derived from ``litellm.callbacks``.
@ -1352,7 +1362,7 @@ class ProxyLogging:
@staticmethod
def _handle_pipeline_result(
result: Any,
result: PipelineExecutionResult,
data: dict,
policy_name: str,
) -> dict:
@ -1367,6 +1377,21 @@ class ProxyLogging:
return data
if result.terminal_action == "block":
original_exception = result.original_exception
if original_exception is not None and not _exception_changes_request_flow(
original_exception
):
blocking_step = result.step_results[-1] if result.step_results else None
if blocking_step is not None:
callback = PipelineExecutor.find_guardrail_callback(
blocking_step.guardrail_name
)
if callback is not None:
_enrich_http_exception_with_guardrail_context(
original_exception, callback
)
raise original_exception
step_results_serializable = [
{
"guardrail": sr.guardrail_name,

View file

@ -99,8 +99,11 @@ class PipelineStepResult(BaseModel):
class PipelineExecutionResult(BaseModel):
"""Result of executing an entire pipeline."""
model_config = ConfigDict(arbitrary_types_allowed=True)
terminal_action: str # block | allow | modify_response
step_results: List[PipelineStepResult]
modified_data: Optional[Dict[str, Any]] = None
error_message: Optional[str] = None
modify_response_message: Optional[str] = None
original_exception: Optional[Exception] = Field(default=None, exclude=True)

View file

@ -209,6 +209,75 @@ async def test_escalation_step1_fails_step2_blocks():
litellm.callbacks = original_callbacks
@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed")
@pytest.mark.asyncio
async def test_block_carries_original_guardrail_exception():
"""A blocking step must expose the guardrail's own raised exception on the
result so the caller can re-raise it verbatim, giving the policy path the
same response/trace as a direct guardrail attachment."""
guard = AlwaysFailGuardrail(guardrail_name="moderation-filter")
pipeline = GuardrailPipeline(
mode="pre_call",
steps=[
PipelineStep(
guardrail="moderation-filter", on_fail="block", on_pass="allow"
)
],
)
original_callbacks = litellm.callbacks.copy()
litellm.callbacks = [guard]
try:
result = await PipelineExecutor.execute_steps(
steps=pipeline.steps,
mode=pipeline.mode,
data={"messages": [{"role": "user", "content": "bad content"}]},
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="content-safety",
)
assert result.terminal_action == "block"
assert isinstance(result.original_exception, HTTPException)
assert result.original_exception.status_code == 400
assert result.original_exception.detail == "Content policy violation"
finally:
litellm.callbacks = original_callbacks
@pytest.mark.asyncio
async def test_unsupported_mode_yields_error_outcome_without_exception():
"""An unexpected hook mode must surface as an error outcome (carrying no
original exception), not crash or run the guardrail."""
guard = AlwaysPassGuardrail(guardrail_name="filter")
original_callbacks = litellm.callbacks.copy()
litellm.callbacks = [guard]
try:
result = await PipelineExecutor.execute_steps(
steps=[PipelineStep(guardrail="filter", on_error="block", on_fail="block")],
mode="during_call",
data={"messages": [{"role": "user", "content": "hi"}]},
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="content-safety",
)
assert guard.calls == 0
assert result.terminal_action == "block"
assert result.step_results[0].outcome == "error"
assert (
"Unsupported pipeline mode: during_call"
in result.step_results[0].error_detail
)
assert result.original_exception is None
finally:
litellm.callbacks = original_callbacks
@pytest.mark.asyncio
async def test_passthrough_guardrail_failure_can_pipeline_block():
"""

View file

@ -17,6 +17,7 @@ import pytest
from fastapi import HTTPException
import litellm
from litellm.exceptions import SensitiveDataRouteException
from litellm.integrations.custom_guardrail import (
CustomGuardrail,
ModifyResponseException,
@ -359,6 +360,7 @@ async def test_maybe_execute_pipelines_blocks_on_block_terminal_action_raises(
fake_result = MagicMock()
fake_result.terminal_action = "block"
fake_result.step_results = []
fake_result.original_exception = None
data = {"metadata": {"_guardrail_pipelines": [("policy-1", pipeline)]}, "messages": [], "model": "m"}
async def fake_execute_steps(**kwargs):
@ -377,6 +379,42 @@ async def test_maybe_execute_pipelines_blocks_on_block_terminal_action_raises(
)
@pytest.mark.asyncio
async def test_maybe_execute_pipelines_reraises_original_guardrail_exception(
proxy_logging, make_user_api_key_auth, monkeypatch
):
"""A policy-wrapped guardrail block must surface the guardrail's own
exception verbatim, identical to the direct-attachment path."""
pipeline = MagicMock()
pipeline.mode = "pre_call"
pipeline.steps = []
original = HTTPException(
status_code=400,
detail={"error": "Violated OpenAI moderation policy", "moderation_result": {"x": 1}},
)
fake_result = MagicMock()
fake_result.terminal_action = "block"
fake_result.step_results = []
fake_result.original_exception = original
data = {"metadata": {"_guardrail_pipelines": [("policy-1", pipeline)]}, "messages": [], "model": "m"}
async def fake_execute_steps(**kwargs):
return fake_result
monkeypatch.setattr(
"litellm.proxy.policy_engine.pipeline_executor.PipelineExecutor.execute_steps",
fake_execute_steps,
)
with pytest.raises(HTTPException) as info:
await proxy_logging._maybe_execute_pipelines(
data=data,
user_api_key_dict=make_user_api_key_auth(),
call_type="completion",
event_hook="pre_call",
)
assert info.value is original
# ---------------------------------------------------------------------------
# _handle_pipeline_result
# ---------------------------------------------------------------------------
@ -391,10 +429,11 @@ def test_handle_pipeline_result_allow_with_modifications():
assert out == {"a": 1, "b": 2, "c": 3}
def test_handle_pipeline_result_block_raises_http_exception():
def test_handle_pipeline_result_block_falls_back_to_generic_when_no_exception():
result = MagicMock()
result.terminal_action = "block"
result.step_results = []
result.original_exception = None
with pytest.raises(HTTPException) as info:
ProxyLogging._handle_pipeline_result(result=result, data={"model": "m"}, policy_name="p")
detail = info.value.detail
@ -410,6 +449,94 @@ def test_handle_pipeline_result_block_raises_http_exception():
}
def test_handle_pipeline_result_block_reraises_original_guardrail_exception():
"""The policy path must re-raise the guardrail's own exception untouched,
not wrap it in a generic ``guardrail_pipeline_error``; this is what makes
the response and trace span identical to the direct-attachment path."""
original = HTTPException(
status_code=400,
detail={
"error": "Violated OpenAI moderation policy",
"moderation_result": {"violated_categories": ["harassment"]},
},
)
result = MagicMock()
result.terminal_action = "block"
result.step_results = []
result.original_exception = original
with pytest.raises(HTTPException) as info:
ProxyLogging._handle_pipeline_result(result=result, data={"model": "m"}, policy_name="p")
assert info.value is original
assert info.value.detail == {
"error": "Violated OpenAI moderation policy",
"moderation_result": {"violated_categories": ["harassment"]},
}
def test_handle_pipeline_result_block_enriches_with_guardrail_name_and_mode():
"""The re-raised exception must gain the blocking guardrail's name and mode,
matching the enrichment the direct-attachment path applies."""
cb = _make_guardrail() # guardrail_name="g", event_hook=pre_call
original = HTTPException(status_code=400, detail={"error": "blocked"})
result = MagicMock()
result.terminal_action = "block"
result.step_results = [MagicMock(guardrail_name="g")]
result.original_exception = original
saved = litellm.callbacks
litellm.callbacks = [cb]
try:
with pytest.raises(HTTPException) as info:
ProxyLogging._handle_pipeline_result(
result=result, data={"model": "m"}, policy_name="p"
)
finally:
litellm.callbacks = saved
assert info.value is original
assert info.value.detail["guardrail_name"] == "g"
assert info.value.detail["guardrail_mode"] == GuardrailEventHooks.pre_call
def test_handle_pipeline_result_block_does_not_reraise_sensitive_data_route():
"""A step configured to block must enforce the block even when the guardrail
raised a reroute exception; re-raising it verbatim would route the request to
an alternate model instead of blocking, bypassing the configured policy."""
original = SensitiveDataRouteException(
route_to_model="on-prem-model",
session_id="sess-1",
guardrail_name="pii-router",
)
result = MagicMock()
result.terminal_action = "block"
result.step_results = [MagicMock(guardrail_name="pii-router")]
result.original_exception = original
with pytest.raises(HTTPException) as info:
ProxyLogging._handle_pipeline_result(result=result, data={"model": "m"}, policy_name="p")
assert info.value.status_code == 400
assert info.value.detail["error"]["type"] == "guardrail_pipeline_error"
def test_handle_pipeline_result_block_does_not_reraise_modify_response():
"""A step configured to block must enforce the block even when the guardrail
raised a passthrough/modify-response exception; re-raising it verbatim would
return the guardrail's synthetic response instead of blocking."""
original = ModifyResponseException(
message="redacted",
model="m",
request_data={"model": "m"},
guardrail_name="masker",
)
result = MagicMock()
result.terminal_action = "block"
result.step_results = [MagicMock(guardrail_name="masker")]
result.original_exception = original
with pytest.raises(HTTPException) as info:
ProxyLogging._handle_pipeline_result(result=result, data={"model": "m"}, policy_name="p")
assert info.value.status_code == 400
assert info.value.detail["error"]["type"] == "guardrail_pipeline_error"
def test_handle_pipeline_result_modify_response_raises_modify_exception():
result = MagicMock()
result.terminal_action = "modify_response"