From 33b123e42c00ee29c2f3a2a539b8022ac2aba369 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 29 Aug 2026 01:46:08 +0000 Subject: [PATCH] fix(policy_engine): propagate post_call pipeline metadata writes to request data execute_steps hands the step hook a shallow copy of the request metadata so the pipeline's guardrail marker cannot poison the caller. On the post_call allow path _handle_pipeline_result then dropped that working copy entirely, so a hook's standard_logging_guardrail_information and applied_guardrails writes never reached the outer request. Deferred async logging built the StandardLoggingPayload from the untouched original, so spend logs, Langfuse, Datadog, and x-litellm-applied-guardrails all looked as if the scanner never ran. Copy the metadata bucket back on the post_call path and restore the should_run_guardrail marker to its pre-step value in the executor so the internal marker still cannot leak. --- .../proxy/policy_engine/pipeline_executor.py | 17 +++-- litellm/proxy/utils.py | 9 ++- .../proxy_logging/test_guardrail_pipeline.py | 67 +++++++++++++++++-- 3 files changed, 81 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 190784a2a60..c387fbb9985 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -177,12 +177,14 @@ class PipelineExecutor: verbose_proxy_logger.warning("Pipeline: guardrail '%s' not found in callbacks", step.guardrail) return ("error", None, f"Guardrail '{step.guardrail}' not found", None) - try: - # Inject guardrail name into metadata so should_run_guardrail() allows it - if "metadata" not in data: - data["metadata"] = {} - data["metadata"]["guardrails"] = [step.guardrail] + if "metadata" not in data: + data["metadata"] = {} + metadata_dict: Final = data["metadata"] + original_guardrails_present: Final = "guardrails" in metadata_dict + original_guardrails: Final = metadata_dict.get("guardrails") + metadata_dict["guardrails"] = [step.guardrail] + try: # A scan_raw_request step evaluates the pristine pre-pipeline # snapshot instead of `data` (which earlier pass_data steps in # this same pipeline may have already rewritten), same reason @@ -245,6 +247,11 @@ class PipelineExecutor: else: verbose_proxy_logger.error("Pipeline: unexpected error from guardrail '%s': %s", step.guardrail, e) return ("error", None, str(e), e) + finally: + if original_guardrails_present: + metadata_dict["guardrails"] = original_guardrails + else: + metadata_dict.pop("guardrails", None) @staticmethod def find_guardrail_callback(guardrail_name: str) -> CustomGuardrail | None: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index d3f2e1d7301..9ffc652eb2a 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1670,8 +1670,13 @@ class ProxyLogging: response carried in ``modified_data`` is adopted by the caller. """ if result.terminal_action == "allow": - if result.modified_data is not None and original_response is None: - data.update(result.modified_data) + if result.modified_data is not None: + if original_response is None: + data.update(result.modified_data) + else: + for key in ("metadata", "litellm_metadata"): + if key in result.modified_data: + data[key] = result.modified_data[key] return data if result.terminal_action == "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 4e1ccf71c5c..0452a1f3730 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 @@ -10,7 +10,7 @@ Covers ``_should_use_guardrail_load_balancing``, ``_execute_guardrail_hook``, from __future__ import annotations import asyncio -from typing import Any, Dict, List +from typing import Any, Dict, Final, List from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -919,6 +919,49 @@ async def test_post_call_success_hook_runs_post_call_pipeline_and_reraises_block assert seen["response"] is response +@pytest.mark.asyncio +async def test_post_call_pipeline_pass_propagates_metadata_writes_to_request_data( + proxy_logging, make_user_api_key_auth, monkeypatch +): + """A post_call pipeline step's writes to the metadata bucket + (standard_logging_guardrail_information for spend logs / Langfuse / Datadog + and applied_guardrails for the x-litellm-applied-guardrails header) must + reach the outer request data. The working-copy the executor gives the hook + is otherwise discarded, so deferred async logging and the response header + look as if the guardrail never ran.""" + + class RecordingGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + from litellm.proxy.common_utils.callback_utils import ( + add_guardrail_to_applied_guardrails_header, + ) + + metadata = data.setdefault("metadata", {}) + metadata.setdefault("standard_logging_guardrail_information", []).append( + {"guardrail_name": "gr-post"} + ) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name="gr-post") + 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() + + await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + metadata = data["metadata"] + assert metadata["applied_guardrails"] == ["gr-post"] + assert metadata["standard_logging_guardrail_information"] == [{"guardrail_name": "gr-post"}] + assert "guardrails" not in metadata + assert "response" not in data + + @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 @@ -1139,18 +1182,32 @@ def test_handle_pipeline_result_modify_response_carries_original_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"]}} +def test_handle_pipeline_result_allow_propagates_metadata_on_post_call(): + """The post_call path must keep the request payload intact (already sent + upstream) but still adopt metadata mutations from the pipeline's working + copy: standard_logging_guardrail_information and applied_guardrails are + written there and must reach deferred async logging plus the + x-litellm-applied-guardrails header.""" + replacement_response: Final = object() + data = {"a": 1, "metadata": {"guardrails": ["other"]}, "litellm_metadata": {"foo": "bar"}} result = MagicMock() result.terminal_action = "allow" - result.modified_data = {"metadata": {"guardrails": ["gr-post"]}, "response": object()} + result.modified_data = { + "metadata": {"guardrails": ["gr-post"], "applied_guardrails": ["gr-post"]}, + "litellm_metadata": {"foo": "bar", "extra": 1}, + "response": replacement_response, + } 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"]}} + assert data == { + "a": 1, + "metadata": {"guardrails": ["gr-post"], "applied_guardrails": ["gr-post"]}, + "litellm_metadata": {"foo": "bar", "extra": 1}, + } @pytest.mark.asyncio