Skip to content

feat: timeout-triggered fallback (timeout + stream_timeout) - #563

Merged
nic-6443 merged 8 commits into
mainfrom
feat/554-timeout-fallback
Jun 8, 2026
Merged

feat: timeout-triggered fallback (timeout + stream_timeout)#563
nic-6443 merged 8 commits into
mainfrom
feat/554-timeout-fallback

Conversation

@jarvis9443

@jarvis9443 jarvis9443 commented Jun 8, 2026

Copy link
Copy Markdown
Contributor

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 retryable BridgeError::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). When stream_timeout is unset, streaming falls back to timeout.

The streaming dispatch paths (/v1/chat/completions, /v1/messages for 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. Shared with_read_timeout / with_read_timeout_bytes combinators and a send_with_deadline connect 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_timeout is a new optional field; existing model configs are unaffected (both fields fold 0/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.ts covers 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 a firstEventDelayMs mock-upstream knob to model "headers fast, first token slow". Docs updated on the Models and Routing pages.

Summary by CodeRabbit

  • New Features

    • Optional per-model streaming timeout plus per-chunk read/connect deadlines to detect stalls and bound upstream calls; peek-first-chunk behavior to enable safe failover.
  • Behavior

    • Unified streaming/non-streaming failover: connect/first-chunk failures may fail over before bytes are sent; mid-stream stalls do not.
    • Model-specific request timeouts applied to upstream calls.
  • Telemetry

    • Winner-scoped routing telemetry and improved mapping of transport errors to timeout semantics.
  • Documentation

    • Docs updated with timeout fields, semantics, and failover behavior.
  • Tests

    • New end-to-end tests and enhanced streaming test harness for timeout/failover scenarios.

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.
@coderabbitai

coderabbitai Bot commented Jun 8, 2026

Copy link
Copy Markdown

Review Change Stack

Note

Reviews paused

It 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 reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

  • @coderabbitai resume to resume automatic reviews.
  • @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

  • ▶️ Resume reviews
  • 🔍 Trigger review

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 032eefd1-f4fc-4c90-85dd-5a001dc7e360

📥 Commits

Reviewing files that changed from the base of the PR and between b2ec424 and 583304c.

📒 Files selected for processing (1)
  • crates/aisix-admin/src/lib.rs

📝 Walkthrough

Walkthrough

This PR adds an optional per-model streaming read timeout (stream_timeout), helpers to fold ms values into Option<Duration> and compute effective streaming timeout, a stream-timeout runtime module, shared reqwest->BridgeError mapping, model-specific deadlines across endpoints, unified streaming/non-streaming per-target failover with first-chunk probing, telemetry winner-scoping, E2E tests, and docs/schema updates.

Changes

Timeout-based streaming failover implementation

Layer / File(s) Summary
Model structure and helpers
crates/aisix-core/src/models/model.rs, crates/aisix-core/src/models/schema.rs, schemas/resources/model.schema.json, docs/configuration/models.md, docs/configuration/routing-and-failover.md
Adds stream_timeout (ms), clarifies timeout semantics, and adds request_timeout(), stream_read_timeout(), and stream_timeout_effective() with unit tests verifying zero/absent handling and precedence.
Stream timeout combinators module
crates/aisix-proxy/src/stream_timeout.rs, crates/aisix-proxy/src/lib.rs
New module exposing with_read_timeout() for typed chunk streams, with_read_timeout_bytes() for raw byte streams, and send_with_deadline() to bound reqwest send()/connect phase; wired into crate root.
Shared reqwest error mapping & endpoint deadline wiring
crates/aisix-proxy/src/dispatch.rs, crates/aisix-proxy/src/completions.rs, crates/aisix-proxy/src/embeddings.rs, crates/aisix-proxy/src/images.rs
Adds reqwest_error_to_bridge() to map reqwest errors (timeouts → BridgeError::Timeout with elapsed_ms, others → BridgeError::Transport) and applies model request_timeout() deadlines where appropriate before bridge calls.
Count-tokens & Rerank send-timeout/error mapping
crates/aisix-proxy/src/count_tokens.rs, crates/aisix-proxy/src/rerank.rs
Build reusable req, apply optional per-model request timeout before send, capture send start time, and map send errors via reqwest_error_to_bridge for cooldown tracking.
Chat / Messages / Responses streaming failover & telemetry
crates/aisix-proxy/src/attempt.rs, crates/aisix-proxy/src/chat.rs, crates/aisix-proxy/src/messages.rs, crates/aisix-proxy/src/responses.rs
Refactors streaming to iterate targets and probe first-chunk/byte under read-timeout before committing success, records per-attempt telemetry, selects winner-scoped telemetry/headers/provider_key_id, removes RoutingTelemetry::single, and enforces connect/read deadlines and per-chunk read timeouts for passthrough and guardrail paths.
Stream read-timeout wiring in cross-provider & Anthropic paths
crates/aisix-proxy/src/messages.rs, crates/aisix-proxy/src/responses.rs
Anthropic passthrough uses send_with_deadline for connect/send; both Anthropic and cross-provider streaming wrap upstream streams with per-chunk read timeouts, peek first chunk for early failover, and re-prepend that chunk for downstream forwarding.
Test harness updates and E2E suite
tests/e2e/src/harness/upstream-openai.ts, tests/e2e/src/cases/timeout-fallback-e2e.test.ts
Adds firstEventDelayMs to mock harness, resilient socket handling, and comprehensive E2E tests validating non-streaming and streaming timeout failover, healthy-primary cases, and mid-stream stall semantics.
Admin test helper adjustments
crates/aisix-admin/src/lib.rs
Increase test body read cap from 64 KiB to 1 MiB across several admin test helpers to avoid truncation when reading merged/openapi embedded payloads.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related issues

Possibly related PRs

  • api7/ai-gateway#378: Overlaps with chat routing/telemetry changes touching /v1/chat/completions.
  • api7/ai-gateway#472: Related changes to routing/streaming candidate resolution and failover iteration.
  • api7/ai-gateway#491: Overlaps with /v1/messages streaming/non-streaming dispatch and output guardrail behavior.
🚥 Pre-merge checks | ✅ 4
✅ Passed checks (4 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The PR title accurately and concisely describes the main feature addition: timeout-triggered failover with two new timeout knobs (timeout and stream_timeout).
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
📝 Generate docstrings
  • Create stacked PR
  • Commit on current branch

Comment @coderabbitai help to get the list of available commands and usage tips.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

📥 Commits

Reviewing files that changed from the base of the PR and between a66c15f and 20cf6ae.

📒 Files selected for processing (19)
  • crates/aisix-core/src/models/model.rs
  • crates/aisix-core/src/models/schema.rs
  • crates/aisix-proxy/src/attempt.rs
  • crates/aisix-proxy/src/chat.rs
  • crates/aisix-proxy/src/completions.rs
  • crates/aisix-proxy/src/count_tokens.rs
  • crates/aisix-proxy/src/dispatch.rs
  • crates/aisix-proxy/src/embeddings.rs
  • crates/aisix-proxy/src/images.rs
  • crates/aisix-proxy/src/lib.rs
  • crates/aisix-proxy/src/messages.rs
  • crates/aisix-proxy/src/rerank.rs
  • crates/aisix-proxy/src/responses.rs
  • crates/aisix-proxy/src/stream_timeout.rs
  • docs/configuration/models.md
  • docs/configuration/routing-and-failover.md
  • schemas/resources/model.schema.json
  • tests/e2e/src/cases/timeout-fallback-e2e.test.ts
  • tests/e2e/src/harness/upstream-openai.ts
💤 Files with no reviewable changes (1)
  • crates/aisix-proxy/src/attempt.rs

Comment thread crates/aisix-core/src/models/model.rs
Comment thread crates/aisix-core/src/models/model.rs
Comment thread crates/aisix-proxy/src/count_tokens.rs
Comment thread crates/aisix-proxy/src/messages.rs Outdated
Comment thread crates/aisix-proxy/src/responses.rs Outdated
Comment thread docs/configuration/models.md Outdated
Comment thread schemas/resources/model.schema.json
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).

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

📥 Commits

Reviewing files that changed from the base of the PR and between df5f9fc and 0732182.

📒 Files selected for processing (5)
  • crates/aisix-core/src/models/model.rs
  • crates/aisix-proxy/src/messages.rs
  • crates/aisix-proxy/src/responses.rs
  • docs/configuration/models.md
  • schemas/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

Comment thread crates/aisix-core/src/models/model.rs
Comment thread crates/aisix-proxy/src/responses.rs Outdated
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.
@nic-6443
nic-6443 merged commit 58464de into main Jun 8, 2026
8 checks passed
@nic-6443
nic-6443 deleted the feat/554-timeout-fallback branch June 8, 2026 10:32
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants