fix(guardrails): drop the steps of later response pipelines a block never reached

When one response pipeline blocks or rewrites the response, the policies after it never run, so
their steps no longer stay listed in the header, the spend log and the trace

litellm/proxy/utils.py

@@ -759,20 +759,27 @@ def _stream_log_waits_for_its_pipelines(data: Mapping[str, object]) -> bool:
     return getattr(data.get("litellm_logging_obj"), "_on_deferred_stream_complete", None) is not None
 
 
+def _pipelines_settled_by(
+    pipelines: Sequence[tuple[str, "GuardrailPipeline"]], index: int, result: PipelineExecutionResult
+) -> Sequence[tuple[str, "GuardrailPipeline"]]:
+    """This pipeline, plus every later one when its verdict ends the request before they run."""
+    return pipelines[index:] if result.terminal_action != "allow" else pipelines[index : index + 1]
+
+
 def _drop_response_steps_that_did_not_run(
     data: dict[str, object],  # mutable-ok: the request's applied guardrails are trimmed in place
-    steps: Sequence["PipelineStep"],
+    pipelines: Sequence[tuple[str, "GuardrailPipeline"]],
     steps_that_ran: frozenset[str],
 ) -> None:
     """
-    Drop from the applied guardrails the steps of a response pipeline that no pipeline ran,
+    Drop from the applied guardrails the steps of these response pipelines that no pipeline ran,
     unless their guardrail also runs outside pipelines. Every step is named up front, so a
     log built before the pipeline finished, such as a cache hit's, still lists the steps
     that ran. The list is trimmed in place, since a route's logging object may hold a
     shallow copy of the metadata.
     """
     pre_call: Final = pipeline_managed_guardrail_names(data, "pre_call")
-    post_call: Final = frozenset(step.guardrail for step in steps)
+    post_call: Final = _pipeline_step_guardrail_names(pipelines)
     not_run: Final = frozenset(
         name for name in post_call - steps_that_ran if _runs_only_in_pipelines(name, data, pre_call, post_call)
     )
@@ -2796,7 +2803,7 @@ class ProxyLogging:
         steps_that_ran: frozenset[str] = frozenset()  # rebind-ok: a later pipeline keeps a step an earlier one ran
         request: Final = cast(dict[str, object], data)  # cast-ok: request payload, updated in place by every pipeline
         log_waits: Final = event_hook == "post_call" and _whole_response_log_waits_for_its_pipelines(request, call_type)
-        for policy_name, pipeline in pipelines:
+        for index, (policy_name, pipeline) in enumerate(pipelines):
             if pipeline.mode != event_hook:
                 continue
 
@@ -2814,7 +2821,9 @@ class ProxyLogging:
 
             steps_that_ran = steps_that_ran | _steps_that_ran(result)
             if log_waits:
-                _drop_response_steps_that_did_not_run(request, pipeline.steps, steps_that_ran)
+                _drop_response_steps_that_did_not_run(
+                    request, _pipelines_settled_by(pipelines, index, result), steps_that_ran
+                )
 
             data = self._handle_pipeline_result(
                 result=result,
@@ -4734,7 +4743,7 @@ class ProxyLogging:
         call_type, endpoint_translation = translation
 
         steps_that_ran: frozenset[str] = frozenset()  # rebind-ok: a later pipeline keeps a step an earlier one ran
-        for policy_name, pipeline in pipelines:
+        for index, (policy_name, pipeline) in enumerate(pipelines):
             result: PipelineExecutionResult = await PipelineExecutor.execute_steps(
                 steps=pipeline.steps,
                 mode="post_call",
@@ -4747,7 +4756,9 @@ class ProxyLogging:
             )
             steps_that_ran = steps_that_ran | _steps_that_ran(result)
             if _stream_log_waits_for_its_pipelines(request_data):
-                _drop_response_steps_that_did_not_run(request_data, pipeline.steps, steps_that_ran)
+                _drop_response_steps_that_did_not_run(
+                    request_data, _pipelines_settled_by(pipelines, index, result), steps_that_ran
+                )
             try:
                 ProxyLogging._handle_pipeline_result(
                     result, data=request_data, policy_name=policy_name, original_response=buffered

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

@@ -12,6 +12,7 @@ from __future__ import annotations
 import asyncio
 import json
 from collections.abc import AsyncGenerator, Iterator
+from contextlib import AbstractContextManager, nullcontext
 from copy import deepcopy
 import logging
 from typing import Any, Callable, Dict, Final, List
@@ -1496,6 +1497,86 @@ async def test_held_stream_pipeline_keeps_a_step_an_earlier_pipeline_ran_that_a_
     assert data["metadata"]["applied_guardrails"] == ["first", "second", "third"]
 
 
+def _ends_before_a_later_pipeline_data(ending: str) -> dict[str, object]:
+    ends: Final = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="first", on_pass=ending)])
+    never_reached: Final = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="second")])
+    return {
+        "model": "m",
+        "messages": [{"role": "user", "content": "hi"}],
+        "metadata": {
+            "guardrails": ["first", "second"],
+            "_guardrail_pipelines": [("ends", ends), ("never-reached", never_reached)],
+            "_pipeline_managed_guardrails": {"first", "second"},
+        },
+    }
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(("ending", "raised"), [("block", HTTPException), ("modify_response", ModifyResponseException)])
+async def test_post_call_pipeline_that_ends_the_request_drops_the_steps_of_a_later_pipeline(
+    proxy_logging: ProxyLogging,
+    make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
+    monkeypatch: pytest.MonkeyPatch,
+    ending: str,
+    raised: type[Exception],
+) -> None:
+    first: Final = _PipelineStepGuardrail("first", "post_call")
+    second: Final = _PipelineStepGuardrail("second", "post_call")
+    monkeypatch.setattr(litellm, "callbacks", [first, second])
+    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
+    data: Final = await proxy_logging.pre_call_hook(
+        user_api_key_dict=make_user_api_key_auth(),
+        data={**_ends_before_a_later_pipeline_data(ending), "litellm_logging_obj": _logging_obj(log_parked=True)},
+        call_type="completion",
+    )
+
+    with pytest.raises(raised):
+        await proxy_logging.post_call_success_hook(
+            data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth()
+        )
+
+    assert (first.calls, second.calls) == (1, 0)
+    assert data["metadata"]["applied_guardrails"] == ["first"]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("ending", "stream_end"), [("block", pytest.raises(HTTPException)), ("modify_response", nullcontext())]
+)
+async def test_held_stream_pipeline_that_ends_the_stream_drops_the_steps_of_a_later_pipeline(
+    proxy_logging: ProxyLogging,
+    make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
+    monkeypatch: pytest.MonkeyPatch,
+    ending: str,
+    stream_end: AbstractContextManager[object],
+) -> None:
+    first: Final = _PipelineStepGuardrail("first", "post_call")
+    second: Final = _PipelineStepGuardrail("second", "post_call")
+    monkeypatch.setattr(litellm, "callbacks", [first, second])
+    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
+    logging_obj: Final = _logging_obj(log_parked=False)
+
+    async def stream_log(*args: object) -> None:
+        return None
+
+    logging_obj._on_deferred_stream_complete = stream_log
+    auth: Final = make_user_api_key_auth(request_route="/v1/chat/completions")
+    data: Final = await proxy_logging.pre_call_hook(
+        user_api_key_dict=auth,
+        data={**_ends_before_a_later_pipeline_data(ending), "stream": True, "litellm_logging_obj": logging_obj},
+        call_type="completion",
+    )
+
+    with stream_end:
+        async for _item in proxy_logging.async_post_call_streaming_iterator_hook(
+            user_api_key_dict=auth, response=_async_chunk_iter(_stream_chunks()), request_data=data
+        ):
+            pass
+
+    assert (first.calls, second.calls) == (1, 0)
+    assert data["metadata"]["applied_guardrails"] == ["first"]
+
+
 @pytest.mark.asyncio
 async def test_post_call_pipeline_keeps_a_step_that_never_ran_whose_guardrail_ran_outside_the_pipeline(
     proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch,