From b63656a5de6a95487d2f4826e2cbe4cccca0d02e Mon Sep 17 00:00:00 2001 From: paveln83 <349044532+paveln83@users.noreply.github.com> Date: Wed, 11 Nov 2026 13:39:19 -0800 Subject: [PATCH] fix(proxy): stream the guardrail block response when the block comes after an SSE keepalive ping (#51114) * fix(proxy): stream the guardrail block response when the block lands after an SSE keepalive ping * docs(proxy): keep the keepalive docstring's reasons together * fix(proxy): stream the /v1/messages block after a keepalive ping and always end a late block stream A passthrough guardrail block on a streamed /v1/messages request that landed after the first keepalive ping still ended as an api_error event. The route now hands its block stream to the keepalive relay, as chat and text completions do. If a route's block stream itself fails after a ping, the relay sends the usual late failure frame instead of raising, so the stream always ends. The passthrough tests now drive the real routes with a real guardrail. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../proxy/anthropic_endpoints/endpoints.py | 56 +++-- litellm/proxy/common_request_processing.py | 26 ++- litellm/proxy/proxy_server.py | 162 ++++++++------- litellm/utils.py | 1 + .../anthropic_endpoints/test_endpoints.py | 81 ++++++-- .../proxy/test_common_request_processing.py | 63 ++++++ ...t_modify_response_streaming_passthrough.py | 196 +++++++++++++++++- 7 files changed, 457 insertions(+), 128 deletions(-) diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 1688992f6fa..636ba987ee8 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -7,7 +7,7 @@ from types import MappingProxyType from typing import Final from fastapi import APIRouter, Depends, HTTPException, Request, Response, status -from fastapi.responses import JSONResponse +from fastapi.responses import JSONResponse, StreamingResponse from pydantic import JsonValue, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict @@ -55,6 +55,7 @@ from litellm.proxy.litellm_pre_call_utils import ( forwarded_anthropic_oauth_token, should_auto_drop_params_for_agentic_cli, ) +from litellm.proxy.utils import ProxyLogging from litellm.router import Router from litellm.types.utils import TokenCountResponse @@ -126,6 +127,32 @@ def _strip_total_tokens_from_anthropic_response(response: object) -> None: usage.pop("total_tokens", None) +async def _record_passthrough_block( + e: ModifyResponseException, proxy_logging_obj: ProxyLogging, user_api_key_dict: UserAPIKeyAuth +) -> None: + await proxy_logging_obj.post_call_failure_hook( + user_api_key_dict=user_api_key_dict, + original_exception=e, + request_data=e.request_data, + ) + + +async def _stream_passthrough_block( + e: ModifyResponseException, + proxy_logging_obj: ProxyLogging, + user_api_key_dict: UserAPIKeyAuth, + headers: Mapping[str, str], +) -> StreamingResponse | JSONResponse: + await _record_passthrough_block(e, proxy_logging_obj, user_api_key_dict) + block_chunks: Final = AnthropicMessagesHandler().build_block_sse_chunks(e) + + async def _blocked_stream() -> AsyncGenerator[str, None]: + for chunk in block_chunks: + yield chunk.decode() + + return await create_response(generator=_blocked_stream(), media_type="text/event-stream", headers=headers) + + @router.post( "/v1/messages", tags=["Anthropic `/v1/messages`"], @@ -174,6 +201,9 @@ async def anthropic_response( user_max_tokens=user_max_tokens, user_api_base=user_api_base, version=version, + on_passthrough_block=lambda block: _stream_passthrough_block( + block, proxy_logging_obj, user_api_key_dict, _NO_HEADERS + ), ) # Optionally strip the non-Anthropic `usage.total_tokens` field # LiteLLM adds internally. Anthropic's official /v1/messages spec @@ -189,30 +219,12 @@ async def anthropic_response( _strip_total_tokens_from_anthropic_response(result) return result except ModifyResponseException as e: - # Guardrail flagged content in passthrough mode - return 200 with violation message - _data: Final = e.request_data - await proxy_logging_obj.post_call_failure_hook( - user_api_key_dict=user_api_key_dict, - original_exception=e, - request_data=_data, - ) - block_headers: Final = litellm_call_id_headers(base_llm_response_processor.litellm_call_id) or _NO_HEADERS - block_handler: Final = AnthropicMessagesHandler() - if data.get("stream") is True: - block_chunks: Final = block_handler.build_block_sse_chunks(e) - - async def _blocked_stream() -> AsyncGenerator[str, None]: - for chunk in block_chunks: - yield chunk.decode() - - return await create_response( - generator=_blocked_stream(), media_type="text/event-stream", headers=block_headers - ) - + return await _stream_passthrough_block(e, proxy_logging_obj, user_api_key_dict, block_headers) + await _record_passthrough_block(e, proxy_logging_obj, user_api_key_dict) fastapi_response.headers.update(block_headers) - return block_handler.build_block_response(e) + return AnthropicMessagesHandler().build_block_response(e) except AnthropicContextManagementError as e: if e.status_code >= 500: # Server-side polyfill failures hit the failure hook for spend/alert diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index c6981015417..1b004423181 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1662,6 +1662,18 @@ async def _relay_late_response(produced: Response, framing: LateSseFraming) -> A await _aclose_late_response(produced) +async def _late_block_or_failure( + exc: Exception, + on_passthrough_block: Callable[[ModifyResponseException], Awaitable[Response]] | None, +) -> Response | Exception: + if not isinstance(exc, ModifyResponseException) or on_passthrough_block is None: + return exc + try: + return await on_passthrough_block(exc) + except Exception as block_failure: # noqa: BLE001 # the client still needs a frame that ends the stream + return block_failure + + async def _sanitized_late_failure( exc: Exception, on_late_failure: "Callable[[Exception], Awaitable[HTTPException | None]] | None", @@ -1692,6 +1704,7 @@ async def open_sse_before_first_byte( media_type: str = "text/event-stream", on_late_failure: Callable[[Exception], Awaitable[HTTPException | None]] | None = None, late_framing: LateSseFraming = OPENAI_LATE_SSE_FRAMING, + on_passthrough_block: Callable[[ModifyResponseException], Awaitable[Response]] | None = None, ) -> _LateResponseT | StreamingResponse: """Write SSE keepalive comments while the upstream LLM call is still in flight. @@ -1708,6 +1721,8 @@ async def open_sse_before_first_byte( the first ping reaches the client as an SSE error frame under a 200 rather than as an HTTP error status, and LiteLLM's own `x-litellm-*` response headers are not yet known. Both are why this stays off until an operator sets an interval. + A passthrough guardrail block is not a failure: when the route passes + `on_passthrough_block`, the stream it builds for the block is relayed. """ interval: Final = coerce_keepalive_interval(ping_interval_seconds) if interval is None: @@ -1730,14 +1745,19 @@ async def open_sse_before_first_byte( try: produced: Final = produce_task.result() except Exception as exc: # noqa: BLE001 # the status line is already sent; surface it as a frame + late: Final = await _late_block_or_failure(exc, on_passthrough_block) + if isinstance(late, Response): + async for chunk in _relay_late_response(late, late_framing): + yield chunk + return verbose_proxy_logger.exception( - "request failed after its SSE keepalive had opened the response: %s", exc + "request failed after its SSE keepalive had opened the response: %s", late ) # The caller's own `except` never sees this, so its failure hook # 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. - for frame in late_framing.failure(await _sanitized_late_failure(exc, on_late_failure)): + for frame in late_framing.failure(await _sanitized_late_failure(late, on_late_failure)): yield frame.encode() return async for chunk in _relay_late_response(produced, late_framing): @@ -3156,6 +3176,7 @@ class ProxyBaseLLMRequestProcessing: contents: list[object] | None = None, skip_pre_call_logic: bool = False, sse_framing: SseFraming = OPENAI_SSE_FRAMING, + on_passthrough_block: Callable[[ModifyResponseException], Awaitable[Response]] | None = None, ) -> object: """Run the request, sending SSE keepalives while the upstream is still silent. @@ -3207,6 +3228,7 @@ class ProxyBaseLLMRequestProcessing: ping_interval_seconds=ttft_keepalive_interval(self.data, llm_router), on_late_failure=_audit_late_failure, late_framing=ANTHROPIC_LATE_SSE_FRAMING if answers_as_anthropic else sse_framing.late, + on_passthrough_block=on_passthrough_block, ) def total_timeout_for( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6fbe562b5d1..8f5db97e7a3 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -210,8 +210,10 @@ from litellm.rust_bridge.catalog import Target from litellm.types.utils import ( PRICING_OVERRIDES_KEY, CallTypes, + Choices, CountContextManagement, LlmProviders, + Message, ModelResponse, ModelResponseStream, TextCompletionResponse, @@ -13263,6 +13265,70 @@ def _blocked_response_usage(original_response: object | None) -> "litellm.Usage" return litellm.Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0) +def _passthrough_block_response(e: ModifyResponseException) -> ModelResponse: + return ModelResponse( + model=e.model, + choices=[Choices(message=Message(content=e.message), finish_reason="content_filter")], + usage=_blocked_response_usage(e.original_response), + ) + + +async def _record_passthrough_block(e: ModifyResponseException, user_api_key_dict: UserAPIKeyAuth) -> None: + await proxy_logging_obj.post_call_failure_hook( + user_api_key_dict=user_api_key_dict, + original_exception=e, + request_data=e.request_data, + ) + + +async def _stream_chat_passthrough_block( + e: ModifyResponseException, user_api_key_dict: UserAPIKeyAuth, request: Request +) -> StreamingResponse: + # Read before post_call_failure_hook, which pops it from the request data. + logging_obj: Final = e.request_data.get("litellm_logging_obj") + await _record_passthrough_block(e, user_api_key_dict) + streaming_response: Final = litellm.CustomStreamWrapper( + completion_stream=litellm.utils.ModelResponseIterator( + model_response=_passthrough_block_response(e), convert_to_delta=True + ), + model=e.model, + custom_llm_provider="cached_response", + logging_obj=logging_obj, + ) + return StreamingResponse( + select_data_generator( + response=streaming_response, + user_api_key_dict=user_api_key_dict, + request_data=e.request_data, + request=request, + ), + media_type="text/event-stream", + status_code=200, + ) + + +async def _stream_text_completion_passthrough_block( + e: ModifyResponseException, user_api_key_dict: UserAPIKeyAuth, request: Request +) -> StreamingResponse: + await _record_passthrough_block(e, user_api_key_dict) + streaming_response: Final = litellm.TextCompletionStreamWrapper( + completion_stream=litellm.utils.ModelResponseIterator( + model_response=_passthrough_block_response(e), convert_to_delta=True + ), + model=e.model, + ) + return StreamingResponse( + select_data_generator( + response=streaming_response, + user_api_key_dict=user_api_key_dict, + request_data=e.request_data, + request=request, + ), + media_type="text/event-stream", + status_code=200, + ) + + @router.post( "/v1/chat/completions", dependencies=[Depends(user_api_key_auth)], @@ -13352,50 +13418,17 @@ async def chat_completion( user_max_tokens=user_max_tokens, user_api_base=user_api_base, version=version, + on_passthrough_block=lambda block: _stream_chat_passthrough_block(block, user_api_key_dict, request), ) if isinstance(result, BaseModel): return model_dump_with_preserved_fields(result, exclude_unset=True) else: return result except ModifyResponseException as e: - # Guardrail flagged content in passthrough mode - return 200 with violation message - _data = e.request_data - # Capture logging_obj before post_call_failure_hook pops it from _data. - _logging_obj: Final = _data.get("litellm_logging_obj") - await proxy_logging_obj.post_call_failure_hook( - user_api_key_dict=user_api_key_dict, - original_exception=e, - request_data=_data, - ) - _chat_response = litellm.ModelResponse() - _chat_response.model = e.model - _chat_response.choices[0].message.content = e.message - _chat_response.choices[0].finish_reason = "content_filter" - # Report the blocked LLM response's real usage (set before the stream - # branch so both paths carry it); zero for pre-call blocks. - _chat_response.usage = _blocked_response_usage(e.original_response) - - if data.get("stream", None) is not None and data["stream"] is True: - _iterator = litellm.utils.ModelResponseIterator(model_response=_chat_response, convert_to_delta=True) - _streaming_response = litellm.CustomStreamWrapper( - completion_stream=_iterator, - model=e.model, - custom_llm_provider="cached_response", - logging_obj=_logging_obj, - ) - selected_data_generator = select_data_generator( - response=_streaming_response, - user_api_key_dict=user_api_key_dict, - request_data=_data, - request=request, - ) - - return StreamingResponse( - selected_data_generator, - media_type="text/event-stream", - status_code=200, # Return 200 for passthrough mode - ) - return _chat_response + if data.get("stream") is True: + return await _stream_chat_passthrough_block(e, user_api_key_dict, request) + await _record_passthrough_block(e, user_api_key_dict) + return _passthrough_block_response(e) except RejectedRequestError as e: _data = e.request_data await proxy_logging_obj.post_call_failure_hook( @@ -13509,49 +13542,20 @@ async def completion( user_max_tokens=user_max_tokens, user_api_base=user_api_base, version=version, + on_passthrough_block=lambda block: _stream_text_completion_passthrough_block( + block, user_api_key_dict, request + ), ) except ModifyResponseException as e: - # Guardrail flagged content in passthrough mode - return 200 with violation message - _data = e.request_data - await proxy_logging_obj.post_call_failure_hook( - user_api_key_dict=user_api_key_dict, - original_exception=e, - request_data=_data, - ) - - if _data.get("stream", None) is not None and _data["stream"] is True: - _text_response: Final = litellm.ModelResponse() - # Set text attribute dynamically for text completion format - setattr(_text_response.choices[0], "text", e.message) - _text_response.model = e.model - _usage = _blocked_response_usage(e.original_response) - # Set usage attribute dynamically (ModelResponse accepts usage in __init__ but it's not in type definition) - setattr(_text_response, "usage", _usage) - _iterator = litellm.utils.ModelResponseIterator(model_response=_text_response, convert_to_delta=True) - _streaming_response = litellm.TextCompletionStreamWrapper( - completion_stream=_iterator, - model=e.model, - ) - - selected_data_generator = select_data_generator( - response=_streaming_response, - user_api_key_dict=user_api_key_dict, - request_data=_data, - request=request, - ) - - return StreamingResponse( - selected_data_generator, - media_type="text/event-stream", - status_code=200, # Return 200 for passthrough mode - ) - else: - _response = litellm.TextCompletionResponse() - _response.choices[0].text = e.message - _response.model = e.model - _usage = _blocked_response_usage(e.original_response) - _response.usage = _usage - return _response + if e.request_data.get("stream") is True: + return await _stream_text_completion_passthrough_block(e, user_api_key_dict, request) + await _record_passthrough_block(e, user_api_key_dict) + text_completion_response: Final = litellm.TextCompletionResponse() + text_completion_response.choices[0].text = e.message + text_completion_response.choices[0].finish_reason = "content_filter" + text_completion_response.model = e.model + text_completion_response.usage = _blocked_response_usage(e.original_response) + return text_completion_response except RejectedRequestError as e: _data = e.request_data await proxy_logging_obj.post_call_failure_hook( diff --git a/litellm/utils.py b/litellm/utils.py index 7fa9d32ae64..fb974b24283 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -9228,6 +9228,7 @@ class ModelResponseIterator: if convert_to_delta is True: _stream_response: Final = ModelResponseStream() _stream_response.choices[0].delta.content = model_response.choices[0].message.content + _stream_response.choices[0].finish_reason = model_response.choices[0].finish_reason self.model_response: ModelResponse | ModelResponseStream = _stream_response else: self.model_response = model_response diff --git a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py index 117d148c287..072de5c0eed 100644 --- a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py +++ b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py @@ -2,6 +2,7 @@ Test for anthropic_endpoints/endpoints.py, focusing on handling dictionary objects in streaming responses """ +import asyncio import datetime import json import logging @@ -85,10 +86,15 @@ _BLOCK_CALL_ID: Final = "messages-block" _BLOCK_TEXT: Final = "blocked by guardrail" -class _BlockEveryRequest(CustomGuardrail): +class _BlockEveryRequestOnceReleased(CustomGuardrail): + def __init__(self, released: asyncio.Event) -> None: + super().__init__(guardrail_name="deny-list", event_hook=GuardrailEventHooks.pre_call, default_on=True) + self.released: Final = released + async def async_pre_call_hook( self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict[str, object], call_type: str ) -> None: + await self.released.wait() self.raise_passthrough_exception(violation_message=_BLOCK_TEXT, request_data=data) @@ -117,21 +123,29 @@ def _messages_request(body: Mapping[str, object]) -> Request: async def _blocked_messages_reply( - body: Mapping[str, object], fastapi_response: Response, rails: tuple[CustomGuardrail, ...] + body: Mapping[str, object], + fastapi_response: Response, + rails: tuple[CustomGuardrail, ...], + keepalive_interval: float | None = None, ) -> tuple[object, str]: import litellm import litellm.proxy.anthropic_endpoints.endpoints as ep import litellm.proxy.proxy_server as proxy_server from litellm.proxy.utils import ProxyLogging - blocker: Final = _BlockEveryRequest( - guardrail_name="deny-list", event_hook=GuardrailEventHooks.pre_call, default_on=True - ) + released: Final = asyncio.Event() + if keepalive_interval is None: + released.set() with ( patch.object( proxy_server, "proxy_logging_obj", new=ProxyLogging(user_api_key_cache=DualCache()) ), # test-quality-ok: module global imported at call time; no injection seam - patch.object(litellm, "callbacks", new=[blocker, *rails]), # test-quality-ok: guardrails register through this global + patch.object( + litellm, "callbacks", new=[_BlockEveryRequestOnceReleased(released), *rails] + ), # test-quality-ok: guardrails register through this global + patch.object( + litellm, "sse_keepalive_ping_interval_seconds", new=keepalive_interval + ), # test-quality-ok: the operator's keepalive interval is read from this global ): response = await ep.anthropic_response( fastapi_response=fastapi_response, @@ -140,12 +154,34 @@ async def _blocked_messages_reply( ) if not isinstance(response, StreamingResponse): return response, "" - streamed: Final = b"".join( - [chunk if isinstance(chunk, bytes) else chunk.encode() async for chunk in response.body_iterator] - ).decode() + + async def chunks() -> AsyncGenerator[bytes, None]: + async for chunk in response.body_iterator: + released.set() + yield chunk if isinstance(chunk, bytes) else str(chunk).encode() + + streamed: Final = b"".join([chunk async for chunk in chunks()]).decode() return response, streamed +def _anthropic_events(streamed: str) -> tuple[tuple[str, Mapping[str, object]], ...]: + frames: Final = tuple(frame for frame in streamed.split("\n\n") if frame.strip() and frame != ": ping") + return tuple( + (frame.split("\n")[0].removeprefix("event: "), json.loads(frame.split("\n")[1].removeprefix("data: "))) + for frame in frames + ) + + +_BLOCK_EVENTS: Final = ( + "message_start", + "content_block_start", + "content_block_delta", + "content_block_stop", + "message_delta", + "message_stop", +) + + class TestBlockedResponseUsage: """Blocked responses report the blocked LLM response's real usage.""" @@ -220,26 +256,27 @@ class TestBlockedResponseUsage: {"model": "claude-haiku-5-5", "stream": True}, Response(), rails ) - frames: Final = tuple(frame for frame in streamed.split("\n\n") if frame.strip()) - events: Final = tuple( - (frame.split("\n")[0].removeprefix("event: "), json.loads(frame.split("\n")[1].removeprefix("data: "))) - for frame in frames - ) + events: Final = _anthropic_events(streamed) assert isinstance(response, StreamingResponse) assert response.media_type == "text/event-stream" assert response.headers["x-litellm-call-id"] == _BLOCK_CALL_ID - assert [name for name, _ in events] == [ - "message_start", - "content_block_start", - "content_block_delta", - "content_block_stop", - "message_delta", - "message_stop", - ] + assert [name for name, _ in events] == list(_BLOCK_EVENTS) assert all(payload["type"] == name for name, payload in events) assert events[2][1]["delta"] == {"type": "text_delta", "text": _BLOCK_TEXT} assert events[4][1]["delta"]["stop_reason"] == "end_turn" + @pytest.mark.asyncio + async def test_a_streaming_block_after_a_keepalive_ping_answers_the_block_events(self) -> None: + _, streamed = await _blocked_messages_reply( + {"model": "claude-haiku-5-5", "stream": True}, Response(), (), keepalive_interval=0.05 + ) + + events: Final = _anthropic_events(streamed) + assert streamed.startswith(": ping\n\n") + assert [name for name, _ in events] == list(_BLOCK_EVENTS) + assert events[2][1]["delta"] == {"type": "text_delta", "text": _BLOCK_TEXT} + assert events[4][1]["delta"]["stop_reason"] == "end_turn" + @pytest.mark.asyncio async def test_json_block_answers_with_the_call_id(self) -> None: fastapi_response: Final = Response() diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index bfe381cb848..ccfac3d4e30 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -11608,6 +11608,69 @@ async def test_a_broken_hook_does_not_replace_the_real_error_with_its_own_bug(): assert "audit backend" not in collected[-2].decode() +async def _block_once_released(release: asyncio.Event) -> Response: + await release.wait() + raise ModifyResponseException( + message="blocked by policy", model="gpt-4o", request_data={"stream": True}, guardrail_name="g" + ) + + +async def _drain_releasing(response: StreamingResponse, release: asyncio.Event) -> tuple[bytes, ...]: + async def chunks() -> AsyncGenerator[bytes, None]: + async for chunk in response.body_iterator: + release.set() + yield chunk if isinstance(chunk, bytes) else str(chunk).encode() + + return tuple([chunk async for chunk in chunks()]) + + +@pytest.mark.asyncio +async def test_a_late_passthrough_block_streams_the_routes_block_response_instead_of_an_error_frame(): + block_frames: Final = ('data: {"choices": [{"delta": {"content": "blocked by policy"}}]}\n\n', "data: [DONE]\n\n") + release: Final = asyncio.Event() + stream_block: Final = AsyncMock(side_effect=lambda _block: _sse_response(block_frames)) + record_failure: Final = AsyncMock(return_value=None) + + response: Final = await open_sse_before_first_byte( + _block_once_released(release), + ping_interval_seconds=0.05, + on_late_failure=record_failure, + on_passthrough_block=stream_block, + ) + collected: Final = await _drain_releasing(response, release) + + assert collected[0] == TTFT_PING + assert [chunk for chunk in collected if chunk != TTFT_PING] == [frame.encode() for frame in block_frames] + assert [call.args[0].message for call in stream_block.await_args_list] == ["blocked by policy"] + record_failure.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_a_late_passthrough_block_on_a_route_without_a_block_stream_still_ends_the_stream(): + release: Final = asyncio.Event() + response: Final = await open_sse_before_first_byte(_block_once_released(release), ping_interval_seconds=0.05) + collected: Final = await _drain_releasing(response, release) + + assert collected[0] == TTFT_PING + assert "error" in json.loads(collected[-2].decode().removeprefix("data: ").strip()) + assert collected[-1] == b"data: [DONE]\n\n" + + +@pytest.mark.asyncio +async def test_a_late_passthrough_block_whose_block_stream_fails_still_ends_the_stream(): + release: Final = asyncio.Event() + broken_block_stream: Final = AsyncMock(side_effect=RuntimeError("block stream unavailable")) + + response: Final = await open_sse_before_first_byte( + _block_once_released(release), ping_interval_seconds=0.05, on_passthrough_block=broken_block_stream + ) + collected: Final = await _drain_releasing(response, release) + + assert collected[0] == TTFT_PING + assert "error" in json.loads(collected[-2].decode().removeprefix("data: ").strip()) + assert collected[-1] == b"data: [DONE]\n\n" + + @pytest.mark.parametrize( "exc,expect_traceback", [ diff --git a/tests/unit/proxy/test_modify_response_streaming_passthrough.py b/tests/unit/proxy/test_modify_response_streaming_passthrough.py index f4d3b659ee1..5c46f1afac4 100644 --- a/tests/unit/proxy/test_modify_response_streaming_passthrough.py +++ b/tests/unit/proxy/test_modify_response_streaming_passthrough.py @@ -14,16 +14,28 @@ Reading it from the outer body passed ``logging_obj=None`` to ``AttributeError: 'NoneType' object has no attribute 'model_call_details'``. """ +import asyncio +import json +from collections.abc import AsyncGenerator, Mapping +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import Request, Response +from fastapi.responses import StreamingResponse +from starlette.types import Message +import litellm +import litellm.proxy.proxy_server as proxy_server +from litellm.caching.dual_cache import DualCache from litellm.exceptions import RejectedRequestError -from litellm.integrations.custom_guardrail import ModifyResponseException +from litellm.integrations.custom_guardrail import CustomGuardrail, ModifyResponseException +from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.proxy_server import chat_completion - +from litellm.proxy.proxy_server import chat_completion, completion +from litellm.proxy.utils import ProxyLogging +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import ModelResponse, ModelResponseStream, TextCompletionResponse async def _run_streaming_block_and_get_wrapper(exception): """Drive chat_completion's streaming guardrail-passthrough handler for the @@ -108,3 +120,181 @@ async def test_streaming_rejected_request_uses_request_data_logging_obj(): mock_csw.assert_called_once() assert mock_csw.call_args.kwargs["logging_obj"] is sentinel_logging_obj + + +_BLOCK_TEXT: Final = "blocked by guardrail" +_PING: Final = ": ping" + + +class _BlockEveryRequestOnceReleased(CustomGuardrail): + def __init__(self, released: asyncio.Event) -> None: + super().__init__(guardrail_name="deny-list", event_hook=GuardrailEventHooks.pre_call, default_on=True) + self.released: Final = released + self.failures: tuple[Exception, ...] = () + + async def async_pre_call_hook( + self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict[str, object], call_type: str + ) -> None: + await self.released.wait() + self.raise_passthrough_exception(violation_message=_BLOCK_TEXT, request_data=data) + + async def async_post_call_failure_hook( + self, + request_data: dict[str, object], + original_exception: Exception, + user_api_key_dict: UserAPIKeyAuth, + traceback_str: str | None = None, + ) -> None: + self.failures = (*self.failures, original_exception) + + +class _RejectEveryRequest(CustomLogger): + async def async_pre_call_hook( + self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict[str, object], call_type: str + ) -> None: + raise RejectedRequestError( + message="rejected by the pre-call hook", model="gpt-4o", llm_provider="openai", request_data=data + ) + + +def _post(path: str, body: Mapping[str, object]) -> Request: + async def receive() -> Message: + return {"type": "http.request", "body": json.dumps(body).encode(), "more_body": False} + + return Request( + { + "type": "http", + "method": "POST", + "path": path, + "query_string": b"", + "headers": [(b"content-type", b"application/json")], + }, + receive, + ) + + +async def _frames_releasing_the_block(response: StreamingResponse, released: asyncio.Event) -> tuple[str, ...]: + async def chunks() -> AsyncGenerator[str, None]: + async for chunk in response.body_iterator: + released.set() + yield chunk if isinstance(chunk, str) else bytes(chunk).decode() + + streamed: Final = "".join([chunk async for chunk in chunks()]) + return tuple(frame for frame in streamed.split("\n\n") if frame) + + +async def _reply( + path: str, + body: Mapping[str, object], + rail: CustomLogger, + released: asyncio.Event, + keepalive_interval: float | None, +) -> tuple[object, tuple[str, ...]]: + endpoint: Final = chat_completion if path == "/v1/chat/completions" else completion + with ( + patch.object( + proxy_server, "proxy_logging_obj", new=ProxyLogging(user_api_key_cache=DualCache()) + ), # test-quality-ok: module global imported at call time; no injection seam + patch.object(litellm, "callbacks", new=[rail]), # test-quality-ok: guardrails register through this global + patch.object( + litellm, "sse_keepalive_ping_interval_seconds", new=keepalive_interval + ), # test-quality-ok: the operator's keepalive interval is read from this global + ): + response: Final = await endpoint( + request=_post(path, body), fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth() + ) + if not isinstance(response, StreamingResponse): + return response, () + return response, await _frames_releasing_the_block(response, released) + + +async def _blocked_reply( + path: str, body: Mapping[str, object], keepalive_interval: float | None +) -> tuple[object, tuple[str, ...], _BlockEveryRequestOnceReleased]: + released: Final = asyncio.Event() + if keepalive_interval is None: + released.set() + blocker: Final = _BlockEveryRequestOnceReleased(released) + response, frames = await _reply(path, body, blocker, released, keepalive_interval) + return response, frames, blocker + + +def _data_frames(frames: tuple[str, ...]) -> tuple[str, ...]: + return tuple(frame.removeprefix("data: ") for frame in frames if frame.startswith("data: {")) + + +KEEPALIVE_INTERVALS: Final = pytest.mark.parametrize( + "keepalive_interval", [0.05, None], ids=["blocked-after-a-keepalive-ping", "keepalive-off"] +) + + +@pytest.mark.asyncio +@KEEPALIVE_INTERVALS +async def test_a_blocked_chat_stream_gets_the_block_message_and_a_content_filter_finish( + keepalive_interval: float | None, +) -> None: + _, frames, blocker = await _blocked_reply( + "/v1/chat/completions", {"model": "gpt-4o", "messages": [], "stream": True}, keepalive_interval + ) + + chunks: Final = tuple(ModelResponseStream.model_validate_json(frame) for frame in _data_frames(frames)) + assert (frames[0] == _PING) is (keepalive_interval is not None) + assert frames[-1] == "data: [DONE]" + assert not any('"error"' in frame for frame in frames) + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == _BLOCK_TEXT + assert chunks[-1].choices[0].finish_reason == "content_filter" + assert [type(failure) for failure in blocker.failures] == [ModifyResponseException] + + +@pytest.mark.asyncio +@KEEPALIVE_INTERVALS +async def test_a_blocked_text_completion_stream_gets_the_block_message(keepalive_interval: float | None) -> None: + _, frames, blocker = await _blocked_reply( + "/v1/completions", {"model": "gpt-4o", "prompt": "hi", "stream": True}, keepalive_interval + ) + + chunks: Final = tuple(TextCompletionResponse.model_validate_json(frame) for frame in _data_frames(frames)) + assert (frames[0] == _PING) is (keepalive_interval is not None) + assert frames[-1] == "data: [DONE]" + assert not any('"error"' in frame for frame in frames) + assert "".join(chunk.choices[0].text or "" for chunk in chunks) == _BLOCK_TEXT + assert chunks[-1].choices[0].finish_reason == "content_filter" + assert [type(failure) for failure in blocker.failures] == [ModifyResponseException] + + +@pytest.mark.asyncio +async def test_a_blocked_chat_request_without_streaming_gets_the_block_message_and_a_content_filter_finish() -> None: + response, _, blocker = await _blocked_reply( + "/v1/chat/completions", {"model": "gpt-4o", "messages": [], "stream": False}, None + ) + + assert isinstance(response, ModelResponse) + assert response.choices[0].message.content == _BLOCK_TEXT + assert response.choices[0].finish_reason == "content_filter" + assert [type(failure) for failure in blocker.failures] == [ModifyResponseException] + + +@pytest.mark.asyncio +async def test_a_blocked_text_completion_without_streaming_gets_the_block_text_and_a_content_filter_finish() -> None: + response, _, blocker = await _blocked_reply( + "/v1/completions", {"model": "gpt-4o", "prompt": "hi", "stream": False}, None + ) + + assert isinstance(response, TextCompletionResponse) + assert response.choices[0].text == _BLOCK_TEXT + assert response.choices[0].finish_reason == "content_filter" + assert [type(failure) for failure in blocker.failures] == [ModifyResponseException] + + +@pytest.mark.asyncio +async def test_a_rejected_text_completion_stream_ends_with_a_stop_finish() -> None: + released: Final = asyncio.Event() + released.set() + + _, frames = await _reply( + "/v1/completions", {"model": "gpt-4o", "prompt": "hi", "stream": True}, _RejectEveryRequest(), released, None + ) + + chunks: Final = tuple(TextCompletionResponse.model_validate_json(frame) for frame in _data_frames(frames)) + assert frames[-1] == "data: [DONE]" + assert chunks[-1].choices[0].finish_reason == "stop"