diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index ad8e29ad4ac..5cd951d19e1 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 36b7a3a4a8a..facf1e3ea8f 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, cast, get_args -from uuid import uuid4 import fastapi from fastapi import APIRouter, Depends, HTTPException, Request, Response @@ -20,8 +18,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 ( @@ -29,7 +28,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, @@ -426,17 +425,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 02b5a921add..ae1fd8a8f9b 100644 --- a/tests/e2e/coverage_registry/guardrail.yaml +++ b/tests/e2e/coverage_registry/guardrail.yaml @@ -36,4 +36,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 007e9392616..da6fb0d410e 100644 --- a/tests/e2e/guardrails/guardrails_client.py +++ b/tests/e2e/guardrails/guardrails_client.py @@ -110,6 +110,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 @@ -117,6 +126,7 @@ GuardrailParamsBody = ( | BlockCodeExecutionParamsBody | PresidioParamsBody | ToolPermissionParamsBody + | CustomCodeParamsBody ) @@ -148,6 +158,7 @@ class _ResponsesGuardrailBody(BaseModel): model: str input: str guardrails: list[str] | None = None + stream: bool | None = None @dataclass(frozen=True, slots=True) @@ -435,6 +446,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..229067919d1 --- /dev/null +++ b/tests/e2e/guardrails/test_responses_pre_call_block_stream_e2e.py @@ -0,0 +1,165 @@ +"""Live e2e: a pre_call guardrail block on POST /v1/responses must answer in the +shape the caller asked for. + +A custom_code guardrail whose apply_guardrail returns block(...) raises +ModifyResponseException before the model is called. On /v1/responses the +caller still gets a 200 carrying the denial, so the response must satisfy the +Responses API contract the request selected: stream=true answers SSE +(text/event-stream) with a typed event sequence ending in exactly one +response.completed whose output[0] is a completed assistant message item with +output_text content, and stream=false/omitted answers application/json with the +same schema-valid output item. Both report zero usage because no provider call +happened. +""" + +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_for_block(call: Callable[[], StreamingResponse]) -> StreamingResponse: + """Retry until the guardrail block lands. A registered guardrail reaches the + data plane on the periodic DB sync, so the first calls can still 404 with + "Guardrail not found" or pass through unblocked; both are retried to the + deadline and the last outcome is what the assertions judge.""" + deadline: Final = time.monotonic() + POLL_TIMEOUT + result = call() + while time.monotonic() < deadline: + if _denial_delivered(result): + return result + if not result.ok and "Guardrail not found" not in result.body and result.status_code not in (-1, 401, 429): + return result + time.sleep(POLL_INTERVAL) + result = call() + return result + + +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/test_litellm/proxy/response_api_endpoints/test_endpoints.py b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py index 4153bf7d7ee..ded7868724e 100644 --- a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py @@ -2193,6 +2193,56 @@ class TestCursorGateRecognizesRoutingGroups: assert "reasoning_effort" not in resolved +BLOCK_MESSAGE = "Content flagged by policy, response withheld" + + +def _post_blocked_responses(original_response, payload=None): + 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, text): + 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): + 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 +2252,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 @@ -2427,3 +2446,63 @@ class TestResponsesInputTokens: assert response.status_code == 429, response.text 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