feat: timeout-triggered fallback (timeout + stream_timeout) - #563
Conversation
Wire the per-model timeout as the non-streaming deadline at every dispatch site, and add stream_timeout as a per-chunk read timeout for streaming with first-chunk fallback (chat, messages, responses). Shared read-timeout combinators + reqwest error mapping; streaming paths now loop over targets.
Add timeout-fallback-e2e covering AC1 (non-stream timeout), AC2 (stream TTFT), AC3 (no interference), and the mid-stream-stall read-timeout semantics; add a firstEventDelayMs harness knob. Document timeout/stream_timeout on the Models page and a Timeout-Triggered Failover section on Routing.
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughThis PR adds an optional per-model streaming read timeout ( ChangesTimeout-based streaming failover implementation
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related issues
Possibly related PRs
🚥 Pre-merge checks | ✅ 4✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
Comment |
There was a problem hiding this comment.
Actionable comments posted: 7
🤖 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-core/src/models/model.rs`:
- Around line 247-253: Field-level docs for the `timeout` field are incomplete:
update the doc comment for `pub timeout: Option<u64>` to mention that when
`stream_timeout` is absent this value also serves as the fallback timeout for
streaming requests, and point readers to the `stream_timeout_effective()` helper
for precise behavior; keep wording concise and consistent with the existing
mention in `stream_timeout_effective()` so callers understand the dual role of
`timeout`.
- Around line 382-420: Add assertions to the existing
deserialises_stream_timeout_and_helpers_fold_zero test to cover
stream_timeout_effective() fallback semantics: verify that when stream_timeout
is present (e.g., 2500) stream_timeout_effective() returns Some(2_500), when
stream_timeout is absent it falls back to timeout (e.g., 30_000 →
Some(Duration::from_millis(30_000))), and when either value is explicitly 0 the
method returns None; reference the Model struct and its
stream_timeout_effective(), request_timeout(), and stream_read_timeout() helpers
when adding these checks.
In `@crates/aisix-proxy/src/count_tokens.rs`:
- Around line 299-313: The send() error handling currently forwards
reqwest::Error::to_string() verbatim into
crate::dispatch::reqwest_error_to_bridge (from the closure in count_tokens.rs)
which can leak URLs/API keys; modify the error path so that before calling
reqwest_error_to_bridge you sanitize/redact any URL, header, or auth info from
the reqwest::Error string (or change reqwest_error_to_bridge signature to accept
&reqwest::Error and perform redaction inside it), and then pass the redacted
message (or a sanitized BridgeError::Transport) into
crate::cooldown::note_failure; keep the existing timeout branch behavior
(e.is_timeout() → BridgeError::Timeout) and only alter the non-timeout transport
message flow to ensure no raw URLs/keys are emitted.
In `@crates/aisix-proxy/src/messages.rs`:
- Around line 1220-1234: The None branch treating upstream.next() EOF as
StreamAborted skips applying cooldown; call crate::cooldown::decide_cooldown
with aisix_gateway::BridgeError::StreamAborted (and model.cooldown.as_ref())
and, if it returns Some((ttl, reason)), invoke
state.runtime_status.mark_cooldown(model_id, ttl, reason) before returning
Err(ProxyError::Bridge(aisix_gateway::BridgeError::StreamAborted)); this mirrors
the error branch handling that uses decide_cooldown(&err, ...) so the dead
target becomes eligible for cooldown.
In `@crates/aisix-proxy/src/responses.rs`:
- Around line 602-625: In the BufferFull path's read loop (the loop that fills
the local `buf` Vec<u8>), treat an EOF that occurs when `buf.is_empty()` as a
first-chunk failure instead of a successful empty 200: when `next` is `None` and
`buf.is_empty()`, call `crate::cooldown::note_failure(&state.runtime_status,
model_id, model.cooldown.as_ref(), aisix_gateway::BridgeError::StreamAborted)`
and return the resulting `ProxyError::Bridge` (same pattern used elsewhere),
otherwise break as before; keep the existing timeout handling via
`model.stream_read_timeout()` and use the same identifiers (`buf`, `model_id`,
`state.runtime_status`, `model.cooldown`, `ProxyError::Bridge`,
`aisix_gateway::BridgeError::StreamAborted`) so behavior matches the verbatim
branch.
In `@docs/configuration/models.md`:
- Line 89: The docs currently say "0 or absent means no timeout" but later state
"When stream_timeout is absent, a streaming request falls back to timeout",
which conflicts; update the wording to explicitly distinguish the two semantics:
clarify that a value of 0 for either timeout or stream_timeout disables timeout
enforcement, whereas absence of stream_timeout does not disable streaming
timeout but instead inherits the numeric value of timeout (if present) —
reference the knobs `stream_timeout` and `timeout` and update the sentences
around both occurrences so readers understand "absent" means "use timeout" while
"0" means "no timeout".
In `@schemas/resources/model.schema.json`:
- Around line 88-96: The "stream_timeout" schema description incorrectly says "0
or absent = no timeout"; update the description for stream_timeout to clarify
semantics: explicitly state that a value of 0 means no timeout, while an
absent/null stream_timeout causes streaming requests to inherit the per-chunk
budget from the global "timeout" setting (i.e., fallback to timeout for
first-chunk and inter-chunk budgeting) and that non-zero integers specify the
millisecond gap timeout; keep references to stream_timeout and its uint64
integer/null typing intact.
🪄 Autofix (Beta)
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: bcdd7a15-754b-4a58-b95a-cda173660b09
📒 Files selected for processing (19)
crates/aisix-core/src/models/model.rscrates/aisix-core/src/models/schema.rscrates/aisix-proxy/src/attempt.rscrates/aisix-proxy/src/chat.rscrates/aisix-proxy/src/completions.rscrates/aisix-proxy/src/count_tokens.rscrates/aisix-proxy/src/dispatch.rscrates/aisix-proxy/src/embeddings.rscrates/aisix-proxy/src/images.rscrates/aisix-proxy/src/lib.rscrates/aisix-proxy/src/messages.rscrates/aisix-proxy/src/rerank.rscrates/aisix-proxy/src/responses.rscrates/aisix-proxy/src/stream_timeout.rsdocs/configuration/models.mddocs/configuration/routing-and-failover.mdschemas/resources/model.schema.jsontests/e2e/src/cases/timeout-fallback-e2e.test.tstests/e2e/src/harness/upstream-openai.ts
💤 Files with no reviewable changes (1)
- crates/aisix-proxy/src/attempt.rs
Unconditionally peeking the first stream chunk converted a streaming response whose first item errors/ends immediately into a pre-commit 502, regressing the committed-200 + in-band-error behavior for models with no stream_timeout. Gate the peek on stream_read_timeout().is_some() so streaming without a configured stream_timeout keeps the prior behavior; the read-timeout wrapper is a no-op when unset.
- Document the timeout/stream_timeout absent-vs-zero semantics (field docs, schema, models.md): 0 disables; absent stream_timeout falls back to timeout. - Cover stream_timeout_effective() cascade in unit tests. - Apply cooldown on first-chunk EOF in the cross-provider stream path, matching the other branches. - Treat a zero-byte buffered /v1/responses stream as a first-chunk failover (gated on stream_timeout, so no-timeout models keep prior behavior).
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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-core/src/models/model.rs`:
- Around line 423-440: The stream timeout helper currently collapses an explicit
Some(0) to None so Model::stream_timeout_effective() falls back to timeout
instead of honoring an explicit "stream_timeout":0; update the helper
stream_read_timeout() (and any callers) so that a stored value of 0 is preserved
as Some(Duration::from_millis(0)) rather than treated as None, and ensure
stream_timeout_effective() uses that preserved value (returning
Some(Duration::from_millis(0)) when stream_timeout is explicitly 0) before
considering the non-stream timeout fallback.
In `@crates/aisix-proxy/src/responses.rs`:
- Around line 602-604: The code currently uses model.stream_read_timeout() for
the first-chunk/inter-chunk deadline (e.g., the read_to binding, the zero-byte
EOF failover block, and calls to with_read_timeout_bytes/first-chunk peek),
which ignores models that only set timeout; change those to use
Model::stream_timeout_effective() (or model.stream_timeout_effective()) so the
effective streaming timeout is used everywhere. Update the read_to assignment,
the zero-byte EOF failover logic, and any
with_read_timeout_bytes/first-chunk-peek calls in the same file (the blocks
around the existing read_to and the other occurrences mentioned) to pass the
effective timeout value instead of stream_read_timeout().
🪄 Autofix (Beta)
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: baa3914a-6be4-48b3-af44-69aea174124c
📒 Files selected for processing (5)
crates/aisix-core/src/models/model.rscrates/aisix-proxy/src/messages.rscrates/aisix-proxy/src/responses.rsdocs/configuration/models.mdschemas/resources/model.schema.json
✅ Files skipped from review due to trivial changes (1)
- docs/configuration/models.md
🚧 Files skipped from review as they are similar to previous changes (2)
- schemas/resources/model.schema.json
- crates/aisix-proxy/src/messages.rs
Apply stream_timeout_effective() (stream_timeout, falling back to timeout) for the per-chunk read-timeout wrapper and the first-chunk peek gate, matching the connect deadline — so a model with only timeout set gets a consistent streaming budget. Clarify that stream_timeout:0 folds to the timeout fallback (matches LiteLLM), document it, and cover it in the unit test.
The merged /admin/openapi.json embeds every resource schema; adding the stream_timeout field + docs pushed it past the tests' 64 KB to_bytes cap on CI (LengthLimitError). Bump the cap to 1 MiB for the self-generated, in-memory spec body.
Problem
Fallback in a routing model (Model Group) only triggered on upstream errors /
5xx. A target that is too slow — a long non-streaming call, or a streaming response that takes too long to emit the first token — kept the request stuck on the primary instead of failing over. Tracks api7/AISIX-Cloud#554.What changed
Two per-model knobs now drive timeout-triggered failover, mirroring LiteLLM's
timeout/stream_timeout:timeout(already in the schema but previously inert on the request path) is now applied as the end-to-end deadline for non-streaming upstream calls at every dispatch site (chat, messages, responses, count_tokens, embeddings, images, completions, rerank). An elapsed deadline is a retryableBridgeError::Timeout(504), so on a routing model it fails over to the next target; on a direct model it returns a clean 504.stream_timeout(new) is a per-chunk read timeout for streaming calls — applied to the first chunk and every inter-chunk gap (resets per chunk). A first-chunk timeout fails over to the next target before any bytes reach the client; a mid-stream stall terminates the stream like any other upstream error (no fallback once the 200 is committed). Whenstream_timeoutis unset, streaming falls back totimeout.The streaming dispatch paths (
/v1/chat/completions,/v1/messagesfor both the cross-provider and native-Anthropic shapes, and/v1/responses) now loop over targets with a first-chunk peek instead of attempting only the first target. Sharedwith_read_timeout/with_read_timeout_bytescombinators and asend_with_deadlineconnect helper keep the chat-stream and raw-passthrough paths consistent. Per-attempt routing telemetry is recorded for the real winning attempt rather than a fabricated single record.Behavior / compatibility
stream_timeoutis a new optional field; existing model configs are unaffected (both fields fold0/absent to "no timeout"). The control-plane + dashboard exposure of these fields is a separate follow-up PR.Tests
New
tests/e2e/src/cases/timeout-fallback-e2e.test.tscovers non-streaming timeout → fallback, streaming first-token timeout → fallback (backup tokens piped, no client error), no-interference on healthy fast calls, and the mid-stream-stall read-timeout semantics. Adds afirstEventDelayMsmock-upstream knob to model "headers fast, first token slow". Docs updated on the Models and Routing pages.Summary by CodeRabbit
New Features
Behavior
Telemetry
Documentation
Tests