From 64eec53fd844684df2f953558c405c66eb27ce2d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 01:25:08 -0700 Subject: [PATCH] fix(guardrails): surface post-flush stream blocks as in-stream error frames and keep guardrail_information in spend logs A guardrail block or failed scan that fires after SSE chunks have been flushed can no longer set an HTTP status, so raising HTTPException there silently truncated the stream. _emit_streaming_http_error now routes post-flush failures through the endpoint translation's build_stream_error_items, emitting the surface-correct error frame on chat completions (data: {error}), /v1/messages (event: error), and /v1/responses (ErrorEvent with the next sequence number). Pre-flush blocks still raise with a real HTTP status. Successful flags-on scans also logged metadata.guardrail_information as null: the chat handler planted litellm_metadata on a route whose bucket is metadata, flipping the bucket for every later write, and responses streams fired their spend log before the eos scan ran. The chat handler now merges user_api_key metadata through get_or_create_metadata_bucket, and deferred stream-complete logging is armed for aresponses like it already was for anthropic_messages. --- .../chat/guardrail_translation/handler.py | 15 ++ .../guardrail_translation/base_translation.py | 48 ++++ .../chat/guardrail_translation/handler.py | 32 +-- .../guardrail_translation/handler.py | 35 +++ litellm/proxy/common_request_processing.py | 14 +- .../guardrail_hooks/bedrock_guardrails.py | 4 +- .../unified_guardrail/unified_guardrail.py | 67 ++++- litellm/responses/streaming_iterator.py | 21 +- .../openai/test_moderations.py | 32 +-- .../test_openai_moderation_streaming.py | 32 ++- .../test_bedrock_guardrails.py | 73 ++++++ .../test_unified_guardrail.py | 235 +++++++++++++++++- .../test_deferred_guardrail_logging.py | 69 +++++ .../proxy/test_common_request_processing.py | 16 +- 14 files changed, 608 insertions(+), 85 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index b9ca18c7843..6cb2e568d0f 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -58,6 +58,8 @@ from litellm.types.utils import ( ) if TYPE_CHECKING: + from fastapi import HTTPException + from litellm.integrations.custom_guardrail import ( CustomGuardrail, ModifyResponseException, @@ -165,6 +167,19 @@ class AnthropicMessagesHandler(BaseTranslation): return self._block_continuation_chunks(exc, responses_so_far or []) return self._standalone_block_chunks(exc) + def build_stream_error_items( + self, + exc: "HTTPException", + responses_so_far: Sequence[Any] | None = None, + ) -> Sequence[Any] | None: + from litellm.proxy.common_request_processing import ( + serialize_http_exception_detail, + ) + from litellm.proxy.guardrails.anthropic_sse import anthropic_sse_error_frames + + message, _ = serialize_http_exception_detail(exc.detail) + return list(anthropic_sse_error_frames(message)) + def _standalone_block_chunks(self, exc: "ModifyResponseException") -> list[bytes]: import uuid diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index ba96ab3dc99..b07a6b986d7 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -1,8 +1,11 @@ from abc import ABC, abstractmethod +from collections.abc import Sequence from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Final, Optional if TYPE_CHECKING: + from fastapi import HTTPException + from litellm.integrations.custom_guardrail import ( CustomGuardrail, ModifyResponseException, @@ -73,6 +76,31 @@ class BaseTranslation(ABC): return transformed + @staticmethod + def merge_user_api_key_metadata_into_request( + request_data: dict[str, Any], # mutable-ok: proxy hooks share and mutate the request payload dict in place + user_api_key_dict: Optional["UserAPIKeyAuth"], + ) -> None: + """ + Add the prefixed ``user_api_key_*`` metadata to the request's resolved + metadata bucket without overwriting existing keys. + + Writes must go through ``get_or_create_metadata_bucket``: creating a + ``litellm_metadata`` key on a route whose bucket is ``metadata`` (chat + completions) flips the bucket for every later metadata write, and spend + logging never sees those writes (e.g. guardrail_information). + """ + from litellm.litellm_core_utils.core_helpers import ( + get_or_create_metadata_bucket, + ) + + user_metadata: Final = BaseTranslation.transform_user_api_key_dict_to_metadata(user_api_key_dict) + if not user_metadata: + return + _, metadata_bucket = get_or_create_metadata_bucket(request_data) + for key, value in user_metadata.items(): + metadata_bucket.setdefault(key, value) + @abstractmethod async def process_input_messages( self, @@ -147,6 +175,26 @@ class BaseTranslation(ABC): """ return None + def build_stream_error_items( + self, + exc: "HTTPException", + responses_so_far: Sequence[Any] | None = None, + ) -> Sequence[Any] | None: + """ + Build the stream items that surface a guardrail HTTPException (a block + with the default exception-on-block config, or a failed scan) after the + response has already started streaming, in this endpoint's wire format. + + Called only once chunks have been sent: the HTTP status is gone, so the + failure must travel as an in-stream error frame. ``responses_so_far`` + holds the chunks the client has already received, for formats whose + error frame continues the stream (e.g. sequence numbers). + + Returns None when the format has no in-stream error frame; the caller + then re-raises ``exc``. Override in endpoint subclasses. + """ + return None + def get_structured_messages(self, data: dict) -> list["AllMessageValues"] | None: """ Convert request data to OpenAI-spec structured messages. diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index e411dc497fc..e61fb719c98 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -14,6 +14,7 @@ Pattern Overview: This pattern can be replicated for other message formats (e.g., Anthropic). """ +from collections.abc import Sequence from typing import TYPE_CHECKING, Any, Final, Union, cast import litellm @@ -46,6 +47,8 @@ from litellm.types.utils import ( ) if TYPE_CHECKING: + from fastapi import HTTPException + from litellm.integrations.custom_guardrail import CustomGuardrail @@ -381,11 +384,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if "response" not in request_data: request_data["response"] = response - # Add user API key metadata with prefixed keys - if "litellm_metadata" not in request_data: - user_metadata: Final = self.transform_user_api_key_dict_to_metadata(user_api_key_dict) - if user_metadata: - request_data["litellm_metadata"] = user_metadata + self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict) inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) if images_to_check: @@ -554,11 +553,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if "responses" not in request_data: request_data["responses"] = responses_so_far - # Add user API key metadata with prefixed keys - if "litellm_metadata" not in request_data: - user_metadata: Final = self.transform_user_api_key_dict_to_metadata(user_api_key_dict) - if user_metadata: - request_data["litellm_metadata"] = user_metadata + self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict) inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) if images_to_check: @@ -590,6 +585,18 @@ class OpenAIChatCompletionsHandler(BaseTranslation): return responses_so_far + def build_stream_error_items( + self, + exc: "HTTPException", + responses_so_far: Sequence[Any] | None = None, + ) -> Sequence[Any] | None: + import json + + from litellm.proxy.common_request_processing import sse_error_payload + + _, error_obj = sse_error_payload(exc) + return [f"data: {json.dumps({'error': error_obj})}\n\n".encode()] + @staticmethod def _accumulate_string_content_by_choice_index( responses_so_far: list["ModelResponseStream"], @@ -652,10 +659,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): request_data = {"responses": responses_so_far} elif "responses" not in request_data: request_data["responses"] = responses_so_far - if "litellm_metadata" not in request_data: - user_metadata: Final = self.transform_user_api_key_dict_to_metadata(user_api_key_dict) - if user_metadata: - request_data["litellm_metadata"] = user_metadata + self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict) inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) if responses_so_far and getattr(responses_so_far[0], "model", None): diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 7c5d8ac99ad..50d7f7452a8 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -48,6 +48,8 @@ from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionToolCallChunk, ChatCompletionToolParam, + ErrorEvent, + ErrorEventError, OpenAIMcpServerTool, ResponsesAPIStreamEvents, ) @@ -59,6 +61,8 @@ from litellm.types.responses.main import ( from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: + from fastapi import HTTPException + from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import UserAPIKeyAuth @@ -80,6 +84,14 @@ class ResponsesStreamChunk(TypedDict, total=False): text: ReadOnly[str] +def _next_stream_sequence_number(responses_so_far: Sequence[Any] | None) -> int: + sequence_numbers: Final = ( + item.get("sequence_number") if isinstance(item, dict) else getattr(item, "sequence_number", None) + for item in reversed(responses_so_far or []) + ) + return next((n + 1 for n in sequence_numbers if isinstance(n, int)), 0) + + class OpenAIResponsesHandler(BaseTranslation): """ Handler for processing OpenAI Responses API with guardrails. @@ -620,6 +632,29 @@ class OpenAIResponsesHandler(BaseTranslation): } return responses_so_far[-1].get("type") in terminal_types + def build_stream_error_items( + self, + exc: "HTTPException", + responses_so_far: Sequence[Any] | None = None, + ) -> Sequence[Any] | None: + from litellm.proxy.common_request_processing import ( + serialize_http_exception_detail, + ) + + message, _ = serialize_http_exception_detail(exc.detail) + return [ + ErrorEvent( + type=ResponsesAPIStreamEvents.ERROR, + sequence_number=_next_stream_sequence_number(responses_so_far), + error=ErrorEventError( + type="guardrail_error", + code=str(exc.status_code), + message=message, + param=None, + ), + ) + ] + def get_streaming_string_so_far(self, responses_so_far: Sequence[ResponsesStreamChunk]) -> str: """ Get the string so far from the responses so far. diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index ff6c8d1b1f8..6c7c77610b3 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -502,7 +502,7 @@ def _as_success_dispatcher(logging_obj: _DispatchesSuccessHandlers) -> _Dispatch return logging_obj -def _serialize_http_exception_detail( +def serialize_http_exception_detail( detail: object, ) -> tuple[str, dict | None]: """ @@ -803,7 +803,7 @@ async def _buffer_first_chunk_honoring_disconnect( raise _ClientDisconnectedBeforeFirstChunk() -def _sse_error_payload(exc: BaseException) -> tuple[int, Mapping[str, object]]: +def sse_error_payload(exc: BaseException) -> tuple[int, Mapping[str, object]]: """Build the ProxyException-shaped ``{"error": ...}`` body used in SSE error frames. Matches ``ProxyException.to_dict()`` so streaming and non-streaming error frames @@ -812,7 +812,7 @@ def _sse_error_payload(exc: BaseException) -> tuple[int, Mapping[str, object]]: # Preserve status code from HTTPException (e.g. guardrail blocks) error_status: Final = getattr(exc, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR) raw_detail: Final = _getattr_object(exc, "detail", "Error processing stream start") - message, structured_fields = _serialize_http_exception_detail(raw_detail) + message, structured_fields = serialize_http_exception_detail(raw_detail) existing_fields: Final = getattr(exc, "provider_specific_fields", None) or {} merged_fields: Final = {**existing_fields, **structured_fields} if structured_fields else (existing_fields or None) @@ -927,7 +927,7 @@ async def create_response( # Unexpected error consuming first chunk. verbose_proxy_logger.exception("Error consuming first chunk from generator: %s", e) - error_status, error_obj = _sse_error_payload(e) + error_status, error_obj = sse_error_payload(e) async def error_gen_message() -> AsyncGenerator[str, None]: for frame in _sse_error_frames(error_obj): @@ -1104,7 +1104,7 @@ async def open_sse_before_first_byte( # would never fire and the failure would go unaudited. The hook # also gets to sanitize what reaches the client, by returning or # raising a replacement, so its answer decides the frame. - _, error_obj = _sse_error_payload(await _sanitized_late_failure(exc, on_late_failure)) + _, error_obj = sse_error_payload(await _sanitized_late_failure(exc, on_late_failure)) for frame in _sse_error_frames(error_obj): yield frame.encode() return @@ -2374,7 +2374,7 @@ class ProxyBaseLLMRequestProcessing: logging_obj._on_deferred_stream_complete = _on_deferred_stream_complete elif ( _post_call_guardrails_active - and route_type == "anthropic_messages" + and route_type in ("anthropic_messages", "aresponses") and self._is_streaming_response(response) ): from litellm.litellm_core_utils.logging_worker import ( @@ -3245,7 +3245,7 @@ class ProxyBaseLLMRequestProcessing: if isinstance(e, HTTPException): raw_detail: Final = _getattr_object(e, "detail", str(e)) - message, structured_fields = _serialize_http_exception_detail(raw_detail) + message, structured_fields = serialize_http_exception_detail(raw_detail) existing_fields: Final = getattr(e, "provider_specific_fields", None) or {} if structured_fields: merged_fields: dict | None = {**existing_fields, **structured_fields} diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index d7f222db35d..951cbc13290 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -42,7 +42,7 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.common_request_processing import _serialize_http_exception_detail +from litellm.proxy.common_request_processing import serialize_http_exception_detail from litellm.proxy.common_utils.sse_keepalive import keepalive_ping_has_fired from litellm.proxy.guardrails.anthropic_sse import ( anthropic_sse_chunks_from_response, @@ -2760,7 +2760,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ) if not raw_sse or (not is_block and not headers_flushed): raise - block_message, _ = _serialize_http_exception_detail(block_detail) + block_message, _ = serialize_http_exception_detail(block_detail) for error_frame in anthropic_sse_error_frames( block_message if is_block else f"{block_exc.status_code}: {block_message}" ): diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index e95e97bfe74..6647ac4c293 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -57,6 +57,9 @@ class _EndpointTranslation(Protocol): @property def build_block_sse_chunks(self) -> "Callable[..., Sequence[bytes] | None]": ... + @property + def build_stream_error_items(self) -> "Callable[..., Sequence[object] | None]": ... + def _as_endpoint_translation(translation: _EndpointTranslation) -> _EndpointTranslation: return translation @@ -408,14 +411,32 @@ class UnifiedLLMGuardrails(CustomLogger): call_type: str | None, responses_so_far: Sequence[object], request_data: dict, + endpoint_translation: _EndpointTranslation | None = None, + stream_started: bool = False, + responses_yielded: Sequence[object] | None = None, ) -> AsyncGenerator[object, None]: - """Surface a mid-stream HTTPException. For A2A call types the response has - already started, so emit an in-stream JSON-RPC error chunk; otherwise - re-raise so the proxy can report it. + """Surface a mid-stream HTTPException (a guardrail block with the default + exception-on-block config, or a failed scan). + + A2A call types emit an in-stream JSON-RPC error chunk. For other call + types, once chunks have already reached the client the HTTP status is + gone, so the failure is delegated to the endpoint translation's + ``build_stream_error_items`` and travels as an in-stream error frame in + that endpoint's wire format. Before the first chunk (or when the format + has no in-stream error frame) the exception is re-raised so the proxy + can report it with a real HTTP status. """ if call_type is not None and CallTypes(call_type) in A2A_CALL_TYPES: yield _a2a_jsonrpc_error_chunk(exc, _get_a2a_request_id(responses_so_far, request_data)) return + if stream_started and endpoint_translation is not None: + error_items: Final = endpoint_translation.build_stream_error_items( + exc, responses_so_far=list(responses_yielded) if responses_yielded is not None else None + ) + if error_items is not None: + for error_item in error_items: + yield error_item + return raise exc def _build_transform_chunk( @@ -586,7 +607,15 @@ class UnifiedLLMGuardrails(CustomLogger): yield block_chunk raise _StreamTerminated() except HTTPException as e: - async for error_item in self._emit_streaming_http_error(e, call_type, responses_so_far, request_data): + async for error_item in self._emit_streaming_http_error( + e, + call_type, + responses_so_far, + request_data, + endpoint_translation=endpoint_translation, + stream_started=bool(responses_yielded), + responses_yielded=responses_yielded, + ): yield error_item raise _StreamTerminated() @@ -1070,11 +1099,17 @@ class UnifiedLLMGuardrails(CustomLogger): return except HTTPException as e: # Response already started (we already yielded chunks); cannot send 400. - # For A2A, yield an in-stream JSON-RPC error so the client sees it. - if call_type is not None and CallTypes(call_type) in A2A_CALL_TYPES: - yield _a2a_jsonrpc_error_chunk(e, _get_a2a_request_id(responses_so_far, request_data)) - return - raise + async for error_item in self._emit_streaming_http_error( + e, + call_type, + responses_so_far, + request_data, + endpoint_translation=endpoint_translation, + stream_started=chunks_yielded, + responses_yielded=responses_yielded, + ): + yield error_item + return chunks_yielded = True responses_yielded.append(original_item) yield original_item @@ -1133,7 +1168,13 @@ class UnifiedLLMGuardrails(CustomLogger): yield block_chunk return except HTTPException as e: - if call_type is not None and CallTypes(call_type) in A2A_CALL_TYPES: - yield _a2a_jsonrpc_error_chunk(e, _get_a2a_request_id(responses_so_far, request_data)) - else: - raise + async for error_item in self._emit_streaming_http_error( + e, + call_type, + responses_so_far, + request_data, + endpoint_translation=endpoint_translation, + stream_started=bool(responses_yielded), + responses_yielded=responses_yielded, + ): + yield error_item diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 368fd481e63..94c93ae7c6b 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -422,15 +422,20 @@ class BaseResponsesAPIStreamingIterator: end_time: Final = datetime.now() if is_async: - asyncio.create_task( - self.logging_obj.dispatch_success_handlers( - logging_response, - start_time=self.start_time, - end_time=end_time, - cache_hit=self._completed_response_cache_hit, - prefer_async_handlers=True, - ) + logging_coroutine: Final = self.logging_obj.dispatch_success_handlers( + logging_response, + start_time=self.start_time, + end_time=end_time, + cache_hit=self._completed_response_cache_hit, + prefer_async_handlers=True, ) + deferred_dispatch_armed: Final = getattr(self.logging_obj, "_on_deferred_stream_complete", None) is not None + if deferred_dispatch_armed: + # End-of-stream guardrail scans write guardrail_information after + # the terminal event; dispatching now would snapshot metadata early. + self.logging_obj._deferred_stream_complete_args = (logging_coroutine,) + else: + asyncio.create_task(logging_coroutine) else: run_async_function( async_function=self.logging_obj.async_success_handler, diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py index 112bc5e6e49..2b43720a126 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py @@ -482,23 +482,25 @@ async def test_openai_moderation_guardrail_streaming_harmful_content(): "metadata": {"guardrails": ["test-openai-moderation"]}, } - # Should raise HTTPException when processing streaming harmful content - from fastapi import HTTPException + # Chunks have already been flushed by end-of-stream moderation, so + # the block surfaces as the in-stream error frame, not a raise. + import json as _json - async def _drain(): - result_chunks = [] - async for chunk in unified_guardrail.async_post_call_streaming_iterator_hook( - user_api_key_dict=user_api_key_dict, - response=mock_stream(), - request_data=request_data, - ): - result_chunks.append(chunk) + result_chunks = [] + async for chunk in unified_guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=mock_stream(), + request_data=request_data, + ): + result_chunks.append(chunk) - with pytest.raises(HTTPException) as exc_info: - await _drain() - - assert exc_info.value.status_code == 400 - assert "Violated OpenAI moderation policy" in str(exc_info.value.detail) + frame = result_chunks[-1] + assert isinstance(frame, bytes) + text = frame.decode() + assert text.startswith("data: ") + assert "Violated OpenAI moderation policy" in text + payload = _json.loads(text[len("data: ") :]) + assert payload["error"]["code"] == "400" @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py index 914af0e2368..476d443d8d8 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py @@ -161,19 +161,27 @@ async def test_openai_moderation_guardrail_streaming_harmful_content(): "metadata": {"guardrails": ["test-openai-moderation"]}, } - # Should raise HTTPException - with pytest.raises(HTTPException) as exc_info: - async for ( - _ - ) in unified_guardrail.async_post_call_streaming_iterator_hook( - user_api_key_dict=user_api_key_dict, - response=mock_stream(), - request_data=request_data, - ): - pass + # Chunks have already been flushed by end-of-stream moderation, so + # the block surfaces as the in-stream error frame, not a raise. + import json as _json - assert exc_info.value.status_code == 400 - assert "Violated OpenAI moderation policy" in str(exc_info.value.detail) + collected = [] + async for ( + chunk + ) in unified_guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=mock_stream(), + request_data=request_data, + ): + collected.append(chunk) + + frame = collected[-1] + assert isinstance(frame, bytes) + text = frame.decode() + assert text.startswith("data: ") + assert "Violated OpenAI moderation policy" in text + payload = _json.loads(text[len("data: ") :]) + assert payload["error"]["code"] == "400" @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index ca4b0a65e0b..235a5c0c09b 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -5517,3 +5517,76 @@ async def test_masking_keeps_buffered_path_even_when_unbuffered_configured(): assert guardrail._streams_incrementally() is False events = await _run_streaming_hook_recording_order(guardrail) assert events[0] == "scan" + + +@pytest.mark.asyncio +async def test_streaming_end_of_stream_block_emits_error_frame_instead_of_truncating(): + """Regression for PR #38722: a topicPolicy DENY caught by the end-of-stream + scan used to raise after SSE headers were flushed, so the client saw a + silently truncated stream. The unified hook must emit the chat in-stream + error frame instead.""" + from litellm.llms import load_guardrail_translation_mappings + from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail import ( + unified_guardrail as unified_module, + ) + from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + UnifiedLLMGuardrails, + ) + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + streaming_end_of_stream_only=True, + streaming_buffer_until_moderated=False, + guardrail_name="bedrock-eos", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + ) + blocked_response = { + "action": "GUARDRAIL_INTERVENED", + "actionReason": "Guardrail blocked.", + "outputs": [{"text": "Sorry, the model cannot answer this question."}], + "assessments": [ + {"topicPolicy": {"topics": [{"name": "Forbidden topic", "type": "DENY", "action": "BLOCKED"}]}} + ], + } + + def _chunk(content, finish_reason=None): + return ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta={"content": content, "role": "assistant"}, + finish_reason=finish_reason, + ) + ], + ) + + async def _mock_stream(): + yield _chunk("the forbidden ") + yield _chunk("topic answer", finish_reason="stop") + + unified_module.endpoint_guardrail_translation_mappings = load_guardrail_translation_mappings() + try: + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.side_effect = guardrail._get_http_exception_for_blocked_guardrail(blocked_response) + + out = [] + async for item in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test", request_route="/v1/chat/completions"), + response=_mock_stream(), + request_data={"guardrail_to_apply": guardrail, "model": "gpt-4"}, + ): + out.append(item) + finally: + unified_module.endpoint_guardrail_translation_mappings = None + + assert len(out) == 3 + assert isinstance(out[0], ModelResponseStream) + frame = out[-1] + assert isinstance(frame, bytes) + payload = json.loads(frame.decode()[len("data: ") :]) + assert payload["error"]["message"] == "Violated guardrail policy" + assert payload["error"]["code"] == "400" + assert payload["error"]["provider_specific_fields"]["guardrailIdentifier"] == "test-guardrail" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 8b9ecfbbeee..0b32558a00a 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -948,19 +948,24 @@ class TestStreamingTransform: assert streamed == "ABCDEFGHIJ" @pytest.mark.asyncio - async def test_incremental_diff_underflow_raises(self): + async def test_incremental_diff_underflow_emits_error_frame(self): """A transform shorter than what was already streamed cannot retract - bytes: it raises HTTPException(stream_transform_underflow).""" + bytes. Chunks have already been flushed by then, so the underflow + surfaces as the in-stream error frame, not an unraisable HTTPException.""" + import json as _json + # First sample emits "ABCDEF" (6 chars); second sample shrinks to 3. guardrail = _StreamingTextGuardrail(shrink_to="ABC", shrink_after=1) chunks = [_stream_chunk("abcdef"), _stream_chunk("ghij")] - with pytest.raises(unified_module.HTTPException) as exc_info: - await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) - assert exc_info.value.status_code == 400 - assert exc_info.value.detail["error"] == "stream_transform_underflow" + frame = out[-1] + assert isinstance(frame, bytes) + payload = _json.loads(frame.decode()[len("data: ") :]) + assert payload["error"]["message"] == "stream_transform_underflow" + assert payload["error"]["code"] == "400" @pytest.mark.asyncio async def test_incremental_diff_final_chunk_preserves_finish_reason(self): @@ -1747,3 +1752,221 @@ class TestAppliedGuardrailsReflectsExecution: async def test_ordinary_guardrail_is_auto_marked_applied(self): data = await self._run(_AutoLoggingGuardrail()) assert "auto-logging" in _applied_guardrails(data) + + +class _EosHttpBlockingGuardrail(CustomGuardrail): + """Raises the bedrock-shaped block HTTPException at end-of-stream scan time.""" + + def __init__(self): + super().__init__(guardrail_name="eos-http-block") + self.streaming_end_of_stream_only = True + + def should_run_guardrail(self, data, event_type): # type: ignore[override] + return True + + async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): + raise unified_module.HTTPException( + status_code=400, + detail={ + "error": "Violated guardrail policy", + "bedrock_guardrail_response": "BLOCKED_TOPIC", + }, + ) + + +def _anthropic_sse_event(event_type, data): + import json as _json + + return f"event: {event_type}\ndata: {_json.dumps(data)}\n\n".encode() + + +def _anthropic_message_chunks(texts): + head = [ + _anthropic_sse_event( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_test", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5", + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 1, "output_tokens": 0}, + }, + }, + ), + _anthropic_sse_event( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + ] + deltas = [ + _anthropic_sse_event( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}, + ) + for text in texts + ] + tail = [ + _anthropic_sse_event("content_block_stop", {"type": "content_block_stop", "index": 0}), + _anthropic_sse_event( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 5}, + }, + ), + _anthropic_sse_event("message_stop", {"type": "message_stop"}), + ] + return head + deltas + tail + + +class TestStreamingHttpErrorFrames: + """A post-flush end-of-stream guardrail block (HTTPException) must surface as + the endpoint's in-stream error frame instead of an unhandled raise that + silently truncates the SSE stream (PR #38722 defect 1).""" + + @pytest.fixture(autouse=True) + def _use_real_mappings(self): + unified_module.endpoint_guardrail_translation_mappings = load_guardrail_translation_mappings() + yield + unified_module.endpoint_guardrail_translation_mappings = None + + @pytest.mark.asyncio + async def test_chat_eos_block_emits_data_error_frame(self): + import json as _json + + guardrail = _EosHttpBlockingGuardrail() + chunks = [_stream_chunk("hello "), _stream_chunk("world", finish_reason="stop")] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert out[:2] == chunks + frame = out[-1] + assert isinstance(frame, bytes) + text = frame.decode() + assert text.startswith("data: ") + payload = _json.loads(text[len("data: ") :]) + assert payload["error"]["message"] == "Violated guardrail policy" + assert payload["error"]["code"] == "400" + + @pytest.mark.asyncio + async def test_messages_eos_block_emits_anthropic_error_event(self): + guardrail = _EosHttpBlockingGuardrail() + chunks = _anthropic_message_chunks(["hello ", "world"]) + + out = await _drive_stream( + UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/messages" + ) + + raw = b"".join(c for c in out if isinstance(c, bytes)).decode() + assert "hello " in raw + assert "event: error" in raw + assert "Violated guardrail policy" in raw + assert "guardrail_error" in raw + + @pytest.mark.asyncio + async def test_responses_eos_block_emits_error_event_with_next_sequence(self): + guardrail = _EosHttpBlockingGuardrail() + chunks = [ + {"type": "response.created", "sequence_number": 0}, + {"type": "response.output_text.delta", "sequence_number": 1, "delta": "hello"}, + { + "type": "response.completed", + "sequence_number": 2, + "response": { + "model": "gpt-4", + "output": [{"type": "message", "content": [{"type": "output_text", "text": "hello"}]}], + }, + }, + ] + + out = await _drive_stream( + UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses" + ) + + assert chunks[0] in out and chunks[1] in out + assert chunks[2] not in out + error_event = out[-1] + assert error_event.type == "error" + assert error_event.sequence_number == 2 + assert error_event.error.message == "Violated guardrail policy" + assert error_event.error.code == "400" + assert error_event.error.type == "guardrail_error" + + @pytest.mark.asyncio + async def test_pre_flush_block_still_raises_http_exception(self): + guardrail = _EosHttpBlockingGuardrail() + guardrail.streaming_buffer_until_moderated = True + chunks = [_stream_chunk("hello "), _stream_chunk("world", finish_reason="stop")] + + with pytest.raises(unified_module.HTTPException) as exc_info: + await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert exc_info.value.status_code == 400 + assert exc_info.value.detail["error"] == "Violated guardrail policy" + + +class _AuditRecordingGuardrail(CustomGuardrail): + """Successful scan that records guardrail_information, like a flags-on audit.""" + + def __init__(self): + super().__init__(guardrail_name="audit-recorder") + self.streaming_end_of_stream_only = True + + def should_run_guardrail(self, data, event_type): # type: ignore[override] + return True + + async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={"action": "NONE"}, + request_data=request_data, + guardrail_status="success", + ) + return inputs + + +class TestStreamingGuardrailInformationBucket: + """guardrail_information written during a chat streaming end-of-stream scan + must land in the request's ``metadata`` bucket that spend logging snapshots. + Regression for PR #38722 defect 2: the chat handler used to plant a + ``litellm_metadata`` key first, flipping the bucket so every later + guardrail_information write was diverted and /spend/logs showed null.""" + + @pytest.fixture(autouse=True) + def _use_real_mappings(self): + unified_module.endpoint_guardrail_translation_mappings = load_guardrail_translation_mappings() + yield + unified_module.endpoint_guardrail_translation_mappings = None + + @pytest.mark.asyncio + async def test_chat_eos_scan_writes_guardrail_information_to_metadata(self): + guardrail = _AuditRecordingGuardrail() + chunks = [_stream_chunk("hello "), _stream_chunk("world", finish_reason="stop")] + + async def _mock_stream(): + for chunk in chunks: + yield chunk + + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key", user_id="user-1", request_route="/v1/chat/completions" + ) + request_data = {"guardrail_to_apply": guardrail, "model": "gpt-4", "metadata": {}} + + out = [] + async for item in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=_mock_stream(), + request_data=request_data, + ): + out.append(item) + + assert "litellm_metadata" not in request_data + recorded = request_data["metadata"]["standard_logging_guardrail_information"] + assert len(recorded) == 1 + assert recorded[0]["guardrail_name"] == "audit-recorder" + assert recorded[0]["guardrail_status"] == "success" + assert request_data["metadata"]["user_api_key_user_id"] == "user-1" diff --git a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py index e70fc61de30..b7317d33b84 100644 --- a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py @@ -1228,3 +1228,72 @@ class TestFireDeferredStreamLogging: assert info is not None, "guardrail_information should be populated" assert len(info) == 1 assert info[0]["guardrail_name"] == "info-writer" + + +class TestResponsesIteratorDeferredLogging: + """Regression for PR #38722 defect 2 on /v1/responses streams: when the + proxy arms _on_deferred_stream_complete, the responses streaming iterator + must store the logging coroutine for ProxyLogging._fire_deferred_stream_logging + (which runs AFTER end-of-stream guardrail scans write guardrail_information) + instead of dispatching immediately with a premature metadata snapshot.""" + + def _iterator(self, logging_obj): + from litellm.responses.streaming_iterator import ( + BaseResponsesAPIStreamingIterator, + ) + + iterator = object.__new__(BaseResponsesAPIStreamingIterator) + iterator.logging_obj = logging_obj + iterator.start_time = None + iterator.completed_response = None + iterator._completed_response_logged = False + iterator._completed_response_cache_hit = None + iterator._persist_completed_response_before_logging = False + return iterator + + def _logging_obj(self): + recorded = {} + + async def dispatch_success_handlers(result=None, **kwargs): + recorded["dispatched"] = True + + logging_obj = MagicMock() + logging_obj.dispatch_success_handlers = dispatch_success_handlers + return logging_obj, recorded + + @pytest.mark.asyncio + async def test_armed_iterator_stores_deferred_coroutine(self): + logging_obj, recorded = self._logging_obj() + logging_obj._on_deferred_stream_complete = MagicMock() + iterator = self._iterator(logging_obj) + + with patch("asyncio.create_task") as mock_create_task: + iterator._log_completed_response(is_async=True) + + mock_create_task.assert_not_called() + args = logging_obj._deferred_stream_complete_args + assert isinstance(args, tuple) and len(args) == 1 + assert "dispatched" not in recorded + await args[0] + assert recorded["dispatched"] is True + + @pytest.mark.asyncio + async def test_unarmed_iterator_dispatches_immediately(self): + logging_obj, recorded = self._logging_obj() + logging_obj._on_deferred_stream_complete = None + iterator = self._iterator(logging_obj) + + created = [] + real_create_task = asyncio.create_task + + def tracking_create_task(coro): + task = real_create_task(coro) + created.append(task) + return task + + with patch("asyncio.create_task", side_effect=tracking_create_task): + iterator._log_completed_response(is_async=True) + + assert len(created) == 1 + await created[0] + assert recorded["dispatched"] is True diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 71d4666416d..99f03d23378 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -1656,32 +1656,32 @@ class TestCommonRequestProcessingHelpers: assert payload["error"]["message"] == "MCP request blocked: no rewritable argument field present" assert payload["error"]["provider_specific_fields"]["error"]["code"] == "panw_prisma_airs_blocked" - async def test_serialize_http_exception_detail_helper(self): + async def testserialize_http_exception_detail_helper(self): """Direct unit coverage for the L1 helper across all branches.""" from litellm.proxy.common_request_processing import ( - _serialize_http_exception_detail, + serialize_http_exception_detail, ) import json as _json - assert _serialize_http_exception_detail("plain") == ("plain", None) + assert serialize_http_exception_detail("plain") == ("plain", None) - msg, fields = _serialize_http_exception_detail({"error": "Violated", "extra": "x"}) + msg, fields = serialize_http_exception_detail({"error": "Violated", "extra": "x"}) assert msg == "Violated" assert fields == {"error": "Violated", "extra": "x"} - msg, fields = _serialize_http_exception_detail({"error": {"message": "blocked", "code": "x"}}) + msg, fields = serialize_http_exception_detail({"error": {"message": "blocked", "code": "x"}}) assert msg == "blocked" assert fields == {"error": {"message": "blocked", "code": "x"}} - msg, fields = _serialize_http_exception_detail({"message": "top-level"}) + msg, fields = serialize_http_exception_detail({"message": "top-level"}) assert msg == "top-level" assert fields == {"message": "top-level"} - msg, fields = _serialize_http_exception_detail({"weird": ["a", "b"]}) + msg, fields = serialize_http_exception_detail({"weird": ["a", "b"]}) assert msg == _json.dumps({"weird": ["a", "b"]}) assert fields == {"weird": ["a", "b"]} - assert _serialize_http_exception_detail(42) == ("42", None) + assert serialize_http_exception_detail(42) == ("42", None) async def test_create_streaming_response_first_chunk_error_string_code(self): """