From 4fcddb0bbd7edbbfbcbd18880ff21ddb10684010 Mon Sep 17 00:00:00 2001 From: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> Date: Mon, 26 Oct 2026 11:55:12 -0700 Subject: [PATCH 01/13] fix(guardrails): name only the pipeline steps that ran in x-litellm-applied-guardrails The pre_call header sweep claimed every policy pipeline step up front, so steps a pipeline never reached showed up in x-litellm-applied-guardrails and in spend logs. The sweep now leaves pipeline steps to the executor, which names each step when its hook runs, whatever it decides. Steps a stream runs after its headers go out (a pipeline the stream cannot hold, or one whose steps all stream live) are still named up front, as before. Resolves LIT-10096 --- litellm/integrations/custom_guardrail.py | 20 +- .../proxy/policy_engine/pipeline_executor.py | 48 ++- litellm/proxy/utils.py | 100 +++++- .../test_guardrail_pipeline_applied_header.py | 314 ++++++++++++++++++ .../policy_engine/test_pipeline_executor.py | 12 +- .../proxy_logging/test_guardrail_pipeline.py | 241 +++++++++++++- .../utils/proxy_logging/test_pre_call_hook.py | 2 +- 7 files changed, 693 insertions(+), 44 deletions(-) create mode 100644 tests/integration/observability/test_guardrail_pipeline_applied_header.py diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 9e2dc902d55..5e1d792718e 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -69,6 +69,7 @@ GUARDRAIL_SESSIONS_TARGET: Final = "guardrail_sessions" # field to suppress a guardrail on the direct-SDK path that never reaches the # proxy's metadata sanitizer. _PRE_CALL_EXECUTED_TOKEN: Final = secrets.token_hex(16) +_PRE_CALL_RAISED_TOKEN: Final = secrets.token_hex(16) _GUARDRAIL_BLOCK_STATUS_CODES: Final = frozenset({400, 403, 422}) @@ -824,13 +825,13 @@ class CustomGuardrail(CustomLogger): return False - def _pre_call_marker(self) -> str | None: + def _pre_call_marker(self, *, raised: bool = False) -> str | None: name: Final = self.guardrail_name if not name: return None - return f"{_PRE_CALL_EXECUTED_TOKEN}:{name}" + return f"{_PRE_CALL_RAISED_TOKEN if raised else _PRE_CALL_EXECUTED_TOKEN}:{name}" - def mark_pre_call_hook_ran(self, data: dict[str, object]) -> None: + def mark_pre_call_hook_ran(self, data: dict[str, object], *, raised: bool = False) -> None: """ Record that this guardrail's ``async_pre_call_hook`` already ran for this request, so the deployment-level hook does not run it a second time. @@ -839,8 +840,11 @@ class CustomGuardrail(CustomLogger): router later spreads a deployment's model-level ``guardrails`` into the top-level request kwargs, which would otherwise re-trigger the same hook from ``async_pre_call_deployment_hook``. + + With ``raised``, records instead that the hook raised as a policy pipeline + step the pipeline then moved past. The deployment-level hook still runs it. """ - marker: Final = self._pre_call_marker() + marker: Final = self._pre_call_marker(raised=raised) if marker is None: return for meta_key in ("metadata", "litellm_metadata"): @@ -855,8 +859,12 @@ class CustomGuardrail(CustomLogger): return data["metadata"] = {PRE_CALL_EXECUTED_GUARDRAILS_KEY: [marker]} - def _pre_call_hook_already_ran(self, data: dict[str, object]) -> bool: - marker: Final = self._pre_call_marker() + def pre_call_hook_ran_or_raised(self, data: Mapping[str, object]) -> bool: + """Whether this guardrail's ``async_pre_call_hook`` ran for this request, whatever it decided.""" + return self._pre_call_hook_already_ran(data) or self._pre_call_hook_already_ran(data, raised=True) + + def _pre_call_hook_already_ran(self, data: Mapping[str, object], *, raised: bool = False) -> bool: + marker: Final = self._pre_call_marker(raised=raised) if marker is None: return False for meta_key in ("metadata", "litellm_metadata"): diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 4ad29bd3478..7f804286282 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -329,6 +329,32 @@ def _prepare_hook_input( return hook_input, scans_raw_request +async def _run_pre_call_hook( + target: CustomLogger, + callback: CustomGuardrail, + data: dict[str, object], # mutable-ok: the request payload the run is recorded on + hook_input: dict[str, object], # mutable-ok: same request-payload shape the hooks mutate + user_api_key_dict: "UserAPIKeyAuth", + call_type: str, +) -> dict[str, object] | None: + try: + response: Final = await target.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=None, + data=hook_input, + call_type=call_type, + ) + except Exception: + callback.mark_pre_call_hook_ran(data, raised=True) + raise + callback.mark_pre_call_hook_ran(data) + if not isinstance(response, dict): + return None + payload: Final = cast(dict[str, object], response) # cast-ok: a pre_call hook returns the request payload + callback.mark_pre_call_hook_ran(payload) + return payload + + def _release_original_chunks( guardrail_name: str, reason: str, @@ -489,8 +515,8 @@ class PipelineExecutor: live, one that drops or adds a tool call, or one the translation or adapter refused with ``UndeliverableStreamRewrite``) is discarded: the buffered chunks go back to the originals and the step passes, so the - client gets the stream the merge base sent, and the guardrail stays out of the - applied-guardrails header since its output never reached the client. The response an + client gets the stream the merge base sent. The guardrail is named in the + applied-guardrails header either way, since it ran. The response an earlier step's translation stored under ``request_data["response"]`` is dropped first, so this step's hook sees the stream as the steps before it left it.""" scanner: Final = ( @@ -525,12 +551,12 @@ class PipelineExecutor: except UndeliverableStreamRewrite as undeliverable: _release_original_chunks(step.guardrail, undeliverable.reason, streaming_chunks, originals) return + finally: + if not callback.records_own_guardrail_information: + add_guardrail_to_applied_guardrails_header(request_data=hook_input, guardrail_name=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) - return - if not callback.records_own_guardrail_information: - add_guardrail_to_applied_guardrails_header(request_data=hook_input, guardrail_name=step.guardrail) @staticmethod async def _run_step( @@ -582,16 +608,7 @@ class PipelineExecutor: try: if mode == "pre_call": - response = await target.async_pre_call_hook( - user_api_key_dict=user_api_key_dict, - cache=None, - data=hook_input, - call_type=call_type, - ) - if isinstance(callback, CustomGuardrail): - callback.mark_pre_call_hook_ran(data) - if isinstance(response, dict): - callback.mark_pre_call_hook_ran(response) + response = await _run_pre_call_hook(target, callback, data, hook_input, user_api_key_dict, call_type) elif mode == "post_call" and streaming_chunks is not None: if endpoint_translation is None: return ( @@ -612,6 +629,7 @@ class PipelineExecutor: ) response = None elif mode == "post_call": + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=step.guardrail) response = await target.async_post_call_success_hook( user_api_key_dict=user_api_key_dict, data=data, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ce7e6afde24..9a960d30d3d 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -706,6 +706,62 @@ def pipeline_managed_guardrail_names( ) +def _ran_or_runs_outside_pipelines( + callback: CustomGuardrail | None, + data: Mapping[str, object], + pipeline_modes: frozenset[str], +) -> bool: + return callback is not None and ( + callback.pre_call_hook_ran_or_raised(data) + or any( + callback._event_hook_is_event_type(hook) # pyright: ignore[reportPrivateUsage] # the only check of a guardrail's hooks + for hook in GuardrailEventHooks + if hook.value not in pipeline_modes + ) + ) + + +# The A2A JSON-RPC methods answered with a stream, as both protocol versions spell them +_A2A_STREAMING_METHODS: Final = frozenset( + ("message/stream", "tasks/resubscribe", "SendStreamingMessage", "SubscribeToTask") +) + + +def _post_call_steps_named_before_headers( + data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth +) -> frozenset[str]: + """ + Post_call pipeline steps the executor names, as it runs them, before the response + headers go out: every step of a whole response, and of a streamed one the steps of + the pipelines that hold the stream until they finish. + """ + if data.get("stream") is not True and data.get("method") not in _A2A_STREAMING_METHODS: + return pipeline_managed_guardrail_names(data, "post_call") + gated: Final = _stream_gated_pipelines(data, user_api_key_dict) + return frozenset() if _pipelines_stream_live(gated) else _pipeline_step_guardrail_names(gated) + + +def _pipeline_steps_that_have_not_run(data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth) -> frozenset[str]: + """ + Guardrails only policy pipelines run whose pre_call hook has not run for this request: + pre_call steps a pipeline skipped and the post_call steps the executor names when it + runs them. A stream runs a pipeline it cannot hold as guardrails on their own, and one + whose steps all stream live after the headers go out, so those steps are named up front + like any other guardrail. + """ + pre_call: Final = pipeline_managed_guardrail_names(data, "pre_call") + post_call: Final = _post_call_steps_named_before_headers(data, user_api_key_dict) + return frozenset( + name + for name in pre_call | post_call + if not _ran_or_runs_outside_pipelines( + PipelineExecutor.find_guardrail_callback(name), + data, + frozenset(mode for mode, names in (("pre_call", pre_call), ("post_call", post_call)) if name in names), + ) + ) + + def _partition_post_call_callbacks() -> tuple[tuple[CustomGuardrail, ...], tuple[CustomLogger, ...]]: resolved: Final = tuple( litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class( @@ -927,19 +983,23 @@ def _streaming_pipeline_translation(user_api_key_dict: UserAPIKeyAuth) -> "BaseT return None if resolved is None else resolved[1] +def _stream_gated_pipelines( + request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth +) -> tuple[tuple[str, "GuardrailPipeline"], ...]: + translation: Final = _streaming_pipeline_translation(user_api_key_dict) + if translation is None: + return () + return tuple( + (policy_name, pipeline) + for policy_name, pipeline in _post_call_pipelines(request_data) + if not _pipeline_unsupported_streaming_guardrails(pipeline, translation) + ) + + def stream_gated_guardrail_names( request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth ) -> frozenset[str]: - translation: Final = _streaming_pipeline_translation(user_api_key_dict) - if translation is None: - return frozenset() - return _pipeline_step_guardrail_names( - tuple( - (policy_name, pipeline) - for policy_name, pipeline in _post_call_pipelines(request_data) - if not _pipeline_unsupported_streaming_guardrails(pipeline, translation) - ) - ) + return _pipeline_step_guardrail_names(_stream_gated_pipelines(request_data, user_api_key_dict)) def _streamable_post_call_pipelines( @@ -2639,7 +2699,7 @@ class ProxyLogging: data.pop("prompt_version", None) data.pop("prompt_environment", None) - def _process_guardrail_metadata(self, data: dict) -> None: + def _process_guardrail_metadata(self, data: dict, user_api_key_dict: UserAPIKeyAuth) -> None: """Process guardrails from metadata and add to applied_guardrails.""" from litellm.proxy.common_utils.callback_utils import ( add_guardrail_to_applied_guardrails_header, @@ -2664,8 +2724,16 @@ class ProxyLogging: if not isinstance(applied_guardrails, list): applied_guardrails = [] + not_run: Final = _pipeline_steps_that_have_not_run( + cast(dict[str, object], data), # cast-ok: request payload + user_api_key_dict, + ) for guardrail_name in guardrails_in_metadata: - if isinstance(guardrail_name, str) and guardrail_name not in applied_guardrails: + if ( + isinstance(guardrail_name, str) + and guardrail_name not in applied_guardrails + and guardrail_name not in not_run + ): add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=guardrail_name) async def _maybe_execute_pipelines( @@ -2952,7 +3020,7 @@ class ProxyLogging: and (guardrails_only or not caps.has_pre_call_override) ): if data is not None: - self._process_guardrail_metadata(data) + self._process_guardrail_metadata(data, user_api_key_dict) return data parallel_guardrails: Final[tuple[CustomGuardrail, ...]] = ( @@ -3056,17 +3124,17 @@ class ProxyLogging: ) if data is not None: - self._process_guardrail_metadata(data) + self._process_guardrail_metadata(data, user_api_key_dict) return data except SensitiveDataRouteException as e: data = await self._handle_sensitive_data_route_exception(e, data, user_api_key_dict) if data is not None: - self._process_guardrail_metadata(data) + self._process_guardrail_metadata(data, user_api_key_dict) return data except Exception: if data is not None: - self._process_guardrail_metadata(data) + self._process_guardrail_metadata(data, user_api_key_dict) raise async def _run_parallel_pre_call_guardrails( diff --git a/tests/integration/observability/test_guardrail_pipeline_applied_header.py b/tests/integration/observability/test_guardrail_pipeline_applied_header.py new file mode 100644 index 00000000000..5d08aa325e7 --- /dev/null +++ b/tests/integration/observability/test_guardrail_pipeline_applied_header.py @@ -0,0 +1,314 @@ +from __future__ import annotations + +import json +import uuid +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml + +from tests.integration._support.client import Gateway, gateway_from_environment +from tests.integration._support.process import owned_proxy +from tests.integration._support.wire import Reply, Request, Wire, wire_server + +_HEADER: Final = "x-litellm-applied-guardrails" +_VENDOR_PATH: Final = "/beta/litellm_basic_guardrail_api" +_PROMPT_SHIELD_PATH: Final = "/contentsafety/text:shieldPrompt" +_PROMPT_SHIELD: Final = "prompt-shield-step" +_REQUEST_FIRST: Final = "request-first-step" +_REQUEST_SECOND: Final = "request-second-step" +_REQUEST_BLOCKER: Final = "request-blocking-step" +_RESPONSE_FIRST: Final = "response-first-step" +_RESPONSE_SECOND: Final = "response-second-step" +_RESPONSE_BLOCKER: Final = "response-blocking-step" +_LIVE_FIRST: Final = "live-first-step" +_LIVE_SECOND: Final = "live-second-step" +_UNGATED_STEP: Final = "ungated-step" +_STOPS_AFTER_FIRST: Final = "stops-after-first" +_RUNS_EVERY_STEP: Final = "runs-every-step" +_BLOCKS_AT_FIRST: Final = "blocks-at-first" +_RESPONSE_STOPS_AFTER_FIRST: Final = "response-stops-after-first" +_SHIELDS_THEN_CHECKS: Final = "shields-then-checks" +_MOVES_PAST_FIRST: Final = "moves-past-first" +_RESPONSE_MOVES_PAST_FIRST: Final = "response-moves-past-first" +_SCANS_AFTER_RELEASE: Final = "scans-after-release" +_CANNOT_HOLD_STREAM: Final = "cannot-hold-stream" + + +@dataclass(frozen=True, slots=True) +class _Rig: + gateway: Gateway + vendor: Wire + upstream: Wire + model: str + + +def _completion_chunk(identity: str, delta: dict[str, str], finish_reason: str | None) -> bytes: + chunk: Final = { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}], + } + return f"data: {json.dumps(chunk)}\n\n".encode() + + +def _completion(request: Request) -> Reply: + assert request.target == "/v1/chat/completions", request.target + identity: Final = f"chatcmpl-{uuid.uuid4().hex}" + if json.loads(request.body).get("stream") is True: + return Reply( + content_type="text/event-stream", + chunks=( + _completion_chunk(identity, {"role": "assistant", "content": "guarded"}, None), + _completion_chunk(identity, {}, "stop"), + b"data: [DONE]\n\n", + ), + ) + return Reply( + body=json.dumps( + { + "id": identity, + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "guarded"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}, + } + ).encode() + ) + + +def _vendor(request: Request) -> Reply: + if request.target.startswith(f"/{_PROMPT_SHIELD}{_PROMPT_SHIELD_PATH}?"): + return Reply( + body=json.dumps({"userPromptAnalysis": {"attackDetected": False}, "documentsAnalysis": []}).encode() + ) + if request.target.startswith(f"/{_UNGATED_STEP}/"): + return Reply(body=json.dumps({"finalAction": "NONE", "appliedPolicies": []}).encode()) + assert request.target.endswith(_VENDOR_PATH), request.target + if request.target.startswith((f"/{_REQUEST_BLOCKER}/", f"/{_RESPONSE_BLOCKER}/")): + return Reply(body=json.dumps({"action": "BLOCKED", "blocked_reason": "first step refused"}).encode()) + return Reply(body=json.dumps({"action": "NONE"}).encode()) + + +def _guardrail(name: str, mode: str, vendor: Wire, **params: object) -> dict[str, object]: + return { + "guardrail_name": name, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": mode, + "default_on": False, + "api_base": f"{vendor.url}/{name}", + "api_key": "synthetic-guardrail-key", + **params, + }, + } + + +def _policy(mode: str, steps: tuple[tuple[str, str, str], ...]) -> dict[str, object]: + return { + "guardrails": {"add": [guardrail for guardrail, _, _ in steps]}, + "pipeline": { + "mode": mode, + "steps": [ + {"guardrail": guardrail, "on_pass": on_pass, "on_fail": on_fail} + for guardrail, on_pass, on_fail in steps + ], + }, + } + + +def _config(directory: Path, vendor: Wire) -> Path: + configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + configuration["guardrails"] = [ + *(_guardrail(name, "pre_call", vendor) for name in (_REQUEST_FIRST, _REQUEST_SECOND, _REQUEST_BLOCKER)), + *(_guardrail(name, "post_call", vendor) for name in (_RESPONSE_FIRST, _RESPONSE_SECOND, _RESPONSE_BLOCKER)), + *( + _guardrail(name, "post_call", vendor, streaming_buffer_until_moderated=False) + for name in (_LIVE_FIRST, _LIVE_SECOND) + ), + { + "guardrail_name": _UNGATED_STEP, + "litellm_params": { + "guardrail": "dynamoai", + "mode": "post_call", + "default_on": False, + "api_base": f"{vendor.url}/{_UNGATED_STEP}", + "api_key": "synthetic-guardrail-key", + }, + }, + { + "guardrail_name": _PROMPT_SHIELD, + "litellm_params": { + "guardrail": "azure/prompt_shield", + "mode": "pre_call", + "default_on": False, + "api_base": f"{vendor.url}/{_PROMPT_SHIELD}", + "api_key": "synthetic-guardrail-key", + }, + }, + ] + configuration["policies"] = { + _STOPS_AFTER_FIRST: _policy( + "pre_call", ((_REQUEST_FIRST, "allow", "block"), (_REQUEST_SECOND, "next", "block")) + ), + _RUNS_EVERY_STEP: _policy("pre_call", ((_REQUEST_FIRST, "next", "block"), (_REQUEST_SECOND, "next", "block"))), + _BLOCKS_AT_FIRST: _policy( + "pre_call", ((_REQUEST_BLOCKER, "next", "block"), (_REQUEST_SECOND, "next", "block")) + ), + _RESPONSE_STOPS_AFTER_FIRST: _policy( + "post_call", ((_RESPONSE_FIRST, "allow", "block"), (_RESPONSE_SECOND, "next", "block")) + ), + _SHIELDS_THEN_CHECKS: _policy( + "pre_call", ((_PROMPT_SHIELD, "next", "block"), (_REQUEST_FIRST, "next", "block")) + ), + _MOVES_PAST_FIRST: _policy( + "pre_call", ((_REQUEST_BLOCKER, "next", "next"), (_REQUEST_SECOND, "next", "block")) + ), + _RESPONSE_MOVES_PAST_FIRST: _policy( + "post_call", ((_RESPONSE_BLOCKER, "next", "next"), (_RESPONSE_SECOND, "next", "block")) + ), + _SCANS_AFTER_RELEASE: _policy("post_call", ((_LIVE_FIRST, "next", "next"), (_LIVE_SECOND, "next", "next"))), + _CANNOT_HOLD_STREAM: _policy( + "post_call", ((_UNGATED_STEP, "next", "block"), (_RESPONSE_SECOND, "next", "block")) + ), + } + configuration["policy_attachments"] = [{"policy": policy, "tags": [policy]} for policy in configuration["policies"]] + path: Final = directory / "guardrail-pipeline-applied-header.yaml" + path.write_text(yaml.safe_dump(configuration)) + return path + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: + directory: Final = tmp_path_factory.mktemp("guardrail-pipeline-applied-header") + with ( + wire_server(_vendor) as vendor, + wire_server(_completion) as upstream, + gateway_from_environment() as shared, + owned_proxy(shared, directory, {}, config=_config(directory, vendor)) as owned, + owned.scenario() as scenario, + ): + yield _Rig(owned, vendor, upstream, scenario.model(api_base=f"{upstream.url}/v1", api_key="sk-fixture")) + + +def _ask_as_team_tagged(rig: _Rig, tag: str, *, stream: bool = False) -> httpx.Response: + rig.vendor.drain() + rig.upstream.drain() + with rig.gateway.scenario() as scenario: + team: Final = scenario.team(metadata={"tags": [tag]}) + key: Final = scenario.key(team_id=team) + return rig.gateway.request( + "POST", + "/v1/chat/completions", + { + "model": rig.model, + "messages": [{"role": "user", "content": f"hello {uuid.uuid4().hex}"}], + "stream": stream, + }, + key=key, + ) + + +def _steps_that_ran(rig: _Rig) -> tuple[str, ...]: + return tuple(request.target.split("/")[1] for request in rig.vendor.drain()) + + +def test_a_request_pipeline_that_allows_at_its_first_step_names_only_that_step(rig: _Rig) -> None: + response: Final = _ask_as_team_tagged(rig, _STOPS_AFTER_FIRST) + assert response.status_code == 200, response.text + assert _steps_that_ran(rig) == (_REQUEST_FIRST,) + assert len(rig.upstream.drain()) == 1 + assert response.headers.get(_HEADER) == _REQUEST_FIRST, dict(response.headers) + + +def test_a_streamed_request_pipeline_that_allows_at_its_first_step_names_only_that_step(rig: _Rig) -> None: + response: Final = _ask_as_team_tagged(rig, _STOPS_AFTER_FIRST, stream=True) + assert response.status_code == 200, response.text + assert response.text.rstrip().endswith("data: [DONE]"), response.text + assert _steps_that_ran(rig) == (_REQUEST_FIRST,) + assert len(rig.upstream.drain()) == 1 + assert response.headers.get(_HEADER) == _REQUEST_FIRST, dict(response.headers) + + +def test_a_request_pipeline_that_runs_every_step_names_every_step_in_order(rig: _Rig) -> None: + response: Final = _ask_as_team_tagged(rig, _RUNS_EVERY_STEP) + assert response.status_code == 200, response.text + assert _steps_that_ran(rig) == (_REQUEST_FIRST, _REQUEST_SECOND) + assert response.headers.get(_HEADER) == f"{_REQUEST_FIRST},{_REQUEST_SECOND}", dict(response.headers) + + +def test_a_request_pipeline_names_a_prompt_shield_step_that_ran_before_a_generic_step(rig: _Rig) -> None: + response: Final = _ask_as_team_tagged(rig, _SHIELDS_THEN_CHECKS) + assert response.status_code == 200, response.text + assert _steps_that_ran(rig) == (_PROMPT_SHIELD, _REQUEST_FIRST) + assert set(response.headers.get(_HEADER, "").split(",")) == {_PROMPT_SHIELD, _REQUEST_FIRST}, dict(response.headers) + + +def test_a_request_pipeline_that_blocks_at_its_first_step_names_only_that_step(rig: _Rig) -> None: + response: Final = _ask_as_team_tagged(rig, _BLOCKS_AT_FIRST) + assert response.status_code == 400, response.text + assert _steps_that_ran(rig) == (_REQUEST_BLOCKER,) + assert rig.upstream.drain() == () + assert response.headers.get(_HEADER) == _REQUEST_BLOCKER, dict(response.headers) + + +def test_a_response_pipeline_that_allows_at_its_first_step_names_only_that_step(rig: _Rig) -> None: + response: Final = _ask_as_team_tagged(rig, _RESPONSE_STOPS_AFTER_FIRST) + assert response.status_code == 200, response.text + assert _steps_that_ran(rig) == (_RESPONSE_FIRST,) + assert response.headers.get(_HEADER) == _RESPONSE_FIRST, dict(response.headers) + + +def test_a_streamed_response_pipeline_that_allows_at_its_first_step_names_only_that_step(rig: _Rig) -> None: + response: Final = _ask_as_team_tagged(rig, _RESPONSE_STOPS_AFTER_FIRST, stream=True) + assert response.status_code == 200, response.text + assert response.text.rstrip().endswith("data: [DONE]"), response.text + assert _steps_that_ran(rig) == (_RESPONSE_FIRST,) + assert response.headers.get(_HEADER) == _RESPONSE_FIRST, dict(response.headers) + + +def test_a_request_pipeline_names_a_step_that_flagged_the_request_before_moving_on(rig: _Rig) -> None: + response: Final = _ask_as_team_tagged(rig, _MOVES_PAST_FIRST) + assert response.status_code == 200, response.text + assert _steps_that_ran(rig) == (_REQUEST_BLOCKER, _REQUEST_SECOND) + assert len(rig.upstream.drain()) == 1 + assert set(response.headers.get(_HEADER, "").split(",")) == {_REQUEST_BLOCKER, _REQUEST_SECOND}, dict( + response.headers + ) + + +@pytest.mark.parametrize("stream", [False, True], ids=["whole", "streamed"]) +def test_a_response_pipeline_names_a_step_that_flagged_the_response_before_moving_on(rig: _Rig, stream: bool) -> None: + response: Final = _ask_as_team_tagged(rig, _RESPONSE_MOVES_PAST_FIRST, stream=stream) + assert response.status_code == 200, response.text + assert _steps_that_ran(rig) == (_RESPONSE_BLOCKER, _RESPONSE_SECOND) + assert set(response.headers.get(_HEADER, "").split(",")) == {_RESPONSE_BLOCKER, _RESPONSE_SECOND}, dict( + response.headers + ) + + +def test_a_stream_released_before_its_pipeline_scans_it_names_the_steps_in_its_headers(rig: _Rig) -> None: + response: Final = _ask_as_team_tagged(rig, _SCANS_AFTER_RELEASE, stream=True) + assert response.status_code == 200, response.text + assert response.text.rstrip().endswith("data: [DONE]"), response.text + assert _steps_that_ran(rig) == (_LIVE_FIRST, _LIVE_SECOND) + assert set(response.headers.get(_HEADER, "").split(",")) == {_LIVE_FIRST, _LIVE_SECOND}, dict(response.headers) + + +def test_a_stream_its_pipeline_cannot_hold_names_the_steps_that_run_on_their_own(rig: _Rig) -> None: + response: Final = _ask_as_team_tagged(rig, _CANNOT_HOLD_STREAM, stream=True) + assert response.status_code == 200, response.text + assert response.text.rstrip().endswith("data: [DONE]"), response.text + assert _RESPONSE_SECOND in _steps_that_ran(rig) + assert set(response.headers.get(_HEADER, "").split(",")) == {_UNGATED_STEP, _RESPONSE_SECOND}, dict( + response.headers + ) diff --git a/tests/unit/proxy/policy_engine/test_pipeline_executor.py b/tests/unit/proxy/policy_engine/test_pipeline_executor.py index c42878bd4f7..eaee18c3929 100644 --- a/tests/unit/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/unit/proxy/policy_engine/test_pipeline_executor.py @@ -1270,7 +1270,7 @@ def _assert_passed_with_discard_warning(result, caplog, reason): assert len(discard_warnings) == 1 assert reason in discard_warnings[0] assert "text rewrites included" in discard_warnings[0] - assert "masker" not in ((result.modified_data or {}).get("metadata") or {}).get("applied_guardrails", []) + assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"] @pytest.mark.asyncio @@ -1433,6 +1433,16 @@ async def test_streaming_step_records_the_guardrail_in_the_applied_guardrails_he assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"] +@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()]) + + result = await _run_streaming_step(_WritingTranslation(), [_chunk()]) + + assert [(step.outcome, step.action_taken) for step in result.step_results] == [("fail", "next")] + assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"] + + @pytest.mark.asyncio async def test_streaming_step_records_guardrail_information_once_on_block(monkeypatch): monkeypatch.setattr(litellm, "callbacks", [_BlockingStreamGuardrail()]) diff --git a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py index acc1463b072..a96c515228f 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -30,12 +30,16 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.prometheus import PrometheusLogger from litellm.llms.base_llm.guardrail_translation.utils import stream_item_field from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header +from litellm.proxy.common_utils.callback_utils import ( + add_guardrail_to_applied_guardrails_header, + get_logging_caching_headers, +) from litellm.llms.openai.chat.guardrail_translation.handler import OpenAIChatCompletionsHandler from litellm.proxy.policy_engine.pipeline_executor import recorded_guardrail_information from litellm.proxy.utils import ProxyLogging, _streamable_post_call_pipelines, stream_gated_guardrail_names from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks +from litellm.types.utils import CallTypes from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.proxy.policy_engine.pipeline_types import ( GuardrailPipeline, @@ -275,7 +279,7 @@ def test_process_guardrail_metadata_calls_header_helper(proxy_logging, monkeypat monkeypatch.setattr(callback_utils, "add_guardrail_to_applied_guardrails_header", fake_add) data = {"metadata": {"guardrails": ["g1", "g2"]}} - proxy_logging._process_guardrail_metadata(data) + proxy_logging._process_guardrail_metadata(data, UserAPIKeyAuth()) snapshot = { "call_count": len(calls), "first_name": calls[0]["name"], @@ -300,7 +304,7 @@ def test_process_guardrail_metadata_skips_already_applied(proxy_logging, monkeyp monkeypatch.setattr(callback_utils, "add_guardrail_to_applied_guardrails_header", fake_add) data = {"metadata": {"guardrails": ["g1", "g2"], "applied_guardrails": ["g1"]}} - proxy_logging._process_guardrail_metadata(data) + proxy_logging._process_guardrail_metadata(data, UserAPIKeyAuth()) assert calls == ["g2"] @@ -312,12 +316,12 @@ def test_process_guardrail_metadata_no_metadata_is_noop(proxy_logging, monkeypat "add_guardrail_to_applied_guardrails_header", MagicMock(side_effect=AssertionError("should not be called")), ) - proxy_logging._process_guardrail_metadata({}) + proxy_logging._process_guardrail_metadata({}, UserAPIKeyAuth()) def test_process_guardrail_metadata_invalid_data_raises(proxy_logging): with pytest.raises(AttributeError): - proxy_logging._process_guardrail_metadata(None) # type: ignore[arg-type] + proxy_logging._process_guardrail_metadata(None, UserAPIKeyAuth()) # type: ignore[arg-type] # --------------------------------------------------------------------------- @@ -1160,6 +1164,233 @@ async def test_pre_call_hook_still_runs_guardrail_managed_only_by_post_call_pipe assert seen["count"] == 1 +class _PipelineStepGuardrail(CustomGuardrail): + def __init__(self, guardrail_name: str, event_hook: str | list[str]) -> None: + super().__init__(guardrail_name=guardrail_name, event_hook=event_hook, default_on=False) + self.calls = 0 + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + self.calls += 1 + return data + + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + self.calls += 1 + + +def _two_step_pipeline_data(mode: str, first_on_pass: str, first_on_fail: str = "block") -> dict[str, Any]: + pipeline = GuardrailPipeline( + mode=mode, + steps=[ + PipelineStep(guardrail="first", on_pass=first_on_pass, on_fail=first_on_fail), + PipelineStep(guardrail="second", on_pass="next", on_fail="block"), + ], + ) + return { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "metadata": { + "guardrails": ["first", "second"], + "_guardrail_pipelines": [("governance", pipeline)], + "_pipeline_managed_guardrails": {"first", "second"}, + }, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("first_on_pass", "expected_calls", "expected_header"), + [("allow", (1, 0), ["first"]), ("next", (1, 1), ["first", "second"])], +) +async def test_pre_call_hook_names_only_the_pre_call_pipeline_steps_that_ran( + proxy_logging, make_user_api_key_auth, monkeypatch, first_on_pass, expected_calls, expected_header +): + first = _PipelineStepGuardrail("first", "pre_call") + second = _PipelineStepGuardrail("second", "pre_call") + monkeypatch.setattr(litellm, "callbacks", [first, second]) + + result = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data=_two_step_pipeline_data("pre_call", first_on_pass), + call_type="completion", + ) + + assert (first.calls, second.calls) == expected_calls + assert result["metadata"]["applied_guardrails"] == expected_header + + +@pytest.mark.asyncio +async def test_pre_call_hook_names_a_skipped_step_whose_guardrail_also_runs_outside_the_pipeline( + proxy_logging, make_user_api_key_auth, monkeypatch +): + first = _PipelineStepGuardrail("first", "pre_call") + second = _PipelineStepGuardrail("second", ["pre_call", "post_call"]) + monkeypatch.setattr(litellm, "callbacks", [first, second]) + + result = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data=_two_step_pipeline_data("pre_call", "allow"), + call_type="completion", + ) + + assert (first.calls, second.calls) == (1, 0) + assert result["metadata"]["applied_guardrails"] == ["first", "second"] + + +class _RaisingPipelineStepGuardrail(_PipelineStepGuardrail): + def __init__(self, guardrail_name: str, error: Exception) -> None: + super().__init__(guardrail_name, "pre_call") + self.error = error + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + self.calls += 1 + raise self.error + + +_PIPELINE_STEP_ERRORS: Final = ( + pytest.param(HTTPException(status_code=400, detail="flagged"), id="flagged"), + pytest.param(RuntimeError("vendor unreachable"), id="errored"), +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("error", _PIPELINE_STEP_ERRORS) +async def test_pre_call_hook_names_a_step_that_raised_before_the_pipeline_moved_on( + proxy_logging, make_user_api_key_auth, monkeypatch, error +): + first = _RaisingPipelineStepGuardrail("first", error) + second = _PipelineStepGuardrail("second", "pre_call") + monkeypatch.setattr(litellm, "callbacks", [first, second]) + + result = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data=_two_step_pipeline_data("pre_call", "next", first_on_fail="next"), + call_type="completion", + ) + + assert (first.calls, second.calls) == (1, 1) + assert result["metadata"]["applied_guardrails"] == ["first", "second"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("error", _PIPELINE_STEP_ERRORS) +async def test_deployment_hook_still_runs_a_pipeline_step_that_raised( + proxy_logging, make_user_api_key_auth, monkeypatch, error +): + first = _RaisingPipelineStepGuardrail("first", error) + second = _PipelineStepGuardrail("second", "pre_call") + monkeypatch.setattr(litellm, "callbacks", [first, second]) + result = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data=_two_step_pipeline_data("pre_call", "next", first_on_fail="next"), + call_type="completion", + ) + + with pytest.raises(type(error)): + await first.async_pre_call_deployment_hook( + kwargs={**result, "guardrails": ["first"]}, call_type=CallTypes.completion + ) + await second.async_pre_call_deployment_hook( + kwargs={**result, "guardrails": ["second"]}, call_type=CallTypes.completion + ) + + assert (first.calls, second.calls) == (2, 1) + + +class _OpensLitellmMetadataGuardrail(_PipelineStepGuardrail): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + self.calls += 1 + data["litellm_metadata"] = {} + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) + return data + + +@pytest.mark.asyncio +async def test_pre_call_hook_names_a_step_that_ran_in_the_litellm_metadata_a_later_step_opened( + proxy_logging, make_user_api_key_auth, monkeypatch +): + first = _PipelineStepGuardrail("first", "pre_call") + second = _OpensLitellmMetadataGuardrail("second", "pre_call") + monkeypatch.setattr(litellm, "callbacks", [first, second]) + + result = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data=_two_step_pipeline_data("pre_call", "next"), + call_type="completion", + ) + + assert (first.calls, second.calls) == (1, 1) + assert "applied_guardrails" not in result["metadata"] + assert result["litellm_metadata"]["applied_guardrails"] == ["second", "first"] + assert set(get_logging_caching_headers(result)["x-litellm-applied-guardrails"].split(",")) == {"first", "second"} + + +@pytest.mark.asyncio +async def test_post_call_pipeline_steps_are_named_only_once_they_run( + proxy_logging, make_user_api_key_auth, monkeypatch +): + first = _PipelineStepGuardrail("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", "allow"), + call_type="completion", + ) + assert "applied_guardrails" not in data["metadata"] + + 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( + ("asks_for", "request_route", "streams_live", "expected_header"), + [ + pytest.param({"stream": False}, "/custom/stream", False, None, id="whole-response"), + pytest.param({"stream": True}, "/v1/chat/completions", False, None, id="stream-held-for-the-pipeline"), + pytest.param({"stream": True}, "/v1/chat/completions", True, ["first", "second"], id="stream-released-live"), + pytest.param( + {"stream": True}, "/custom/stream", False, ["first", "second"], id="stream-runs-steps-on-their-own" + ), + pytest.param({"method": "message/send"}, "/a2a/agent", False, None, id="a2a-whole-response"), + pytest.param({"method": "message/stream"}, "/a2a/agent", False, ["first", "second"], id="a2a-stream"), + pytest.param({"method": "SubscribeToTask"}, "/a2a/agent", False, ["first", "second"], id="a2a-v1-resubscribe"), + ], +) +async def test_pre_call_hook_names_up_front_the_post_call_steps_a_stream_runs_after_its_headers( + proxy_logging, make_user_api_key_auth, monkeypatch, asks_for, request_route, streams_live, expected_header +): + guardrails = [_PipelineStepGuardrail(name, "post_call") for name in ("first", "second")] + for guardrail in guardrails: + guardrail.streaming_buffer_until_moderated = not streams_live + monkeypatch.setattr(litellm, "callbacks", guardrails) + data = _two_step_pipeline_data("post_call", "next") + data["metadata"]["_guardrail_pipelines"] = [ + ( + "governance", + GuardrailPipeline( + mode="post_call", + steps=[PipelineStep(guardrail=name, on_pass="next", on_fail="next") for name in ("first", "second")], + ), + ) + ] + data.update(asks_for) + + result = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(request_route=request_route), + data=data, + call_type="completion", + ) + + assert result["metadata"].get("applied_guardrails") == expected_header + + @pytest.mark.asyncio async def test_post_call_pipeline_replacement_response_reaches_caller( proxy_logging, make_user_api_key_auth, monkeypatch diff --git a/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py b/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py index b7fc8c6dbeb..615e39c495d 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py +++ b/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py @@ -170,7 +170,7 @@ async def test_pre_call_hook_processes_guardrail_metadata_when_no_overrides(prox proxy_logging.slack_alerting_instance = MagicMock(alerting=None) invoked = {} - def fake_process(d): + def fake_process(d, user_api_key_dict): invoked["data"] = d proxy_logging._process_guardrail_metadata = fake_process # type: ignore[assignment] From 043593cf7bbc7c3db047b4d57e1d4bd471f772fb Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 26 Oct 2026 15:26:53 -0700 Subject: [PATCH 02/13] test(guardrails): pin applied guardrails on every endpoint, client, spend log and trace Covers /v1/responses and /v1/messages, the OpenAI and Anthropic SDKs sync and async, the spend log and OTel span, a step whose service errored, forged names and a concurrent burst across two workers --- .../test_guardrail_pipeline_applied_header.py | 381 +++++++++++++++++- 1 file changed, 362 insertions(+), 19 deletions(-) diff --git a/tests/integration/observability/test_guardrail_pipeline_applied_header.py b/tests/integration/observability/test_guardrail_pipeline_applied_header.py index 5d08aa325e7..0ac8d0342f9 100644 --- a/tests/integration/observability/test_guardrail_pipeline_applied_header.py +++ b/tests/integration/observability/test_guardrail_pipeline_applied_header.py @@ -1,17 +1,24 @@ from __future__ import annotations +import ast +import asyncio +import itertools import json import uuid -from collections.abc import Iterator +from collections.abc import Callable, Iterator, Mapping from dataclasses import dataclass from pathlib import Path +from types import MappingProxyType from typing import Final +import anthropic import httpx +import openai import pytest import yaml -from tests.integration._support.client import Gateway, gateway_from_environment +from tests.integration._support.client import Gateway, eventually, gateway_from_environment, object_value +from tests.integration._support.otlp_sink import Span, SpanSinks, recorded_spans from tests.integration._support.process import owned_proxy from tests.integration._support.wire import Reply, Request, Wire, wire_server @@ -22,6 +29,7 @@ _PROMPT_SHIELD: Final = "prompt-shield-step" _REQUEST_FIRST: Final = "request-first-step" _REQUEST_SECOND: Final = "request-second-step" _REQUEST_BLOCKER: Final = "request-blocking-step" +_REQUEST_ERRORING: Final = "request-erroring-step" _RESPONSE_FIRST: Final = "response-first-step" _RESPONSE_SECOND: Final = "response-second-step" _RESPONSE_BLOCKER: Final = "response-blocking-step" @@ -37,6 +45,8 @@ _MOVES_PAST_FIRST: Final = "moves-past-first" _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" +_ENDPOINTS: Final = ("/v1/chat/completions", "/v1/responses", "/v1/messages") @dataclass(frozen=True, slots=True) @@ -45,6 +55,7 @@ class _Rig: vendor: Wire upstream: Wire model: str + traces: str def _completion_chunk(identity: str, delta: dict[str, str], finish_reason: str | None) -> bytes: @@ -58,7 +69,63 @@ def _completion_chunk(identity: str, delta: dict[str, str], finish_reason: str | return f"data: {json.dumps(chunk)}\n\n".encode() +def _event(event: dict[str, object]) -> bytes: + return f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() + + +def _response(identity: str) -> dict[str, object]: + return { + "id": identity, + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [ + { + "id": f"msg_{identity}", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "guarded", "annotations": []}], + } + ], + "usage": {"input_tokens": 3, "output_tokens": 1, "total_tokens": 4}, + } + + +def _responses_reply(request: Request) -> Reply: + identity: Final = f"resp_{uuid.uuid4().hex}" + response: Final = _response(identity) + if json.loads(request.body).get("stream") is not True: + return Reply(body=json.dumps(response).encode()) + return Reply( + content_type="text/event-stream", + chunks=( + _event( + { + "type": "response.created", + "sequence_number": 0, + "response": {**response, "status": "in_progress", "output": []}, + } + ), + _event( + { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": f"msg_{identity}", + "output_index": 0, + "content_index": 0, + "delta": "guarded", + } + ), + _event({"type": "response.completed", "sequence_number": 2, "response": response}), + ), + ) + + def _completion(request: Request) -> Reply: + if request.target == "/v1/responses": + return _responses_reply(request) assert request.target == "/v1/chat/completions", request.target identity: Final = f"chatcmpl-{uuid.uuid4().hex}" if json.loads(request.body).get("stream") is True: @@ -94,6 +161,8 @@ def _vendor(request: Request) -> Reply: if request.target.startswith(f"/{_UNGATED_STEP}/"): return Reply(body=json.dumps({"finalAction": "NONE", "appliedPolicies": []}).encode()) assert request.target.endswith(_VENDOR_PATH), request.target + if request.target.startswith(f"/{_REQUEST_ERRORING}/"): + return Reply(status=503, body=json.dumps({"error": "guardrail service unavailable"}).encode()) if request.target.startswith((f"/{_REQUEST_BLOCKER}/", f"/{_RESPONSE_BLOCKER}/")): return Reply(body=json.dumps({"action": "BLOCKED", "blocked_reason": "first step refused"}).encode()) return Reply(body=json.dumps({"action": "NONE"}).encode()) @@ -113,13 +182,13 @@ def _guardrail(name: str, mode: str, vendor: Wire, **params: object) -> dict[str } -def _policy(mode: str, steps: tuple[tuple[str, str, str], ...]) -> dict[str, object]: +def _policy(mode: str, steps: tuple[tuple[str, str, str], ...], *, on_error: str | None = None) -> dict[str, object]: return { "guardrails": {"add": [guardrail for guardrail, _, _ in steps]}, "pipeline": { "mode": mode, "steps": [ - {"guardrail": guardrail, "on_pass": on_pass, "on_fail": on_fail} + {"guardrail": guardrail, "on_pass": on_pass, "on_fail": on_fail, "on_error": on_error} for guardrail, on_pass, on_fail in steps ], }, @@ -128,8 +197,12 @@ def _policy(mode: str, steps: tuple[tuple[str, str, str], ...]) -> dict[str, obj 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["guardrails"] = [ - *(_guardrail(name, "pre_call", vendor) for name in (_REQUEST_FIRST, _REQUEST_SECOND, _REQUEST_BLOCKER)), + *( + _guardrail(name, "pre_call", vendor) + for name in (_REQUEST_FIRST, _REQUEST_SECOND, _REQUEST_BLOCKER, _REQUEST_ERRORING) + ), *(_guardrail(name, "post_call", vendor) for name in (_RESPONSE_FIRST, _RESPONSE_SECOND, _RESPONSE_BLOCKER)), *( _guardrail(name, "post_call", vendor, streaming_buffer_until_moderated=False) @@ -180,6 +253,9 @@ def _config(directory: Path, vendor: Wire) -> Path: _CANNOT_HOLD_STREAM: _policy( "post_call", ((_UNGATED_STEP, "next", "block"), (_RESPONSE_SECOND, "next", "block")) ), + _ERRORS_PAST_FIRST: _policy( + "pre_call", ((_REQUEST_ERRORING, "next", "block"), (_REQUEST_SECOND, "next", "block")), on_error="next" + ), } configuration["policy_attachments"] = [{"policy": policy, "tags": [policy]} for policy in configuration["policies"]] path: Final = directory / "guardrail-pipeline-applied-header.yaml" @@ -188,34 +264,60 @@ def _config(directory: Path, vendor: Wire) -> Path: @pytest.fixture(scope="module") -def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: +def rig(tmp_path_factory: pytest.TempPathFactory, audit_sinks: SpanSinks) -> Iterator[_Rig]: directory: Final = tmp_path_factory.mktemp("guardrail-pipeline-applied-header") with ( wire_server(_vendor) as vendor, wire_server(_completion) as upstream, gateway_from_environment() as shared, - owned_proxy(shared, directory, {}, config=_config(directory, vendor)) as owned, + owned_proxy( + shared, + directory, + {"OTEL_EXPORTER": "http/json", "OTEL_ENDPOINT": audit_sinks.operator, "OTEL_BSP_SCHEDULE_DELAY": "100"}, + config=_config(directory, vendor), + workers=2, + ) as owned, owned.scenario() as scenario, ): - yield _Rig(owned, vendor, upstream, scenario.model(api_base=f"{upstream.url}/v1", api_key="sk-fixture")) + model: Final = scenario.model(api_base=f"{upstream.url}/v1", api_key="sk-fixture") + yield _Rig(owned, vendor, upstream, model, audit_sinks.operator) -def _ask_as_team_tagged(rig: _Rig, tag: str, *, stream: bool = False) -> httpx.Response: +def _body(path: str, model: str, stream: bool) -> dict[str, object]: + prompt: Final = f"hello {uuid.uuid4().hex}" + if path == "/v1/responses": + return {"model": model, "input": prompt, "stream": stream} + messages: Final = [{"role": "user", "content": prompt}] + if path == "/v1/messages": + return {"model": model, "max_tokens": 20, "messages": messages, "stream": stream} + return {"model": model, "messages": messages, "stream": stream} + + +def _ask_as_new_team( + rig: _Rig, + tag: str, + *, + stream: bool = False, + path: str = "/v1/chat/completions", + extra: Mapping[str, object] = MappingProxyType({}), +) -> tuple[str, httpx.Response]: rig.vendor.drain() rig.upstream.drain() with rig.gateway.scenario() as scenario: team: Final = scenario.team(metadata={"tags": [tag]}) key: Final = scenario.key(team_id=team) - return rig.gateway.request( - "POST", - "/v1/chat/completions", - { - "model": rig.model, - "messages": [{"role": "user", "content": f"hello {uuid.uuid4().hex}"}], - "stream": stream, - }, - key=key, - ) + return team, rig.gateway.request("POST", path, {**_body(path, rig.model, stream), **extra}, key=key) + + +def _ask_as_team_tagged( + rig: _Rig, + tag: str, + *, + stream: bool = False, + path: str = "/v1/chat/completions", + extra: Mapping[str, object] = MappingProxyType({}), +) -> httpx.Response: + return _ask_as_new_team(rig, tag, stream=stream, path=path, extra=extra)[1] def _steps_that_ran(rig: _Rig) -> tuple[str, ...]: @@ -312,3 +414,244 @@ def test_a_stream_its_pipeline_cannot_hold_names_the_steps_that_run_on_their_own assert set(response.headers.get(_HEADER, "").split(",")) == {_UNGATED_STEP, _RESPONSE_SECOND}, dict( response.headers ) + + +@pytest.mark.parametrize("stream", [False, True], ids=["whole", "streamed"]) +@pytest.mark.parametrize("path", ["/v1/responses", "/v1/messages"]) +def test_a_request_pipeline_that_allows_at_its_first_step_names_only_that_step_on_every_endpoint( + rig: _Rig, path: str, stream: bool +) -> None: + response: Final = _ask_as_team_tagged(rig, _STOPS_AFTER_FIRST, stream=stream, path=path) + assert response.status_code == 200, response.text + assert _steps_that_ran(rig) == (_REQUEST_FIRST,) + assert len(rig.upstream.drain()) == 1 + assert response.headers.get(_HEADER) == _REQUEST_FIRST, dict(response.headers) + + +@pytest.mark.parametrize("stream", [False, True], ids=["whole", "streamed"]) +@pytest.mark.parametrize("path", ["/v1/responses", "/v1/messages"]) +def test_a_response_pipeline_that_allows_at_its_first_step_names_only_that_step_on_every_endpoint( + rig: _Rig, path: str, stream: bool +) -> None: + response: Final = _ask_as_team_tagged(rig, _RESPONSE_STOPS_AFTER_FIRST, stream=stream, path=path) + assert response.status_code == 200, response.text + assert _steps_that_ran(rig) == (_RESPONSE_FIRST,) + assert response.headers.get(_HEADER) == _RESPONSE_FIRST, dict(response.headers) + + +def test_a_request_pipeline_names_a_step_whose_service_errored_before_moving_on(rig: _Rig) -> None: + response: Final = _ask_as_team_tagged(rig, _ERRORS_PAST_FIRST) + assert response.status_code == 200, response.text + assert _steps_that_ran(rig) == (_REQUEST_ERRORING, _REQUEST_SECOND) + assert len(rig.upstream.drain()) == 1 + assert set(response.headers.get(_HEADER, "").split(",")) == {_REQUEST_ERRORING, _REQUEST_SECOND}, dict( + response.headers + ) + + +@pytest.mark.parametrize( + "forged", + [ + pytest.param( + { + "metadata": { + "applied_guardrails": [_REQUEST_SECOND], + "_pre_call_executed_guardrails": [f"deadbeef:{_REQUEST_SECOND}", _REQUEST_SECOND], + } + }, + id="names-in-its-metadata", + ), + pytest.param({"guardrails": [_REQUEST_SECOND]}, id="asks-for-the-skipped-step"), + ], +) +def test_a_request_cannot_get_a_step_that_never_ran_named(rig: _Rig, forged: Mapping[str, object]) -> None: + response: Final = _ask_as_team_tagged(rig, _STOPS_AFTER_FIRST, extra=forged) + assert response.status_code == 200, response.text + assert _steps_that_ran(rig) == (_REQUEST_FIRST,) + assert response.headers.get(_HEADER) == _REQUEST_FIRST, dict(response.headers) + + +def _spans_naming_guardrails(rig: _Rig, since: int, team: str) -> tuple[Span, ...]: + spans: Final = recorded_spans(rig.traces, since)[1] + return tuple( + span + for span in spans + if span["attributes"].get("metadata.user_api_key_team_id") == team + and "metadata.applied_guardrails" in span["attributes"] + ) + + +@pytest.mark.parametrize( + ("tag", "path", "stream", "expected"), + [ + pytest.param(_STOPS_AFTER_FIRST, "/v1/chat/completions", False, [_REQUEST_FIRST], id="request-pipeline"), + pytest.param(_STOPS_AFTER_FIRST, "/v1/chat/completions", True, [_REQUEST_FIRST], id="request-pipeline-stream"), + pytest.param(_STOPS_AFTER_FIRST, "/v1/responses", True, [_REQUEST_FIRST], id="request-pipeline-responses"), + pytest.param( + _RESPONSE_STOPS_AFTER_FIRST, "/v1/chat/completions", False, [_RESPONSE_FIRST], id="response-pipeline" + ), + pytest.param( + _RESPONSE_STOPS_AFTER_FIRST, "/v1/chat/completions", True, [_RESPONSE_FIRST], id="response-pipeline-stream" + ), + pytest.param( + _RESPONSE_STOPS_AFTER_FIRST, "/v1/messages", True, [_RESPONSE_FIRST], id="response-pipeline-messages" + ), + ], +) +def test_the_spend_log_and_the_trace_name_only_the_steps_that_ran( + rig: _Rig, tag: str, path: str, stream: bool, expected: list[str] +) -> None: + since: Final = recorded_spans(rig.traces)[0] + team, response = _ask_as_new_team(rig, tag, stream=stream, path=path) + assert response.status_code == 200, response.text + 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 object_value(rows[0]["metadata"]).get("applied_guardrails") == expected, rows + assert [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, + openai.OpenAI(api_key=key, base_url=f"{base}/v1", max_retries=0, http_client=transport) as client, + ): + response: Final = client.chat.completions.with_raw_response.create( + model=model, messages=[{"role": "user", "content": f"hello {uuid.uuid4().hex}"}] + ) + return response.headers.get(_HEADER) + + +def _openai_responses(base: str, key: str, model: str) -> str | None: + with ( + httpx.Client(timeout=30, trust_env=False) as transport, + openai.OpenAI(api_key=key, base_url=f"{base}/v1", max_retries=0, http_client=transport) as client, + ): + return client.responses.with_raw_response.create(model=model, input=f"hello {uuid.uuid4().hex}").headers.get( + _HEADER + ) + + +def _async_openai_streamed_chat(base: str, key: str, model: str) -> str | None: + async def run() -> str | None: + async with ( + httpx.AsyncClient(timeout=30, trust_env=False) as transport, + openai.AsyncOpenAI(api_key=key, base_url=f"{base}/v1", max_retries=0, http_client=transport) as client, + client.chat.completions.with_streaming_response.create( + model=model, messages=[{"role": "user", "content": f"hello {uuid.uuid4().hex}"}], stream=True + ) as response, + ): + lines: Final = [line async for line in response.iter_lines()] + assert "data: [DONE]" in lines, lines + return response.headers.get(_HEADER) + + return asyncio.run(run()) + + +def _async_openai_streamed_responses(base: str, key: str, model: str) -> str | None: + async def run() -> str | None: + async with ( + httpx.AsyncClient(timeout=30, trust_env=False) as transport, + openai.AsyncOpenAI(api_key=key, base_url=f"{base}/v1", max_retries=0, http_client=transport) as client, + client.responses.with_streaming_response.create( + model=model, input=f"hello {uuid.uuid4().hex}", stream=True + ) as response, + ): + lines: Final = [line async for line in response.iter_lines()] + assert any("response.completed" in line for line in lines), lines + return response.headers.get(_HEADER) + + return asyncio.run(run()) + + +def _anthropic_streamed_messages(base: str, key: str, model: str) -> str | None: + with ( + httpx.Client(timeout=30, trust_env=False) as transport, + anthropic.Anthropic(api_key=key, base_url=base, max_retries=0, http_client=transport) as client, + client.messages.with_streaming_response.create( + model=model, + max_tokens=20, + messages=[{"role": "user", "content": f"hello {uuid.uuid4().hex}"}], + stream=True, + ) as response, + ): + lines: Final = tuple(response.iter_lines()) + assert any("message_stop" in line for line in lines), lines + return response.headers.get(_HEADER) + + +def _async_anthropic_messages(base: str, key: str, model: str) -> str | None: + async def run() -> str | None: + async with ( + httpx.AsyncClient(timeout=30, trust_env=False) as transport, + anthropic.AsyncAnthropic(api_key=key, base_url=base, max_retries=0, http_client=transport) as client, + ): + response: Final = await client.messages.with_raw_response.create( + model=model, max_tokens=20, messages=[{"role": "user", "content": f"hello {uuid.uuid4().hex}"}] + ) + return response.headers.get(_HEADER) + + return asyncio.run(run()) + + +@pytest.mark.parametrize( + "client", + [ + pytest.param(_openai_chat, id="openai-chat"), + pytest.param(_async_openai_streamed_chat, id="async-openai-chat-stream"), + pytest.param(_openai_responses, id="openai-responses"), + pytest.param(_async_openai_streamed_responses, id="async-openai-responses-stream"), + pytest.param(_anthropic_streamed_messages, id="anthropic-messages-stream"), + pytest.param(_async_anthropic_messages, id="async-anthropic-messages"), + ], +) +def test_every_sdk_client_sees_only_the_step_that_ran(rig: _Rig, client: Callable[[str, str, str], str | None]) -> None: + rig.vendor.drain() + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(team_id=scenario.team(metadata={"tags": [_STOPS_AFTER_FIRST]})) + header: Final = client(str(rig.gateway.client.base_url).rstrip("/"), key, rig.model) + assert _steps_that_ran(rig) == (_REQUEST_FIRST,) + assert header == _REQUEST_FIRST + + +async def _send_concurrently( + rig: _Rig, keys: Mapping[str, str], cases: tuple[tuple[str, str, bool], ...] +) -> tuple[httpx.Response, ...]: + async with httpx.AsyncClient(base_url=rig.gateway.client.base_url, timeout=60, trust_env=False) as client: + return tuple( + await asyncio.gather( + *( + client.post( + path, json=_body(path, rig.model, stream), headers={"Authorization": f"Bearer {keys[tag]}"} + ) + for tag, path, stream in cases + ) + ) + ) + + +def test_a_concurrent_burst_names_only_each_requests_own_steps(rig: _Rig) -> None: + expected: Final = MappingProxyType( + { + _STOPS_AFTER_FIRST: _REQUEST_FIRST, + _RUNS_EVERY_STEP: f"{_REQUEST_FIRST},{_REQUEST_SECOND}", + _RESPONSE_STOPS_AFTER_FIRST: _RESPONSE_FIRST, + } + ) + cases: Final = tuple(itertools.product(expected, _ENDPOINTS, (False, True))) * 2 + with rig.gateway.scenario() as scenario: + keys: Final = MappingProxyType( + {tag: scenario.key(team_id=scenario.team(metadata={"tags": [tag]})) for tag in expected} + ) + responses: Final = asyncio.run(_send_concurrently(rig, keys, cases)) + rig.vendor.drain() + rig.upstream.drain() + assert [ + (tag, path, stream, response.status_code, response.headers.get(_HEADER)) + for (tag, path, stream), response in zip(cases, responses, strict=True) + ] == [(tag, path, stream, 200, expected[tag]) for tag, path, stream in cases] From 0c91ebf606dc8d5f305ec9250a32f3b584d29210 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 26 Oct 2026 17:25:07 -0700 Subject: [PATCH 03/13] 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 --- .../proxy/policy_engine/pipeline_executor.py | 31 ++++++- .../test_guardrail_pipeline_applied_header.py | 84 +++++++++++++++++++ .../policy_engine/test_pipeline_executor.py | 22 +++++ .../proxy_logging/test_guardrail_pipeline.py | 28 +++++++ 4 files changed, 163 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 7f804286282..c5213f4c138 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/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, diff --git a/tests/integration/observability/test_guardrail_pipeline_applied_header.py b/tests/integration/observability/test_guardrail_pipeline_applied_header.py index 0ac8d0342f9..e0fb2905a06 100644 --- a/tests/integration/observability/test_guardrail_pipeline_applied_header.py +++ b/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, diff --git a/tests/unit/proxy/policy_engine/test_pipeline_executor.py b/tests/unit/proxy/policy_engine/test_pipeline_executor.py index eaee18c3929..e981b7bdeb0 100644 --- a/tests/unit/proxy/policy_engine/test_pipeline_executor.py +++ b/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()]) diff --git a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py index a96c515228f..c8d19c465f8 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/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"), From 0eb21e4b403d067b5db773e348767c5966be516e Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 26 Oct 2026 17:33:00 -0700 Subject: [PATCH 04/13] 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 --- .../proxy/policy_engine/pipeline_executor.py | 3 +- .../test_guardrail_pipeline_applied_header.py | 56 +++++++++++++++++++ .../policy_engine/test_pipeline_executor.py | 13 +++++ 3 files changed, 70 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index c5213f4c138..c8d5e4c3b48 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/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) diff --git a/tests/integration/observability/test_guardrail_pipeline_applied_header.py b/tests/integration/observability/test_guardrail_pipeline_applied_header.py index e0fb2905a06..58e46141cac 100644 --- a/tests/integration/observability/test_guardrail_pipeline_applied_header.py +++ b/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", diff --git a/tests/unit/proxy/policy_engine/test_pipeline_executor.py b/tests/unit/proxy/policy_engine/test_pipeline_executor.py index e981b7bdeb0..068102aac08 100644 --- a/tests/unit/proxy/policy_engine/test_pipeline_executor.py +++ b/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.""" From 5bd75f813e72b5a634b3e169f974aae8242567b8 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 26 Oct 2026 17:44:03 -0700 Subject: [PATCH 05/13] test(guardrails): type the pipeline step test helpers --- .../policy_engine/test_pipeline_executor.py | 13 ++++-- .../proxy_logging/test_guardrail_pipeline.py | 40 ++++++++++--------- 2 files changed, 32 insertions(+), 21 deletions(-) diff --git a/tests/unit/proxy/policy_engine/test_pipeline_executor.py b/tests/unit/proxy/policy_engine/test_pipeline_executor.py index 068102aac08..bd8b9eb24d7 100644 --- a/tests/unit/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/unit/proxy/policy_engine/test_pipeline_executor.py @@ -15,6 +15,7 @@ import pytest import litellm from litellm.caching.dual_cache import DualCache from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import ( @@ -30,7 +31,7 @@ from litellm.types.proxy.policy_engine.pipeline_types import ( GuardrailPipeline, PipelineStep, ) -from litellm.types.utils import CallTypesLiteral +from litellm.types.utils import CallTypesLiteral, GenericGuardrailAPIInputs try: from fastapi.exceptions import HTTPException @@ -1450,10 +1451,16 @@ async def test_streaming_step_whose_guardrail_records_its_own_results_is_named(m class _LitellmMetadataOpeningGuardrail(CustomGuardrail): """Opens ``litellm_metadata`` on the request while scanning, as the A2A translation does.""" - def __init__(self): + def __init__(self) -> None: 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): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: request_data["litellm_metadata"] = {"user_api_key_hash": "hashed"} return inputs diff --git a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py index c8d19c465f8..29438c63a56 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -39,7 +39,8 @@ from litellm.proxy.policy_engine.pipeline_executor import recorded_guardrail_inf from litellm.proxy.utils import ProxyLogging, _streamable_post_call_pipelines, stream_gated_guardrail_names from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks -from litellm.types.utils import CallTypes +from litellm.caching.dual_cache import DualCache +from litellm.types.utils import CallTypes, CallTypesLiteral, LLMResponseTypes from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.proxy.policy_engine.pipeline_types import ( GuardrailPipeline, @@ -1169,20 +1170,26 @@ class _PipelineStepGuardrail(CustomGuardrail): super().__init__(guardrail_name=guardrail_name, event_hook=event_hook, default_on=False) self.calls = 0 - async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + async def async_pre_call_hook( + self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral + ) -> dict: self.calls += 1 return data - async def async_post_call_success_hook(self, data, user_api_key_dict, response): + async def async_post_call_success_hook( + self, data: dict, user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes + ) -> None: self.calls += 1 -def _two_step_pipeline_data(mode: str, first_on_pass: str, first_on_fail: str = "block") -> dict[str, Any]: +def _two_step_pipeline_data( + mode: str, first_on_pass: str, first_on_fail: str = "block", second_on_fail: str = "block" +) -> dict[str, object]: pipeline = GuardrailPipeline( mode=mode, steps=[ PipelineStep(guardrail="first", on_pass=first_on_pass, on_fail=first_on_fail), - PipelineStep(guardrail="second", on_pass="next", on_fail="block"), + PipelineStep(guardrail="second", on_pass="next", on_fail=second_on_fail), ], ) return { @@ -1241,7 +1248,9 @@ class _RaisingPipelineStepGuardrail(_PipelineStepGuardrail): super().__init__(guardrail_name, "pre_call") self.error = error - async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + async def async_pre_call_hook( + self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral + ) -> dict: self.calls += 1 raise self.error @@ -1297,7 +1306,9 @@ async def test_deployment_hook_still_runs_a_pipeline_step_that_raised( class _OpensLitellmMetadataGuardrail(_PipelineStepGuardrail): - async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + async def async_pre_call_hook( + self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral + ) -> dict: self.calls += 1 data["litellm_metadata"] = {} add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) @@ -1349,7 +1360,9 @@ async def test_post_call_pipeline_steps_are_named_only_once_they_run( class _OpensLitellmMetadataOnResponseGuardrail(_PipelineStepGuardrail): - async def async_post_call_success_hook(self, data, user_api_key_dict, response): + async def async_post_call_success_hook( + self, data: dict, user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes + ) -> None: self.calls += 1 data["litellm_metadata"] = {"user_api_key_hash": "hashed"} @@ -1398,16 +1411,7 @@ async def test_pre_call_hook_names_up_front_the_post_call_steps_a_stream_runs_af for guardrail in guardrails: guardrail.streaming_buffer_until_moderated = not streams_live monkeypatch.setattr(litellm, "callbacks", guardrails) - data = _two_step_pipeline_data("post_call", "next") - data["metadata"]["_guardrail_pipelines"] = [ - ( - "governance", - GuardrailPipeline( - mode="post_call", - steps=[PipelineStep(guardrail=name, on_pass="next", on_fail="next") for name in ("first", "second")], - ), - ) - ] + data = _two_step_pipeline_data("post_call", "next", first_on_fail="next", second_on_fail="next") data.update(asks_for) result = await proxy_logging.pre_call_hook( From e2304e63defea14446e871f5f49e80d98a7cf7d8 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 26 Oct 2026 17:55:23 -0700 Subject: [PATCH 06/13] test(guardrails): type the new pipeline test signatures --- .../policy_engine/test_pipeline_executor.py | 10 +++-- .../proxy_logging/test_guardrail_pipeline.py | 41 +++++++++++-------- .../utils/proxy_logging/test_pre_call_hook.py | 2 +- 3 files changed, 33 insertions(+), 20 deletions(-) diff --git a/tests/unit/proxy/policy_engine/test_pipeline_executor.py b/tests/unit/proxy/policy_engine/test_pipeline_executor.py index bd8b9eb24d7..c56b6364bb8 100644 --- a/tests/unit/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/unit/proxy/policy_engine/test_pipeline_executor.py @@ -1440,7 +1440,7 @@ class _SelfRecordingGuardrail(_TextReturningGuardrail): @pytest.mark.asyncio -async def test_streaming_step_whose_guardrail_records_its_own_results_is_named(monkeypatch): +async def test_streaming_step_whose_guardrail_records_its_own_results_is_named(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "callbacks", [_SelfRecordingGuardrail(["hello world"])]) result = await _run_streaming_step(_WritingTranslation(), [_chunk()]) @@ -1466,7 +1466,9 @@ class _LitellmMetadataOpeningGuardrail(CustomGuardrail): @pytest.mark.asyncio -async def test_streaming_step_is_named_in_the_metadata_the_request_started_with(monkeypatch): +async def test_streaming_step_is_named_in_the_metadata_the_request_started_with( + monkeypatch: pytest.MonkeyPatch, +) -> None: monkeypatch.setattr(litellm, "callbacks", [_LitellmMetadataOpeningGuardrail()]) result = await _run_streaming_step(_WritingTranslation(), [_chunk()]) @@ -1476,7 +1478,9 @@ async def test_streaming_step_is_named_in_the_metadata_the_request_started_with( @pytest.mark.asyncio -async def test_streaming_step_that_flags_the_stream_is_named_when_the_pipeline_moves_on(monkeypatch): +async def test_streaming_step_that_flags_the_stream_is_named_when_the_pipeline_moves_on( + monkeypatch: pytest.MonkeyPatch, +) -> None: monkeypatch.setattr(litellm, "callbacks", [_BlockingStreamGuardrail()]) result = await _run_streaming_step(_WritingTranslation(), [_chunk()]) diff --git a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 29438c63a56..2426ee8b7f0 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -1209,8 +1209,11 @@ def _two_step_pipeline_data( [("allow", (1, 0), ["first"]), ("next", (1, 1), ["first", "second"])], ) async def test_pre_call_hook_names_only_the_pre_call_pipeline_steps_that_ran( - proxy_logging, make_user_api_key_auth, monkeypatch, first_on_pass, expected_calls, expected_header -): + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, + first_on_pass: str, + expected_calls: tuple[int, int], + expected_header: list[str], +) -> None: first = _PipelineStepGuardrail("first", "pre_call") second = _PipelineStepGuardrail("second", "pre_call") monkeypatch.setattr(litellm, "callbacks", [first, second]) @@ -1227,8 +1230,8 @@ async def test_pre_call_hook_names_only_the_pre_call_pipeline_steps_that_ran( @pytest.mark.asyncio async def test_pre_call_hook_names_a_skipped_step_whose_guardrail_also_runs_outside_the_pipeline( - proxy_logging, make_user_api_key_auth, monkeypatch -): + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, +) -> None: first = _PipelineStepGuardrail("first", "pre_call") second = _PipelineStepGuardrail("second", ["pre_call", "post_call"]) monkeypatch.setattr(litellm, "callbacks", [first, second]) @@ -1264,8 +1267,9 @@ _PIPELINE_STEP_ERRORS: Final = ( @pytest.mark.asyncio @pytest.mark.parametrize("error", _PIPELINE_STEP_ERRORS) async def test_pre_call_hook_names_a_step_that_raised_before_the_pipeline_moved_on( - proxy_logging, make_user_api_key_auth, monkeypatch, error -): + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, + error: Exception, +) -> None: first = _RaisingPipelineStepGuardrail("first", error) second = _PipelineStepGuardrail("second", "pre_call") monkeypatch.setattr(litellm, "callbacks", [first, second]) @@ -1283,8 +1287,9 @@ async def test_pre_call_hook_names_a_step_that_raised_before_the_pipeline_moved_ @pytest.mark.asyncio @pytest.mark.parametrize("error", _PIPELINE_STEP_ERRORS) async def test_deployment_hook_still_runs_a_pipeline_step_that_raised( - proxy_logging, make_user_api_key_auth, monkeypatch, error -): + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, + error: Exception, +) -> None: first = _RaisingPipelineStepGuardrail("first", error) second = _PipelineStepGuardrail("second", "pre_call") monkeypatch.setattr(litellm, "callbacks", [first, second]) @@ -1317,8 +1322,8 @@ class _OpensLitellmMetadataGuardrail(_PipelineStepGuardrail): @pytest.mark.asyncio async def test_pre_call_hook_names_a_step_that_ran_in_the_litellm_metadata_a_later_step_opened( - proxy_logging, make_user_api_key_auth, monkeypatch -): + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, +) -> None: first = _PipelineStepGuardrail("first", "pre_call") second = _OpensLitellmMetadataGuardrail("second", "pre_call") monkeypatch.setattr(litellm, "callbacks", [first, second]) @@ -1337,8 +1342,8 @@ async def test_pre_call_hook_names_a_step_that_ran_in_the_litellm_metadata_a_lat @pytest.mark.asyncio async def test_post_call_pipeline_steps_are_named_only_once_they_run( - proxy_logging, make_user_api_key_auth, monkeypatch -): + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, +) -> None: first = _PipelineStepGuardrail("first", "post_call") second = _PipelineStepGuardrail("second", "post_call") monkeypatch.setattr(litellm, "callbacks", [first, second]) @@ -1369,8 +1374,8 @@ class _OpensLitellmMetadataOnResponseGuardrail(_PipelineStepGuardrail): @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 -): + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, +) -> None: first = _OpensLitellmMetadataOnResponseGuardrail("first", "post_call") second = _PipelineStepGuardrail("second", "post_call") monkeypatch.setattr(litellm, "callbacks", [first, second]) @@ -1405,8 +1410,12 @@ async def test_post_call_pipeline_names_every_step_that_ran_in_metadata_after_a_ ], ) async def test_pre_call_hook_names_up_front_the_post_call_steps_a_stream_runs_after_its_headers( - proxy_logging, make_user_api_key_auth, monkeypatch, asks_for, request_route, streams_live, expected_header -): + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, + asks_for: dict[str, object], + request_route: str, + streams_live: bool, + expected_header: list[str] | None, +) -> None: guardrails = [_PipelineStepGuardrail(name, "post_call") for name in ("first", "second")] for guardrail in guardrails: guardrail.streaming_buffer_until_moderated = not streams_live diff --git a/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py b/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py index 615e39c495d..48dac33d010 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py +++ b/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py @@ -170,7 +170,7 @@ async def test_pre_call_hook_processes_guardrail_metadata_when_no_overrides(prox proxy_logging.slack_alerting_instance = MagicMock(alerting=None) invoked = {} - def fake_process(d, user_api_key_dict): + def fake_process(d: dict[str, object], user_api_key_dict: UserAPIKeyAuth) -> None: invoked["data"] = d proxy_logging._process_guardrail_metadata = fake_process # type: ignore[assignment] From 0676e6a460fdbf54492910ebca228cdd64dea0e9 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 26 Oct 2026 18:42:19 -0700 Subject: [PATCH 07/13] test(guardrails): key the test guardrail hook payloads by str --- .../policy_engine/test_pipeline_executor.py | 2 +- .../proxy_logging/test_guardrail_pipeline.py | 16 ++++++++-------- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/tests/unit/proxy/policy_engine/test_pipeline_executor.py b/tests/unit/proxy/policy_engine/test_pipeline_executor.py index c56b6364bb8..c5cb7b3444f 100644 --- a/tests/unit/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/unit/proxy/policy_engine/test_pipeline_executor.py @@ -1457,7 +1457,7 @@ class _LitellmMetadataOpeningGuardrail(CustomGuardrail): async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, - request_data: dict, + request_data: dict[str, object], input_type: Literal["request", "response"], logging_obj: LiteLLMLoggingObj | None = None, ) -> GenericGuardrailAPIInputs: diff --git a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 2426ee8b7f0..bd0480182b7 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -1171,13 +1171,13 @@ class _PipelineStepGuardrail(CustomGuardrail): self.calls = 0 async def async_pre_call_hook( - self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral - ) -> dict: + self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict[str, object], call_type: CallTypesLiteral + ) -> dict[str, object]: self.calls += 1 return data async def async_post_call_success_hook( - self, data: dict, user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes + self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes ) -> None: self.calls += 1 @@ -1252,8 +1252,8 @@ class _RaisingPipelineStepGuardrail(_PipelineStepGuardrail): self.error = error async def async_pre_call_hook( - self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral - ) -> dict: + self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict[str, object], call_type: CallTypesLiteral + ) -> dict[str, object]: self.calls += 1 raise self.error @@ -1312,8 +1312,8 @@ async def test_deployment_hook_still_runs_a_pipeline_step_that_raised( class _OpensLitellmMetadataGuardrail(_PipelineStepGuardrail): async def async_pre_call_hook( - self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral - ) -> dict: + self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict[str, object], call_type: CallTypesLiteral + ) -> dict[str, object]: self.calls += 1 data["litellm_metadata"] = {} add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) @@ -1366,7 +1366,7 @@ async def test_post_call_pipeline_steps_are_named_only_once_they_run( class _OpensLitellmMetadataOnResponseGuardrail(_PipelineStepGuardrail): async def async_post_call_success_hook( - self, data: dict, user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes + self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes ) -> None: self.calls += 1 data["litellm_metadata"] = {"user_api_key_hash": "hashed"} From 80b95eb8c0910fe834c59d1fb765e15fcee17a6a Mon Sep 17 00:00:00 2001 From: yucheng Date: Tue, 27 Oct 2026 16:35:41 -0700 Subject: [PATCH 08/13] refactor(a2a): share the A2A method map with the applied-guardrails stream check The applied-guardrails header check now reads the A2A streaming methods from the same map the A2A endpoint uses to translate method names, instead of its own copy --- litellm/proxy/a2a/version_convert.py | 23 ++++++++++++++++++- .../proxy/agent_endpoints/a2a_endpoints.py | 20 +++------------- litellm/proxy/utils.py | 9 ++------ 3 files changed, 27 insertions(+), 25 deletions(-) diff --git a/litellm/proxy/a2a/version_convert.py b/litellm/proxy/a2a/version_convert.py index bde2ff45a88..eb55bff50bc 100644 --- a/litellm/proxy/a2a/version_convert.py +++ b/litellm/proxy/a2a/version_convert.py @@ -24,7 +24,8 @@ The two wire shapes: ``Task`` result is a bare object without ``kind``. """ -from collections.abc import Callable +from collections.abc import Callable, Mapping +from types import MappingProxyType from typing import Final, Literal from pydantic import BaseModel @@ -39,6 +40,26 @@ JsonDict = dict[str, object] _V1_SEND_ENVELOPE_KEYS: Final = frozenset({"message", "task"}) _V1_STREAM_ENVELOPE_KEYS: Final = frozenset({"message", "task", "statusUpdate", "artifactUpdate"}) +A2A_PASCAL_TO_WIRE: Final[Mapping[str, str]] = MappingProxyType( + { + "SendMessage": "message/send", + "SendStreamingMessage": "message/stream", + "GetTask": "tasks/get", + "ListTasks": "tasks/list", + "CancelTask": "tasks/cancel", + "SubscribeToTask": "tasks/resubscribe", + "CreateTaskPushNotificationConfig": "tasks/pushNotificationConfig/set", + "GetTaskPushNotificationConfig": "tasks/pushNotificationConfig/get", + "ListTaskPushNotificationConfigs": "tasks/pushNotificationConfig/list", + "DeleteTaskPushNotificationConfig": "tasks/pushNotificationConfig/delete", + "GetExtendedAgentCard": "agent/getAuthenticatedExtendedCard", + } +) +_STREAMING_WIRE_METHODS: Final = frozenset({"message/stream", "tasks/resubscribe"}) +A2A_STREAMING_METHODS: Final = _STREAMING_WIRE_METHODS | frozenset( + pascal for pascal, wire in A2A_PASCAL_TO_WIRE.items() if wire in _STREAMING_WIRE_METHODS +) + def _dump_03(model: BaseModel) -> JsonDict: """Dump a compat (0.3) pydantic model to its camelCase wire dict.""" diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 4691c7b0bc5..85a5ca30816 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -28,6 +28,7 @@ from litellm.litellm_core_utils.url_utils import SSRFError, validate_url from litellm.llms.a2a.common_utils import resolve_a2a_hop_auth_header from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.a2a.version_convert import ( + A2A_PASCAL_TO_WIRE, A2AVersion, normalize_agent_card, normalize_jsonrpc_response, @@ -65,21 +66,6 @@ _SSE_KEEPALIVE_HEADERS: Final[Mapping[str, str]] = MappingProxyType( } ) -_PASCAL_TO_WIRE: Final[Mapping[str, str]] = { - "SendMessage": "message/send", - "SendStreamingMessage": "message/stream", - "GetTask": "tasks/get", - "ListTasks": "tasks/list", - "CancelTask": "tasks/cancel", - "SubscribeToTask": "tasks/resubscribe", - "CreateTaskPushNotificationConfig": "tasks/pushNotificationConfig/set", - "GetTaskPushNotificationConfig": "tasks/pushNotificationConfig/get", - "ListTaskPushNotificationConfigs": "tasks/pushNotificationConfig/list", - "DeleteTaskPushNotificationConfig": "tasks/pushNotificationConfig/delete", - "GetExtendedAgentCard": "agent/getAuthenticatedExtendedCard", -} - - _DECODED_JSON: Final = TypeAdapter(object) @@ -132,7 +118,7 @@ def _served_version(agent: "AgentResponse", request: Request, original_method: s configured: Final = (agent.agent_card_params or {}).get("protocolVersion") if configured in ("0.3", "1.0"): return configured - if original_method in _PASCAL_TO_WIRE: + if original_method in A2A_PASCAL_TO_WIRE: return "1.0" return "1.0" if request.headers.get("a2a-version", "").startswith("1.") else "0.3" @@ -699,7 +685,7 @@ async def invoke_agent_a2a( params = body.get("params", {}) if method: - method = _PASCAL_TO_WIRE.get(method, method) + method = A2A_PASCAL_TO_WIRE.get(method, method) if isinstance(params, dict): # extract any litellm params from the params - eg. 'guardrails' diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 9a960d30d3d..36e5c95c79c 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -179,6 +179,7 @@ from litellm.proxy._types import ( SpecialModelNames, UserAPIKeyAuth, ) +from litellm.proxy.a2a.version_convert import A2A_STREAMING_METHODS from litellm.proxy.agent_endpoints.auth.agent_access_groups import CeilingResolver, resolve_agent_access_group_ceiling from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header @@ -721,12 +722,6 @@ def _ran_or_runs_outside_pipelines( ) -# The A2A JSON-RPC methods answered with a stream, as both protocol versions spell them -_A2A_STREAMING_METHODS: Final = frozenset( - ("message/stream", "tasks/resubscribe", "SendStreamingMessage", "SubscribeToTask") -) - - def _post_call_steps_named_before_headers( data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth ) -> frozenset[str]: @@ -735,7 +730,7 @@ def _post_call_steps_named_before_headers( headers go out: every step of a whole response, and of a streamed one the steps of the pipelines that hold the stream until they finish. """ - if data.get("stream") is not True and data.get("method") not in _A2A_STREAMING_METHODS: + if data.get("stream") is not True and data.get("method") not in A2A_STREAMING_METHODS: return pipeline_managed_guardrail_names(data, "post_call") gated: Final = _stream_gated_pipelines(data, user_api_key_dict) return frozenset() if _pipelines_stream_live(gated) else _pipeline_step_guardrail_names(gated) From 5434ce4675f7950e05d8ea7b0e4fa670fdaa4efe Mon Sep 17 00:00:00 2001 From: yucheng Date: Tue, 27 Oct 2026 16:54:52 -0700 Subject: [PATCH 09/13] fix(guardrails): name the LLM pass-through route's response pipeline up front The LLM pass-through route builds its response headers before the response pipeline runs, so the pre_call sweep claims those steps as on main. Adds Final to the new test locals. --- litellm/proxy/utils.py | 29 +++-- .../test_guardrail_pipeline_applied_header.py | 107 +++++++++++++++++- .../proxy_logging/test_guardrail_pipeline.py | 80 ++++++++----- .../utils/proxy_logging/test_pre_call_hook.py | 2 +- 4 files changed, 175 insertions(+), 43 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 36e5c95c79c..bfc7175cc65 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -723,29 +723,35 @@ def _ran_or_runs_outside_pipelines( def _post_call_steps_named_before_headers( - data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth + data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, call_type: str ) -> frozenset[str]: """ Post_call pipeline steps the executor names, as it runs them, before the response headers go out: every step of a whole response, and of a streamed one the steps of - the pipelines that hold the stream until they finish. + the pipelines that hold the stream until they finish. The LLM pass-through route builds + its headers before any response step runs, so none of its steps. """ + if call_type == CallTypes.allm_passthrough_route.value: + return frozenset() if data.get("stream") is not True and data.get("method") not in A2A_STREAMING_METHODS: return pipeline_managed_guardrail_names(data, "post_call") gated: Final = _stream_gated_pipelines(data, user_api_key_dict) return frozenset() if _pipelines_stream_live(gated) else _pipeline_step_guardrail_names(gated) -def _pipeline_steps_that_have_not_run(data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth) -> frozenset[str]: +def _pipeline_steps_that_have_not_run( + data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, call_type: str +) -> frozenset[str]: """ Guardrails only policy pipelines run whose pre_call hook has not run for this request: pre_call steps a pipeline skipped and the post_call steps the executor names when it runs them. A stream runs a pipeline it cannot hold as guardrails on their own, and one - whose steps all stream live after the headers go out, so those steps are named up front - like any other guardrail. + whose steps all stream live after the headers go out, and the LLM pass-through route runs + its response steps after its headers are built, so those steps are named up front like + any other guardrail. """ pre_call: Final = pipeline_managed_guardrail_names(data, "pre_call") - post_call: Final = _post_call_steps_named_before_headers(data, user_api_key_dict) + post_call: Final = _post_call_steps_named_before_headers(data, user_api_key_dict, call_type) return frozenset( name for name in pre_call | post_call @@ -2694,7 +2700,7 @@ class ProxyLogging: data.pop("prompt_version", None) data.pop("prompt_environment", None) - def _process_guardrail_metadata(self, data: dict, user_api_key_dict: UserAPIKeyAuth) -> None: + def _process_guardrail_metadata(self, data: dict, user_api_key_dict: UserAPIKeyAuth, call_type: str) -> None: """Process guardrails from metadata and add to applied_guardrails.""" from litellm.proxy.common_utils.callback_utils import ( add_guardrail_to_applied_guardrails_header, @@ -2722,6 +2728,7 @@ class ProxyLogging: not_run: Final = _pipeline_steps_that_have_not_run( cast(dict[str, object], data), # cast-ok: request payload user_api_key_dict, + call_type, ) for guardrail_name in guardrails_in_metadata: if ( @@ -3015,7 +3022,7 @@ class ProxyLogging: and (guardrails_only or not caps.has_pre_call_override) ): if data is not None: - self._process_guardrail_metadata(data, user_api_key_dict) + self._process_guardrail_metadata(data, user_api_key_dict, call_type) return data parallel_guardrails: Final[tuple[CustomGuardrail, ...]] = ( @@ -3119,17 +3126,17 @@ class ProxyLogging: ) if data is not None: - self._process_guardrail_metadata(data, user_api_key_dict) + self._process_guardrail_metadata(data, user_api_key_dict, call_type) return data except SensitiveDataRouteException as e: data = await self._handle_sensitive_data_route_exception(e, data, user_api_key_dict) if data is not None: - self._process_guardrail_metadata(data, user_api_key_dict) + self._process_guardrail_metadata(data, user_api_key_dict, call_type) return data except Exception: if data is not None: - self._process_guardrail_metadata(data, user_api_key_dict) + self._process_guardrail_metadata(data, user_api_key_dict, call_type) raise async def _run_parallel_pre_call_guardrails( diff --git a/tests/integration/observability/test_guardrail_pipeline_applied_header.py b/tests/integration/observability/test_guardrail_pipeline_applied_header.py index 58e46141cac..174a88ac306 100644 --- a/tests/integration/observability/test_guardrail_pipeline_applied_header.py +++ b/tests/integration/observability/test_guardrail_pipeline_applied_header.py @@ -20,6 +20,7 @@ import yaml from tests.integration._support.client import Gateway, eventually, gateway_from_environment, object_value from tests.integration._support.otlp_sink import Span, SpanSinks, recorded_spans from tests.integration._support.process import owned_proxy +from tests.integration._support.upstream import aws_event_stream_frame from tests.integration._support.wire import Reply, Request, Wire, wire_server _HEADER: Final = "x-litellm-applied-guardrails" @@ -50,6 +51,8 @@ _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") +_BEDROCK_MODEL: Final = "pipeline-bedrock" +_BEDROCK_MODEL_ID: Final = "anthropic.claude-sonnet-5-v1:0" @dataclass(frozen=True, slots=True) @@ -126,7 +129,58 @@ def _responses_reply(request: Request) -> Reply: ) +def _bedrock_event(event_type: str, payload: dict[str, object]) -> bytes: + return aws_event_stream_frame( + {":event-type": event_type, ":content-type": "application/json", ":message-type": "event"}, + json.dumps(payload).encode(), + ) + + +def _bedrock_reply(request: Request) -> Reply: + action: Final = request.target.rsplit("/", 1)[1] + if action == "converse-stream": + return Reply( + content_type="application/vnd.amazon.eventstream", + chunks=( + _bedrock_event("messageStart", {"role": "assistant"}), + _bedrock_event("contentBlockDelta", {"contentBlockIndex": 0, "delta": {"text": "guarded"}}), + _bedrock_event("contentBlockStop", {"contentBlockIndex": 0}), + _bedrock_event("messageStop", {"stopReason": "end_turn"}), + _bedrock_event( + "metadata", {"usage": {"inputTokens": 3, "outputTokens": 1, "totalTokens": 4}, "metrics": {}} + ), + ), + ) + if action == "invoke": + return Reply( + body=json.dumps( + { + "id": f"msg_{uuid.uuid4().hex}", + "type": "message", + "role": "assistant", + "model": _BEDROCK_MODEL_ID, + "content": [{"type": "text", "text": "guarded"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 3, "output_tokens": 1}, + } + ).encode() + ) + assert action == "converse", request.target + return Reply( + body=json.dumps( + { + "output": {"message": {"role": "assistant", "content": [{"text": "guarded"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 3, "outputTokens": 1, "totalTokens": 4}, + "metrics": {"latencyMs": 1}, + } + ).encode() + ) + + def _completion(request: Request) -> Reply: + if request.target.startswith("/model/"): + return _bedrock_reply(request) if request.target == "/v1/responses": return _responses_reply(request) assert request.target == "/v1/chat/completions", request.target @@ -208,11 +262,21 @@ def _policy(mode: str, steps: tuple[tuple[str, str, str], ...], *, on_error: str } -def _config(directory: Path, vendor: Wire) -> Path: +def _config(directory: Path, vendor: Wire, upstream: 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": _BEDROCK_MODEL, + "litellm_params": { + "model": f"bedrock/{_BEDROCK_MODEL_ID}", + "api_base": upstream.url, + "aws_access_key_id": "AKIASYNTHETICPROVIDER", + "aws_secret_access_key": "synthetic-provider-secret", + "aws_region_name": "us-east-1", + }, + }, { "model_name": _JUDGE_MODEL, "litellm_params": { @@ -316,7 +380,7 @@ def rig(tmp_path_factory: pytest.TempPathFactory, audit_sinks: SpanSinks) -> Ite shared, directory, {"OTEL_EXPORTER": "http/json", "OTEL_ENDPOINT": audit_sinks.operator, "OTEL_BSP_SCHEDULE_DELAY": "100"}, - config=_config(directory, vendor), + config=_config(directory, vendor, upstream), workers=2, ) as owned, owned.scenario() as scenario, @@ -572,6 +636,45 @@ def test_a_response_pipeline_names_a_judge_step_that_records_its_own_results(rig assert sorted(object_value(rows[0]["metadata"]).get("applied_guardrails") or []) == [_JUDGE, _RESPONSE_SECOND] +def _bedrock_body(action: str) -> dict[str, object]: + prompt: Final = f"hello {uuid.uuid4().hex}" + if action == "invoke": + return { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 20, + "messages": [{"role": "user", "content": [{"type": "text", "text": prompt}]}], + } + return {"messages": [{"role": "user", "content": [{"text": prompt}]}]} + + +@pytest.mark.parametrize("action", ["converse", "invoke", "converse-stream"]) +def test_a_pass_through_route_names_its_response_pipeline_up_front_because_its_headers_are_built_first( + rig: _Rig, action: str +) -> None: + rig.vendor.drain() + rig.upstream.drain() + with rig.gateway.scenario() as scenario: + team: Final = scenario.team(metadata={"tags": [_RESPONSE_STOPS_AFTER_FIRST]}) + response: Final = rig.gateway.request( + "POST", f"/bedrock/model/{_BEDROCK_MODEL}/{action}", _bedrock_body(action), key=scenario.key(team_id=team) + ) + assert response.status_code == 200, response.text + assert _steps_that_ran(rig) == (_RESPONSE_FIRST,) + assert [request.target.rsplit("/", 1)[1] for request in rig.upstream.drain()] == [action] + assert sorted(response.headers.get(_HEADER, "").split(",")) == [_RESPONSE_FIRST, _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 []) == [ + _RESPONSE_FIRST, + _RESPONSE_SECOND, + ], rows + + def _agent_card(name: str, url: str) -> dict[str, object]: return { "protocolVersion": "0.3", diff --git a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py index bd0480182b7..ddeaf45d32b 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -280,7 +280,7 @@ def test_process_guardrail_metadata_calls_header_helper(proxy_logging, monkeypat monkeypatch.setattr(callback_utils, "add_guardrail_to_applied_guardrails_header", fake_add) data = {"metadata": {"guardrails": ["g1", "g2"]}} - proxy_logging._process_guardrail_metadata(data, UserAPIKeyAuth()) + proxy_logging._process_guardrail_metadata(data, UserAPIKeyAuth(), "completion") snapshot = { "call_count": len(calls), "first_name": calls[0]["name"], @@ -305,7 +305,7 @@ def test_process_guardrail_metadata_skips_already_applied(proxy_logging, monkeyp monkeypatch.setattr(callback_utils, "add_guardrail_to_applied_guardrails_header", fake_add) data = {"metadata": {"guardrails": ["g1", "g2"], "applied_guardrails": ["g1"]}} - proxy_logging._process_guardrail_metadata(data, UserAPIKeyAuth()) + proxy_logging._process_guardrail_metadata(data, UserAPIKeyAuth(), "completion") assert calls == ["g2"] @@ -317,12 +317,12 @@ def test_process_guardrail_metadata_no_metadata_is_noop(proxy_logging, monkeypat "add_guardrail_to_applied_guardrails_header", MagicMock(side_effect=AssertionError("should not be called")), ) - proxy_logging._process_guardrail_metadata({}, UserAPIKeyAuth()) + proxy_logging._process_guardrail_metadata({}, UserAPIKeyAuth(), "completion") def test_process_guardrail_metadata_invalid_data_raises(proxy_logging): with pytest.raises(AttributeError): - proxy_logging._process_guardrail_metadata(None, UserAPIKeyAuth()) # type: ignore[arg-type] + proxy_logging._process_guardrail_metadata(None, UserAPIKeyAuth(), "completion") # type: ignore[arg-type] # --------------------------------------------------------------------------- @@ -1185,7 +1185,7 @@ class _PipelineStepGuardrail(CustomGuardrail): def _two_step_pipeline_data( mode: str, first_on_pass: str, first_on_fail: str = "block", second_on_fail: str = "block" ) -> dict[str, object]: - pipeline = GuardrailPipeline( + pipeline: Final = GuardrailPipeline( mode=mode, steps=[ PipelineStep(guardrail="first", on_pass=first_on_pass, on_fail=first_on_fail), @@ -1214,11 +1214,11 @@ async def test_pre_call_hook_names_only_the_pre_call_pipeline_steps_that_ran( expected_calls: tuple[int, int], expected_header: list[str], ) -> None: - first = _PipelineStepGuardrail("first", "pre_call") - second = _PipelineStepGuardrail("second", "pre_call") + first: Final = _PipelineStepGuardrail("first", "pre_call") + second: Final = _PipelineStepGuardrail("second", "pre_call") monkeypatch.setattr(litellm, "callbacks", [first, second]) - result = await proxy_logging.pre_call_hook( + result: Final = await proxy_logging.pre_call_hook( user_api_key_dict=make_user_api_key_auth(), data=_two_step_pipeline_data("pre_call", first_on_pass), call_type="completion", @@ -1232,11 +1232,11 @@ async def test_pre_call_hook_names_only_the_pre_call_pipeline_steps_that_ran( async def test_pre_call_hook_names_a_skipped_step_whose_guardrail_also_runs_outside_the_pipeline( proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, ) -> None: - first = _PipelineStepGuardrail("first", "pre_call") - second = _PipelineStepGuardrail("second", ["pre_call", "post_call"]) + first: Final = _PipelineStepGuardrail("first", "pre_call") + second: Final = _PipelineStepGuardrail("second", ["pre_call", "post_call"]) monkeypatch.setattr(litellm, "callbacks", [first, second]) - result = await proxy_logging.pre_call_hook( + result: Final = await proxy_logging.pre_call_hook( user_api_key_dict=make_user_api_key_auth(), data=_two_step_pipeline_data("pre_call", "allow"), call_type="completion", @@ -1270,11 +1270,11 @@ async def test_pre_call_hook_names_a_step_that_raised_before_the_pipeline_moved_ proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, error: Exception, ) -> None: - first = _RaisingPipelineStepGuardrail("first", error) - second = _PipelineStepGuardrail("second", "pre_call") + first: Final = _RaisingPipelineStepGuardrail("first", error) + second: Final = _PipelineStepGuardrail("second", "pre_call") monkeypatch.setattr(litellm, "callbacks", [first, second]) - result = await proxy_logging.pre_call_hook( + result: Final = await proxy_logging.pre_call_hook( user_api_key_dict=make_user_api_key_auth(), data=_two_step_pipeline_data("pre_call", "next", first_on_fail="next"), call_type="completion", @@ -1290,10 +1290,10 @@ async def test_deployment_hook_still_runs_a_pipeline_step_that_raised( proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, error: Exception, ) -> None: - first = _RaisingPipelineStepGuardrail("first", error) - second = _PipelineStepGuardrail("second", "pre_call") + first: Final = _RaisingPipelineStepGuardrail("first", error) + second: Final = _PipelineStepGuardrail("second", "pre_call") monkeypatch.setattr(litellm, "callbacks", [first, second]) - result = await proxy_logging.pre_call_hook( + result: Final = await proxy_logging.pre_call_hook( user_api_key_dict=make_user_api_key_auth(), data=_two_step_pipeline_data("pre_call", "next", first_on_fail="next"), call_type="completion", @@ -1324,11 +1324,11 @@ class _OpensLitellmMetadataGuardrail(_PipelineStepGuardrail): async def test_pre_call_hook_names_a_step_that_ran_in_the_litellm_metadata_a_later_step_opened( proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, ) -> None: - first = _PipelineStepGuardrail("first", "pre_call") - second = _OpensLitellmMetadataGuardrail("second", "pre_call") + first: Final = _PipelineStepGuardrail("first", "pre_call") + second: Final = _OpensLitellmMetadataGuardrail("second", "pre_call") monkeypatch.setattr(litellm, "callbacks", [first, second]) - result = await proxy_logging.pre_call_hook( + result: Final = await proxy_logging.pre_call_hook( user_api_key_dict=make_user_api_key_auth(), data=_two_step_pipeline_data("pre_call", "next"), call_type="completion", @@ -1344,12 +1344,12 @@ async def test_pre_call_hook_names_a_step_that_ran_in_the_litellm_metadata_a_lat async def test_post_call_pipeline_steps_are_named_only_once_they_run( proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, ) -> None: - first = _PipelineStepGuardrail("first", "post_call") - second = _PipelineStepGuardrail("second", "post_call") + 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 = await proxy_logging.pre_call_hook( + data: Final = await proxy_logging.pre_call_hook( user_api_key_dict=make_user_api_key_auth(), data=_two_step_pipeline_data("post_call", "allow"), call_type="completion", @@ -1376,12 +1376,12 @@ class _OpensLitellmMetadataOnResponseGuardrail(_PipelineStepGuardrail): async def test_post_call_pipeline_names_every_step_that_ran_in_metadata_after_a_step_opens_litellm_metadata( proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, ) -> None: - first = _OpensLitellmMetadataOnResponseGuardrail("first", "post_call") - second = _PipelineStepGuardrail("second", "post_call") + first: Final = _OpensLitellmMetadataOnResponseGuardrail("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 = await proxy_logging.pre_call_hook( + data: Final = 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", @@ -1416,14 +1416,14 @@ async def test_pre_call_hook_names_up_front_the_post_call_steps_a_stream_runs_af streams_live: bool, expected_header: list[str] | None, ) -> None: - guardrails = [_PipelineStepGuardrail(name, "post_call") for name in ("first", "second")] + guardrails: Final = [_PipelineStepGuardrail(name, "post_call") for name in ("first", "second")] for guardrail in guardrails: guardrail.streaming_buffer_until_moderated = not streams_live monkeypatch.setattr(litellm, "callbacks", guardrails) - data = _two_step_pipeline_data("post_call", "next", first_on_fail="next", second_on_fail="next") + data: Final = _two_step_pipeline_data("post_call", "next", first_on_fail="next", second_on_fail="next") data.update(asks_for) - result = await proxy_logging.pre_call_hook( + result: Final = await proxy_logging.pre_call_hook( user_api_key_dict=make_user_api_key_auth(request_route=request_route), data=data, call_type="completion", @@ -1432,6 +1432,28 @@ async def test_pre_call_hook_names_up_front_the_post_call_steps_a_stream_runs_af assert result["metadata"].get("applied_guardrails") == expected_header +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True], ids=["whole-response", "stream"]) +async def test_pre_call_hook_names_up_front_the_post_call_steps_the_llm_pass_through_route_runs_after_its_headers( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + stream: bool, +) -> None: + guardrails: Final = [_PipelineStepGuardrail(name, "post_call") for name in ("first", "second")] + monkeypatch.setattr(litellm, "callbacks", guardrails) + data: Final = _two_step_pipeline_data("post_call", "allow") + data["stream"] = stream + + result: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/bedrock/model/m/converse"), + data=data, + call_type="allm_passthrough_route", + ) + + assert result["metadata"].get("applied_guardrails") == ["first", "second"] + + @pytest.mark.asyncio async def test_post_call_pipeline_replacement_response_reaches_caller( proxy_logging, make_user_api_key_auth, monkeypatch diff --git a/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py b/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py index 48dac33d010..89a59d44380 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py +++ b/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py @@ -170,7 +170,7 @@ async def test_pre_call_hook_processes_guardrail_metadata_when_no_overrides(prox proxy_logging.slack_alerting_instance = MagicMock(alerting=None) invoked = {} - def fake_process(d: dict[str, object], user_api_key_dict: UserAPIKeyAuth) -> None: + def fake_process(d: dict[str, object], user_api_key_dict: UserAPIKeyAuth, call_type: str) -> None: invoked["data"] = d proxy_logging._process_guardrail_metadata = fake_process # type: ignore[assignment] From 747b6afc5c750de843006890bad813ee293b1744 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 28 Oct 2026 10:31:17 -0700 Subject: [PATCH 10/13] fix(guardrails): name response steps up front on routes that log before they run Image generation, files, batches and fine-tuning build their spend log before the response pipeline runs, so leaving its steps for the executor to name dropped them from the spend row and OTel. Leave them only when the success log waits for the response hooks, which the shared LLM request path and A2A message/send now mark before pre_call_hook. --- .../proxy/agent_endpoints/a2a_endpoints.py | 1 + litellm/proxy/common_request_processing.py | 20 ++++++--- litellm/proxy/utils.py | 23 +++++++--- .../test_guardrail_pipeline_applied_header.py | 33 ++++++++++++++ .../proxy/test_common_request_processing.py | 41 ++++++++++++++++++ .../proxy_logging/test_guardrail_pipeline.py | 43 +++++++++++++++++-- 6 files changed, 146 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 85a5ca30816..612988a74e9 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -788,6 +788,7 @@ async def invoke_agent_a2a( proxy_config=proxy_config, route_type="asend_message", version=version, + defer_async_logging=method == "message/send", ) request_data = data diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index cec176ac559..6579afb90ef 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2279,6 +2279,7 @@ class ProxyBaseLLMRequestProcessing: rate_limited_model: str | None = None, skip_guardrails: bool = False, logging_auth: UserAPIKeyAuth | None = None, + defer_async_logging: bool = False, ) -> tuple[dict, LiteLLMLoggingObj]: start_time: Final = datetime.now() # start before calling guardrail hooks @@ -2444,6 +2445,10 @@ class ProxyBaseLLMRequestProcessing: ) self.data["litellm_logging_obj"] = logging_obj + # Set before pre_call_hook: its guardrail sweep leaves a whole response's pipeline + # steps for the executor to name only when the success log waits for them. + if defer_async_logging: + logging_obj.defer_async_logging = True # Merge model-level guardrails before pre_call_hook so DB/UI-configured # guardrails actually execute on pre_call. Without this, guardrails set @@ -2534,6 +2539,7 @@ class ProxyBaseLLMRequestProcessing: model: str | None, route_type: str, llm_router: Router | None, + defer_async_logging: bool = False, ) -> tuple[dict, LiteLLMLoggingObj]: from litellm.proxy.common_utils.proxy_rate_limit_error import ( PER_MODEL_RATE_LIMIT_DESCRIPTOR_KEYS, @@ -2574,6 +2580,7 @@ class ProxyBaseLLMRequestProcessing: model=model, route_type=route_type, llm_router=llm_router, + defer_async_logging=defer_async_logging, ) except ProxyRateLimitError as original_exc: rate_limited_data: Final = self.data @@ -2630,6 +2637,7 @@ class ProxyBaseLLMRequestProcessing: route_type=route_type, llm_router=llm_router, rate_limited_model=original_model, + defer_async_logging=defer_async_logging, ) except ProxyRateLimitError as fallback_exc: if is_hard_per_model_limit(fallback_exc): @@ -2860,6 +2868,11 @@ class ProxyBaseLLMRequestProcessing: requested_model_from_client: Final[str | None] = client_model if isinstance(client_model, str) else None self._debug_log_request_payload() + # Defer async logging when post-call guardrails are configured so the + # StandardLoggingPayload is built after guardrails write to metadata. + # Cache the result to avoid scanning litellm.callbacks twice. + _post_call_guardrails_active: Final = self._has_post_call_guardrails() + if skip_pre_call_logic: logging_obj = self.data.get("litellm_logging_obj") if logging_obj is None: @@ -2884,14 +2897,11 @@ class ProxyBaseLLMRequestProcessing: model=model, route_type=route_type, llm_router=llm_router, + defer_async_logging=_post_call_guardrails_active + and not self._is_streaming_request(data=self.data, is_streaming_request=is_streaming_request), # pyright: ignore[reportUnknownMemberType] # self.data is the untyped request dict ) phase_event("litellm.request.pre_call_completed") - # Defer async logging when post-call guardrails are configured so the - # StandardLoggingPayload is built after guardrails write to metadata. - # Cache the result to avoid scanning litellm.callbacks twice. - _post_call_guardrails_active: Final = self._has_post_call_guardrails() - # Non-streaming: defer the create_task in wrapper_async so the # SLP is built after guardrails write to metadata. Streaming # uses a separate closure mechanism (see below). diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index bfc7175cc65..8e4ba76381d 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -722,19 +722,29 @@ def _ran_or_runs_outside_pipelines( ) +def _success_log_waits_for_response_hooks(data: Mapping[str, object]) -> bool: + logging_obj: Final = data.get("litellm_logging_obj") + return isinstance(logging_obj, Logging) and logging_obj.defer_async_logging + + def _post_call_steps_named_before_headers( data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, call_type: str ) -> frozenset[str]: """ Post_call pipeline steps the executor names, as it runs them, before the response - headers go out: every step of a whole response, and of a streamed one the steps of - the pipelines that hold the stream until they finish. The LLM pass-through route builds + headers go out and the success log is built: every step of a whole response whose + success log waits for the response hooks, and of a streamed one the steps of the + pipelines that hold the stream until they finish. The LLM pass-through route builds its headers before any response step runs, so none of its steps. """ if call_type == CallTypes.allm_passthrough_route.value: return frozenset() if data.get("stream") is not True and data.get("method") not in A2A_STREAMING_METHODS: - return pipeline_managed_guardrail_names(data, "post_call") + return ( + pipeline_managed_guardrail_names(data, "post_call") + if _success_log_waits_for_response_hooks(data) + else frozenset() + ) gated: Final = _stream_gated_pipelines(data, user_api_key_dict) return frozenset() if _pipelines_stream_live(gated) else _pipeline_step_guardrail_names(gated) @@ -746,9 +756,10 @@ def _pipeline_steps_that_have_not_run( Guardrails only policy pipelines run whose pre_call hook has not run for this request: pre_call steps a pipeline skipped and the post_call steps the executor names when it runs them. A stream runs a pipeline it cannot hold as guardrails on their own, and one - whose steps all stream live after the headers go out, and the LLM pass-through route runs - its response steps after its headers are built, so those steps are named up front like - any other guardrail. + whose steps all stream live after the headers go out, the LLM pass-through route runs + its response steps after its headers are built, and a route whose success log does not + wait for the response hooks builds that log first, so those steps are named up front + like any other guardrail. """ pre_call: Final = pipeline_managed_guardrail_names(data, "pre_call") post_call: Final = _post_call_steps_named_before_headers(data, user_api_key_dict, call_type) diff --git a/tests/integration/observability/test_guardrail_pipeline_applied_header.py b/tests/integration/observability/test_guardrail_pipeline_applied_header.py index 174a88ac306..6e408d24e9a 100644 --- a/tests/integration/observability/test_guardrail_pipeline_applied_header.py +++ b/tests/integration/observability/test_guardrail_pipeline_applied_header.py @@ -183,6 +183,8 @@ def _completion(request: Request) -> Reply: return _bedrock_reply(request) if request.target == "/v1/responses": return _responses_reply(request) + if request.target == "/v1/images/generations": + return Reply(body=json.dumps({"created": 1, "data": [{"b64_json": "aW1n", "revised_prompt": None}]}).encode()) assert request.target == "/v1/chat/completions", request.target identity: Final = f"chatcmpl-{uuid.uuid4().hex}" if json.loads(request.body).get("stream") is True: @@ -675,6 +677,37 @@ def test_a_pass_through_route_names_its_response_pipeline_up_front_because_its_h ], rows +def test_an_image_generation_names_its_response_pipeline_up_front_because_its_spend_log_is_built_first( + rig: _Rig, +) -> None: + rig.upstream.drain() + with rig.gateway.scenario() as scenario: + image_model: Final = scenario.model( + model="openai/dall-e-3", api_base=f"{rig.upstream.url}/v1", api_key="sk-fixture" + ) + team: Final = scenario.team(metadata={"tags": [_RESPONSE_STOPS_AFTER_FIRST]}) + response: Final = rig.gateway.request( + "POST", + "/v1/images/generations", + {"model": image_model, "prompt": f"a cat {uuid.uuid4().hex}", "size": "1024x1024"}, + key=scenario.key(team_id=team), + ) + assert response.status_code == 200, response.text + assert [request.target for request in rig.upstream.drain()] == ["/v1/images/generations"] + assert sorted(response.headers.get(_HEADER, "").split(",")) == [_RESPONSE_FIRST, _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 []) == [ + _RESPONSE_FIRST, + _RESPONSE_SECOND, + ], rows + + def _agent_card(name: str, url: str) -> dict[str, object]: return { "protocolVersion": "0.3", diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 82dd3479321..7a2fe623032 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -483,6 +483,47 @@ class TestProxyBaseLLMRequestProcessing: pytest.fail("litellm_call_id is not a valid UUID") assert data_passed["litellm_call_id"] == returned_data["litellm_call_id"] + @pytest.mark.asyncio + @pytest.mark.parametrize("defer_async_logging", [False, True]) + async def test_common_processing_pre_call_logic_defers_the_success_log_before_pre_call_hook_runs( + self, monkeypatch: pytest.MonkeyPatch, defer_async_logging: bool + ) -> None: + processing_obj: Final = ProxyBaseLLMRequestProcessing(data={}) + mock_request: Final = MagicMock(spec=Request) + mock_request.headers = {} + deferred_when_hook_ran: Final[list[bool]] = [] + + async def mock_add_litellm_data_to_request(*args: object, **kwargs: object) -> dict[str, object]: + return {} + + async def mock_pre_call_hook( + user_api_key_dict: UserAPIKeyAuth, data: dict[str, object], call_type: str, skip_guardrails: bool = False + ) -> dict[str, object]: + logging_obj: Final = data["litellm_logging_obj"] + assert isinstance(logging_obj, LiteLLMLoggingObj) + deferred_when_hook_ran.append(logging_obj.defer_async_logging) + return data + + mock_proxy_logging_obj: Final = MagicMock(spec=ProxyLogging) + mock_proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=mock_pre_call_hook) + monkeypatch.setattr( + litellm.proxy.common_request_processing, + "add_litellm_data_to_request", + mock_add_litellm_data_to_request, + ) + + await processing_obj.common_processing_pre_call_logic( + request=mock_request, + general_settings={}, + user_api_key_dict=MagicMock(spec=UserAPIKeyAuth), + proxy_logging_obj=mock_proxy_logging_obj, + proxy_config=MagicMock(spec=ProxyConfig), + route_type="acompletion", + defer_async_logging=defer_async_logging, + ) + + assert deferred_when_hook_ran == [defer_async_logging] + @pytest.mark.asyncio @pytest.mark.parametrize("requested_model", [{"bad": "value"}, ["gpt-5.2"], 1]) async def test_common_processing_pre_call_logic_rejects_a_non_string_model_with_400( diff --git a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py index ddeaf45d32b..5839bdcda4d 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -28,6 +28,7 @@ from litellm.integrations.custom_guardrail import ( ) from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.prometheus import PrometheusLogger +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.guardrail_translation.utils import stream_item_field from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.callback_utils import ( @@ -1203,6 +1204,24 @@ def _two_step_pipeline_data( } +def _logging_obj(*, defer: bool) -> LiteLLMLoggingObj: + logging_obj: Final = LiteLLMLoggingObj( + model="m", + messages=[], + stream=False, + call_type="acompletion", + start_time=None, + litellm_call_id="call-1", + function_id="fn-1", + ) + logging_obj.defer_async_logging = defer + return logging_obj + + +def _logged_after_response_hooks(data: dict[str, object]) -> dict[str, object]: + return {**data, "litellm_logging_obj": _logging_obj(defer=True)} + + @pytest.mark.asyncio @pytest.mark.parametrize( ("first_on_pass", "expected_calls", "expected_header"), @@ -1351,7 +1370,7 @@ async def test_post_call_pipeline_steps_are_named_only_once_they_run( data: Final = await proxy_logging.pre_call_hook( user_api_key_dict=make_user_api_key_auth(), - data=_two_step_pipeline_data("post_call", "allow"), + data=_logged_after_response_hooks(_two_step_pipeline_data("post_call", "allow")), call_type="completion", ) assert "applied_guardrails" not in data["metadata"] @@ -1383,7 +1402,7 @@ async def test_post_call_pipeline_names_every_step_that_ran_in_metadata_after_a_ data: Final = await proxy_logging.pre_call_hook( user_api_key_dict=make_user_api_key_auth(), - data=_two_step_pipeline_data("post_call", "next"), + data=_logged_after_response_hooks(_two_step_pipeline_data("post_call", "next")), call_type="completion", ) await proxy_logging.post_call_success_hook( @@ -1399,6 +1418,20 @@ async def test_post_call_pipeline_names_every_step_that_ran_in_metadata_after_a_ ("asks_for", "request_route", "streams_live", "expected_header"), [ pytest.param({"stream": False}, "/custom/stream", False, None, id="whole-response"), + pytest.param( + {"litellm_logging_obj": None}, + "/v1/images/generations", + False, + ["first", "second"], + id="whole-response-logged-before-it-runs", + ), + pytest.param( + {"litellm_logging_obj": _logging_obj(defer=False)}, + "/v1/chat/completions", + False, + ["first", "second"], + id="whole-response-log-not-deferred", + ), pytest.param({"stream": True}, "/v1/chat/completions", False, None, id="stream-held-for-the-pipeline"), pytest.param({"stream": True}, "/v1/chat/completions", True, ["first", "second"], id="stream-released-live"), pytest.param( @@ -1420,7 +1453,9 @@ async def test_pre_call_hook_names_up_front_the_post_call_steps_a_stream_runs_af for guardrail in guardrails: guardrail.streaming_buffer_until_moderated = not streams_live monkeypatch.setattr(litellm, "callbacks", guardrails) - data: Final = _two_step_pipeline_data("post_call", "next", first_on_fail="next", second_on_fail="next") + data: Final = _logged_after_response_hooks( + _two_step_pipeline_data("post_call", "next", first_on_fail="next", second_on_fail="next") + ) data.update(asks_for) result: Final = await proxy_logging.pre_call_hook( @@ -1442,7 +1477,7 @@ async def test_pre_call_hook_names_up_front_the_post_call_steps_the_llm_pass_thr ) -> None: guardrails: Final = [_PipelineStepGuardrail(name, "post_call") for name in ("first", "second")] monkeypatch.setattr(litellm, "callbacks", guardrails) - data: Final = _two_step_pipeline_data("post_call", "allow") + data: Final = _logged_after_response_hooks(_two_step_pipeline_data("post_call", "allow")) data["stream"] = stream result: Final = await proxy_logging.pre_call_hook( From 967c4f336743c6e641e599774e08f4e5ab05621e Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 28 Oct 2026 11:10:58 -0700 Subject: [PATCH 11/13] test(guardrails): pin the deferred-log flag through rate-limit fallback and an A2A step that allows --- .../observability/test_guardrail_pipeline_applied_header.py | 1 + tests/unit/proxy/test_common_request_processing.py | 4 ++++ 2 files changed, 5 insertions(+) diff --git a/tests/integration/observability/test_guardrail_pipeline_applied_header.py b/tests/integration/observability/test_guardrail_pipeline_applied_header.py index 6e408d24e9a..55f6f0764c3 100644 --- a/tests/integration/observability/test_guardrail_pipeline_applied_header.py +++ b/tests/integration/observability/test_guardrail_pipeline_applied_header.py @@ -756,6 +756,7 @@ def agent(rig: _Rig) -> Iterator[str]: pytest.param( _RESPONSE_MOVES_PAST_FIRST, [_RESPONSE_BLOCKER, _RESPONSE_SECOND], id="first-step-flags-and-moves-on" ), + pytest.param(_RESPONSE_STOPS_AFTER_FIRST, [_RESPONSE_FIRST], id="first-step-allows"), ], ) def test_an_agent_message_names_every_response_step_that_ran_in_the_spend_log_and_the_trace( diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 7a2fe623032..7d091327613 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -7542,10 +7542,12 @@ class TestPreCallWithFallbacksOnLocalRateLimit: processor = ProxyBaseLLMRequestProcessing(data={"model": primary_model}) call_count = 0 + deferred_per_call = [] async def mock_pre_call_logic(**kwargs): nonlocal call_count call_count += 1 + deferred_per_call.append(kwargs["defer_async_logging"]) model_in_data = processor.data.get("model") if model_in_data == primary_model: raise ProxyRateLimitError( @@ -7578,10 +7580,12 @@ class TestPreCallWithFallbacksOnLocalRateLimit: model=primary_model, route_type="acompletion", llm_router=mock_router, + defer_async_logging=True, ) assert processor.data["model"] == fallback_model assert call_count == 2 + assert deferred_per_call == [True, True] @pytest.mark.asyncio async def test_raises_when_no_fallbacks_configured(self): From 27cb942c58cd990db40541555053fe64911f8846 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 28 Oct 2026 12:08:40 -0700 Subject: [PATCH 12/13] test(guardrails): drive the deferred-log flag through a real pre-call hook --- .../proxy/test_common_request_processing.py | 59 ++++++++++--------- 1 file changed, 30 insertions(+), 29 deletions(-) diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 7d091327613..0737e668a7a 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -25,7 +25,7 @@ from litellm.litellm_core_utils.bug_report import ( build_bug_report, ) from litellm.types.router import RouterRateLimitError, RouterRateLimitErrorBasic -from litellm.types.utils import CachingDetails, Delta, ModelResponseStream, StreamingChoices +from litellm.types.utils import CachingDetails, CallTypesLiteral, Delta, ModelResponseStream, StreamingChoices from litellm.constants import ( CLIENT_REQUESTED_MODEL_SCOPE_KEY, MAX_LITELLM_CALL_ID_LENGTH, @@ -207,6 +207,24 @@ def test_include_guardrail_response_requested_is_false_without_exact_true(): assert include_guardrail_response_requested({}) is False +class _RecordsDeferral(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.deferred: list[bool] = [] + + async def async_pre_call_hook( + self, + user_api_key_dict: ProxyUserAPIKeyAuth, + cache: DualCache, + data: dict, + call_type: CallTypesLiteral, + ) -> dict: + logging_obj: Final = data["litellm_logging_obj"] + assert isinstance(logging_obj, LiteLLMLoggingObj) + self.deferred.append(logging_obj.defer_async_logging) + return data + + class TestProxyBaseLLMRequestProcessing: @pytest.mark.asyncio async def test_base_passthrough_process_llm_request_preserves_litellm_headers_for_non_streaming_response( @@ -488,41 +506,24 @@ class TestProxyBaseLLMRequestProcessing: async def test_common_processing_pre_call_logic_defers_the_success_log_before_pre_call_hook_runs( self, monkeypatch: pytest.MonkeyPatch, defer_async_logging: bool ) -> None: - processing_obj: Final = ProxyBaseLLMRequestProcessing(data={}) - mock_request: Final = MagicMock(spec=Request) - mock_request.headers = {} - deferred_when_hook_ran: Final[list[bool]] = [] + from litellm.proxy.proxy_server import ProxyConfig as RealProxyConfig - async def mock_add_litellm_data_to_request(*args: object, **kwargs: object) -> dict[str, object]: - return {} + recorder: Final = _RecordsDeferral() + monkeypatch.setattr(litellm, "callbacks", [recorder]) - async def mock_pre_call_hook( - user_api_key_dict: UserAPIKeyAuth, data: dict[str, object], call_type: str, skip_guardrails: bool = False - ) -> dict[str, object]: - logging_obj: Final = data["litellm_logging_obj"] - assert isinstance(logging_obj, LiteLLMLoggingObj) - deferred_when_hook_ran.append(logging_obj.defer_async_logging) - return data - - mock_proxy_logging_obj: Final = MagicMock(spec=ProxyLogging) - mock_proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=mock_pre_call_hook) - monkeypatch.setattr( - litellm.proxy.common_request_processing, - "add_litellm_data_to_request", - mock_add_litellm_data_to_request, - ) - - await processing_obj.common_processing_pre_call_logic( - request=mock_request, + await ProxyBaseLLMRequestProcessing( + data={"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hi"}]} + ).common_processing_pre_call_logic( + request=Request(scope={"type": "http", "method": "POST", "path": "/v1/chat/completions", "headers": []}), general_settings={}, - user_api_key_dict=MagicMock(spec=UserAPIKeyAuth), - proxy_logging_obj=mock_proxy_logging_obj, - proxy_config=MagicMock(spec=ProxyConfig), + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), + proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()), + proxy_config=RealProxyConfig(), route_type="acompletion", defer_async_logging=defer_async_logging, ) - assert deferred_when_hook_ran == [defer_async_logging] + assert recorder.deferred == [defer_async_logging] @pytest.mark.asyncio @pytest.mark.parametrize("requested_model", [{"bad": "value"}, ["gpt-5.2"], 1]) From 076e3e374c1089f9f8403a306a090e1b0c4d9fd5 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 28 Oct 2026 12:53:53 -0700 Subject: [PATCH 13/13] test(guardrails): type the deferral recorder's request dict and fallback list --- tests/unit/proxy/test_common_request_processing.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 0737e668a7a..8d01edb0096 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -216,9 +216,9 @@ class _RecordsDeferral(CustomLogger): self, user_api_key_dict: ProxyUserAPIKeyAuth, cache: DualCache, - data: dict, + data: dict[str, object], call_type: CallTypesLiteral, - ) -> dict: + ) -> dict[str, object]: logging_obj: Final = data["litellm_logging_obj"] assert isinstance(logging_obj, LiteLLMLoggingObj) self.deferred.append(logging_obj.defer_async_logging) @@ -7543,7 +7543,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit: processor = ProxyBaseLLMRequestProcessing(data={"model": primary_model}) call_count = 0 - deferred_per_call = [] + deferred_per_call: Final[list[bool]] = [] async def mock_pre_call_logic(**kwargs): nonlocal call_count