fix(anthropic_messages): preserve provider error semantics on upstream stream failure

The detached pump previously caught every upstream exception (Bedrock read,
decode, provider-response, or chunk-conversion error) and terminated the
client stream normally, masking the original provider exception and its
status so downstream failure handling never ran.

Now, when the upstream fails while the client is still connected, forward the
original exception through the queue so the client-facing generator re-raises
it and the proxy's failure handling (status code, post_call_failure_hook)
runs unchanged. Only when the client has already disconnected, where there is
no one to propagate to and no failure hook will fire, fall back to salvaging
partial spend from the collected chunks.
This commit is contained in:
nuernber 2026-08-05 17:35:26 -07:00 committed by nuernber
parent 1b401af716
commit ce25486702
2 changed files with 128 additions and 26 deletions

View file

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

View file

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