feat(retries): move the retry budget onto the model and apply it everywhere#816
feat(retries): move the retry budget onto the model and apply it everywhere#816jarvis9443 wants to merge 2 commits into
Conversation
…ywhere
`retries` lived only on `Routing`, so a direct model — one referenced
without a model group — ran the dispatch loop with the budget pinned to
zero and had no field to change it. A single transient upstream fault
went straight back to the caller. Seven endpoints had no retry path at
all, and the ensemble judge carried a hardcoded retry that ignored
configuration.
Add `Model.retries` and resolve the budget across three levels:
attempt.model.retries # this upstream's own budget
?? routing.retries # group-wide default (unchanged semantics)
?? config.upstream.retries # deployment default, built-in 2
Per-target beats per-group because a routing target is a Model: how many
times an upstream may be re-hit is a property of that upstream, and one
target tolerating three retries says nothing about its neighbour. An
explicit 0 at any level is an opt-out and does not fall through.
The deployment default additionally defers when another candidate target
is queued behind the current one: with a fallback available, prefer
failing over to grinding a failing upstream. Without it a two-target
group whose first target times out would burn three timeouts before
moving on, tripling the latency of timeout-driven fail-over (#554). An
explicitly configured budget is always honoured as written, so this only
gates the default. This also tracks what LiteLLM actually does, which is
easy to misread: its `num_retries` does not re-hit one deployment — each
retry re-enters deployment selection and the failed deployment has
meanwhile been cooled down, so a retry inside a multi-deployment group
lands elsewhere. Same-target grinding is what it does only for a
single-deployment group, which is the case the default is kept for.
The budget now applies on every endpoint. `chat`, `messages`, `responses`
and `count_tokens` keep their own loops (they also walk fall-over targets
and emit per-attempt telemetry); `count_tokens` gains the same-target
retry it never had. The single-model endpoints — embeddings, rerank,
completions, audio transcription and speech, images, videos, passthrough
— share a new `routing::retrying_dispatch`. Semantic-routing models are
fixed by the same change: they read the budget from `routing`, which they
do not have, so they were silently pinned at zero. Ensemble panel members
and the judge retry through `ProxyModelCaller`, using each member model's
own budget, which replaces the judge's hardcoded single retry so the two
layers cannot multiply.
Two scoping notes. Passthrough relays upstream statuses verbatim, so a
5xx never becomes an error there and only transport/decode faults are
retried — re-sending a relayed 5xx would break that contract and stack on
the provider edge's own retries. Audio's multipart form is rebuilt per
attempt, which is only sound because every part is an in-memory `Bytes`.
Backoff now honours an upstream `Retry-After` when the hint is within
5s, falling back to the existing exponential term otherwise. LiteLLM
honours up to 60s; an inline proxy cannot, because the wait is spent
inside the caller's own latency budget.
The default of 2 matches `openai.DEFAULT_MAX_RETRIES`, which is also
where the LiteLLM router lands when nothing is configured. This is a
behaviour change on upgrade: a request against a single upstream now
makes up to three attempts instead of one. Deployments that need the old
behaviour set `upstream.retries: 0`.
Refs api7/AISIX-Cloud#1132
|
Warning Review limit reachedYou’ve reached a temporary PR review limit under our Fair Usage Limits Policy. Next review available in: 19 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 (2)
📝 WalkthroughWalkthroughThis PR adds deployment-wide and per-model retry configuration, resolves effective per-target retry budgets, honors bounded upstream ChangesRetry configuration and state wiring
Retry engine and routing
Endpoint dispatch integration
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant ProxyRouter
participant UpstreamTarget
Client->>ProxyRouter: API request
ProxyRouter->>UpstreamTarget: initial request
UpstreamTarget-->>ProxyRouter: retryable failure with optional Retry-After
ProxyRouter->>ProxyRouter: resolve budget and calculate backoff
ProxyRouter->>UpstreamTarget: retry same target
UpstreamTarget-->>ProxyRouter: response
ProxyRouter-->>Client: upstream response or failover result
Possibly related issues
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
…t Transport `audio.rs` and `passthrough.rs` mapped every reqwest send error to `BridgeError::Transport`, discarding the timeout signal that `reqwest_error_to_bridge` preserves for the other direct-upstream paths. That was inert while those endpoints had no retry loop. With one, it was not: `RetryBudget::covers` keys off `BridgeError::Timeout` to decide whether an unconfigured budget is spent, so a model's own `timeout` was being retried. `audio-timeout-e2e` (#911 [22]) caught it on CI — a 400ms timeout took ~2.2s, and the assertion had only ~30ms of headroom left. It now completes in ~450ms. Refs api7/AISIX-Cloud#1132
There was a problem hiding this comment.
Actionable comments posted: 5
🤖 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/audio.rs`:
- Line 626: Update multipart_dispatch to derive a static retry endpoint label
from upstream_path before calling retrying_dispatch: use
"/v1/audio/translations" for translation requests and "/v1/audio/transcriptions"
otherwise, then pass that label instead of the hardcoded transcription path.
- Around line 639-648: Route reqwest send failures through the shared error
classifier instead of constructing BridgeError::Transport directly, preserving
timeout classification as 504 and its appropriate retry budget. Apply this
change to the send error handlers at crates/aisix-proxy/src/audio.rs:639-648 and
1011-1020, and crates/aisix-proxy/src/passthrough.rs:515-524; keep the existing
cooldown::note_failure integration.
In `@crates/aisix-proxy/src/completions.rs`:
- Around line 327-331: Move the cooldown tracker’s note_failure call into the
retry closure passed to routing::retrying_dispatch in the completions flow, so
every failed bridge.complete attempt records a cooldown before retrying.
Preserve the existing final dispatch result handling and match the established
pattern used by rerank.rs, audio.rs, and videos.rs.
In `@crates/aisix-proxy/src/count_tokens.rs`:
- Around line 239-244: Update the retry-budget calculation in the loop
containing effective_retries so has_fallback_targets reflects only remaining
usable Anthropic targets, excluding entries skipped by continue (such as
non-Anthropic targets). Do not derive it from attempt_models.len() or target_idx
alone; preserve same-target retries when no eligible fallback remains.
In `@crates/aisix-proxy/src/passthrough.rs`:
- Around line 528-535: The retry flow must not replay non-idempotent upstream
writes after a response-body read failure. In passthrough.rs, update the
handling around upstream_resp.bytes() and retrying_dispatch so UpstreamDecode is
not retried for non-idempotent methods, while preserving retries for safe
idempotent requests. In videos.rs, apply the same protection to the submit POST
path, or add an idempotency safeguard before allowing retries.
🪄 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: d9106941-a65e-400d-8377-18c04395f17c
📒 Files selected for processing (25)
config.example.yamlcrates/aisix-core/src/config.rscrates/aisix-core/src/models/model.rscrates/aisix-core/src/models/routing.rscrates/aisix-proxy/src/audio.rscrates/aisix-proxy/src/chat.rscrates/aisix-proxy/src/completions.rscrates/aisix-proxy/src/count_tokens.rscrates/aisix-proxy/src/embeddings.rscrates/aisix-proxy/src/ensemble.rscrates/aisix-proxy/src/images.rscrates/aisix-proxy/src/lib.rscrates/aisix-proxy/src/messages.rscrates/aisix-proxy/src/passthrough.rscrates/aisix-proxy/src/rerank.rscrates/aisix-proxy/src/responses.rscrates/aisix-proxy/src/routing.rscrates/aisix-proxy/src/state.rscrates/aisix-proxy/src/videos.rscrates/aisix-server/src/main.rsschemas/resources/model.schema.jsonschemas/resources/routing.schema.jsontests/e2e/src/cases/cooldown-contract-e2e.test.tstests/e2e/src/cases/direct-model-retries-e2e.test.tstests/e2e/src/cases/fallback-edges-e2e.test.ts
| // Send, check the status, and read the body as one retryable unit. See | ||
| // the same shape in rerank.rs for why `note_failure` stays per attempt. | ||
| let (upstream_headers, body_bytes) = | ||
| match crate::routing::retrying_dispatch(state, model, "/v1/audio/transcriptions", || { |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Retry endpoint label is hardcoded to transcriptions.
multipart_dispatch serves both /v1/audio/transcriptions and /v1/audio/translations (via upstream_path), but the retrying_dispatch label is fixed to "/v1/audio/transcriptions". Translation retries will be traced/logged under the wrong endpoint in retrying_dispatch's debug!/warn! output, misleading operators debugging translation retry behavior.
Since the endpoint parameter is &'static str, map upstream_path to a static label:
Proposed fix
let retry_endpoint = if upstream_path == "/audio/translations" {
"/v1/audio/translations"
} else {
"/v1/audio/transcriptions"
};🤖 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/audio.rs` at line 626, Update multipart_dispatch to
derive a static retry endpoint label from upstream_path before calling
retrying_dispatch: use "/v1/audio/translations" for translation requests and
"/v1/audio/transcriptions" otherwise, then pass that label instead of the
hardcoded transcription path.
| match crate::routing::retrying_dispatch(state, model, "/v1/completions", || { | ||
| bridge.complete(&body, &ctx) | ||
| }) | ||
| .await | ||
| { |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Inspect note_failure's semantics and how other retry-wrapped endpoints call it.
rg -n -B3 -A20 'fn note_failure' crates/aisix-proxy/src/cooldown.rs
# Compare call-site placement (inside vs outside retrying_dispatch closures)
# across all endpoints wrapped in this PR.
for f in audio.rs embeddings.rs rerank.rs completions.rs videos.rs; do
echo "=== crates/aisix-proxy/src/$f ==="
rg -n -B5 -A2 'note_failure|retrying_dispatch' "crates/aisix-proxy/src/$f"
doneRepository: api7/aisix
Length of output: 12162
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== retrying_dispatch definition ==="
rg -n -A40 -B10 'fn retrying_dispatch|retrying_dispatch<' crates/aisix-proxy/src
echo
echo "=== completions.rs around the retry block ==="
sed -n '300,540p' crates/aisix-proxy/src/completions.rs
echo
echo "=== bridge.complete references ==="
rg -n -A6 -B6 'fn complete|complete\(' crates/aisix-proxy/srcRepository: api7/aisix
Length of output: 20961
Record cooldown on each retry attempt
This only calls note_failure after retrying_dispatch gives up, so transient failures that succeed on a later retry never reach the cooldown tracker. Move it into the retry closure to match rerank.rs/audio.rs/videos.rs.
🤖 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/completions.rs` around lines 327 - 331, Move the
cooldown tracker’s note_failure call into the retry closure passed to
routing::retrying_dispatch in the completions flow, so every failed
bridge.complete attempt records a cooldown before retrying. Preserve the
existing final dispatch result handling and match the established pattern used
by rerank.rs, audio.rs, and videos.rs.
| let budget = crate::routing::effective_retries( | ||
| &target.model, | ||
| &target.id, | ||
| request_id, | ||
| ) | ||
| .await | ||
| { | ||
| Ok(resp) => return Ok((resp, "anthropic".to_string())), | ||
| Err(e) => { | ||
| let retryable = matches!( | ||
| &e, | ||
| ProxyError::Bridge(be) if crate::routing::is_retryable(be, retry_on_429, fallback_statuses) | ||
| ); | ||
| last_err = Some(e); | ||
| if !retryable { | ||
| break; | ||
| model_entry.value.routing.as_ref(), | ||
| state.default_retries, | ||
| target_idx + 1 < attempt_models.len(), | ||
| ); |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
has_fallback_targets should ignore skipped non-Anthropic targets.
attempt_models.len() includes targets this loop continues past, so a mixed group like [anthropic, openai] makes the default retry budget think there is a fallback when there isn’t one. That suppresses same-target retries on the only usable Anthropic target and makes transient failures fail too early.
Fix
+ let has_anthropic_fallback = attempt_models[target_idx + 1..]
+ .iter()
+ .any(|t| t.model.provider.as_deref() == Some("anthropic"));
let budget = crate::routing::effective_retries(
&target.model,
model_entry.value.routing.as_ref(),
state.default_retries,
- target_idx + 1 < attempt_models.len(),
+ has_anthropic_fallback,
);📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| let budget = crate::routing::effective_retries( | |
| &target.model, | |
| &target.id, | |
| request_id, | |
| ) | |
| .await | |
| { | |
| Ok(resp) => return Ok((resp, "anthropic".to_string())), | |
| Err(e) => { | |
| let retryable = matches!( | |
| &e, | |
| ProxyError::Bridge(be) if crate::routing::is_retryable(be, retry_on_429, fallback_statuses) | |
| ); | |
| last_err = Some(e); | |
| if !retryable { | |
| break; | |
| model_entry.value.routing.as_ref(), | |
| state.default_retries, | |
| target_idx + 1 < attempt_models.len(), | |
| ); | |
| let has_anthropic_fallback = attempt_models[target_idx + 1..] | |
| .iter() | |
| .any(|t| t.model.provider.as_deref() == Some("anthropic")); | |
| let budget = crate::routing::effective_retries( | |
| &target.model, | |
| model_entry.value.routing.as_ref(), | |
| state.default_retries, | |
| has_anthropic_fallback, | |
| ); |
🤖 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/count_tokens.rs` around lines 239 - 244, Update the
retry-budget calculation in the loop containing effective_retries so
has_fallback_targets reflects only remaining usable Anthropic targets, excluding
entries skipped by continue (such as non-Anthropic targets). Do not derive it
from attempt_models.len() or target_idx alone; preserve same-target retries when
no eligible fallback remains.
| let resp_body = upstream_resp.bytes().await.map_err(|e| { | ||
| crate::cooldown::note_failure( | ||
| tracker, | ||
| model_id, | ||
| cooldown_cfg, | ||
| aisix_gateway::BridgeError::UpstreamDecode(e.to_string()), | ||
| ) | ||
| })?; |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
rg -nP -C6 'fn is_retryable\b' crates/aisix-proxy/src/routing.rsRepository: api7/aisix
Length of output: 830
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- routing.rs: is_retryable + retrying_dispatch ---'
rg -n -C 20 'pub fn is_retryable\b|fn retrying_dispatch\b|retrying_dispatch\(' crates/aisix-proxy/src/routing.rs
echo
echo '--- passthrough.rs: retry wrapper context ---'
rg -n -C 12 'retrying_dispatch\(|UpstreamDecode|bytes\(\)\.await' crates/aisix-proxy/src/passthrough.rs
echo
echo '--- videos.rs: retry wrapper context and HTTP method ---'
rg -n -C 16 'retrying_dispatch\(|method\.clone\(|UpstreamDecode|bytes\(\)\.await' crates/aisix-proxy/src/videos.rsRepository: api7/aisix
Length of output: 15990
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- method / idempotency gating around retries ---'
rg -n -C 4 'idempot|request_id|x-aisix-request-id|retry_on_429|fallback_on_statuses|covers\(' crates/aisix-proxy/src
echo
echo '--- video route entry points and methods ---'
rg -n -C 6 'provider_call\(|method\s*:\s*Method|Method::POST|Method::GET|submit|poll|fetch' crates/aisix-proxy/src/videos.rsRepository: api7/aisix
Length of output: 50367
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- video submit route ---'
sed -n '1088,1115p' crates/aisix-proxy/src/videos.rs
echo
echo '--- retrying_dispatch docs/body ---'
sed -n '226,320p' crates/aisix-proxy/src/routing.rs
echo
echo '--- passthrough retry block ---'
sed -n '462,540p' crates/aisix-proxy/src/passthrough.rsRepository: api7/aisix
Length of output: 8582
Don’t retry non-idempotent upstream writes
retrying_dispatch treats UpstreamDecode as retryable, so a body-read failure after the upstream already accepted the request can replay the same POST and duplicate side effects. passthrough.rs should stop retrying this case, or retries should be limited to idempotent methods; videos.rs should likewise avoid retrying the submit POST unless it has an idempotency safeguard.
📍 Affects 2 files
crates/aisix-proxy/src/passthrough.rs#L528-L535(this comment)crates/aisix-proxy/src/videos.rs#L744-L813
🤖 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/passthrough.rs` around lines 528 - 535, The retry flow
must not replay non-idempotent upstream writes after a response-body read
failure. In passthrough.rs, update the handling around upstream_resp.bytes() and
retrying_dispatch so UpstreamDecode is not retried for non-idempotent methods,
while preserving retries for safe idempotent requests. In videos.rs, apply the
same protection to the submit POST path, or add an idempotency safeguard before
allowing retries.
Problem
retrieslived only onRouting, so a direct model — one referenced without a model group — ran the dispatch loop with the budget pinned to zero and had no field to change it. A single transient upstream fault went straight back to the caller.The machinery was already there: a direct model runs the same loop with a one-element target list. Only the parameter was hardcoded:
Beyond that, seven endpoints had no retry path at all, semantic-routing models read the budget from a
routingblock they don't have, and the ensemble judge carried a hardcoded retry that ignored configuration.Change
Model.retries, resolved across three levels:Per-target beats per-group because a routing target is a Model: how many times an upstream may be re-hit is a property of that upstream. An explicit
0at any level is an opt-out and does not fall through.The deployment default additionally defers when another candidate target is queued, and does not spend itself on a timeout. Both only gate the default — an explicitly configured budget is honoured as written.
Those two rules exist because the first version of this change broke things, and the E2E suite caught it:
timeout-fallback-e2e, /v1/completions: add output guardrails on generated choices[].text #554).timeout: 700waited ~2.1s instead of ~0.7s — the opposite of what setting a timeout asks for (audio-timeout-e2e, #911 [22]).They also track what LiteLLM actually does, which is easy to misread. Its
num_retriesdoes not re-hit one deployment: each retry re-enters deployment selection, and the failed deployment has meanwhile been cooled down, so a retry inside a multi-deployment group lands elsewhere. Same-target grinding is what LiteLLM does only for a single-deployment group — exactly the case the default is kept for.Coverage
All 13 endpoints now apply the budget:
chat,messages,responses,count_tokenskeep their own loops (they also walk fall-over targets and emit per-attempt telemetry).count_tokensgains the same-target retry it never had.routing::retrying_dispatch.ProxyModelCallerusing each member model's own budget. The judge's hardcoded single retry is removed so the two layers can't multiply into six upstream calls.Two scoping decisions:
Bytes— noted in the code, since a streamed part could not be replayed.Backoff now honours an upstream
Retry-Afterwithin 5s, falling back to the existing exponential term otherwise. LiteLLM honours up to 60s; an inline proxy can't, because the wait is spent inside the caller's own latency budget — the same reasoning that already tightened the exponential bounds relative to LiteLLM's defaults.Behaviour change
The default of 2 matches
openai.DEFAULT_MAX_RETRIES, where the LiteLLM router also lands when nothing is configured. On upgrade, a request against a single upstream makes up to three attempts instead of one, so a failing upstream sees up to 3x the load and a request that ultimately fails takes longer (it now includes the backoff waits).upstream.retries: 0restores the previous behaviour deployment-wide.Groups with a configured
routing.retriesare unaffected. Groups without one keep failing over on the first failure as before — the fallback rule is what preserves that.Tests
routing::tests— the three-level precedence, explicit-zero at each level, the fallback-deferral rule, and that a default budget refuses a timeout while a configured one accepts it.direct-model-retries-e2e(new) — a direct model spends the default budget (3 attempts),retries: 0makes exactly one, and the budget reaches/v1/embeddings, which had no retry loop. Verified to fail without the change: reverting the default to 0 givesexpected 1 to be 3on the first and third.responses::upstream_5xx_emits_zero_token_failed_attempt_eventnow pins theinitial, retry, retryattempt sequence a direct model produces.fallback-edges-e2edocuments the asymmetry the fallback rule creates (1 attempt on the target with a fallback behind it, 3 on the last).cargo test --workspacegreen (47 binaries),cargo fmt/clippy -D warningsclean, and the full 157-file E2E suite passes locally.Known limitation
retriescan be set and changed but not cleared back to "inherit" via the API —*int32can't distinguish an absent field from an explicitnull. Setting an explicit value is the workaround; worth a follow-up if it comes up.Refs api7/AISIX-Cloud#1132
Summary by CodeRabbit
New Features
Retry-Afterguidance.Bug Fixes
Tests