fix(azure): apply stream_timeout per chunk, not as a whole-response ceiling - #809
Conversation
…eiling `Model::stream_timeout` is documented as "maximum gap in milliseconds between upstream streaming chunks", and that is how the proxy's `with_read_timeout` wrapper and the OpenAI bridge treat it. The Azure bridge instead computed `started + d` once and passed that single absolute instant to `timeout_at` on every chunk wait, turning the budget into a ceiling on the entire response. A long but perfectly healthy stream was therefore cut off with `BridgeError::Timeout` as soon as its total duration passed one gap's budget — the more output a model produced, the more likely it was to be killed. With `stream_timeout` unset the budget falls back to `timeout`, so a deployment that only set `timeout` was affected too. Switching to `tokio::time::timeout(d, …)` restarts the budget on each chunk, which is what the field means and what every other streaming path already did. `elapsed_ms` on the timeout error now reports the gap budget, matching `with_read_timeout`. Test: three chunks 80ms apart under a 250ms budget — every gap inside the budget, the ~320ms total past it. The whole stream must arrive. Fails with the absolute deadline, passes per-chunk. Fixes api7/AISIX-Cloud#1122
|
Warning Review limit reachedYou’ve reached a temporary PR review limit under our Fair Usage Limits Policy. Next review available in: 55 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)
📝 WalkthroughWalkthroughAzure OpenAI streaming timeout handling now treats ChangesAzure streaming timeout
Estimated code review effort: 2 (Simple) | ~10 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
|
@coderabbitai review |
|
@coderabbitai full review |
✅ Action performedFull review finished. Your included review limit is currently reached under our Fair Usage Limits Policy. Your recent PR review activity is in the 95th percentile or higher among CodeRabbit users, so adaptive limits apply. This review may still proceed through usage-based billing if eligible. Your next included review will be available in 42 minutes. |
✅ Action performedReview finished.
|
|
@coderabbitai full review |
✅ Action performedFull review finished. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
crates/aisix-provider-azure-openai/src/bridge.rs (1)
2020-2068: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy liftExercise this behavior through
chat_stream.This helper-level test still passes if
chat_streamstops forwardingctx.deadlineintobuild_chunk_stream. Add a source-blind streaming test that configures the 250ms budget, emits delayed SSE frames, and asserts all content arrives after total elapsed time exceeds the budget.As per coding guidelines, “Prioritize end-to-end coverage over unit or integration coverage when coverage is limited,” and “Treat E2E tests as source-blind and verify observable user contracts rather than implementation details.”
🤖 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-provider-azure-openai/src/bridge.rs` around lines 2020 - 2068, The existing test exercises build_chunk_stream directly, so it cannot verify that chat_stream forwards the configured deadline. Add a source-blind async test through chat_stream that configures a 250ms budget, supplies delayed SSE frames, asserts all expected content is received, and verifies total elapsed time exceeds the single-gap budget. Use only observable chat_stream behavior rather than referencing build_chunk_stream internals.Source: Coding guidelines
🤖 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.
Nitpick comments:
In `@crates/aisix-provider-azure-openai/src/bridge.rs`:
- Around line 2020-2068: The existing test exercises build_chunk_stream
directly, so it cannot verify that chat_stream forwards the configured deadline.
Add a source-blind async test through chat_stream that configures a 250ms
budget, supplies delayed SSE frames, asserts all expected content is received,
and verifies total elapsed time exceeds the single-gap budget. Use only
observable chat_stream behavior rather than referencing build_chunk_stream
internals.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 24904c7a-4b7e-430f-8a53-f2972a6f8751
📒 Files selected for processing (1)
crates/aisix-provider-azure-openai/src/bridge.rs
CodeRabbit review on #809: the semantics test drove `build_chunk_stream` directly, so it would still pass if `chat_stream` stopped forwarding `ctx.deadline` into the stream. That is true, but the suggested fix does not close it: a delivery test ("gaps within budget, total beyond it, everything arrives") passes just as well when no deadline is forwarded at all, because then nothing can time out. Closing it needs the opposite assertion, so both are added: - `chat_stream_delivers_a_long_stream_whose_gaps_stay_within_budget` — the suggested end-to-end shape, pinning per-chunk semantics through the public entry point. - `chat_stream_times_out_when_one_gap_exceeds_the_budget` — headers flush immediately so the connect-phase `with_deadline` is already satisfied, and only then does one gap exceed the budget. The per-chunk timeout is the only thing that can fire, so this is what fails if the wiring is dropped. Verified by mutation, and the two failure modes are caught by different tests: absolute deadline (the original bug) → helper + delivery test fail ctx.deadline not forwarded → only the timeout test fails Both need an upstream that emits frames on a schedule, which wiremock cannot express (`set_delay` delays the response as a whole — precisely the distinction under test), so they hand-roll a one-shot SSE server: no content-length, `connection: close`, body framed by EOF. That needs tokio's `net` + `io-util` in dev-dependencies.
|
Valid point, addressed in 0b82b50 — with one correction to the proposed shape. You're right that the helper-level test cannot see whether Closing it needs the opposite assertion, so both are now there:
Verified by mutation rather than by inspection, and the two failure modes turn out to be caught by different tests:
The middle column is the empirical version of the point above — the suggested test alone would not have caught the regression it was aimed at. One implementation note: both need an upstream that emits frames on a schedule, which wiremock cannot express ( |
Conflict in the Azure per-chunk timeout, where #808 landed `BridgeError::Timeout`'s new `cause` field on the same lines this branch rewrote. Resolved by keeping both: this branch's `d.as_millis()` (the per-chunk gap budget — the whole point of the change, and what `with_read_timeout` already reports) plus `cause: String::new()`, since an elapsed gateway-owned deadline has no transport-layer cause to name.
1 participant
Problem
Model::stream_timeoutis defined as "Maximum gap in milliseconds between upstream streaming chunks" (crates/aisix-core/src/models/model.rs). The proxy'swith_read_timeoutwrapper and the OpenAI bridge both implement it that way — the budget bounds eachnext()and resets after every successful read.The Azure bridge did something different:
Computing the instant once and reusing it turns the budget into a ceiling on the entire response. A long but perfectly healthy stream was cut off with
BridgeError::Timeoutas soon as its total duration passed one gap's budget — so the more output a model produced, the more likely it was to be killed mid-answer. Reasoning models and long completions are the worst case.This is not limited to deployments that set
stream_timeout:stream_timeout_effective()falls back totimeout, so a model configured with onlytimeout: 30000had its Azure streams capped at 30s of total wall-clock.Fix
Pass the budget as a
Durationand usetokio::time::timeout(d, …), which restarts on every chunk.elapsed_mson the timeout error now reports the gap budget, matching whatwith_read_timeoutreports on the same condition.The Audit H2 protection this code was added for is unchanged — a hung upstream body still can't wedge the connection after headers arrive, because a single gap exceeding the budget still fails.
Test
stream_budget_bounds_each_chunk_gap_not_the_whole_responsedrivesbuild_chunk_streamwith three chunks 80ms apart under a 250ms budget: every gap is comfortably inside the budget while the ~320ms total is past it, so the semantics are distinguishable. The whole stream must be delivered.Verified it fails under the old absolute deadline (
Timeoutpartway) and passes per-chunk.The existing
chat_stream_enforces_per_chunk_deadlinetest uses wiremock'sset_delay, which delays the response as a whole and so can't tell the two semantics apart — it passes either way. No E2E case is added: reproducing this needs precise per-chunk delivery timing from the upstream, which the unit test drives directly and an E2E mock would only approximate.Fixes api7/AISIX-Cloud#1122
Summary by CodeRabbit