mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
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:
parent
1b401af716
commit
ce25486702
2 changed files with 128 additions and 26 deletions
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue