From eea1d0f2696d6ab6b67e8b208c85bae9fa624e1e Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 26 Sep 2026 18:05:21 -0700 Subject: [PATCH] fix(responses): stream guardrail pre-call block as SSE with a typed output item (#42507) * fix(responses): stream guardrail pre-call block as SSE with a typed output item Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(responses): import blocked usage helper from the guardrail utils module Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(e2e): drop narrating docstrings and poll without rebinding Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): cover pre-call guardrail block on /v1/responses stream and json Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): audit cells for responses guardrail block contract Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): observe upstream on the recorded chat route for responses denial cells Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): tidy responses denial audit cells Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): wait for worker count to recover after SIGKILL Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): require a replacement worker after SIGKILL Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(responses): type the blocked response test helpers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: yucheng --- .../guardrail_translation/handler.py | 6 +- .../proxy/response_api_endpoints/endpoints.py | 32 +- tests/e2e/coverage_registry/guardrail.yaml | 1 + tests/e2e/guardrails/guardrails_client.py | 29 + ...est_responses_pre_call_block_stream_e2e.py | 154 ++++ .../observability/test_guardrail_effects.py | 695 +++++++++++++++++- .../response_api_endpoints/test_endpoints.py | 148 +++- .../proxy/test_blocked_response_usage.py | 20 +- 8 files changed, 1022 insertions(+), 63 deletions(-) create mode 100644 tests/e2e/guardrails/test_responses_pre_call_block_stream_e2e.py diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 66cebe0175d..d6d68e0607a 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -1540,7 +1540,7 @@ class OpenAIResponsesHandler(BaseTranslation): from litellm.responses.streaming_iterator import build_synthetic_response_events return build_synthetic_response_events( - transformed=_blocked_response(exc, response_id=f"resp_{uuid.uuid4()}", model=exc.model), + transformed=build_blocked_response(exc), logging_obj=None, chunk_size=max(len(exc.message), 1), ) @@ -1648,6 +1648,10 @@ def _blocked_output_item(exc: "ModifyResponseException") -> GenericResponseOutpu return GenericResponseOutputItem.model_validate(payload) +def build_blocked_response(exc: "ModifyResponseException") -> ResponsesAPIResponse: + return _blocked_response(exc, response_id=f"resp_{uuid.uuid4()}", model=exc.model) + + def _blocked_response( exc: "ModifyResponseException", response_id: str, diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 75eefb2e73b..c5d702ad65a 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -1,13 +1,11 @@ import asyncio import contextlib import json -import time -from collections.abc import AsyncIterator, Awaitable, Mapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Mapping, Sequence from enum import Enum from functools import partial from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, NamedTuple, Protocol, TypeAlias, cast, get_args -from uuid import uuid4 import fastapi from fastapi import APIRouter, Depends, HTTPException, Request, Response @@ -21,8 +19,9 @@ from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.constants import EMPTY_MAPPING from litellm.integrations.custom_guardrail import ModifyResponseException -from litellm.llms.base_llm.guardrail_translation.utils import ( - blocked_responses_api_usage as _blocked_responses_api_usage, +from litellm.llms.openai.responses.guardrail_translation.handler import ( + OpenAIResponsesHandler, + build_blocked_response, ) from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import ( @@ -30,7 +29,7 @@ from litellm.proxy.auth.user_api_key_auth import ( user_api_key_auth, user_api_key_auth_websocket, ) -from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing, create_response from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, _safe_set_request_parsed_body, @@ -440,17 +439,16 @@ async def responses_api( request_data=_data, ) - violation_text: Final = e.message - response_obj: Final = ResponsesAPIResponse( - id=f"resp_{uuid4()}", - object="response", - created_at=int(time.time()), - model=e.model or data.get("model"), - output=cast(Any, [{"content": [{"type": "text", "text": violation_text}]}]), - status="completed", - usage=_blocked_responses_api_usage(e.original_response), - ) - return response_obj + if data.get("stream") is True: + block_chunks: Final = OpenAIResponsesHandler().build_block_sse_chunks(e) + + async def _blocked_stream() -> AsyncGenerator[str, None]: + for chunk in block_chunks: + yield chunk.decode() + yield "data: [DONE]\n\n" + + return await create_response(generator=_blocked_stream(), media_type="text/event-stream", headers={}) + return build_blocked_response(e) except Exception as e: raise await processor._handle_llm_api_exception( e=e, diff --git a/tests/e2e/coverage_registry/guardrail.yaml b/tests/e2e/coverage_registry/guardrail.yaml index f49568c883b..920a288aea6 100644 --- a/tests/e2e/coverage_registry/guardrail.yaml +++ b/tests/e2e/coverage_registry/guardrail.yaml @@ -37,4 +37,5 @@ - {id: guardrail.mcp_security.pre_call.blocks, module: guardrail, tier: P2, hook_point: pre_call, assertions: [blocks], exercised_on: [mcp_operations], source: "guardrail_hooks/mcp_security", rationale: "MCP protocol security"} - {id: guardrail.llm_as_a_judge.pre_call.blocks, module: guardrail, tier: P2, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/llm_as_a_judge", rationale: "LLM-based judgment guardrail"} - {id: guardrail.litellm_content_filter.pre_mcp_call.blocks, module: guardrail, tier: P1, hook_point: pre_mcp_call, assertions: [blocks], exercised_on: [mcp_operations], source: "guardrail_hooks/litellm_content_filter/content_filter.py:_scan_mcp_tool_call_arguments", rationale: "A general content-filter guardrail configured mode=pre_mcp_call blocks a banned keyword in an MCP tool call's arguments before it reaches the upstream MCP server; a clean argument passes"} +- {id: guardrail.custom_code.pre_call.blocks, module: guardrail, tier: P1, hook_point: pre_call, assertions: [blocks], exercised_on: [responses], source: "response_api_endpoints/endpoints.py ModifyResponseException handler", rationale: "A pre_call custom_code block on /v1/responses must answer in the requested shape: SSE response.completed with a completed assistant output_text message item when stream=true, schema-valid JSON when not, both with zero usage"} - {id: guardrail.dispatch.pre_call.rejects_unknown_name, module: guardrail, tier: P1, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "proxy guardrail dispatch (per-request `guardrails` selector)", rationale: "A request naming a guardrail this proxy does not serve must fail closed with a 4xx; today it is silently served unguarded, so a typo'd name drops the protection the caller asked for"} diff --git a/tests/e2e/guardrails/guardrails_client.py b/tests/e2e/guardrails/guardrails_client.py index 17223dc36fa..1f4fc43355b 100644 --- a/tests/e2e/guardrails/guardrails_client.py +++ b/tests/e2e/guardrails/guardrails_client.py @@ -111,6 +111,15 @@ class ToolPermissionParamsBody(GuardrailParamsBase): on_disallowed_action: Literal["block", "rewrite"] = "block" +class CustomCodeParamsBody(GuardrailParamsBase): + """Custom-code guardrail params: `custom_code` is the sandboxed source the + proxy compiles, which must define `apply_guardrail(inputs, request_data, + input_type)` returning `allow()` or `block(reason)`.""" + + guardrail: Literal["custom_code"] = "custom_code" + custom_code: str + + GuardrailParamsBody = ( ContentFilterParamsBody | BedrockGuardrailParamsBody @@ -118,6 +127,7 @@ GuardrailParamsBody = ( | BlockCodeExecutionParamsBody | PresidioParamsBody | ToolPermissionParamsBody + | CustomCodeParamsBody ) @@ -174,6 +184,7 @@ class _ResponsesGuardrailBody(BaseModel): model: str input: str guardrails: list[str] | None = None + stream: bool | None = None @dataclass(frozen=True, slots=True) @@ -509,6 +520,24 @@ class GuardrailsClient: json=_ResponsesGuardrailBody(model=model, input=text, guardrails=guardrails), ) + def responses_stream_raw( + self, + key: str, + model: str, + text: str, + *, + guardrails: list[str] | None = None, + ) -> StreamingResponse: + """Drive /v1/responses with stream=true, returning the raw HTTP outcome: + a streamed block is judged on status, content-type, and the SSE event + sequence, not a typed JSON body.""" + return self.proxy.transport.send( + "/v1/responses", + headers=self.proxy.transport.bearer(key), + json=_ResponsesGuardrailBody(model=model, input=text, guardrails=guardrails, stream=True), + stream=True, + ) + def apply_guardrail(self, key: str, *, name: str, text: str) -> Result[ApplyGuardrailResponse]: return self.proxy.transport.post( "/guardrails/apply_guardrail", diff --git a/tests/e2e/guardrails/test_responses_pre_call_block_stream_e2e.py b/tests/e2e/guardrails/test_responses_pre_call_block_stream_e2e.py new file mode 100644 index 00000000000..93512e2a64c --- /dev/null +++ b/tests/e2e/guardrails/test_responses_pre_call_block_stream_e2e.py @@ -0,0 +1,154 @@ +from __future__ import annotations + +import time +from collections.abc import Callable +from typing import Final + +import pytest +from e2e_config import POLL_INTERVAL, POLL_TIMEOUT, unique_marker +from e2e_http import StreamingResponse +from guardrails_client import CustomCodeParamsBody, GuardrailsClient +from lifecycle import ResourceManager +from pydantic import BaseModel, TypeAdapter + +pytestmark = pytest.mark.e2e + +DENIAL: Final = "This model is not currently available. Please contact support if you think this is a mistake." + +CUSTOM_CODE: Final = f''' +def apply_guardrail(inputs, request_data, input_type): + return block("{DENIAL}") +''' + + +class _ContentPart(BaseModel): + type: str + text: str | None = None + + +class _OutputItem(BaseModel): + type: str | None = None + id: str | None = None + role: str | None = None + status: str | None = None + content: list[_ContentPart] = [] + + +class _Usage(BaseModel): + total_tokens: int = 0 + + +class _ResponseBody(BaseModel): + output: list[_OutputItem] = [] + usage: _Usage | None = None + + +class _EventHead(BaseModel): + type: str + + +class _CompletedEvent(BaseModel): + type: str + response: _ResponseBody + + +_EVENT_HEAD: Final = TypeAdapter(_EventHead) + + +def _denial_delivered(result: StreamingResponse) -> bool: + if not result.ok: + return False + if DENIAL in result.body: + return True + return any(DENIAL in event for event in result.stream_events) + + +def _poll_terminal(result: StreamingResponse) -> bool: + if _denial_delivered(result): + return True + if result.ok: + return False + return "Guardrail not found" not in result.body and result.status_code not in (-1, 401, 429) + + +def _poll_attempt(call: Callable[[], StreamingResponse], deadline: float) -> StreamingResponse: + result: Final = call() + if _poll_terminal(result) or time.monotonic() >= deadline: + return result + time.sleep(POLL_INTERVAL) + return _poll_attempt(call, deadline) + + +def _poll_for_block(call: Callable[[], StreamingResponse]) -> StreamingResponse: + return _poll_attempt(call, time.monotonic() + POLL_TIMEOUT) + + +def _assert_blocked_response(response: _ResponseBody) -> None: + item = next(iter(response.output), None) + assert item is not None, f"blocked response carried no output item: {response.output!r}" + assert item.type == "message", f"output[0] must be a message item, got {item.type!r}: {item!r}" + assert item.role == "assistant", f"output[0] role must be assistant, got {item.role!r}" + assert item.status == "completed", f"output[0] status must be completed, got {item.status!r}" + part = next(iter(item.content), None) + assert part is not None, f"output[0] carried no content part: {item!r}" + assert part.type == "output_text", f"content[0] must be output_text, got {part.type!r}" + assert part.text == DENIAL, f"content[0] text must be the denial, got {part.text!r}" + assert response.usage is not None and response.usage.total_tokens == 0, ( + f"a blocked response never reached a provider, usage must be zero: {response.usage!r}" + ) + + +class TestResponsesPreCallBlock: + def _register_block(self, client: GuardrailsClient, resources: ResourceManager) -> str: + name: Final = f"e2e-custom-code-responses-block-{unique_marker()}" + guardrail_id: Final = client.register( + name, + CustomCodeParamsBody(mode="pre_call", default_on=False, custom_code=CUSTOM_CODE), + ) + resources.defer(lambda: client.delete_guardrail(guardrail_id)) + return name + + @pytest.mark.covers("guardrail.custom_code.pre_call.blocks", exercised_on=["responses"]) + def test_stream_block_is_sse_with_completed_assistant_message( + self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str + ) -> None: + name: Final = self._register_block(client, resources) + model: Final = client.create_backend_model( + resources, prefix="e2e-responses-block", backend="openai/gpt-4.1-mini", api_key="os.environ/OPENAI_API_KEY" + ) + + result: Final = _poll_for_block( + lambda: client.responses_stream_raw(scoped_key, model, "say hi", guardrails=[name]) + ) + + assert result.status_code == 200, f"a pre_call block answers 200, got {result.status_code}: {result.body[:400]}" + assert (result.content_type or "").startswith("text/event-stream"), ( + f"stream=true must answer SSE, got content-type {result.content_type!r}: {result.body[:400]}" + ) + events: Final = tuple(_EVENT_HEAD.validate_json(payload).type for payload in result.stream_events) + completed: Final = tuple( + _CompletedEvent.model_validate_json(payload) + for payload, event_type in zip(result.stream_events, events) + if event_type == "response.completed" + ) + assert len(completed) == 1, ( + f"the denial stream must end in exactly one response.completed event, got events {events!r}" + ) + _assert_blocked_response(completed[0].response) + + @pytest.mark.covers("guardrail.custom_code.pre_call.blocks", exercised_on=["responses"]) + def test_non_stream_block_is_schema_valid_json( + self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str + ) -> None: + name: Final = self._register_block(client, resources) + model: Final = client.create_backend_model( + resources, prefix="e2e-responses-block", backend="openai/gpt-4.1-mini", api_key="os.environ/OPENAI_API_KEY" + ) + + result: Final = _poll_for_block(lambda: client.responses(scoped_key, model, "say hi", guardrails=[name])) + + assert result.status_code == 200, f"a pre_call block answers 200, got {result.status_code}: {result.body[:400]}" + assert (result.content_type or "").startswith("application/json"), ( + f"a non-streaming block answers JSON, got content-type {result.content_type!r}" + ) + _assert_blocked_response(_ResponseBody.model_validate_json(result.body)) diff --git a/tests/integration/observability/test_guardrail_effects.py b/tests/integration/observability/test_guardrail_effects.py index 4fac42a796d..9f5f3da4302 100644 --- a/tests/integration/observability/test_guardrail_effects.py +++ b/tests/integration/observability/test_guardrail_effects.py @@ -1,15 +1,22 @@ import json +import os +import signal +import socket import uuid +from concurrent.futures import ThreadPoolExecutor from pathlib import Path from typing import Final +import httpx +import psutil import pytest import yaml from integration._support.client import Gateway, eventually, object_value from integration._support.database import read_rows from integration._support.mcp import mcp_peer, register_mcp, tool_names -from integration._support.process import owned_proxy +from integration._support.process import group_members, owned_proxy, owned_proxy_process from integration._support.wire import Reply, Request, wire_server +from openai import AsyncOpenAI, OpenAI @pytest.mark.covers("other.observability.guardrails.rewrite_reaches_correct_anthropic_positions") @@ -208,8 +215,6 @@ def test_guardrail_denial_prevents_provider_and_preserves_allowed_control(gatewa with owned_proxy(gateway, tmp_path, {}, config=path) as candidate, candidate.scenario() as scenario: model: Final = scenario.model() key: Final = scenario.key(models=[model]) - import httpx - with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as observed: observed.get("/__observations") denied: Final = candidate.request( @@ -500,3 +505,687 @@ def test_request_selected_mcp_guardrail_blocks_direct_and_virtual_calls(gateway: assert len(calls) == 1 assert calls[0]["body"]["params"]["name"] == tool assert calls[0]["body"]["params"]["arguments"] == arguments + + +_RESPONSES_DENIAL: Final = "This model is not currently available." + + +def _deny_guardrail(name: str, denial: str = _RESPONSES_DENIAL) -> dict[str, object]: + return { + "guardrail_name": name, + "litellm_params": { + "guardrail": "custom_code", + "mode": "pre_call", + "default_on": False, + "custom_code": (f"def apply_guardrail(inputs, request_data, input_type):\n return block({denial!r})\n"), + }, + } + + +def _responses_denial_config(tmp_path: Path, identity: str, denial: str = _RESPONSES_DENIAL) -> Path: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["guardrails"] = [_deny_guardrail(identity, denial)] + path: Final = tmp_path / "responses-deny.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +def _assert_blocked_message_item(item: dict[str, object], response: dict[str, object]) -> None: + assert item["type"] == "message", item + assert item["role"] == "assistant", item + assert item["status"] == "completed", item + assert str(item["id"]).startswith("msg_"), item + assert item["content"] == [{"type": "output_text", "text": _RESPONSES_DENIAL, "annotations": []}], item + assert response["status"] == "completed", response + usage: Final = response["usage"] + assert isinstance(usage, dict), response + assert (usage["input_tokens"], usage["output_tokens"], usage["total_tokens"]) == (0, 0, 0), usage + + +def _response_id(index: int, response: httpx.Response) -> str: + assert response.status_code == 200, (index, response.text) + if index % 3 == 0: + assert response.headers["content-type"].startswith("text/event-stream"), response.text + return str(_blocked_stream_events(response.text)[-1]["response"]["id"]) + if index % 3 == 1: + assert response.headers["content-type"].startswith("text/event-stream"), response.text + blocked: Final = _blocked_stream_events(response.text)[-1]["response"] + _assert_blocked_message_item(blocked["output"][0], blocked) + return str(blocked["id"]) + assert response.headers["content-type"].startswith("application/json"), response.text + body: Final = response.json() + _assert_blocked_message_item(body["output"][0], body) + return str(body["id"]) + + +@pytest.mark.covers("other.observability.guardrails.responses_pre_call_denial_streams_typed_message") +def test_responses_pre_call_denial_streams_sse_with_typed_message_item(gateway: Gateway, tmp_path: Path) -> None: + identity: Final = "guardrail" + uuid.uuid4().hex + config: Final = _responses_denial_config(tmp_path, identity) + with owned_proxy(gateway, tmp_path, {}, config=config) as candidate, candidate.scenario() as scenario: + model: Final = scenario.model(use_chat_completions_api=True) + with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as observed: + observed.get("/__observations") + response: Final = candidate.request( + "POST", + "/v1/responses", + {"model": model, "input": "say hi", "stream": True, "guardrails": [identity]}, + ) + assert response.status_code == 200, response.text + assert response.headers["content-type"].startswith("text/event-stream"), ( + response.headers["content-type"], + response.text, + ) + lines: Final = tuple(line for line in response.text.split("\n") if line.startswith("data: ")) + assert lines[-1] == "data: [DONE]", response.text + events: Final = tuple(json.loads(line.removeprefix("data: ")) for line in lines[:-1]) + kinds: Final = tuple(event["type"] for event in events) + assert tuple(kind for kind in kinds if kind != "response.output_text.delta") == ( + "response.created", + "response.in_progress", + "response.output_item.added", + "response.content_part.added", + "response.output_text.done", + "response.content_part.done", + "response.output_item.done", + "response.completed", + ), kinds + assert kinds.index("response.output_text.delta") == kinds.index("response.content_part.added") + 1, kinds + assert "".join(event["delta"] for event in events if event["type"] == "response.output_text.delta") == ( + _RESPONSES_DENIAL + ) + completed: Final = events[-1]["response"] + assert completed["output"] == [events[-2]["item"]], (completed, events[-2]) + _assert_blocked_message_item(completed["output"][0], completed) + assert observed.get("/__observations").json()["requests"] == [] + + +@pytest.mark.covers("other.observability.guardrails.responses_pre_call_denial_returns_typed_message") +def test_responses_pre_call_denial_returns_json_with_typed_message_item(gateway: Gateway, tmp_path: Path) -> None: + identity: Final = "guardrail" + uuid.uuid4().hex + config: Final = _responses_denial_config(tmp_path, identity) + with owned_proxy(gateway, tmp_path, {}, config=config) as candidate, candidate.scenario() as scenario: + model: Final = scenario.model(use_chat_completions_api=True) + with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as observed: + observed.get("/__observations") + response: Final = candidate.request( + "POST", "/v1/responses", {"model": model, "input": "say hi", "guardrails": [identity]} + ) + assert response.status_code == 200, response.text + assert response.headers["content-type"].startswith("application/json"), response.headers["content-type"] + body: Final = response.json() + assert body["object"] == "response", body + assert len(body["output"]) == 1, body + _assert_blocked_message_item(body["output"][0], body) + assert observed.get("/__observations").json()["requests"] == [] + + +_RESPONSES_OUTPUT_DENIAL: Final = "Output withheld by policy." +_UPSTREAM_INPUT_TOKENS: Final = 20 +_UPSTREAM_OUTPUT_TOKENS: Final = 20 +_UPSTREAM_TOTAL_TOKENS: Final = 40 + + +def _responses_output_denial_config(tmp_path: Path, identity: str, model: str) -> Path: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["guardrails"] = [ + { + "guardrail_name": identity, + "litellm_params": { + "guardrail": "custom_code", + "mode": "post_call", + "default_on": False, + "custom_code": ( + "def apply_guardrail(inputs, request_data, input_type):\n" + f" return block({_RESPONSES_OUTPUT_DENIAL!r})\n" + ), + }, + } + ] + config["policies"] = { + f"{identity}-pipeline": { + "guardrails": {"add": [identity]}, + "pipeline": { + "mode": "post_call", + "steps": [ + { + "guardrail": identity, + "on_pass": "allow", + "on_fail": "modify_response", + "modify_response_message": _RESPONSES_OUTPUT_DENIAL, + } + ], + }, + } + } + config["policy_attachments"] = [{"policy": f"{identity}-pipeline", "models": [model]}] + path: Final = tmp_path / "responses-output-deny.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +def _blocked_stream_events(text: str) -> tuple[dict[str, object], ...]: + lines: Final = tuple(line for line in text.split("\n") if line.startswith("data: ")) + assert lines[-1] == "data: [DONE]", text + return tuple(json.loads(line.removeprefix("data: ")) for line in lines[:-1]) + + +def _dead_api_base() -> str: + with socket.socket() as reserve: + reserve.bind(("127.0.0.1", 0)) + port: Final = reserve.getsockname()[1] + return f"http://127.0.0.1:{port}/v1" + + +@pytest.mark.covers("other.observability.guardrails.responses_pre_call_denial_openai_sdk_streams_typed_message") +def test_responses_pre_call_denial_openai_sdk_streams_typed_message(gateway: Gateway, tmp_path: Path) -> None: + identity: Final = "guardrail" + uuid.uuid4().hex + config: Final = _responses_denial_config(tmp_path, identity) + with owned_proxy(gateway, tmp_path, {}, config=config) as candidate, candidate.scenario() as scenario: + model: Final = scenario.model(use_chat_completions_api=True) + client: Final = OpenAI( + base_url=f"{candidate.client.base_url}/v1", api_key=candidate.key, max_retries=0, timeout=15 + ) + events: Final = tuple( + client.responses.create(model=model, input="say hi", stream=True, extra_body={"guardrails": [identity]}) + ) + assert events[-1].type == "response.completed", [event.type for event in events] + completed: Final = events[-1].response + assert completed is not None and len(completed.output) == 1, completed + item: Final = completed.output[0] + assert item.type == "message", item + assert item.role == "assistant" and item.status == "completed", item + assert item.content[0].type == "output_text" and item.content[0].text == _RESPONSES_DENIAL, item.content + assert completed.usage is not None and completed.usage.total_tokens == 0, completed.usage + + +@pytest.mark.covers("other.observability.guardrails.responses_pre_call_denial_openai_async_sdk_streams_typed_message") +async def test_responses_pre_call_denial_openai_async_sdk_streams_typed_message( + gateway: Gateway, tmp_path: Path +) -> None: + identity: Final = "guardrail" + uuid.uuid4().hex + config: Final = _responses_denial_config(tmp_path, identity) + with owned_proxy(gateway, tmp_path, {}, config=config) as candidate, candidate.scenario() as scenario: + model: Final = scenario.model(use_chat_completions_api=True) + client: Final = AsyncOpenAI( + base_url=f"{candidate.client.base_url}/v1", api_key=candidate.key, max_retries=0, timeout=15 + ) + stream: Final = await client.responses.create( + model=model, input="say hi", stream=True, extra_body={"guardrails": [identity]} + ) + kinds: Final = [event.type async for event in stream] + assert kinds[-1] == "response.completed", kinds + assert "response.output_text.delta" in kinds, kinds + assert "response.in_progress" in kinds, kinds + + +@pytest.mark.covers("other.observability.guardrails.responses_pre_call_denial_openai_sdk_returns_typed_message") +def test_responses_pre_call_denial_openai_sdk_returns_typed_message(gateway: Gateway, tmp_path: Path) -> None: + identity: Final = "guardrail" + uuid.uuid4().hex + config: Final = _responses_denial_config(tmp_path, identity) + with owned_proxy(gateway, tmp_path, {}, config=config) as candidate, candidate.scenario() as scenario: + model: Final = scenario.model(use_chat_completions_api=True) + client: Final = OpenAI( + base_url=f"{candidate.client.base_url}/v1", api_key=candidate.key, max_retries=0, timeout=15 + ) + body: Final = client.responses.create(model=model, input="say hi", extra_body={"guardrails": [identity]}) + assert body.object == "response" and body.status == "completed", body + assert len(body.output) == 1, body.output + item: Final = body.output[0] + assert item.type == "message" and item.role == "assistant", item + assert item.content[0].type == "output_text" and item.content[0].text == _RESPONSES_DENIAL, item.content + assert body.output_text == _RESPONSES_DENIAL, body + assert body.usage is not None and body.usage.total_tokens == 0, body.usage + + +@pytest.mark.covers("other.observability.guardrails.responses_pre_call_denial_stream_false_returns_json") +def test_responses_pre_call_denial_stream_false_returns_json(gateway: Gateway, tmp_path: Path) -> None: + identity: Final = "guardrail" + uuid.uuid4().hex + config: Final = _responses_denial_config(tmp_path, identity) + with owned_proxy(gateway, tmp_path, {}, config=config) as candidate, candidate.scenario() as scenario: + model: Final = scenario.model(use_chat_completions_api=True) + response: Final = candidate.request( + "POST", "/v1/responses", {"model": model, "input": "say hi", "stream": False, "guardrails": [identity]} + ) + assert response.status_code == 200, response.text + assert response.headers["content-type"].startswith("application/json"), response.text + body: Final = response.json() + _assert_blocked_message_item(body["output"][0], body) + + +@pytest.mark.covers("other.observability.guardrails.responses_pre_call_denial_stream_string_true_returns_json") +def test_responses_pre_call_denial_stream_string_true_returns_json(gateway: Gateway, tmp_path: Path) -> None: + identity: Final = "guardrail" + uuid.uuid4().hex + config: Final = _responses_denial_config(tmp_path, identity) + with owned_proxy(gateway, tmp_path, {}, config=config) as candidate, candidate.scenario() as scenario: + model: Final = scenario.model(use_chat_completions_api=True) + response: Final = candidate.request( + "POST", "/v1/responses", {"model": model, "input": "say hi", "stream": "true", "guardrails": [identity]} + ) + assert response.status_code == 200, response.text + assert response.headers["content-type"].startswith("application/json"), ( + response.headers["content-type"], + response.text, + ) + body: Final = response.json() + _assert_blocked_message_item(body["output"][0], body) + + +@pytest.mark.covers("other.observability.guardrails.responses_pre_call_denial_stream_event_vocabulary") +def test_responses_pre_call_denial_stream_event_vocabulary(gateway: Gateway, tmp_path: Path) -> None: + identity: Final = "guardrail" + uuid.uuid4().hex + second: Final = "guardrail-2-" + uuid.uuid4().hex + config: Final = _responses_denial_config(tmp_path, identity) + loaded: Final = yaml.safe_load(config.read_text()) + loaded["guardrails"].append(_deny_guardrail(second)) + config.write_text(yaml.safe_dump(loaded)) + with owned_proxy(gateway, tmp_path, {}, config=config) as candidate, candidate.scenario() as scenario: + model: Final = scenario.model(use_chat_completions_api=True) + with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as observed: + observed.get("/__observations") + response: Final = candidate.request( + "POST", + "/v1/responses", + {"model": model, "input": "say hi", "stream": True, "guardrails": [identity, second]}, + ) + assert response.status_code == 200, response.text + assert response.headers["content-type"].startswith("text/event-stream"), response.text + events: Final = _blocked_stream_events(response.text) + kinds: Final = {event["type"] for event in events} + assert kinds == { + "response.created", + "response.in_progress", + "response.output_item.added", + "response.content_part.added", + "response.output_text.delta", + "response.output_text.done", + "response.content_part.done", + "response.output_item.done", + "response.completed", + }, kinds + item_done: Final = tuple(event for event in events if event["type"] == "response.output_item.done") + assert len(item_done) == 1, events + assert len(events[-1]["response"]["output"]) == 1, events[-1] + assert observed.get("/__observations").json()["requests"] == [] + + +@pytest.mark.covers("other.observability.guardrails.responses_pre_call_denial_stream_large_denial_text") +def test_responses_pre_call_denial_stream_large_denial_text(gateway: Gateway, tmp_path: Path) -> None: + identity: Final = "guardrail" + uuid.uuid4().hex + denial: Final = ("Denied: " + "mixed ascii and unicode text " * 200 + "fin")[:5000] + config: Final = _responses_denial_config(tmp_path, identity, denial) + with owned_proxy(gateway, tmp_path, {}, config=config) as candidate, candidate.scenario() as scenario: + model: Final = scenario.model(use_chat_completions_api=True) + response: Final = candidate.request( + "POST", + "/v1/responses", + {"model": model, "input": "say hi", "stream": True, "guardrails": [identity]}, + ) + assert response.status_code == 200, response.text + assert response.headers["content-type"].startswith("text/event-stream"), response.text + events: Final = _blocked_stream_events(response.text) + assert "".join(event["delta"] for event in events if event["type"] == "response.output_text.delta") == denial + done: Final = next(event for event in events if event["type"] == "response.output_text.done") + assert done["text"] == denial, done + completed: Final = events[-1]["response"] + assert completed["output"][0]["content"][0]["text"] == denial, completed + + +@pytest.mark.covers("other.observability.guardrails.responses_pre_call_denial_stream_requests_have_distinct_ids") +def test_responses_pre_call_denial_stream_requests_have_distinct_ids(gateway: Gateway, tmp_path: Path) -> None: + identity: Final = "guardrail" + uuid.uuid4().hex + config: Final = _responses_denial_config(tmp_path, identity) + with owned_proxy(gateway, tmp_path, {}, config=config) as candidate, candidate.scenario() as scenario: + model: Final = scenario.model(use_chat_completions_api=True) + with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as observed: + observed.get("/__observations") + responses: Final = tuple( + candidate.request( + "POST", + "/v1/responses", + {"model": model, "input": "say hi", "stream": True, "guardrails": [identity]}, + ) + for _ in range(2) + ) + completed: Final = tuple(_blocked_stream_events(response.text)[-1]["response"] for response in responses) + for response in responses: + assert response.status_code == 200, response.text + assert response.headers["content-type"].startswith("text/event-stream"), response.text + assert completed[0]["id"] != completed[1]["id"], completed + assert completed[0]["output"][0]["id"] != completed[1]["output"][0]["id"], completed + assert observed.get("/__observations").json()["requests"] == [] + + +def _register_named_model(candidate: Gateway, name: str, api_base: str | None = None, **parameters: object) -> str: + created: Final = candidate.post( + "/model/new", + { + "model_name": name, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "integration-provider-key", + "api_base": api_base or f"{candidate.upstream_url}/v1", + **parameters, + }, + }, + ) + return str(created["model_info"]["id"]) + + +@pytest.mark.covers("other.observability.guardrails.responses_post_call_pipeline_denial_streams_real_usage") +def test_responses_post_call_pipeline_denial_streams_real_usage(gateway: Gateway, tmp_path: Path) -> None: + identity: Final = "guardrail" + uuid.uuid4().hex + model: Final = f"integration-{uuid.uuid4().hex}" + config: Final = _responses_output_denial_config(tmp_path, identity, model) + with owned_proxy(gateway, tmp_path, {}, config=config) as candidate: + model_id: Final = _register_named_model(candidate, model, use_chat_completions_api=True) + try: + response: Final = candidate.request( + "POST", "/v1/responses", {"model": model, "input": f"say hi {uuid.uuid4().hex}", "stream": True} + ) + assert response.status_code == 200, response.text + assert response.headers["content-type"].startswith("text/event-stream"), response.text + events: Final = _blocked_stream_events(response.text) + assert events[-1]["type"] == "response.completed", events + completed: Final = events[-1]["response"] + item: Final = completed["output"][0] + assert item["type"] == "message" and item["role"] == "assistant", item + assert item["content"][0]["type"] == "output_text", item + assert item["content"][0]["text"] == _RESPONSES_OUTPUT_DENIAL, item + usage: Final = completed["usage"] + assert ( + usage["input_tokens"], + usage["output_tokens"], + usage["total_tokens"], + ) == (_UPSTREAM_INPUT_TOKENS, _UPSTREAM_OUTPUT_TOKENS, _UPSTREAM_TOTAL_TOKENS), usage + finally: + candidate.post("/model/delete", {"id": model_id}) + + +@pytest.mark.covers("other.observability.guardrails.responses_post_call_pipeline_denial_returns_real_usage") +def test_responses_post_call_pipeline_denial_returns_real_usage(gateway: Gateway, tmp_path: Path) -> None: + identity: Final = "guardrail" + uuid.uuid4().hex + model: Final = f"integration-{uuid.uuid4().hex}" + config: Final = _responses_output_denial_config(tmp_path, identity, model) + with owned_proxy(gateway, tmp_path, {}, config=config) as candidate: + model_id: Final = _register_named_model(candidate, model, use_chat_completions_api=True) + try: + response: Final = candidate.request( + "POST", "/v1/responses", {"model": model, "input": f"say hi {uuid.uuid4().hex}"} + ) + assert response.status_code == 200, response.text + body: Final = response.json() + item: Final = body["output"][0] + assert item["type"] == "message" and item["role"] == "assistant", item + assert item["content"][0]["type"] == "output_text", item + assert item["content"][0]["text"] == _RESPONSES_OUTPUT_DENIAL, item + usage: Final = body["usage"] + assert ( + usage["input_tokens"], + usage["output_tokens"], + usage["total_tokens"], + ) == (_UPSTREAM_INPUT_TOKENS, _UPSTREAM_OUTPUT_TOKENS, _UPSTREAM_TOTAL_TOKENS), usage + finally: + candidate.post("/model/delete", {"id": model_id}) + + +@pytest.mark.covers("other.observability.guardrails.responses_denial_requires_authentication") +def test_responses_denial_requires_authentication(gateway: Gateway, tmp_path: Path) -> None: + identity: Final = "guardrail" + uuid.uuid4().hex + config: Final = _responses_denial_config(tmp_path, identity) + with owned_proxy(gateway, tmp_path, {}, config=config) as candidate, candidate.scenario() as scenario: + model: Final = scenario.model() + response: Final = candidate.request( + "POST", "/v1/responses", {"model": model, "input": "say hi", "guardrails": [identity]}, key="sk-invalid" + ) + assert response.status_code == 401, (response.status_code, response.text) + assert response.json()["error"]["type"] == "token_not_found_in_db", response.text + + +@pytest.mark.covers("other.observability.guardrails.responses_pre_call_denial_stream_does_not_reach_upstream") +def test_responses_pre_call_denial_stream_does_not_reach_upstream(gateway: Gateway, tmp_path: Path) -> None: + identity: Final = "guardrail" + uuid.uuid4().hex + config: Final = _responses_denial_config(tmp_path, identity) + with owned_proxy(gateway, tmp_path, {}, config=config) as candidate, candidate.scenario() as scenario: + model: Final = scenario.model(api_base=_dead_api_base()) + response: Final = candidate.request( + "POST", + "/v1/responses", + {"model": model, "input": "say hi", "stream": True, "guardrails": [identity]}, + ) + assert response.status_code == 200, response.text + assert response.headers["content-type"].startswith("text/event-stream"), response.text + events: Final = _blocked_stream_events(response.text) + completed: Final = events[-1]["response"] + _assert_blocked_message_item(completed["output"][0], completed) + + +@pytest.mark.covers("other.observability.guardrails.responses_unguarded_stream_reaches_upstream") +def test_responses_unguarded_stream_reaches_upstream(gateway: Gateway, tmp_path: Path) -> None: + identity: Final = "guardrail" + uuid.uuid4().hex + config: Final = _responses_denial_config(tmp_path, identity) + with owned_proxy(gateway, tmp_path, {}, config=config) as candidate, candidate.scenario() as scenario: + dead: Final = scenario.model(api_base=_dead_api_base()) + denied: Final = candidate.request( + "POST", + "/v1/responses", + {"model": dead, "input": "say hi", "stream": True, "guardrails": [identity]}, + ) + assert denied.status_code == 200, denied.text + model: Final = scenario.model(use_chat_completions_api=True) + with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as observed: + observed.get("/__observations") + response: Final = candidate.request( + "POST", + "/v1/responses", + {"model": model, "input": f"say hi {uuid.uuid4().hex}", "stream": True, "guardrails": []}, + ) + assert response.status_code == 200, response.text + assert response.headers["content-type"].startswith("text/event-stream"), response.text + assert "response.completed" in response.text, response.text + requests: Final = eventually( + lambda: observed.get("/__observations").json()["requests"], + lambda values: len(values) >= 1, + seconds=30, + ) + assert len(requests) == 1, requests + assert requests[0]["path"] == "/v1/chat/completions", requests + + +@pytest.mark.covers("other.observability.guardrails.chat_pre_call_denial_streams_content_filter") +def test_chat_pre_call_denial_streams_content_filter(gateway: Gateway, tmp_path: Path) -> None: + identity: Final = "guardrail" + uuid.uuid4().hex + config: Final = _responses_denial_config(tmp_path, identity) + with owned_proxy(gateway, tmp_path, {}, config=config) as candidate, candidate.scenario() as scenario: + model: Final = scenario.model() + response: Final = candidate.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": "say hi"}], + "stream": True, + "guardrails": [identity], + }, + ) + assert response.status_code == 200, response.text + assert response.headers["content-type"].startswith("text/event-stream"), response.text + lines: Final = tuple(line for line in response.text.split("\n") if line.startswith("data: ")) + assert lines[-1] == "data: [DONE]", response.text + chunks: Final = tuple(json.loads(line.removeprefix("data: ")) for line in lines[:-1]) + assert chunks[0]["choices"][0]["delta"]["content"] == _RESPONSES_DENIAL, chunks + assert chunks[-1]["choices"][0]["finish_reason"] == "stop", chunks + + +@pytest.mark.covers("other.observability.guardrails.chat_pre_call_denial_returns_content_filter") +def test_chat_pre_call_denial_returns_content_filter(gateway: Gateway, tmp_path: Path) -> None: + identity: Final = "guardrail" + uuid.uuid4().hex + config: Final = _responses_denial_config(tmp_path, identity) + with owned_proxy(gateway, tmp_path, {}, config=config) as candidate, candidate.scenario() as scenario: + model: Final = scenario.model() + response: Final = candidate.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": "say hi"}], + "guardrails": [identity], + }, + ) + assert response.status_code == 200, response.text + body: Final = response.json() + choice: Final = body["choices"][0] + assert choice["finish_reason"] == "content_filter", body + assert choice["message"]["content"] == _RESPONSES_DENIAL, body + assert ( + body["usage"]["prompt_tokens"], + body["usage"]["completion_tokens"], + body["usage"]["total_tokens"], + ) == (0, 0, 0), body["usage"] + + +@pytest.mark.covers("other.observability.guardrails.messages_pre_call_denial_returns_message") +def test_messages_pre_call_denial_returns_message(gateway: Gateway, tmp_path: Path) -> None: + identity: Final = "guardrail" + uuid.uuid4().hex + config: Final = _responses_denial_config(tmp_path, identity) + with owned_proxy(gateway, tmp_path, {}, config=config) as candidate, candidate.scenario() as scenario: + model: Final = scenario.model() + response: Final = candidate.request( + "POST", + "/v1/messages", + { + "model": model, + "messages": [{"role": "user", "content": "say hi"}], + "max_tokens": 16, + "guardrails": [identity], + }, + ) + assert response.status_code == 200, response.text + body: Final = response.json() + assert body["type"] == "message" and body["role"] == "assistant", body + assert body["content"] == [{"type": "text", "text": _RESPONSES_DENIAL}], body + assert body["stop_reason"] == "end_turn", body + assert (body["usage"]["input_tokens"], body["usage"]["output_tokens"]) == (0, 0), body + + +@pytest.mark.covers("other.observability.guardrails.messages_pre_call_denial_streams_message") +def test_messages_pre_call_denial_streams_message(gateway: Gateway, tmp_path: Path) -> None: + identity: Final = "guardrail" + uuid.uuid4().hex + config: Final = _responses_denial_config(tmp_path, identity) + with owned_proxy(gateway, tmp_path, {}, config=config) as candidate, candidate.scenario() as scenario: + model: Final = scenario.model() + response: Final = candidate.request( + "POST", + "/v1/messages", + { + "model": model, + "messages": [{"role": "user", "content": "say hi"}], + "max_tokens": 16, + "stream": True, + "guardrails": [identity], + }, + ) + assert response.status_code == 200, response.text + assert response.headers["content-type"].startswith("text/event-stream"), response.text + lines: Final = tuple(line for line in response.text.split("\n") if line.startswith("data: ")) + assert len(lines) == 1, response.text + body: Final = json.loads(lines[0].removeprefix("data: ")) + assert body["type"] == "message" and body["role"] == "assistant", body + assert body["content"] == [{"type": "text", "text": _RESPONSES_DENIAL}], body + + +@pytest.mark.covers("other.observability.guardrails.responses_pre_call_denial_writes_zero_spend_row") +def test_responses_pre_call_denial_writes_zero_spend_row(gateway: Gateway, tmp_path: Path) -> None: + identity: Final = "guardrail" + uuid.uuid4().hex + config: Final = _responses_denial_config(tmp_path, identity) + with owned_proxy(gateway, tmp_path, {}, config=config) as candidate, candidate.scenario() as scenario: + model: Final = scenario.model() + response: Final = candidate.request( + "POST", "/v1/responses", {"model": model, "input": "say hi", "guardrails": [identity]} + ) + assert response.status_code == 200, response.text + rows: Final = eventually( + lambda: read_rows( + 'SELECT spend, total_tokens FROM "LiteLLM_SpendLogs" WHERE model=%s AND call_type=%s', + (model, "aresponses"), + ), + lambda values: len(values) == 1, + seconds=70, + ) + assert float(rows[0]["spend"]) == 0, rows + assert rows[0]["total_tokens"] == 0, rows + + +@pytest.mark.covers("other.observability.guardrails.responses_pre_call_denial_stream_survives_worker_burst") +def test_responses_pre_call_denial_stream_survives_worker_burst(gateway: Gateway, tmp_path: Path) -> None: + identity: Final = "guardrail" + uuid.uuid4().hex + config: Final = _responses_denial_config(tmp_path, identity) + with owned_proxy(gateway, tmp_path, {}, config=config, workers=2) as candidate, candidate.scenario() as scenario: + model: Final = scenario.model(api_base=_dead_api_base()) + healthy: Final = scenario.model(use_chat_completions_api=True) + with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as observed: + observed.get("/__observations") + + def burst(index: int) -> httpx.Response: + if index % 3 == 0: + return candidate.request( + "POST", + "/v1/responses", + {"model": healthy, "input": f"say hi {uuid.uuid4().hex} {index}", "stream": True}, + ) + stream: Final = index % 3 == 1 + return candidate.request( + "POST", + "/v1/responses", + {"model": model, "input": f"say hi {index}", "stream": stream, "guardrails": [identity]}, + ) + + with ThreadPoolExecutor(max_workers=8) as pool: + responses: Final = tuple(pool.map(burst, range(30))) + response_ids: Final = frozenset(_response_id(index, response) for index, response in enumerate(responses)) + assert len(response_ids) == 30, response_ids + assert len(observed.get("/__observations").json()["requests"]) == 10 + + +@pytest.mark.covers("other.observability.guardrails.responses_pre_call_denial_stream_survives_worker_kill") +def test_responses_pre_call_denial_stream_survives_worker_kill(gateway: Gateway, tmp_path: Path) -> None: + identity: Final = "guardrail" + uuid.uuid4().hex + config: Final = _responses_denial_config(tmp_path, identity) + with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned: + candidate: Final = owned.gateway + with candidate.scenario() as scenario: + model: Final = scenario.model(api_base=_dead_api_base()) + members: Final = tuple( + member for member in group_members(owned.process.pid) if member.pid != owned.process.pid + ) + children: Final = tuple(member.pid for member in members) + workers: Final = tuple( + member.pid for member in members if any("spawn_main" in part for part in member.cmdline()) + ) + assert len(workers) >= 2, workers + os.kill(workers[0], signal.SIGKILL) + expected: Final = len(children) + eventually( + lambda: tuple( + member.pid + for member in group_members(owned.process.pid) + if member.pid != owned.process.pid + and member.is_running() + and member.status() != psutil.STATUS_ZOMBIE + ), + lambda pids: len(pids) >= expected and any(pid not in children for pid in pids), + seconds=30, + ) + + def burst(index: int) -> httpx.Response: + return candidate.request( + "POST", + "/v1/responses", + {"model": model, "input": f"say hi {index}", "stream": True, "guardrails": [identity]}, + ) + + with ThreadPoolExecutor(max_workers=5) as pool: + responses: Final = tuple(pool.map(burst, range(10))) + for response in responses: + assert response.status_code == 200, response.text + assert response.headers["content-type"].startswith("text/event-stream"), response.text diff --git a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py index e684aa55b33..656dc33e88c 100644 --- a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py @@ -3,6 +3,7 @@ Test for response_api_endpoints/endpoints.py """ import unittest +from collections.abc import Mapping from typing import Any, Final, Literal from unittest.mock import AsyncMock, MagicMock, patch @@ -14,6 +15,7 @@ from httpx import Response import litellm from litellm.proxy.proxy_server import app +from litellm.types.llms.openai import ResponsesAPIResponse @pytest.mark.asyncio @@ -2193,6 +2195,59 @@ class TestCursorGateRecognizesRoutingGroups: assert "reasoning_effort" not in resolved +BLOCK_MESSAGE = "Content flagged by policy, response withheld" + + +def _post_blocked_responses( + original_response: ResponsesAPIResponse | litellm.ModelResponse | None, + payload: Mapping[str, object] | None = None, +) -> httpx.Response: + from litellm.integrations.custom_guardrail import ModifyResponseException + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + exc = ModifyResponseException( + message=BLOCK_MESSAGE, + model="gpt-4o-mini", + request_data={"model": "gpt-4o-mini", "input": "hi"}, + guardrail_name="zero-usage-regression", + original_response=original_response, + ) + mock_proxy_logging = MagicMock() + mock_proxy_logging.post_call_failure_hook = AsyncMock() + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + api_key="sk-test", request_route="/v1/responses" + ) + body = {"model": "gpt-4o-mini", "input": "Write a haiku about token accounting"} + if payload: + body.update(payload) + try: + with ( + patch( + "litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing.base_process_llm_request", + new=AsyncMock(side_effect=exc), + ), + patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging), + ): + client = TestClient(app) + return client.post("/v1/responses", json=body, headers={"Authorization": "Bearer sk-1234"}) + finally: + app.dependency_overrides.pop(user_api_key_auth, None) + + +def _assert_blocked_output_item(item: Mapping[str, object], text: str) -> None: + assert item["type"] == "message" + assert item["id"].startswith("msg_") + assert item["role"] == "assistant" + assert item["status"] == "completed" + assert item["content"][0]["type"] == "output_text" + assert item["content"][0]["text"] == text + + +def _sse_data_frames(text: str) -> list[str]: + return [line.removeprefix("data: ").strip() for line in text.splitlines() if line.startswith("data: ")] + + class TestGuardrailBlockedResponsesUsage: """Regression tests for https://github.com/BerriAI/litellm/issues/36880. @@ -2202,38 +2257,7 @@ class TestGuardrailBlockedResponsesUsage: e.original_response, exactly like /v1/chat/completions already does.""" def _post_blocked_responses(self, original_response): - from litellm.integrations.custom_guardrail import ModifyResponseException - from litellm.proxy._types import UserAPIKeyAuth - from litellm.proxy.auth.user_api_key_auth import user_api_key_auth - - exc = ModifyResponseException( - message="Content flagged by policy, response withheld", - model="gpt-4o-mini", - request_data={"model": "gpt-4o-mini", "input": "hi"}, - guardrail_name="zero-usage-regression", - original_response=original_response, - ) - mock_proxy_logging = MagicMock() - mock_proxy_logging.post_call_failure_hook = AsyncMock() - app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( - api_key="sk-test", request_route="/v1/responses" - ) - try: - with ( - patch( - "litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing.base_process_llm_request", - new=AsyncMock(side_effect=exc), - ), - patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging), - ): - client = TestClient(app) - return client.post( - "/v1/responses", - json={"model": "gpt-4o-mini", "input": "Write a haiku about token accounting"}, - headers={"Authorization": "Bearer sk-1234"}, - ) - finally: - app.dependency_overrides.pop(user_api_key_auth, None) + return _post_blocked_responses(original_response) def test_post_call_block_reports_real_upstream_usage(self): from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse @@ -2429,6 +2453,66 @@ class TestResponsesInputTokens: assert response.json()["error"]["message"] == "rate limited" +class TestGuardrailBlockedResponsesShape: + """A pre_call block raises ModifyResponseException before any provider call. + + The reply must satisfy the Responses API contract the request selected: + stream=true answers SSE ending in one response.completed whose output[0] is + a completed assistant message item with output_text content, and a plain + POST answers JSON with the same item, both with the usage the blocked call + consumed (zero for pre_call).""" + + def test_non_stream_block_is_a_completed_assistant_message(self): + response = _post_blocked_responses(None) + + assert response.status_code == 200, response.text + assert response.headers["content-type"].startswith("application/json") + body = response.json() + _assert_blocked_output_item(body["output"][0], BLOCK_MESSAGE) + assert body["usage"]["total_tokens"] == 0 + + def test_stream_block_answers_sse_with_completed_event(self): + response = _post_blocked_responses(None, payload={"stream": True}) + + assert response.status_code == 200, response.text + assert response.headers["content-type"].startswith("text/event-stream") + frames = _sse_data_frames(response.text) + assert frames[-1] == "[DONE]" + events = [json.loads(frame) for frame in frames[:-1]] + types = [event["type"] for event in events] + assert "response.created" in types + completed = [event for event in events if event["type"] == "response.completed"] + assert len(completed) == 1 + completed_response = completed[0]["response"] + _assert_blocked_output_item(completed_response["output"][0], BLOCK_MESSAGE) + assert completed_response["usage"]["total_tokens"] == 0 + delta_text = "".join(event["delta"] for event in events if event["type"] == "response.output_text.delta") + assert delta_text == BLOCK_MESSAGE + + def test_stream_block_keeps_upstream_usage(self): + from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse + + original = ResponsesAPIResponse( + id="resp_upstream", + created_at=1, + model="gpt-4o-mini", + object="response", + output=[], + status="completed", + usage=ResponseAPIUsage(input_tokens=14, output_tokens=20, total_tokens=34), + ) + + response = _post_blocked_responses(original, payload={"stream": True}) + + assert response.status_code == 200, response.text + frames = _sse_data_frames(response.text) + completed = [json.loads(frame) for frame in frames[:-1] if json.loads(frame)["type"] == "response.completed"] + usage = completed[0]["response"]["usage"] + assert usage["input_tokens"] == 14 + assert usage["output_tokens"] == 20 + assert usage["total_tokens"] == 34 + + def test_responses_routes_document_response_models_in_openapi_schema(): from typing import cast diff --git a/tests/test_litellm/proxy/test_blocked_response_usage.py b/tests/test_litellm/proxy/test_blocked_response_usage.py index 4f20f35e94b..90d861be8e0 100644 --- a/tests/test_litellm/proxy/test_blocked_response_usage.py +++ b/tests/test_litellm/proxy/test_blocked_response_usage.py @@ -4,7 +4,7 @@ proxy endpoints (/v1/chat/completions, /v1/completions, and /v1/responses). A post-call block replaces the LLM response with the violation message, but the upstream call already consumed tokens. `_blocked_response_usage` (and its -Responses API counterpart `_blocked_responses_api_usage`) reports that real +Responses API counterpart `blocked_responses_api_usage`) reports that real usage (carried on `ModifyResponseException.original_response`) rather than zero; a pre-call block never invoked the LLM, so usage is zero. """ @@ -91,8 +91,8 @@ def test_responses_api_blocked_reply_carries_real_usage(): """ import time - from litellm.proxy.response_api_endpoints.endpoints import ( - _blocked_responses_api_usage, + from litellm.llms.base_llm.guardrail_translation.utils import ( + blocked_responses_api_usage, ) original_response = ResponsesAPIResponse( @@ -105,7 +105,7 @@ def test_responses_api_blocked_reply_carries_real_usage(): usage=ResponseAPIUsage(input_tokens=14, output_tokens=20, total_tokens=34), ) - usage = _blocked_responses_api_usage(original_response) + usage = blocked_responses_api_usage(original_response) assert usage.input_tokens == 14 assert usage.output_tokens == 20 @@ -114,11 +114,11 @@ def test_responses_api_blocked_reply_carries_real_usage(): def test_responses_api_blocked_reply_zero_usage_when_no_original_response(): """Pre-call block has no original_response, so usage must be zero.""" - from litellm.proxy.response_api_endpoints.endpoints import ( - _blocked_responses_api_usage, + from litellm.llms.base_llm.guardrail_translation.utils import ( + blocked_responses_api_usage, ) - usage = _blocked_responses_api_usage(None) + usage = blocked_responses_api_usage(None) assert usage.input_tokens == 0 assert usage.output_tokens == 0 @@ -128,14 +128,14 @@ def test_responses_api_blocked_reply_zero_usage_when_no_original_response(): def test_responses_api_blocked_reply_maps_bridged_chat_usage(): """A chat model bridged through /v1/responses blocks with a ModelResponse whose Usage fields must map prompt_tokens -> input_tokens and completion_tokens -> output_tokens.""" - from litellm.proxy.response_api_endpoints.endpoints import ( - _blocked_responses_api_usage, + from litellm.llms.base_llm.guardrail_translation.utils import ( + blocked_responses_api_usage, ) resp = litellm.ModelResponse() resp.usage = litellm.Usage(prompt_tokens=14, completion_tokens=18, total_tokens=32) - usage = _blocked_responses_api_usage(resp) + usage = blocked_responses_api_usage(resp) assert usage.input_tokens == 14 assert usage.output_tokens == 18