feat(routing): mid-stream fallback — resume a committed stream on fallback targets - #882
Merged
Conversation
…lback targets
Once a streaming response's first chunk is delivered the 200 is
committed and the pre-stream retry/failover loop is out of reach: a
mid-generation transport break, inter-chunk stall, decode failure, or
provider in-band error could only terminate the stream (in-band error
frame, no [DONE]). Long generations died with no recovery even when
healthy fallback targets existed (AISIX-Cloud#1222).
Add routing.stream_failure — explicit opt-in, default terminate:
stream_failure:
mode: continue # terminate (default) | continue
on: [transport_error, read_timeout,
upstream_decode_error, upstream_in_band_error] # default: all
max_fallbacks: 1
With continue, a qualifying mid-stream error hands the SAME client
stream to the remaining targets (strategy order, runtime state
re-checked, cooldown/health recorded). The fallback request carries
the original messages plus LiteLLM's verbatim continuation
instruction and the partial text as an assistant message — native
prefill on Anthropic-wire targets. Retryability reuses is_retryable
(non-429 in-band 4xx never triggers; retry_on_429 and
fallback_on_statuses apply).
Safety boundaries (deliberate divergence from LiteLLM, which falls
back with no output-shape guards): streams that emitted tool-call or
reasoning deltas, structured-output requests
(response_format json_object/json_schema), and partials past the 1MiB
continuation cap keep the terminate behavior.
Telemetry: the failed attempt emits its own per-attempt UsageEvent
with estimated partial spend (prompt + delivered partial, #655/#1074
contracts); the terminal event is attributed to the SERVING target
(attempt_kind mid_stream_fallback) with the terminal error class when
the stream still dies; new request-level stream_outcome
(success | partial_failed | partial_recovered) separates 'HTTP 200'
from 'stream completed' — partial_failed also lands on plain
terminate-mode failures. Client-facing usage frames fold the partial
estimate in (LiteLLM's merge semantics). New counter
aisix_mid_stream_fallbacks_total{model,outcome}.
Client-cancel safety is structural: the failover combinator advances
only when the client pulls, so an abandoned stream can never dispatch
a ghost fallback request (#1094 family).
First phase covers /v1/chat/completions; /v1/responses and
/v1/messages keep terminate semantics pending protocol-specific
continuation design.
|
Warning Review limit reachedYou’ve reached a temporary PR review limit under our Fair Usage Limits Policy. Next review available in: 31 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (6)
📝 WalkthroughWalkthroughThe PR adds configurable mid-stream failure handling. Streaming requests can terminate or continue on fallback targets. The proxy preserves partial output, emits continuation requests, tracks partial usage, and records stream outcomes and fallback metrics. ChangesMid-stream failover
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant Proxy
participant Primary
participant Fallback
Client->>Proxy: Start streaming chat completion
Proxy->>Primary: Forward request
Primary-->>Proxy: Partial chunks and mid-stream failure
Proxy->>Fallback: Send continuation request with partial output
Fallback-->>Proxy: Remaining completion chunks
Proxy-->>Client: Combined stream and one [DONE]
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 7
🧹 Nitpick comments (4)
crates/aisix-core/src/models/routing.rs (2)
195-199: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRemove the internal tracker id from the public model comment.
The doc comment on
StreamFailureis rendered intoschemas/resources/routing.schema.json(line 178) and into the Admin API OpenAPI description.AISIX-Cloud#1222is internal shorthand and has no meaning for an API consumer. Sibling fields inRoutingcarry no tracker ids. Keep the issue reference in the PR description or in a non-doc comment.♻️ Proposed comment rewrite
-/// Mid-stream failure policy for streaming responses (AISIX-Cloud#1222). +/// Mid-stream failure policy for streaming responses. /// /// Applies only to failures that occur after the response head (and /// possibly some chunks) reached the client; failures before the first /// chunk keep using the regular retry/failover loop.Regenerate the schemas after the edit:
cargo run -p aisix-core --bin dump-schemaAs per coding guidelines: "Write model comments as public API reference text, avoid internal shorthand" and "Regenerate resource schemas after changing model comments with
cargo run -p aisix-core --bin dump-schema".🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/aisix-core/src/models/routing.rs` around lines 195 - 199, Remove the internal “AISIX-Cloud#1222” tracker reference from the public doc comment on StreamFailure while preserving the behavioral description of mid-stream failures. Regenerate the resource schemas using the aisix-core dump-schema command so the generated routing schema and Admin API documentation reflect the updated comment.Source: Coding guidelines
220-241: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd unit tests for the new accessors and the serialized wire values.
This test module pins the wire contract for every other routing knob:
sticky_parses_and_defaults_false,when_all_unavailable_parses_try_anyway, andwhen_all_unavailable_rejects_unknown_value.stream_failurehas no equivalent here. The proxy tests exercise behavior, but nothing in this crate pins thatmode: "continue"andon: ["read_timeout"]deserialize, that an unknown mode is rejected, or thatmax_fallbacks_or_default()returns 1 andon_or_default()returns all four triggers when unset.💚 Proposed test
#[test] fn stream_failure_parses_and_defaults() { let r: Routing = serde_json::from_str( r#"{"targets":[{"model":"a"},{"model":"b"}], "stream_failure":{"mode":"continue","on":["read_timeout"]}}"#, ) .unwrap(); let sf = r.stream_failure.as_ref().unwrap(); assert_eq!(sf.mode_or_default(), StreamFailureMode::Continue); assert_eq!(sf.on_or_default(), &[StreamFailureTrigger::ReadTimeout]); // Unset max_fallbacks defaults to one mid-stream attempt. assert_eq!(sf.max_fallbacks_or_default(), 1); // Omitted `on` selects every trigger class. let all: Routing = serde_json::from_str( r#"{"targets":[{"model":"a"}],"stream_failure":{}}"#, ) .unwrap(); let sf = all.stream_failure.as_ref().unwrap(); assert_eq!(sf.mode_or_default(), StreamFailureMode::Terminate); assert_eq!(sf.on_or_default().len(), 4); // Unknown mode is rejected by the strict enum. assert!(serde_json::from_str::<Routing>( r#"{"targets":[{"model":"a"}],"stream_failure":{"mode":"resume"}}"#, ) .is_err()); }As per coding guidelines: "Prioritize end-to-end coverage over unit or integration coverage when coverage is limited, and never skip, disable, or use
.onlyto silence a failing test."🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/aisix-core/src/models/routing.rs` around lines 220 - 241, Add a routing test covering StreamFailure deserialization and accessor defaults. In the existing routing test module, verify mode "continue", on ["read_timeout"], and max_fallbacks_or_default() returning 1; also verify omitted on defaults to all four triggers and the default mode is Terminate, and assert that an unknown mode such as "resume" is rejected.Source: Coding guidelines
crates/aisix-obs/src/metrics.rs (1)
1190-1203: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
aisix_mid_stream_fallbacks_totalships with no test coverage at any level. The PR adds the metric, its recorder, and bothoutcomevalues, but no test asserts the counter, itsmodellabel, or either outcome. Ifrecord_mid_stream_fallbackis never called, or is called with the wrong label, the whole change set still passes. Every comparable counter incrates/aisix-obs/src/metrics.rshas a rendering test, and the proxy suite already scrapesstate.metrics.render()elsewhere.
crates/aisix-obs/src/metrics.rs#L1190-L1203: add a unit test that callsrecord_mid_stream_fallbackonce withrecovered = trueand once withrecovered = false, then asserts the rendered exposition carriesaisix_mid_stream_fallbacks_totalwithoutcome="recovered"andoutcome="failed"under the expectedmodellabel, each at 1.crates/aisix-proxy/src/lib.rs#L4771-L4788: keep theMetricshandle frombuild_stateinmid_stream_failure_continues_on_fallback_target_in_same_stream, then assert the scrape containsmodel="smart"withoutcome="recovered".crates/aisix-proxy/src/lib.rs#L4889-L4899: apply the same handle-and-scrape pattern inmid_stream_fallback_exhaustion_surfaces_error_without_done, assertingoutcome="failed".🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/aisix-obs/src/metrics.rs` around lines 1190 - 1203, The new mid-stream fallback metric lacks coverage. In crates/aisix-obs/src/metrics.rs:1190-1203, add a unit test for record_mid_stream_fallback that records both recovered and failed outcomes and verifies rendered aisix_mid_stream_fallbacks_total samples have model and outcome labels with value 1. In crates/aisix-proxy/src/lib.rs:4771-4788, retain the Metrics handle from build_state in mid_stream_failure_continues_on_fallback_target_in_same_stream and assert the scrape includes model="smart", outcome="recovered"; apply the same pattern at crates/aisix-proxy/src/lib.rs:4889-4899 for mid_stream_fallback_exhaustion_surfaces_error_without_done, asserting outcome="failed".crates/aisix-proxy/src/stream_failover.rs (1)
346-357: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winLog the configuration-error skips.
The runtime-state skip at lines 339-343 emits a
debugevent naming the target and the status. The three misconfiguration skips here emit nothing. If every remaining candidate is misconfigured, the client receiveslast_err— the original upstream error — and the operator gets no signal that the fallback chain was abandoned for a configuration reason.resolve_provider_keyalready builds a descriptive message (seecrates/aisix-proxy/src/dispatch.rslines 75-91) that is discarded here.♻️ Proposed fix
let model = &attempt.model; - let Ok(provider) = crate::dispatch::require_provider(model) else { - continue; - }; + let provider = match crate::dispatch::require_provider(model) { + Ok(p) => p, + Err(err) => { + tracing::warn!( + request_id = %plan.request_id, + target = %model.display_name, + error = %err, + "skipping mid-stream fallback candidate (misconfigured target)", + ); + continue; + } + }; let provider = provider.to_ascii_lowercase(); let snapshot = plan.state.snapshot.load(); - let Ok(pk_entry) = crate::dispatch::resolve_provider_key(&snapshot, model) else { - continue; - }; + let pk_entry = match crate::dispatch::resolve_provider_key(&snapshot, model) { + Ok(pk) => pk, + Err(err) => { + tracing::warn!( + request_id = %plan.request_id, + target = %model.display_name, + error = %err, + "skipping mid-stream fallback candidate (provider key unresolved)", + ); + continue; + } + }; let Some(bridge) = crate::dispatch::resolve_bridge(&plan.state.hub, &pk_entry.value) else { + tracing::warn!( + request_id = %plan.request_id, + target = %model.display_name, + "skipping mid-stream fallback candidate (no registered bridge)", + ); continue; };🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/aisix-proxy/src/stream_failover.rs` around lines 346 - 357, Update the fallback candidate handling around require_provider, resolve_provider_key, and resolve_bridge to emit debug events before each configuration-error continue, including the target/model and the descriptive error where available. Reuse the error message produced by resolve_provider_key, and clearly identify provider-resolution and bridge-resolution skips while preserving the existing continue behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@crates/aisix-proxy/src/chat.rs`:
- Around line 1732-1746: Update the record_mid_stream_fallback call to derive
recovered from the mid_stream_fallbacks recovery state rather than comparing
stream_outcome to "partial_recovered". Preserve the existing stream_outcome
classification while ensuring any stream with a successful mid-stream fallback
is recorded as recovered, including client-aborted or guardrail-blocked cases.
- Around line 1655-1679: Update stream_failover::acquire_fallback_stream to call
reserve_routing_target before dispatching each fallback target through
bridge.chat_stream. Skip targets whose reservation is rejected, and retain the
successful reservation so its usage is included in post_stream_keys and billed
to that fallback target. Ensure the MidStreamPlan fallback flow carries this
reservation through post-stream accounting.
- Around line 4705-4715: Update the usage handling in build_sse_stream so
mid_stream_extra is folded into only one terminal usage frame after
stream_failover::wrap completes, rather than every usage-bearing chunk. Preserve
the serving-attempt usage captured by comp and ensure earlier usage chunks
remain unmodified while the final client-facing usage includes the accumulated
failed-attempt estimate once.
In `@crates/aisix-proxy/src/stream_failover.rs`:
- Around line 358-377: Reserve the fallback target’s per-model quota before
dispatching each continuation in the loop containing bridge.chat_stream. Resolve
the gate from attempt.model using the existing quota::reserve_model_only flow,
skip the candidate and continue when reservation is refused, and commit or drop
the reservation along the existing success and failure paths.
- Around line 151-213: Update stream_failover::wrap so the
Box::pin(async_stream::stream! { ... }) generator is passed through
crate::request_id::in_request_span before returning. Preserve the existing
generator body and ChatChunkStream behavior while ensuring tracing and outbound
requests execute within the request span.
In `@schemas/resources/model.schema.json`:
- Around line 669-715: Update the doc comments for StreamFailure and
StreamFailureMode in routing.rs to remove the internal AISIX-Cloud tracker
reference and replace the unresolved StreamFailure::mode rustdoc link with a
standalone public description of the enum. Regenerate the schema using
dump-schema so the StreamFailure and StreamFailureMode descriptions contain only
consumer-facing API text.
In `@tests/e2e/src/cases/mid-stream-fallback-e2e.test.ts`:
- Around line 119-195: Set an explicit 30-second timeout on the test named
“connection drop mid-stream continues on the fallback target in the same
stream,” matching the timeout used by its sibling tests. Keep the test logic and
assertions unchanged.
---
Nitpick comments:
In `@crates/aisix-core/src/models/routing.rs`:
- Around line 195-199: Remove the internal “AISIX-Cloud#1222” tracker reference
from the public doc comment on StreamFailure while preserving the behavioral
description of mid-stream failures. Regenerate the resource schemas using the
aisix-core dump-schema command so the generated routing schema and Admin API
documentation reflect the updated comment.
- Around line 220-241: Add a routing test covering StreamFailure deserialization
and accessor defaults. In the existing routing test module, verify mode
"continue", on ["read_timeout"], and max_fallbacks_or_default() returning 1;
also verify omitted on defaults to all four triggers and the default mode is
Terminate, and assert that an unknown mode such as "resume" is rejected.
In `@crates/aisix-obs/src/metrics.rs`:
- Around line 1190-1203: The new mid-stream fallback metric lacks coverage. In
crates/aisix-obs/src/metrics.rs:1190-1203, add a unit test for
record_mid_stream_fallback that records both recovered and failed outcomes and
verifies rendered aisix_mid_stream_fallbacks_total samples have model and
outcome labels with value 1. In crates/aisix-proxy/src/lib.rs:4771-4788, retain
the Metrics handle from build_state in
mid_stream_failure_continues_on_fallback_target_in_same_stream and assert the
scrape includes model="smart", outcome="recovered"; apply the same pattern at
crates/aisix-proxy/src/lib.rs:4889-4899 for
mid_stream_fallback_exhaustion_surfaces_error_without_done, asserting
outcome="failed".
In `@crates/aisix-proxy/src/stream_failover.rs`:
- Around line 346-357: Update the fallback candidate handling around
require_provider, resolve_provider_key, and resolve_bridge to emit debug events
before each configuration-error continue, including the target/model and the
descriptive error where available. Reuse the error message produced by
resolve_provider_key, and clearly identify provider-resolution and
bridge-resolution skips while preserving the existing continue behavior.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: e52bccc6-9943-4b2f-b539-20bcb7b7a2fc
📒 Files selected for processing (13)
crates/aisix-admin/src/openapi.rscrates/aisix-core/src/lib.rscrates/aisix-core/src/models/mod.rscrates/aisix-core/src/models/routing.rscrates/aisix-obs/src/metrics.rscrates/aisix-obs/src/usage.rscrates/aisix-proxy/src/chat.rscrates/aisix-proxy/src/lib.rscrates/aisix-proxy/src/routing.rscrates/aisix-proxy/src/stream_failover.rsschemas/resources/model.schema.jsonschemas/resources/routing.schema.jsontests/e2e/src/cases/mid-stream-fallback-e2e.test.ts
…mpt usage isolation, honest metric - Reserve the fallback target's own rate-limit layers before the continuation dispatch (reserve_routing_target, AISIX-Cloud#1087 parity): refused reservation skips the candidate as a recorded 429 attempt without burning fallback budget; a granted one converts to a stream-lifetime concurrency hold (released on switch/stream end) and registers its keys for post-stream TPM accounting in the completion closure. - Reset the pump's usage accumulators when the serving attempt switches (shared attempt_seq): a per-chunk-usage provider (Gemini) would otherwise max-wins-mix the failed attempt's partial counters into the serving attempt's terminal event and double-fold with the estimated partial. - The mid-stream fallback metric now reports recovered = no terminal stream error, so a client that disconnects after a successful switch (or a guardrail block on recovered content) no longer counts as a failed failover. - Schema polish: drop the issue-tracker shorthand and the rustdoc intra-doc link from the public StreamFailure/StreamFailureMode descriptions (regenerated). - e2e: explicit 30s timeout on the connection-drop test, matching its siblings.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
Once a streaming response's first chunk reaches the client the 200 is committed and the pre-first-chunk failover of #554 can no longer help. A mid-generation transport break / inter-chunk stall / decode failure / provider in-band error today can only terminate the stream (in-band error frame, no
[DONE]) — long generations die with no recovery even when healthy fallback targets exist, and by HTTP status alone the request still looks like a success.Change
New
routing.stream_failureblock — explicit opt-in, omitted keeps today's behavior exactly:With
continue, a qualifying mid-stream error hands the same client stream to the remaining targets (strategy order, runtime health/cooldown re-checked at switch time, failed target cooled down):is_pre_first_chunkbranch).is_retryable: a non-429 in-band 4xx never triggers;retry_on_429/fallback_on_statusesapply.[DONE]. On exhaustion: in-band error frame, no fabricated[DONE].stream_failover.rs) only advances when the client pulls, so an abandoned stream can never dispatch a ghost fallback request.Safety boundaries (deliberate divergence from LiteLLM, which has no output-shape guards and will splice a re-emitted tool call after half-delivered arguments): streams that already emitted tool-call or reasoning deltas, structured-output requests (
response_format: json_object|json_schema), and partials past the 1 MiB continuation cap keep the terminate behavior.Telemetry (#655 / #1074 contracts)
usage_estimated), error class/message,attempt_kindpreserved.attempt_kind: mid_stream_fallback, its latency clock), not the pre-stream winner; when the stream still dies it now carries the terminalerror_class/error_message.UsageEvent.stream_outcome:success/partial_failed/partial_recovered— the HTTP status froze at 200, this is the signal that separates "200" from "stream actually completed".partial_failedalso lands on plain terminate-mode failures, so the observability gain applies with the feature off.aisix_mid_stream_fallbacks_total{model, outcome=recovered|failed}.Scope
First phase covers
/v1/chat/completions(parsed-SSE)./v1/responsesand/v1/messageskeep terminate semantics pending protocol-specific continuation design (issue scope agrees).served_by_targetin the access log still names the pre-stream winner (the log line is emitted at head commit); usage events carry the real serving attempts.Tests
[DONE]); tool-call safety gate;ontrigger-list narrowing. Unit: trigger taxonomy, continuation builder, structured-output detection, schema round-trips.aisix+ etcd + mock upstream): real mid-body connection drop → seamless continuation (client sees one logical answer, fallback got the continuation shape); real inter-chunk stall paststream_timeout→read_timeouttrigger; client abort mid-stream → fallback target never called.Fixes api7/AISIX-Cloud#1222
Summary by CodeRabbit
New Features
Bug Fixes