scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 79

fix(guardrails): surface post-flush stream blocks as in-stream error frames and keep guardrail_information in spend logs

Browse source 
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.
This commit is contained in:
mateo-berri 2026-08-29 01:25:08 -07:00
parent e486e43d0f
commit 64eec53fd8
14 changed files with 608 additions and 85 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
litellm/llms/anthropic/chat/guardrail_translation/handler.py
View file

@ -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

48
litellm/llms/base_llm/guardrail_translation/base_translation.py
View file

@ -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.

32
litellm/llms/openai/chat/guardrail_translation/handler.py
View file

@ -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):

35
litellm/llms/openai/responses/guardrail_translation/handler.py
View file

@ -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.

14
litellm/proxy/common_request_processing.py
View file

@ -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}

4
litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
View file

@ -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}"
):

67
litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py
View file

@ -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

21
litellm/responses/streaming_iterator.py
View file

@ -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,

32
tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py
View file

@ -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

32
tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py
View file

@ -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

73
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py
View file

@ -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"

235
tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py
View file

@ -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"

69
tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py
View file

@ -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

16
tests/test_litellm/proxy/test_common_request_processing.py
View file

@ -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):
"""