diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py index abcd817ea18..3458705c7f8 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py @@ -209,37 +209,24 @@ class BaseAnthropicMessagesStreamingIterator: once the client goes away the pump only buffers for billing so the queue can't grow unbounded. + An upstream failure (Bedrock read / decode / chunk-conversion error) + that happens while the client is still connected is forwarded through + the queue and re-raised here, so the original provider exception (and + its status) reaches the proxy's failure handling unchanged rather than + being masked by a generic incomplete-stream event. + This method provides the common logic for both Anthropic and Bedrock implementations. """ from litellm._logging import verbose_proxy_logger - queue: Final[asyncio.Queue[bytes | None]] = asyncio.Queue() + queue: Final[asyncio.Queue[bytes | None | BaseException]] = asyncio.Queue() client_detached: Final = asyncio.Event() async def _pump_upstream() -> None: collected_chunks: Final[list[bytes]] = [] saw_terminal_event = False # rebind-ok: accumulates across the upstream loop - try: - async for chunk in completion_stream: - if self.completion_start_time is None: - self.completion_start_time = datetime.now() - saw_terminal_event = saw_terminal_event or _is_terminal_stream_chunk(chunk) - encoded_chunk = self._convert_chunk_to_sse_format(chunk) - collected_chunks.append(encoded_chunk) - if not client_detached.is_set(): - queue.put_nowait(encoded_chunk) - except Exception as exc: # noqa: BLE001 # must still flush partial usage in finally, not crash the pump - verbose_proxy_logger.warning( - "async_sse_wrapper upstream pump stopped after %d chunks: %s(%s)", - len(collected_chunks), - type(exc).__name__, - exc, - ) - finally: - if not client_detached.is_set(): - if not saw_terminal_event: - queue.put_nowait(_incomplete_stream_error_sse_event()) - queue.put_nowait(None) + + async def _bill() -> None: try: await self._handle_streaming_logging(collected_chunks) except Exception as exc: # noqa: BLE001 # billing is best-effort; never crash the pump @@ -250,6 +237,42 @@ class BaseAnthropicMessagesStreamingIterator: exc, ) + try: + async for chunk in completion_stream: + if self.completion_start_time is None: + self.completion_start_time = datetime.now() + saw_terminal_event = saw_terminal_event or _is_terminal_stream_chunk(chunk) + encoded_chunk = self._convert_chunk_to_sse_format(chunk) + collected_chunks.append(encoded_chunk) + if not client_detached.is_set(): + queue.put_nowait(encoded_chunk) + except Exception as exc: # noqa: BLE001 # forward the provider error to a live client, else salvage spend + if not client_detached.is_set(): + # Preserve the provider-specific failure: hand the original + # exception to the client-facing generator so it re-raises + # and the proxy's failure handling (status code, + # post_call_failure_hook) runs. The failure path owns + # logging here, so don't also success-bill. + queue.put_nowait(exc) + return + # Client already disconnected: nothing to propagate to and no + # failure hook will run, so salvage the partial spend instead + # of dropping the request entirely. + verbose_proxy_logger.warning( + "async_sse_wrapper upstream pump failed after client disconnect (%d chunks): %s(%s)", + len(collected_chunks), + type(exc).__name__, + exc, + ) + await _bill() + return + + if not client_detached.is_set(): + if not saw_terminal_event: + queue.put_nowait(_incomplete_stream_error_sse_event()) + queue.put_nowait(None) + await _bill() + pump_task: Final = asyncio.create_task(_pump_upstream()) _UPSTREAM_PUMP_TASKS.add(pump_task) pump_task.add_done_callback(_UPSTREAM_PUMP_TASKS.discard) @@ -259,10 +282,13 @@ class BaseAnthropicMessagesStreamingIterator: item = await queue.get() if item is None: break + if isinstance(item, BaseException): + raise item yield item finally: - # Client-facing generator is being torn down (normal end or a - # disconnect GeneratorExit). Signal the pump to stop enqueueing so - # the queue can't grow unbounded, but let it keep draining upstream - # to its terminal usage event for accurate billing. + # Client-facing generator is being torn down (normal end, a + # re-raised upstream error, or a disconnect GeneratorExit). Signal + # the pump to stop enqueueing so the queue can't grow unbounded, but + # let it keep draining upstream to its terminal usage event for + # accurate billing. client_detached.set() diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py index 6bb061c9048..21eeb514094 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py @@ -364,3 +364,79 @@ async def test_async_sse_wrapper_bills_full_stream_when_client_reads_all(): assert len(client_chunks) == len(_STREAM_PREFIX) + len(_STREAM_TAIL) assert _output_tokens_from_logged_chunks(iterator.logged_chunks) == 64 assert not any(c.startswith(b"event: error\n") for c in iterator.logged_chunks) + + +class _ProviderStreamError(Exception): + """Stand-in for a provider-specific streaming failure carrying a status code.""" + + def __init__(self, message: str, status_code: int): + super().__init__(message) + self.status_code = status_code + + +@pytest.mark.asyncio +async def test_async_sse_wrapper_reraises_upstream_error_to_connected_client(): + """ + Regression: an upstream failure (Bedrock read / decode / chunk-conversion) + before message_stop must propagate the ORIGINAL provider exception to a + still-connected client, so the proxy's failure handling keeps the + provider-specific status. The pump must not swallow it into a generic + api_error event + normal termination. + """ + + async def _failing_stream(): + yield {"type": "message_start", "message": {"id": "msg_1", "usage": {"input_tokens": 52, "output_tokens": 1}}} + yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "partial"}} + raise _ProviderStreamError("bedrock stream blew up", status_code=529) + + iterator = _RecordingLoggingIterator( + litellm_logging_obj=_make_logging_obj("test_reraises_upstream_error"), + request_body={}, + ) + + received = [] + with pytest.raises(_ProviderStreamError) as excinfo: + async for chunk in iterator.async_sse_wrapper(_failing_stream()): + received.append(chunk) + + # Original exception + status preserved, not masked by a synthetic api_error. + assert excinfo.value.status_code == 529 + assert received # the client still got the pre-error chunks + assert not any(c.startswith(b"event: error\n") for c in received) + # On the failure path we do NOT success-bill (failure handling owns logging). + assert iterator.logged_chunks == [] + + +@pytest.mark.asyncio +async def test_async_sse_wrapper_salvages_partial_spend_on_upstream_error_after_disconnect(): + """ + When the upstream errors AFTER the client has already disconnected there is + no live client to re-raise to and no failure hook will run, so the pump + salvages partial spend from what it collected instead of dropping the row. + """ + tail_gated = asyncio.Event() + + async def _gated_failing_stream(): + yield {"type": "message_start", "message": {"id": "msg_1", "usage": {"input_tokens": 52, "output_tokens": 1}}} + yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "partial"}} + await tail_gated.wait() + raise _ProviderStreamError("late failure", status_code=500) + + iterator = _RecordingLoggingIterator( + litellm_logging_obj=_make_logging_obj("test_salvage_partial_on_late_error"), + request_body={}, + ) + + gen = iterator.async_sse_wrapper(_gated_failing_stream()) + received = [await gen.__anext__(), await gen.__anext__()] + await gen.aclose() # client disconnects before the upstream error + + tail_gated.set() # let the upstream raise now, after disconnect + for _ in range(100): + if iterator.logged_chunks: + break + await asyncio.sleep(0.01) + + assert len(received) == 2 + # Partial spend was still recorded rather than the whole request being dropped. + assert iterator.logged_chunks == received