From 98e459cdcd02edb708a43410b1748e039d9599b2 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Fri, 26 Jun 2026 09:56:14 +0300 Subject: [PATCH] fix(guardrails): match policy-pipeline block response to direct guardrail attachment 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). A guardrail can also raise a control-flow exception that the proxy turns into an alternate request flow rather than a block: SensitiveDataRouteException reroutes to another model and ModifyResponseException returns a 200 passthrough response. Re-raising those verbatim under a step configured on_fail=block would convert the configured block into the guardrail's alternate behavior, letting a user bypass the block. Those two exceptions now fall through to the generic pipeline block so the policy's block is honored; only exceptions that already represent a block (content-policy HTTP 400, GuardrailRaisedException, BlockedPiiEntityError) are re-raised verbatim. Resolves LIT-4041 --- .../proxy/policy_engine/pipeline_executor.py | 36 +++-- litellm/proxy/utils.py | 27 +++- .../proxy/policy_engine/pipeline_types.py | 3 + .../policy_engine/test_pipeline_executor.py | 69 ++++++++++ .../proxy_logging/test_guardrail_pipeline.py | 129 +++++++++++++++++- 5 files changed, 251 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index e46e3e1dc9f..1fde9109e19 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -6,7 +6,7 @@ """ 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 @@ async def execute_steps( 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 @@ async def execute_steps( terminal_action="block", step_results=step_results, error_message=error_detail, + original_exception=original_exception, ) if action == "modify_response": @@ -134,22 +140,30 @@ async def _run_step( 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 @@ async def _run_step( 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): diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 781f0e9f301..2f8b12e6529 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -363,6 +363,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``. @@ -1356,7 +1366,7 @@ async def _maybe_execute_pipelines( @staticmethod def _handle_pipeline_result( - result: Any, + result: PipelineExecutionResult, data: dict, policy_name: str, ) -> dict: @@ -1371,6 +1381,21 @@ def _handle_pipeline_result( 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, diff --git a/litellm/types/proxy/policy_engine/pipeline_types.py b/litellm/types/proxy/policy_engine/pipeline_types.py index abbb127cd7a..b7b82bff723 100644 --- a/litellm/types/proxy/policy_engine/pipeline_types.py +++ b/litellm/types/proxy/policy_engine/pipeline_types.py @@ -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) 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 058d5b0283b..adf4e8d47c6 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -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(): """ diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 1ff9fbf8d83..9ecfd390dd0 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -17,6 +17,7 @@ from fastapi import HTTPException import litellm +from litellm.exceptions import SensitiveDataRouteException from litellm.integrations.custom_guardrail import ( CustomGuardrail, ModifyResponseException, @@ -358,6 +359,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): @@ -376,6 +378,42 @@ async def fake_execute_steps(**kwargs): ) +@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 # --------------------------------------------------------------------------- @@ -390,10 +428,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 @@ -409,6 +448,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"