feat(router): stream chat completions through the Rust backend

The call driver reads an attempt's stream until its first content chunk and hands back a
ReplayStream that replays what it buffered. A MidStreamFallbackError before content goes
straight to the fallback chain without retrying the group, as Python's chat stream wrapper
does; any other error before content is raised as is, and after content a mid-stream
fallback error surfaces as its original exception.

A native test runs the same failing-before-first-chunk fallback chain through both backends
and gets the same content from the same sequence of model groups.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Aw1xktpUnix5tWhvSzey8t

litellm-rust/crates/router/POC.md

@@ -107,6 +107,8 @@ Design notes made while implementing:
 
 - Step 5 (in progress): `litellm/router_backends/selection.py::pinned_backend` fixes the backend decision; `tests/unit/router_backends/backends.py::router_backend` runs an opted-in module on both backends (Rust pinned to `RUST_REQUIRED`), with known Rust-side failures as strict xfails in `RUST_GAPS`. Opted in so far: `tests/unit/test_router_per_deployment_num_retries.py`, `tests/unit/test_router_exception_redaction.py`. `tests/test_litellm_rust/router/test_redis_parity.py` runs both backends against a live `redis-server` and compares every key, raw value (timestamp masked) and TTL: cooldowns, `allowed_fails` counters and failure RPM keys match byte for byte, and a Rust router honors cooldowns a Python router wrote
 
-Next: opt more suites into `router_backend` (`tests/unit/test_router_order_fallback.py` and the fallback/cooldown tests in `tests/unit/router_utils` that go through `Router(...)`), then streaming (the driver buffers to the first content chunk and reports `StreamFailedBeforeContent`; decide whether chat keeps Python's fallback-only rule), then the generic operations (`aresponses`, `aanthropic_messages`, `aimage_generation`), then `upsert_deployment`/`delete_deployment`/`update_settings` through `Registry::replace`
+Streaming (async chat): `RoutedCall` reads the stream until the first chunk with content (Python's `_stream_chunks_have_generated_content`) and returns a `ReplayStream` (a `FallbackAwareStreamWrapper`) that replays the buffered chunks. A `MidStreamFallbackError` before content is classified `stream_failure: before_content` and goes straight to the fallback chain without in-group retries, matching Python's chat wrapper (spec section 9) rather than retrying in the group; any other pre-content stream error is `terminal` and raised as is; after content a mid-stream fallback error surfaces as its original exception. Not done: the sync stream path, and combining the failed attempt's partial usage into the fallback's
 
-Known gaps to close or decline: chat streaming pre-first-chunk failures (Python goes straight to fallbacks, no in-group retry), `include_fallback_errors`, request-level `model_group_retry_policy`/`specific_deployment`/`priority`, CustomLogger fallback-event hooks, `router_cooldown_event_callback`, the 200-key `InMemoryCache` eviction, fallback keys containing `/` (provider-prefixed matching needs the cost map)
+Next: opt more suites into `router_backend` (`tests/unit/test_router_order_fallback.py` and the fallback/cooldown tests in `tests/unit/router_utils` that go through `Router(...)`), then sync streaming, then the generic operations (`aresponses`, `aanthropic_messages`, `aimage_generation`), then `upsert_deployment`/`delete_deployment`/`update_settings` through `Registry::replace`
+
+Known gaps to close or decline: `include_fallback_errors`, request-level `model_group_retry_policy`/`specific_deployment`/`priority`, CustomLogger fallback-event hooks, `router_cooldown_event_callback`, the 200-key `InMemoryCache` eviction, fallback keys containing `/` (provider-prefixed matching needs the cost map)

litellm-rust/crates/router/src/engine.rs

@@ -10,7 +10,7 @@ use std::{
 
 use crate::{
     cooldown::{self, status_is_retryable},
-    failure::{Classified, ExceptionClass, Failure, Raised},
+    failure::{Classified, ExceptionClass, Failure, Raised, StreamFailure},
     fallback::{chain_for_groups, generic_targets},
     host::{
         Attempt, HopStamp, Invoked, MockFailure, Op, RetryStamp, RouterHost, Target, TypedFallback,
@@ -308,7 +308,8 @@ impl<'a, H: RouterHost> Run<'a, H> {
             Ok(success) => return Ok(success),
             Err(stop) => self.failed(stop).map_err(Stop::Host)?,
         };
-        if failure.classified.guardrail_intervention {
+        if failure.classified.guardrail_intervention || failure.classified.stream_failure.is_some()
+        {
             return Err(Stop::Failed(failure));
         }
         if self.call.num_retries.is_none()
@@ -350,6 +351,9 @@ impl<'a, H: RouterHost> Run<'a, H> {
                 }
                 Err(stop) => self.failed(stop).map_err(Stop::Host)?,
             };
+            if failure.classified.stream_failure.is_some() {
+                return Err(Stop::Failed(failure));
+            }
             self.log_retry(hop, &failure);
             let remaining = num_retries - attempt - 1;
             healthy = self.healthy_counts(&hop.group).await.0;
@@ -530,7 +534,10 @@ impl<'a, H: RouterHost> Run<'a, H> {
         hop: &Hop,
         failure: Box<Failure<H::Error>>,
     ) -> Attempted<H::Response, H::Error, H::Fault> {
-        if self.call.disable_fallbacks || failure.classified.guardrail_intervention {
+        if self.call.disable_fallbacks
+            || failure.classified.guardrail_intervention
+            || failure.classified.stream_failure == Some(StreamFailure::Terminal)
+        {
             return Err(Stop::Failed(failure));
         }
         let top_level = hop.depth == 0;

litellm-rust/crates/router/src/failure.rs

@@ -19,6 +19,17 @@ pub enum ExceptionClass {
     ApiConnection,
 }
 
+/// How a streamed attempt failed before its first content chunk reached the caller.
+#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq)]
+#[serde(rename_all = "snake_case")]
+pub enum StreamFailure {
+    /// A failure the stream can fall back from (`MidStreamFallbackError`). Chat goes straight
+    /// to the fallback chain, without retrying in the group, as Python's stream wrapper does.
+    BeforeContent,
+    /// Any other stream error, which reaches the caller as raised.
+    Terminal,
+}
+
 /// What the host learned about an exception an attempt raised. The host computes the facts
 /// that depend on Python objects (header parsing, logging state); routing reads only these.
 #[derive(Clone, Debug, Default, Deserialize, PartialEq)]
@@ -45,6 +56,8 @@ pub struct Classified {
     pub exempt_from_cooldown: bool,
     /// `type(e) in litellm.LITELLM_EXCEPTION_TYPES`, which gates the exhausted-retry stamp.
     pub exact_litellm_type: bool,
+    #[serde(default)]
+    pub stream_failure: Option<StreamFailure>,
 }
 
 impl Classified {

litellm-rust/crates/router/tests/engine.rs

@@ -5,7 +5,7 @@ use std::{
 
 use litellm_router::{
     engine::{Engine, Override, RouteError, RouterCall},
-    failure::{Classified, ExceptionClass, Raised},
+    failure::{Classified, ExceptionClass, Raised, StreamFailure},
     host::{Attempt, Invoked, Op, RouterHost, Target},
     operation::Operation,
     pyrepr::PyNumber,
@@ -20,6 +20,7 @@ use rstest::rstest;
 enum Scripted {
     Ok,
     Fail(i64, ExceptionClass),
+    Stream(StreamFailure),
 }
 
 /// Answers each deployment's attempts from its own script; an unscripted attempt succeeds.
@@ -113,6 +114,13 @@ impl RouterHost for ScriptedHost {
                 error: format!("{id}#{index}"),
                 classified: classified(status, class),
             },
+            Scripted::Stream(failure) => Invoked::Failure {
+                error: format!("{id}#{index}"),
+                classified: Classified {
+                    stream_failure: Some(failure),
+                    ..classified(503, ExceptionClass::ServiceUnavailable)
+                },
+            },
         })
     }
 
@@ -421,3 +429,46 @@ async fn exhausted_retries_stamp_the_latest_error() {
             .any(|op| op == format!("stamp({}, 2)", raised(&error)))
     );
 }
+
+#[rstest]
+#[tokio::test]
+async fn a_stream_failing_before_content_falls_back_without_retrying_the_group() {
+    let mut settings = settings();
+    settings.fallbacks = Some(vec![chain("g", &["h"])]);
+    let engine = engine(
+        vec![
+            deployment("a", "g"),
+            deployment("b", "g"),
+            deployment("c", "h"),
+        ],
+        settings,
+    );
+    let before_content = Scripted::Stream(StreamFailure::BeforeContent);
+    let host = ScriptedHost::default()
+        .script("a", &[before_content.clone()])
+        .script("b", &[before_content]);
+
+    let routed = engine.route(&host, call("g")).await.ok().unwrap();
+
+    assert_eq!(host.attempts().len(), 2);
+    assert_eq!(host.attempts()[1], "c");
+    assert_eq!(routed.outcome.model_group, "h");
+    assert!(host.sleeps.lock().unwrap().is_empty());
+}
+
+#[rstest]
+#[tokio::test]
+async fn a_terminal_stream_failure_is_raised_without_retries_or_fallbacks() {
+    let mut settings = settings();
+    settings.fallbacks = Some(vec![chain("g", &["h"])]);
+    let engine = engine(vec![deployment("a", "g"), deployment("c", "h")], settings);
+    let host = ScriptedHost::default().script("a", &[Scripted::Stream(StreamFailure::Terminal)]);
+
+    let Err(RouteError::Failed(failed)) = engine.route(&host, call("g")).await else {
+        panic!("a terminal stream failure is raised");
+    };
+
+    assert_eq!(host.attempts(), ["a"]);
+    assert_eq!(raised(&failed.error), "a#1");
+    assert!(failed.ops.is_empty());
+}

litellm/router_backends/rust_call.py

@@ -17,12 +17,18 @@ from typing import Final, Literal, Protocol, TypeAlias, cast
 from pydantic import BaseModel, ConfigDict, TypeAdapter
 
 import litellm
+from litellm.exceptions import MidStreamFallbackError
 from litellm.integrations.custom_guardrail import is_guardrail_intervention
 from litellm.litellm_core_utils.exception_mapping_utils import (
     _get_response_headers,  # pyright: ignore[reportPrivateUsage]  # the header reader the cooldown callback uses
 )
 from litellm.litellm_core_utils.secret_redaction import redact_string
 from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_structure
+from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
+from litellm.router_backends.python_router import (
+    FallbackAwareStreamWrapper,
+    _stream_chunks_have_generated_content,  # pyright: ignore[reportPrivateUsage]  # the content rule Python's stream wrapper applies
+)
 from litellm.router_strategy.complexity_router.context_compaction import compact_to_fit
 from litellm.router_utils.add_retry_fallback_headers import (
     add_fallback_headers_to_response,
@@ -39,7 +45,7 @@ from litellm.router_utils.cooldown_handlers import (
     is_caller_timeout_408,
 )
 from litellm.types.router import RouterErrors, RouterRateLimitError
-from litellm.types.utils import ModelResponse
+from litellm.types.utils import ModelResponse, ModelResponseStream
 from litellm.utils import (
     _get_retry_after_from_exception_header,  # pyright: ignore[reportPrivateUsage]  # the Retry-After parser retries and cooldowns use
 )
@@ -200,6 +206,73 @@ def classify(error: BaseException, kwargs: Mapping[str, object], callbacks_ran:
     }
 
 
+class _StreamFailed(Exception):
+    """A stream that failed before its first content chunk."""
+
+    def __init__(self, error: BaseException, kind: Literal["before_content", "terminal"]) -> None:
+        super().__init__(str(error))
+        self.error: Final = error
+        self.kind: Final = kind
+
+
+class ReplayStream(FallbackAwareStreamWrapper):
+    """The attempt's stream, replaying the chunks read while waiting for its first content.
+
+    After content has reached the caller the stream cannot fall back, so a mid-stream
+    fallback error surfaces as the error that caused it, as the Python router's wrapper does.
+    """
+
+    def __init__(self, inner: CustomStreamWrapper, buffered: Sequence[ModelResponseStream]) -> None:
+        super().__init__(  # pyright: ignore[reportUnknownMemberType]  # CustomStreamWrapper's untyped constructor
+            completion_stream=cast(object, inner.completion_stream),  # cast-ok: untyped attribute
+            model=cast(object, inner.model),  # cast-ok: untyped attribute
+            custom_llm_provider=inner.custom_llm_provider,
+            logging_obj=inner.logging_obj,
+            _response_headers=getattr(inner, "_response_headers", None),
+        )
+        self._inner: Final = inner
+        self._replay: Final = iter(tuple(buffered))
+        self.chunks = cast(list[ModelResponseStream], inner.chunks)  # cast-ok: untyped attribute
+        self._hidden_params = dict(cast(Mapping[str, object], inner._hidden_params))  # pyright: ignore[reportPrivateUsage]  # cast-ok: the stream's own untyped metadata
+
+    def __aiter__(self) -> ReplayStream:
+        return self
+
+    async def __anext__(self) -> ModelResponseStream:
+        buffered: Final = next(self._replay, None)
+        if buffered is not None:
+            return buffered
+        try:
+            return await self._inner.__anext__()
+        except MidStreamFallbackError as error:
+            if error.original_exception is not None:
+                raise error.original_exception from error
+            raise
+
+    async def aclose(self) -> None:
+        await self._inner.aclose()
+
+
+def _not_fetched(stream: CustomStreamWrapper) -> bool:
+    return cast(object, stream.completion_stream) is None and stream.make_call is not None  # cast-ok: untyped attribute
+
+
+async def _replay_from_first_content(stream: CustomStreamWrapper) -> ReplayStream:
+    """Reads `stream` until a chunk carries content, so a failure before then can still go to
+    another deployment; the chunks read so far replay ahead of the rest."""
+    buffered: Final[list[ModelResponseStream]] = []  # mutable-ok: chunks read before the first content
+    try:
+        async for chunk in stream:
+            buffered.append(chunk)
+            if _stream_chunks_have_generated_content([chunk]):
+                break
+    except MidStreamFallbackError as error:
+        raise _StreamFailed(error, "before_content") from error
+    except Exception as error:
+        raise _StreamFailed(error, "terminal") from error
+    return ReplayStream(stream, buffered)
+
+
 class _AttemptFailed(Exception):
     """An attempt's error that litellm did not raise, so its failure callbacks never ran for it."""
 
@@ -244,6 +317,12 @@ class RoutedCall:
             return self._mock(parsed, kwargs)
         try:
             response: Final = await self._acompletion(parsed, kwargs)
+        except _StreamFailed as failed:
+            return (
+                "error",
+                failed.error,
+                {**classify(failed.error, kwargs, callbacks_ran=True), "stream_failure": failed.kind},
+            )
         except _AttemptFailed as failed:
             return ("error", failed.error, classify(failed.error, kwargs, callbacks_ran=False))
         except Exception as error:
@@ -382,8 +461,14 @@ class RoutedCall:
                         message="Response output was blocked.", model=attempt.model_group, llm_provider=""
                     )
                 )
+            if isinstance(response, CustomStreamWrapper) and _not_fetched(response):
+                await response.fetch_stream()
             normalizer.success_calls[model_name] += 1
+            if isinstance(response, CustomStreamWrapper):
+                return await _replay_from_first_content(response)
             return response
+        except _StreamFailed:
+            raise
         except _AttemptFailed as failed:
             self._stamp(failed.error, deployment, kwargs, model_name)
             raise

litellm/router_backends/rust_router.py

@@ -271,8 +271,8 @@ class RustRouter:
         unsupported: Final = [name for name in _UNSUPPORTED_REQUEST_KWARGS if name in kwargs]
         if unsupported:
             raise NotImplementedError(f"the Rust router backend does not support request option {unsupported[0]!r} yet")
-        if kwargs.get("stream") is True:
-            raise NotImplementedError("the Rust router backend does not support streaming yet")
+        if kwargs.get("stream") is True and operation == "completion":
+            raise NotImplementedError("the Rust router backend does not support sync streaming yet")
         clientside: Final = _IS_CLIENTSIDE_CREDENTIAL(kwargs)
         if clientside:
             raise NotImplementedError("the Rust router backend does not support client-side credentials yet")

tests/test_litellm_rust/router/test_rust_router.py

@@ -6,6 +6,8 @@ import random
 from collections.abc import Callable, Mapping, Sequence
 from typing import Final, Protocol
 
+from unittest.mock import MagicMock, patch
+
 import pytest
 
 import litellm
@@ -15,7 +17,9 @@ from litellm.router_backends.selection import select_backend
 from litellm.rust_bridge.catalog import Route, RouteRule
 from litellm.rust_bridge.configuration import Rollout
 from litellm.types.router import RouterRateLimitError
-from litellm.types.utils import ModelResponse
+from litellm.exceptions import MidStreamFallbackError
+from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
+from litellm.types.utils import ModelResponse, ModelResponseStream
 
 pytestmark = pytest.mark.requires_rust_extension
 
@@ -178,3 +182,69 @@ async def test_the_facade_serves_a_supported_config_from_the_rust_backend() -> N
     assert backend.get_model_names() == ["g"]
     with pytest.raises(NotImplementedError, match=r"Router\.upsert_deployment"):
         _ = backend.upsert_deployment
+
+
+class _FailingStream(CustomStreamWrapper):
+    """A provider stream that dies before its first chunk, as a 500 mid-stream does."""
+
+    def __init__(self, model: str) -> None:
+        super().__init__(completion_stream=object(), model=model, custom_llm_provider="openai", logging_obj=MagicMock())
+
+    def __aiter__(self) -> _FailingStream:
+        return self
+
+    async def __anext__(self) -> ModelResponseStream:
+        raise MidStreamFallbackError(
+            message=f"provider 500 from {self.model}",
+            model=str(self.model),
+            llm_provider="openai",
+            is_pre_first_chunk=True,
+            original_exception=litellm.InternalServerError(
+                message=f"provider 500 from {self.model}", model=str(self.model), llm_provider="openai"
+            ),
+        )
+
+
+class _OkStream(_FailingStream):
+    def __init__(self, model: str) -> None:
+        super().__init__(model)
+        self._replies: Final = iter(
+            (
+                ModelResponseStream(choices=[{"index": 0, "delta": {"role": "assistant"}}]),
+                ModelResponseStream(choices=[{"index": 0, "delta": {"content": f"ok-from-{model}"}}]),
+            )
+        )
+
+    async def __anext__(self) -> ModelResponseStream:
+        reply: Final = next(self._replies, None)
+        if reply is None:
+            raise StopAsyncIteration
+        return reply
+
+
+async def _fake_stream(**kwargs: object) -> CustomStreamWrapper:  # kwargs-ok: litellm.acompletion's surface
+    model: Final = str(kwargs["model"])
+    return _OkStream(model) if "fb2" in model else _FailingStream(model)
+
+
+async def test_a_stream_failing_before_content_walks_the_fallback_chain_on_both_backends() -> None:
+    arguments: Final = {
+        "model_list": (
+            {"model_name": "primary", "litellm_params": {"model": "openai/primary-model", "api_key": "k"}},
+            {"model_name": "fb1", "litellm_params": {"model": "openai/fb1-model", "api_key": "k"}},
+            {"model_name": "fb2", "litellm_params": {"model": "openai/fb2-model", "api_key": "k"}},
+        ),
+        "fallbacks": [{"primary": ["fb1", "fb2"]}],
+        "num_retries": 2,
+    }
+    observed: Final[list[tuple[str, tuple[str, ...]]]] = []  # mutable-ok: one observation per backend
+    for build in BACKENDS:
+        backend = build(arguments, 1)
+        with patch("litellm.acompletion", side_effect=_fake_stream) as acompletion:
+            response = await backend.acompletion("primary", MESSAGES, stream=True)
+            content = "".join([chunk.choices[0].delta.content or "" async for chunk in response])  # pyright: ignore[reportAttributeAccessIssue,reportUnknownMemberType,reportUnknownVariableType]  # stream chunks
+        groups = tuple(str(call.kwargs["metadata"]["model_group"]) for call in acompletion.call_args_list)
+        observed.append((content, groups))
+
+    assert observed[0] == observed[1]
+    assert observed[1] == ("ok-from-openai/fb2-model", ("primary", "fb1", "fb2"))