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"