scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 109

fix(router): fall over on raised mid-stream errors in /v1/messages streams

Browse source
This commit is contained in:
mateo-berri 2026-08-27 15:29:50 -07:00
parent 7083c47998
commit 9f290d8b99
2 changed files with 207 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

101
litellm/router.py
View file

@ -420,6 +420,46 @@ def _anthropic_stream_should_decline_fallback(has_generated_content: bool, error
return has_generated_content or not error.is_pre_first_chunk
def _anthropic_stream_raised_error_status(error: Exception) -> int | None:
raw_status: Final = getattr(error, "status_code", None)
if isinstance(raw_status, int):
return raw_status
if isinstance(raw_status, str) and raw_status.isdigit():
return int(raw_status)
response_status: Final = getattr(getattr(error, "response", None), "status_code", None)
return response_status if isinstance(response_status, int) else None
def _anthropic_stream_fallback_error_for_raised(
error: Exception, model: str, llm_provider: str, has_generated_content: bool
) -> "MidStreamFallbackError | None":
"""
A provider iterator that fails mid-stream by raising (Bedrock surfaces
its event-stream exception frames as a BedrockError, a transport drop
raises httpx's error) never produces the Anthropic SSE `event: error`
frame the wrapper detects, so the raise is converted into the same
MidStreamFallbackError a detected error event gets, under the same gate:
only before real content reached the caller and only for a retriable
status (429, 5xx, or none at all for a transport failure), mirroring
CustomStreamWrapper._handle_stream_fallback_error on /chat/completions.
None means the exception propagates to the caller unchanged.
"""
from litellm.exceptions import MidStreamFallbackError
if has_generated_content:
return None
status_code: Final = _anthropic_stream_raised_error_status(error)
if status_code is not None and not _is_retriable_anthropic_status(status_code):
return None
return MidStreamFallbackError(
message=str(error),
model=model,
llm_provider=llm_provider,
original_exception=error,
is_pre_first_chunk=True,
)
def _anthropic_stream_commits_now(chunk: object, has_generated_content: bool, buffered_chunk_count: int) -> bool:
"""
Whether `chunk` should make Router._aanthropic_messages_streaming_iterator
@ -5019,6 +5059,8 @@ class Router:
has_generated_content = False # rebind-ok: set once real content is seen, or the buffer cap is hit
buffered_lifecycle_chunks: tuple[bytes, ...] = () # rebind-ok: flushed once committed or on decline
model: Final = cast(str, initial_kwargs.get("model")) # cast-ok: kwargs always carries the model group
custom_llm_provider: Final = initial_kwargs.get("custom_llm_provider")
llm_provider: Final = custom_llm_provider if isinstance(custom_llm_provider, str) else "anthropic"
try:
async for chunk in source_iterator:
if _anthropic_stream_forwards_ping_live(
@ -5061,14 +5103,16 @@ class Router:
yield chunk
for buffered_chunk in buffered_lifecycle_chunks:
yield buffered_chunk
except MidStreamFallbackError as e:
if _anthropic_stream_should_decline_fallback(has_generated_content, e):
for buffered_chunk in buffered_lifecycle_chunks:
yield buffered_chunk
if e.original_exception is not None:
raise e.original_exception from e
raise
async for item in self._aanthropic_messages_fallback_attempt(e, initial_kwargs, wrapper):
except Exception as stream_error: # noqa: BLE001 # any raised provider error must reach the fallback gate, like CustomStreamWrapper
async for item in self._aanthropic_messages_recover_stream_error(
stream_error,
has_generated_content,
buffered_lifecycle_chunks,
model,
llm_provider,
initial_kwargs,
wrapper,
):
yield item
finally:
with anyio.CancelScope(shield=True), contextlib.suppress(BaseException):
@ -5080,6 +5124,47 @@ class Router:
wrapper: Final = FallbackAwareAnthropicMessagesStream(stream_with_fallbacks(), source_iterator)
return wrapper
async def _aanthropic_messages_recover_stream_error(
self,
stream_error: Exception,
has_generated_content: bool,
buffered_lifecycle_chunks: tuple[bytes, ...],
model: str,
llm_provider: str,
initial_kwargs: dict[str, Any], # mutable-ok: handed to _aanthropic_messages_fallback_attempt, which mutates it
wrapper: "FallbackAwareAnthropicMessagesStream",
) -> AsyncGenerator[bytes, None]:
"""
Decides what a source-iterator failure in
Router._aanthropic_messages_streaming_iterator turns into: a fallback
attempt, or the error reaching the caller. A MidStreamFallbackError
(completion-bridge path, or the wrapper's own SSE error-event
detection) is declined per _anthropic_stream_should_decline_fallback
with the held-back lifecycle frames flushed first; any other raise is
converted per _anthropic_stream_fallback_error_for_raised and, when
not convertible, propagates untouched so the caller still gets a
clean error response.
"""
from litellm.exceptions import MidStreamFallbackError
if isinstance(stream_error, MidStreamFallbackError) and _anthropic_stream_should_decline_fallback(
has_generated_content, stream_error
):
for buffered_chunk in buffered_lifecycle_chunks:
yield buffered_chunk
if stream_error.original_exception is not None:
raise stream_error.original_exception from stream_error
raise stream_error
fallback_error: Final = (
stream_error
if isinstance(stream_error, MidStreamFallbackError)
else _anthropic_stream_fallback_error_for_raised(stream_error, model, llm_provider, has_generated_content)
)
if fallback_error is None:
raise stream_error
async for item in self._aanthropic_messages_fallback_attempt(fallback_error, initial_kwargs, wrapper):
yield item
async def _aanthropic_messages_fallback_attempt(
self,
e: "MidStreamFallbackError",

114
tests/test_litellm/test_router.py
View file

@ -15,6 +15,7 @@ import litellm
from litellm import Router
from litellm.exceptions import MidStreamFallbackError
from litellm.integrations.custom_logger import CustomLogger
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import (
SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES,
)
@ -10222,6 +10223,119 @@ async def test_anthropic_messages_fallback_also_catches_raised_midstream_error()
assert mock_fallback.await_args.kwargs["e"] is raised_error
@pytest.mark.asyncio
@pytest.mark.parametrize(
"raised_error",
[
BedrockError(status_code=503, message='serviceUnavailableException {"message": "Service unavailable"}'),
BedrockError(status_code=500, message='internalServerException {"message": "Internal error"}'),
BedrockError(status_code=429, message='throttlingException {"message": "Too many requests"}'),
httpx.ReadError("connection reset by upstream"),
],
ids=["503", "500", "429", "transport-drop"],
)
async def test_anthropic_messages_raised_provider_error_before_content_triggers_fallback(raised_error):
"""Bedrock surfaces a mid-stream exception frame by raising BedrockError
out of its iterator rather than yielding an Anthropic SSE error event, so
the wrapper must convert a retriable pre-content raise into a fallback
attempt exactly like a detected error event (parity with
CustomStreamWrapper._handle_stream_fallback_error on /chat/completions)."""
router = _anthropic_messages_make_router()
source = _AnthropicMessagesRaisingByteStream([_anthropic_messages_message_start_chunk()], raised_error)
fallback_stream = _AnthropicMessagesFallbackByteStream([_anthropic_messages_content_chunk("fallback answer")])
with patch.object(
router,
"async_function_with_fallbacks_common_utils",
new=AsyncMock(return_value=fallback_stream),
) as mock_fallback:
wrapped = await router._aanthropic_messages_streaming_iterator(
response=source,
initial_kwargs={"model": "primary"},
)
collected = [chunk async for chunk in wrapped]
assert collected == [_anthropic_messages_content_chunk("fallback answer")]
mock_fallback.assert_awaited_once()
converted = mock_fallback.await_args.kwargs["e"]
assert isinstance(converted, MidStreamFallbackError)
assert converted.original_exception is raised_error
assert converted.is_pre_first_chunk is True
assert source.closed is True
@pytest.mark.asyncio
@pytest.mark.parametrize(
"raised_error",
[
BedrockError(status_code=400, message='validationException {"message": "Malformed input"}'),
BedrockError(status_code=424, message='modelStreamErrorException {"message": "Model stream error"}'),
],
ids=["400", "424"],
)
async def test_anthropic_messages_raised_non_retriable_provider_error_propagates_unchanged(raised_error):
"""A raised 4xx (other than 429) is a client error no other deployment can
fix: it reaches the caller as the very same exception, with no fallback
attempt and nothing flushed, so the proxy still answers a clean 4xx."""
router = _anthropic_messages_make_router()
source = _AnthropicMessagesRaisingByteStream([_anthropic_messages_message_start_chunk()], raised_error)
with patch.object(
router,
"async_function_with_fallbacks_common_utils",
new=AsyncMock(),
) as mock_fallback:
wrapped = await router._aanthropic_messages_streaming_iterator(
response=source,
initial_kwargs={"model": "primary"},
)
collected = []
async def _consume():
async for chunk in wrapped:
collected.append(chunk)
with pytest.raises(BedrockError) as exc_info:
await _consume()
assert collected == []
assert exc_info.value is raised_error
mock_fallback.assert_not_awaited()
@pytest.mark.asyncio
async def test_anthropic_messages_raised_provider_error_after_content_propagates_unchanged():
"""Once real content reached the caller a fallback would append a second
message lifecycle to the same SSE stream, so a raised provider error after
content propagates as-is even when its status is retriable."""
router = _anthropic_messages_make_router()
content = _anthropic_messages_content_chunk("partial answer")
raised_error = BedrockError(status_code=503, message='serviceUnavailableException {"message": "Service unavailable"}')
source = _AnthropicMessagesRaisingByteStream([content], raised_error)
with patch.object(
router,
"async_function_with_fallbacks_common_utils",
new=AsyncMock(),
) as mock_fallback:
wrapped = await router._aanthropic_messages_streaming_iterator(
response=source,
initial_kwargs={"model": "primary"},
)
collected = []
async def _consume():
async for chunk in wrapped:
collected.append(chunk)
with pytest.raises(BedrockError) as exc_info:
await _consume()
assert collected == [content]
assert exc_info.value is raised_error
mock_fallback.assert_not_awaited()
@pytest.mark.asyncio
async def test_anthropic_messages_non_retriable_client_error_skips_fallback():
"""A 4xx (non-429) error type (e.g. invalid_request_error) is a client