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 1/9] 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 2/9] 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 3/9] 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 4/9] 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 5/9] 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 6/9] 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 7/9] 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 8/9] 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 9/9] 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]