fix(guardrails): name every pipeline step that ran in the metadata the request started with

An A2A response step opens litellm_metadata on the request, so the next step's
name landed there while the spend log and the trace read metadata and lost it

litellm/proxy/policy_engine/pipeline_executor.py

@@ -355,6 +355,22 @@ async def _run_pre_call_hook(
     return payload
 
 
+def _name_step_that_ran(
+    request_data: dict[str, object],  # mutable-ok: the request payload the step's name is recorded on
+    metadata_key: Literal["metadata", "litellm_metadata"],
+    guardrail_name: str,
+) -> None:
+    """Name a step that ran in the applied guardrails of the metadata bucket the pipeline's request started with."""
+    existing: Final = request_data.get(metadata_key)
+    bucket: Final[dict[str, object]] = (  # mutable-ok: the request's own metadata bucket, written in place
+        cast("dict[str, object]", existing)  # cast-ok: metadata buckets are str-keyed
+        if isinstance(existing, dict)
+        else {}
+    )
+    request_data[metadata_key] = bucket  # rebind-ok: the step's name is recorded on the request payload
+    add_guardrail_to_applied_guardrails_header(request_data={metadata_key: bucket}, guardrail_name=guardrail_name)
+
+
 def _release_original_chunks(
     guardrail_name: str,
     reason: str,
@@ -415,6 +431,9 @@ class PipelineExecutor:
             PipelineExecutionResult with terminal action and step results
         """
         step_results: Final[list[PipelineStepResult]] = []
+        metadata_key: Final = get_metadata_variable_name_from_kwargs(
+            cast(Mapping[str, object], data)  # cast-ok: the request payload is str-keyed
+        )
         working_data = data.copy()
         if "metadata" in working_data:
             working_data["metadata"] = working_data["metadata"].copy()
@@ -433,6 +452,7 @@ class PipelineExecutor:
                 data=working_data,
                 user_api_key_dict=user_api_key_dict,
                 call_type=call_type,
+                metadata_key=metadata_key,
                 raw_request_snapshot=raw_request_snapshot,
                 streaming_chunks=streaming_chunks,
                 endpoint_translation=endpoint_translation,
@@ -506,6 +526,7 @@ class PipelineExecutor:
         user_api_key_dict: "UserAPIKeyAuth",
         litellm_logging_obj: "LiteLLMLoggingObj | None",
         stream_already_sent: bool,
+        metadata_key: Literal["metadata", "litellm_metadata"],
     ) -> None:
         """Run one streaming post_call step through the endpoint translation, delivering
         text and tool-call rewrites on translations that support ended-stream write-back. A
@@ -553,7 +574,7 @@ class PipelineExecutor:
             return
         finally:
             if not callback.records_own_guardrail_information:
-                add_guardrail_to_applied_guardrails_header(request_data=hook_input, guardrail_name=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)
@@ -565,6 +586,7 @@ class PipelineExecutor:
         data: dict,
         user_api_key_dict: "UserAPIKeyAuth",
         call_type: str,
+        metadata_key: Literal["metadata", "litellm_metadata"],
         raw_request_snapshot: dict | None = None,  # mutable-ok: same request-payload shape as data
         streaming_chunks: list[object] | None = None,  # mutable-ok: shared buffered-stream chunks, read per step
         endpoint_translation: "BaseTranslation | None" = None,
@@ -626,10 +648,15 @@ class PipelineExecutor:
                     user_api_key_dict=user_api_key_dict,
                     litellm_logging_obj=data.get("litellm_logging_obj"),
                     stream_already_sent=stream_already_sent,
+                    metadata_key=metadata_key,
                 )
                 response = None
             elif mode == "post_call":
-                add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=step.guardrail)
+                _name_step_that_ran(
+                    cast(dict[str, object], data),  # cast-ok: the request payload is str-keyed
+                    metadata_key,
+                    step.guardrail,
+                )
                 response = await target.async_post_call_success_hook(
                     user_api_key_dict=user_api_key_dict,
                     data=data,

tests/integration/observability/test_guardrail_pipeline_applied_header.py

@@ -516,6 +516,90 @@ def test_the_spend_log_and_the_trace_name_only_the_steps_that_ran(
     ] * len(spans), spans
 
 
+def _agent_card(name: str, url: str) -> dict[str, object]:
+    return {
+        "protocolVersion": "0.3",
+        "name": name,
+        "description": "Scripted agent that answers every message with one text part",
+        "version": "1.0.0",
+        "url": url,
+        "capabilities": {"streaming": True},
+        "defaultInputModes": ["text"],
+        "defaultOutputModes": ["text"],
+        "skills": [],
+    }
+
+
+@pytest.fixture(scope="module")
+def agent(rig: _Rig) -> Iterator[str]:
+    name: Final = f"pipeline-agent-{uuid.uuid4().hex}"
+
+    def answer(request: Request) -> Reply:
+        if request.method == "GET":
+            return Reply(body=json.dumps(_agent_card(name, f"{peer.url}/")).encode())
+        message: Final = {
+            "kind": "message",
+            "role": "agent",
+            "messageId": uuid.uuid4().hex,
+            "parts": [{"kind": "text", "text": "guarded"}],
+        }
+        return Reply(
+            body=json.dumps({"jsonrpc": "2.0", "id": json.loads(request.body)["id"], "result": message}).encode()
+        )
+
+    with wire_server(answer) as peer:
+        identity: Final = str(
+            rig.gateway.post(
+                "/v1/agents", {"agent_name": name, "agent_card_params": _agent_card(name, f"{peer.url}/")}
+            )["agent_id"]
+        )
+        yield identity
+        rig.gateway.request("DELETE", f"/v1/agents/{identity}")
+
+
+@pytest.mark.parametrize(
+    ("tag", "expected"),
+    [
+        pytest.param(_SCANS_AFTER_RELEASE, [_LIVE_FIRST, _LIVE_SECOND], id="every-step-passes"),
+        pytest.param(
+            _RESPONSE_MOVES_PAST_FIRST, [_RESPONSE_BLOCKER, _RESPONSE_SECOND], id="first-step-flags-and-moves-on"
+        ),
+    ],
+)
+def test_an_agent_message_names_every_response_step_that_ran_in_the_spend_log_and_the_trace(
+    rig: _Rig, agent: str, tag: str, expected: list[str]
+) -> None:
+    since: Final = recorded_spans(rig.traces)[0]
+    rig.vendor.drain()
+    message: Final = {
+        "kind": "message",
+        "role": "user",
+        "messageId": uuid.uuid4().hex,
+        "parts": [{"kind": "text", "text": f"hello {uuid.uuid4().hex}"}],
+    }
+    with rig.gateway.scenario() as scenario:
+        team: Final = scenario.team(metadata={"tags": [tag]})
+        response: Final = rig.gateway.request(
+            "POST",
+            f"/a2a/{agent}",
+            {"jsonrpc": "2.0", "id": uuid.uuid4().hex, "method": "message/send", "params": {"message": message}},
+            key=scenario.key(team_id=team),
+            headers={"Accept": "application/json"},
+        )
+    assert response.status_code == 200, response.text
+    assert list(_steps_that_ran(rig)) == expected
+    rows: Final = eventually(
+        lambda: rig.gateway.rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE team_id = %s', (team,)),
+        lambda found: len(found) == 1,
+        seconds=70,
+    )
+    spans: Final = eventually(lambda: _spans_naming_guardrails(rig, since, team), bool, seconds=30)
+    assert sorted(object_value(rows[0]["metadata"]).get("applied_guardrails") or []) == expected, rows
+    assert [sorted(ast.literal_eval(str(span["attributes"]["metadata.applied_guardrails"]))) for span in spans] == [
+        expected
+    ] * len(spans), spans
+
+
 def _openai_chat(base: str, key: str, model: str) -> str | None:
     with (
         httpx.Client(timeout=30, trust_env=False) as transport,

tests/unit/proxy/policy_engine/test_pipeline_executor.py

@@ -1157,6 +1157,7 @@ async def test_pipeline_step_keeps_native_hook_when_opted_out(monkeypatch):
         data=data,
         user_api_key_dict=MagicMock(),
         call_type="completion",
+        metadata_key="metadata",
     )
 
     assert outcome == "pass"
@@ -1433,6 +1434,27 @@ async def test_streaming_step_records_the_guardrail_in_the_applied_guardrails_he
     assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"]
 
 
+class _LitellmMetadataOpeningGuardrail(CustomGuardrail):
+    """Opens ``litellm_metadata`` on the request while scanning, as the A2A translation does."""
+
+    def __init__(self):
+        super().__init__(guardrail_name="masker", event_hook="post_call", default_on=True)
+
+    async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
+        request_data["litellm_metadata"] = {"user_api_key_hash": "hashed"}
+        return inputs
+
+
+@pytest.mark.asyncio
+async def test_streaming_step_is_named_in_the_metadata_the_request_started_with(monkeypatch):
+    monkeypatch.setattr(litellm, "callbacks", [_LitellmMetadataOpeningGuardrail()])
+
+    result = await _run_streaming_step(_WritingTranslation(), [_chunk()])
+
+    assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"]
+    assert "applied_guardrails" not in result.modified_data["litellm_metadata"]
+
+
 @pytest.mark.asyncio
 async def test_streaming_step_that_flags_the_stream_is_named_when_the_pipeline_moves_on(monkeypatch):
     monkeypatch.setattr(litellm, "callbacks", [_BlockingStreamGuardrail()])

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

@@ -1348,6 +1348,34 @@ async def test_post_call_pipeline_steps_are_named_only_once_they_run(
     assert data["metadata"]["applied_guardrails"] == ["first"]
 
 
+class _OpensLitellmMetadataOnResponseGuardrail(_PipelineStepGuardrail):
+    async def async_post_call_success_hook(self, data, user_api_key_dict, response):
+        self.calls += 1
+        data["litellm_metadata"] = {"user_api_key_hash": "hashed"}
+
+
+@pytest.mark.asyncio
+async def test_post_call_pipeline_names_every_step_that_ran_in_metadata_after_a_step_opens_litellm_metadata(
+    proxy_logging, make_user_api_key_auth, monkeypatch
+):
+    first = _OpensLitellmMetadataOnResponseGuardrail("first", "post_call")
+    second = _PipelineStepGuardrail("second", "post_call")
+    monkeypatch.setattr(litellm, "callbacks", [first, second])
+    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
+
+    data = await proxy_logging.pre_call_hook(
+        user_api_key_dict=make_user_api_key_auth(),
+        data=_two_step_pipeline_data("post_call", "next"),
+        call_type="completion",
+    )
+    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, 1)
+    assert data["metadata"]["applied_guardrails"] == ["first", "second"]
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     ("asks_for", "request_route", "streams_live", "expected_header"),