fix(guardrails): name a streamed pipeline step whose guardrail records its own results

A decision_model step on a held stream ran but was left out of the header and
the spend log, where main listed it

litellm/proxy/policy_engine/pipeline_executor.py

@@ -573,8 +573,7 @@ class PipelineExecutor:
             _release_original_chunks(step.guardrail, undeliverable.reason, streaming_chunks, originals)
             return
         finally:
-            if not callback.records_own_guardrail_information:
-                _name_step_that_ran(hook_input, metadata_key, step.guardrail)
+            _name_step_that_ran(hook_input, metadata_key, step.guardrail)
         discard_reason: Final = observer.discard_reason(rewrite_undeliverable_reason)
         if discard_reason is not None:
             _release_original_chunks(step.guardrail, discard_reason, streaming_chunks, originals)

tests/integration/observability/test_guardrail_pipeline_applied_header.py

@@ -36,6 +36,8 @@ _RESPONSE_BLOCKER: Final = "response-blocking-step"
 _LIVE_FIRST: Final = "live-first-step"
 _LIVE_SECOND: Final = "live-second-step"
 _UNGATED_STEP: Final = "ungated-step"
+_JUDGE: Final = "judge-step"
+_JUDGE_MODEL: Final = "pipeline-judge"
 _STOPS_AFTER_FIRST: Final = "stops-after-first"
 _RUNS_EVERY_STEP: Final = "runs-every-step"
 _BLOCKS_AT_FIRST: Final = "blocks-at-first"
@@ -46,6 +48,7 @@ _RESPONSE_MOVES_PAST_FIRST: Final = "response-moves-past-first"
 _SCANS_AFTER_RELEASE: Final = "scans-after-release"
 _CANNOT_HOLD_STREAM: Final = "cannot-hold-stream"
 _ERRORS_PAST_FIRST: Final = "errors-past-first"
+_JUDGES_THEN_CHECKS: Final = "judges-then-checks"
 _ENDPOINTS: Final = ("/v1/chat/completions", "/v1/responses", "/v1/messages")
 
 
@@ -154,6 +157,16 @@ def _completion(request: Request) -> Reply:
 
 
 def _vendor(request: Request) -> Reply:
+    if request.target == f"/{_JUDGE}/v1/decisions":
+        return Reply(
+            body=json.dumps(
+                {
+                    "model": "gpt-6-luna",
+                    "answers": [{"type": "predicate", "name": "unsafe_output", "probability": 0.01}],
+                    "usage": {"input_tokens": 10, "output_tokens": 1, "total_tokens": 11},
+                }
+            ).encode()
+        )
     if request.target.startswith(f"/{_PROMPT_SHIELD}{_PROMPT_SHIELD_PATH}?"):
         return Reply(
             body=json.dumps({"userPromptAnalysis": {"attackDetected": False}, "documentsAnalysis": []}).encode()
@@ -198,6 +211,17 @@ def _policy(mode: str, steps: tuple[tuple[str, str, str], ...], *, on_error: str
 def _config(directory: Path, vendor: Wire) -> Path:
     configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
     configuration["litellm_settings"] = {**configuration.get("litellm_settings", {}), "callbacks": ["otel"]}
+    configuration["model_list"] = [
+        *configuration.get("model_list", []),
+        {
+            "model_name": _JUDGE_MODEL,
+            "litellm_params": {
+                "model": "openai/gpt-6-luna",
+                "api_base": f"{vendor.url}/{_JUDGE}/v1",
+                "api_key": "synthetic-judge-key",
+            },
+        },
+    ]
     configuration["guardrails"] = [
         *(
             _guardrail(name, "pre_call", vendor)
@@ -218,6 +242,23 @@ def _config(directory: Path, vendor: Wire) -> Path:
                 "api_key": "synthetic-guardrail-key",
             },
         },
+        {
+            "guardrail_name": _JUDGE,
+            "litellm_params": {
+                "guardrail": "decision_model",
+                "mode": "post_call",
+                "default_on": False,
+                "decision_model": _JUDGE_MODEL,
+                "checks": [
+                    {
+                        "name": "unsafe_output",
+                        "instructions": "Does the response say anything unsafe?",
+                        "action": "block",
+                        "threshold": 0.5,
+                    }
+                ],
+            },
+        },
         {
             "guardrail_name": _PROMPT_SHIELD,
             "litellm_params": {
@@ -253,6 +294,7 @@ def _config(directory: Path, vendor: Wire) -> Path:
         _CANNOT_HOLD_STREAM: _policy(
             "post_call", ((_UNGATED_STEP, "next", "block"), (_RESPONSE_SECOND, "next", "block"))
         ),
+        _JUDGES_THEN_CHECKS: _policy("post_call", ((_JUDGE, "next", "block"), (_RESPONSE_SECOND, "next", "block"))),
         _ERRORS_PAST_FIRST: _policy(
             "pre_call", ((_REQUEST_ERRORING, "next", "block"), (_REQUEST_SECOND, "next", "block")), on_error="next"
         ),
@@ -516,6 +558,20 @@ def test_the_spend_log_and_the_trace_name_only_the_steps_that_ran(
     ] * len(spans), spans
 
 
+@pytest.mark.parametrize("stream", [False, True], ids=["whole", "streamed"])
+def test_a_response_pipeline_names_a_judge_step_that_records_its_own_results(rig: _Rig, stream: bool) -> None:
+    team, response = _ask_as_new_team(rig, _JUDGES_THEN_CHECKS, stream=stream)
+    assert response.status_code == 200, response.text
+    assert _steps_that_ran(rig) == (_JUDGE, _RESPONSE_SECOND)
+    assert sorted(response.headers.get(_HEADER, "").split(",")) == [_JUDGE, _RESPONSE_SECOND], dict(response.headers)
+    rows: Final = eventually(
+        lambda: rig.gateway.rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE team_id = %s', (team,)),
+        lambda found: len(found) == 1,
+        seconds=70,
+    )
+    assert sorted(object_value(rows[0]["metadata"]).get("applied_guardrails") or []) == [_JUDGE, _RESPONSE_SECOND]
+
+
 def _agent_card(name: str, url: str) -> dict[str, object]:
     return {
         "protocolVersion": "0.3",

tests/unit/proxy/policy_engine/test_pipeline_executor.py

@@ -1434,6 +1434,19 @@ async def test_streaming_step_records_the_guardrail_in_the_applied_guardrails_he
     assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"]
 
 
+class _SelfRecordingGuardrail(_TextReturningGuardrail):
+    records_own_guardrail_information = True
+
+
+@pytest.mark.asyncio
+async def test_streaming_step_whose_guardrail_records_its_own_results_is_named(monkeypatch):
+    monkeypatch.setattr(litellm, "callbacks", [_SelfRecordingGuardrail(["hello world"])])
+
+    result = await _run_streaming_step(_WritingTranslation(), [_chunk()])
+
+    assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"]
+
+
 class _LitellmMetadataOpeningGuardrail(CustomGuardrail):
     """Opens ``litellm_metadata`` on the request while scanning, as the A2A translation does."""