fix(policy_engine): execute post_call guardrail pipelines on responses

litellm/proxy/utils.py

@@ -455,6 +455,30 @@ def _pipeline_managed_guardrail_names(data: Mapping[str, object]) -> frozenset[s
     )
 
 
+def _raise_for_streaming_post_call_pipelines(data: Mapping[str, object]) -> None:
+    if data.get("stream") is not True:
+        return
+    post_call_policies: Final = tuple(
+        policy_name for policy_name, pipeline in _policy_pipelines(data) if pipeline.mode == "post_call"
+    )
+    if not post_call_policies:
+        return
+    raise HTTPException(
+        status_code=400,
+        detail={
+            "error": {
+                "message": (
+                    "Policies with post_call guardrail pipelines cannot govern streaming responses yet: "
+                    f"{', '.join(post_call_policies)}. Retry with stream=false, or move these policies' output "
+                    "guardrails from pipeline steps to guardrails.add, which scans streamed output."
+                ),
+                "type": "guardrail_pipeline_error",
+                "policies": list(post_call_policies),
+            }
+        },
+    )
+
+
 def _prompt_block_text(block: object) -> str:
     if isinstance(block, str):
         return block
@@ -1578,6 +1602,7 @@ class ProxyLogging:
         call_type: str,
         event_hook: str,
         raw_request_snapshot: dict | None = None,  # mutable-ok: same request-payload shape as data
+        response: LLMResponseTypes | None = None,
     ) -> dict:
         """
         Execute guardrail pipelines if any are configured for this request.
@@ -1596,6 +1621,8 @@ class ProxyLogging:
         if not pipelines:
             return data
 
+        step_input: Final = {**data, "response": response} if response is not None else data
+
         for policy_name, pipeline in pipelines:
             if pipeline.mode != event_hook:
                 continue
@@ -1603,7 +1630,7 @@ class ProxyLogging:
             result: PipelineExecutionResult = await PipelineExecutor.execute_steps(
                 steps=pipeline.steps,
                 mode=pipeline.mode,
-                data=data,
+                data=step_input,
                 user_api_key_dict=user_api_key_dict,
                 call_type=call_type,
                 policy_name=policy_name,
@@ -1614,6 +1641,7 @@ class ProxyLogging:
                 result=result,
                 data=data,
                 policy_name=policy_name,
+                original_response=response,
             )
 
         return data
@@ -1623,14 +1651,18 @@ class ProxyLogging:
         result: PipelineExecutionResult,
         data: dict,
         policy_name: str,
+        original_response: LLMResponseTypes | None = None,
     ) -> dict:
         """
         Handle a PipelineExecutionResult — allow, block, or modify_response.
 
         Returns data dict if allowed, raises on block/modify_response.
+        ``original_response`` is set on the post_call path, where allowed
+        modifications land on the response object in place, so the request
+        payload (already sent upstream) is left untouched.
         """
         if result.terminal_action == "allow":
-            if result.modified_data is not None:
+            if result.modified_data is not None and original_response is None:
                 data.update(result.modified_data)
             return data
 
@@ -1671,6 +1703,7 @@ class ProxyLogging:
                 request_data=data,
                 guardrail_name=f"pipeline:{policy_name}",
                 detection_info=None,
+                original_response=original_response,
             )
 
         return data
@@ -1786,6 +1819,8 @@ class ProxyLogging:
         )
 
         try:
+            _raise_for_streaming_post_call_pipelines(data)
+
             # Execute guardrail pipelines before the normal callback loop
             data = await self._maybe_execute_pipelines(
                 data=data,
@@ -2774,6 +2809,14 @@ class ProxyLogging:
         from litellm.proxy.proxy_server import llm_router
         from litellm.types.guardrails import GuardrailEventHooks
 
+        await self._maybe_execute_pipelines(
+            data=data,
+            user_api_key_dict=user_api_key_dict,
+            call_type=getattr(data.get("litellm_logging_obj"), "call_type", None) or "acompletion",
+            event_hook="post_call",
+            response=response,
+        )
+
         guardrail_callbacks: Final[list[CustomGuardrail]] = []
         other_callbacks: Final[list[CustomLogger]] = []
         try:

tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py

@@ -23,7 +23,7 @@ from litellm.integrations.custom_guardrail import (
     ModifyResponseException,
 )
 from litellm.integrations.prometheus import PrometheusLogger
-from litellm.proxy.utils import ProxyLogging
+from litellm.proxy.utils import ProxyLogging, _raise_for_streaming_post_call_pipelines
 from litellm.types.guardrails import GuardrailEventHooks
 from litellm.types.proxy.policy_engine.pipeline_types import (
     GuardrailPipeline,
@@ -865,3 +865,153 @@ async def test_process_prompt_template_aresponses_swaps_model_and_merges_input(p
     hook_kwargs = logging_obj.async_get_chat_completion_prompt.await_args.kwargs
     assert hook_kwargs["messages"] == [{"role": "user", "content": "Who are you?"}]
     assert hook_kwargs["prompt_spec"] is prompt_spec
+
+
+# ---------------------------------------------------------------------------
+# post_call pipeline execution (LIT-6410)
+# ---------------------------------------------------------------------------
+
+
+def _post_call_pipeline_data(guardrail: str = "gr-post", **extra: Any) -> Dict[str, Any]:
+    pipeline = GuardrailPipeline(
+        mode="post_call",
+        steps=[PipelineStep(guardrail=guardrail, on_pass="allow", on_fail="block")],
+    )
+    return {
+        "model": "m",
+        "messages": [{"role": "user", "content": "hi"}],
+        "metadata": {
+            "_guardrail_pipelines": [("response-governance", pipeline)],
+            "_pipeline_managed_guardrails": {guardrail},
+        },
+        **extra,
+    }
+
+
+@pytest.mark.asyncio
+async def test_post_call_success_hook_runs_post_call_pipeline_and_reraises_block(
+    proxy_logging, make_user_api_key_auth, monkeypatch
+):
+    seen: Dict[str, Any] = {}
+
+    class OutputBlockingGuardrail(CustomGuardrail):
+        async def async_post_call_success_hook(self, data, user_api_key_dict, response):
+            seen["response"] = response
+            raise HTTPException(status_code=400, detail={"error": "output blocked"})
+
+    monkeypatch.setattr(
+        litellm,
+        "callbacks",
+        [OutputBlockingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)],
+    )
+    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
+    data = _post_call_pipeline_data()
+    response = litellm.ModelResponse()
+
+    with pytest.raises(HTTPException) as info:
+        await proxy_logging.post_call_success_hook(
+            data=data, response=response, user_api_key_dict=make_user_api_key_auth()
+        )
+
+    assert info.value.detail["error"] == "output blocked"
+    assert seen["response"] is response
+
+
+@pytest.mark.asyncio
+async def test_post_call_pipeline_pass_runs_once_and_leaves_request_data_untouched(
+    proxy_logging, make_user_api_key_auth, monkeypatch
+):
+    seen: Dict[str, Any] = {"count": 0}
+
+    class RecordingGuardrail(CustomGuardrail):
+        async def async_post_call_success_hook(self, data, user_api_key_dict, response):
+            seen["count"] += 1
+            seen["response"] = response
+            return None
+
+    monkeypatch.setattr(
+        litellm,
+        "callbacks",
+        [RecordingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)],
+    )
+    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
+    data = _post_call_pipeline_data()
+    response = litellm.ModelResponse()
+
+    out = await proxy_logging.post_call_success_hook(
+        data=data, response=response, user_api_key_dict=make_user_api_key_auth()
+    )
+
+    assert out is response
+    assert seen["response"] is response
+    assert seen["count"] == 1
+    assert "response" not in data
+    assert "guardrails" not in data["metadata"]
+
+
+def test_handle_pipeline_result_modify_response_carries_original_response():
+    result = MagicMock()
+    result.terminal_action = "modify_response"
+    result.modify_response_message = "filtered"
+    response = litellm.ModelResponse()
+
+    with pytest.raises(ModifyResponseException) as info:
+        ProxyLogging._handle_pipeline_result(
+            result=result, data={"model": "m"}, policy_name="p", original_response=response
+        )
+
+    assert info.value.original_response is response
+
+
+def test_handle_pipeline_result_allow_discards_modifications_on_post_call():
+    data = {"a": 1, "metadata": {"guardrails": ["other"]}}
+    result = MagicMock()
+    result.terminal_action = "allow"
+    result.modified_data = {"metadata": {"guardrails": ["gr-post"]}, "response": object()}
+
+    out = ProxyLogging._handle_pipeline_result(
+        result=result, data=data, policy_name="p", original_response=litellm.ModelResponse()
+    )
+
+    assert out is data
+    assert data == {"a": 1, "metadata": {"guardrails": ["other"]}}
+
+
+@pytest.mark.asyncio
+async def test_pre_call_hook_rejects_streaming_request_with_post_call_pipeline(
+    proxy_logging, make_user_api_key_auth, monkeypatch
+):
+    monkeypatch.setattr(litellm, "callbacks", [])
+    data = _post_call_pipeline_data(stream=True)
+
+    with pytest.raises(HTTPException) as info:
+        await proxy_logging.pre_call_hook(
+            user_api_key_dict=make_user_api_key_auth(),
+            data=data,
+            call_type="completion",
+            guardrails_only=True,
+        )
+
+    assert info.value.status_code == 400
+    assert info.value.detail["error"]["policies"] == ["response-governance"]
+    assert "stream=false" in info.value.detail["error"]["message"]
+
+
+def test_raise_for_streaming_post_call_pipelines_ignores_non_streaming_and_pre_call():
+    post_call = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="g", on_fail="block")])
+    pre_call = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="g", on_fail="block")])
+
+    assert (
+        _raise_for_streaming_post_call_pipelines(
+            {"stream": False, "metadata": {"_guardrail_pipelines": [("p", post_call)]}}
+        )
+        is None
+    )
+    assert _raise_for_streaming_post_call_pipelines({"metadata": {"_guardrail_pipelines": [("p", post_call)]}}) is None
+    assert (
+        _raise_for_streaming_post_call_pipelines(
+            {"stream": True, "metadata": {"_guardrail_pipelines": [("p", pre_call)]}}
+        )
+        is None
+    )
+    assert _raise_for_streaming_post_call_pipelines({"stream": True}) is None