feat(proxy): mid-stream fallback for /v1/messages — both dispatch legs - #893
Draft
jarvis9443 wants to merge 2 commits into
Draft
feat(proxy): mid-stream fallback for /v1/messages — both dispatch legs#893jarvis9443 wants to merge 2 commits into
jarvis9443 wants to merge 2 commits into
Conversation
Extends routing.stream_failure (AISIX-Cloud#1222) beyond /v1/chat/completions: the cross-provider leg wraps its ChatChunkStream with the existing combinator (the AnthropicSseEncoder outside it keeps the envelope seamless), and the Anthropic passthrough leg gains a frame-granular byte combinator that withholds in-band error frames, owns the read timeout when armed, and resumes the client's committed message envelope from fallback chunks via AnthropicSseEncoder::resume. Failed serving attempts emit per-attempt events stamped for this endpoint; terminal events gain stream_outcome + serving-target attribution; the continuation prompt moves to LiteLLM's current (Responses-surface) instruction. Ref api7/AISIX-Cloud#1222
📝 WalkthroughWalkthroughThe proxy adds ChangesMessages stream failover
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant MessagesDispatch
participant PrimaryUpstream
participant FallbackTarget
participant AnthropicSseEncoder
Client->>MessagesDispatch: Request /v1/messages stream
MessagesDispatch->>PrimaryUpstream: Start streaming attempt
PrimaryUpstream-->>MessagesDispatch: Partial SSE output and failure
MessagesDispatch->>FallbackTarget: Send continuation request
FallbackTarget-->>AnthropicSseEncoder: Return fallback stream
AnthropicSseEncoder-->>Client: Continue the existing message envelope
Possibly related PRs
Important Pre-merge checks failedPlease resolve all errors before merging. Addressing warnings is optional. ❌ Failed checks (1 error, 1 inconclusive)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
Contributor
Author
|
@coderabbitai review |
✅ Action performedReview finished.
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (6)
crates/aisix-proxy/src/messages.rs (2)
3766-3775: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winReuse
anthropic_error_framein the typed pump.Lines 2758-2762 build the identical frame with an inline
format!. Two copies of the same wire format can drift. Call this helper from the typed pump'sErrarm instead.♻️ Proposed change at lines 2756-2763
// Hold-back: the held (unscanned) chunks are dropped — // fail closed; only the error frame reaches the client. - let frame = format!( - "event: error\ndata: {{\"type\":\"error\",\"error\":{{\"type\":\"{}\",\"message\":{}}}}}\n\n", - e.error_type(), - serde_json::to_string(&e.to_string()).unwrap_or_else(|_| "\"error\"".into()), - ); - yield Ok(bytes::Bytes::from(frame)); + yield Ok(anthropic_error_frame(&e)); return;🤖 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/messages.rs` around lines 3766 - 3775, Update the typed pump’s Err arm to call the existing anthropic_error_frame helper instead of constructing the Anthropic error frame with an inline format!, preserving the current BridgeError value and response behavior while eliminating the duplicated wire-format logic.
959-961: 📐 Maintainability & Code Quality | 🔵 TrivialDocument the deferred
/v1/responseswork and link its follow-up issue.This commit wires
routing.stream_failurefor chat completions and both messages dispatch legs.crates/aisix-proxy/src/responses.rshas no corresponding wiring. If/v1/responsesis intentionally deferred, state the reason in the PR description and link the follow-up issue. TheAISIX-Cloud#1222reference alone does not identify that follow-up.🤖 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/messages.rs` around lines 959 - 961, Document in the PR description that routing.stream_failure remains unwired for /v1/responses, state the reason it is intentionally deferred, and link the specific follow-up issue; retain the existing AISIX-Cloud#1222 reference for the implemented mid-stream failover context.Source: Coding guidelines
tests/e2e/src/cases/mid-stream-fallback-messages-e2e.test.ts (1)
101-107: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider a provider-neutral tracking helper.
anthropicUpstreamonly starts an upstream and registers it for cleanup; nothing in it is Anthropic-specific. The cross-protocol test at lines 338-341 therefore callsstartOpenAiUpstreamand pushes toupstreamsby hand. Rename the helper (for exampletrackedUpstream) and use it in both places so no call site can forget the cleanup push.🤖 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 `@tests/e2e/src/cases/mid-stream-fallback-messages-e2e.test.ts` around lines 101 - 107, Rename anthropicUpstream to a provider-neutral helper such as trackedUpstream, preserving its start-and-register cleanup behavior. Update both the Anthropic setup and the cross-protocol test call site to use this helper, removing the duplicated direct startOpenAiUpstream call and manual upstreams.push.crates/aisix-proxy/src/lib.rs (2)
5200-5203: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winTighten the usage assertion so it cannot pass on the primary frame.
MSG_MID_STREAM_HEAD_SSEalready contains"output_tokens":1inside themessage_startusage object. The substring check therefore succeeds even if the fallback never emits the closingmessage_delta. Assert on themessage_deltaframe to pin the intended contract.♻️ Proposed assertion change
- assert!( - wire.contains("\"output_tokens\""), - "closing message_delta carries usage:\n{wire}" - ); + assert_eq!( + wire.matches("event: message_delta").count(), + 1, + "exactly one closing message_delta:\n{wire}" + ); + assert!( + wire.contains("\"output_tokens\":7"), + "closing message_delta carries the fallback's usage:\n{wire}" + );🤖 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/lib.rs` around lines 5200 - 5203, Update the assertion in the stream test around MSG_MID_STREAM_HEAD_SSE so it inspects the closing message_delta frame rather than the entire wire output. Assert that the message_delta-specific content includes "output_tokens", ensuring the check cannot pass solely because message_start already contains that field.
5438-5446: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider pinning the event count.
The loop reads exactly two events and stops. A regression that emits a third event for the same request would not fail this test. The sibling test at lines 5964-5975 drains with a short timeout and panics on an extra event. Apply the same pattern here.
🤖 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/lib.rs` around lines 5438 - 5446, Update the event collection in this test to drain the receiver with a short timeout after the expected events, and fail if an additional event is received, matching the sibling test’s pattern. Preserve the existing assertions for the two expected events and the sorting by attempt_index.tests/e2e/src/harness/upstream-openai.ts (1)
218-227: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winRelease the stall timer when the socket closes.
await sleep(600_000)holds a pending promise and an active Node timer for 10 minutes.server.close()inafterAlldoes not destroy in-flight sockets, so this handler keeps awaiting after the suite ends and can delay worker teardown. Resolve the wait as soon as the connection goes away.♻️ Proposed fix
if ( step.stallAfterEvents !== undefined && i >= step.stallAfterEvents ) { // Hang without closing: the socket stays open until the // gateway abandons it (read timeout) or the test ends. - await sleep(600_000); + await new Promise<void>((resolve) => { + const timer = setTimeout(resolve, 600_000); + timer.unref(); + res.on("close", () => { + clearTimeout(timer); + resolve(); + }); + }); return; }🤖 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 `@tests/e2e/src/harness/upstream-openai.ts` around lines 218 - 227, Update the stall branch in the upstream request handler to wait on a promise that resolves when the response socket closes, and clean up the associated listener/timer when either event occurs. Replace the fixed 10-minute sleep while preserving the intentional open-connection behavior until the client disconnects.
🤖 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/messages.rs`:
- Around line 1479-1499: Update the verbatim-phase ineligible in-band error
branch, near the existing upstream error-frame return, to record the failure in
the shared terminal_failure slot before returning. Reuse the same failure
representation used by handle_passthrough_failure and the ineligible
mid-continuation branch, so the downstream stream_outcome derivation reports
partial_failed and preserves the existing error-frame behavior.
- Around line 2743-2755: Align terminal mid-stream error handling across the
equivalent error arms in messages.rs, chat.rs, and responses.rs: record the same
reached_end, stream_failed, error class, and error message state so all
endpoints produce identical status and stream_outcome telemetry. Update the
shared endpoint-specific logic rather than changing unrelated EOF handling, and
add coverage verifying mid-stream errors emit the unified contract.
- Around line 1213-1258: The mid-stream continuation currently parses the
post-override body, causing primary ProviderKey request overrides to leak into
fallback requests. Preserve a post-redaction, pre-override request body for
continuation construction, then have the mid_stream_ctx block parse that body
while retaining the existing model restoration and translation behavior. Add
coverage using different primary and fallback param_renames to verify fallback
continuations use the fallback mapping.
In `@tests/e2e/src/cases/mid-stream-fallback-messages-e2e.test.ts`:
- Around line 227-231: Update the primary upstream setup in the mid-stream
fallback test to append a fourth raw stream frame after the existing three,
ensuring the harness reaches index 3 and triggers disconnectAfterEvents. Keep
the existing delay and disconnect configuration unchanged so the test exercises
an actual mid-body connection drop.
---
Nitpick comments:
In `@crates/aisix-proxy/src/lib.rs`:
- Around line 5200-5203: Update the assertion in the stream test around
MSG_MID_STREAM_HEAD_SSE so it inspects the closing message_delta frame rather
than the entire wire output. Assert that the message_delta-specific content
includes "output_tokens", ensuring the check cannot pass solely because
message_start already contains that field.
- Around line 5438-5446: Update the event collection in this test to drain the
receiver with a short timeout after the expected events, and fail if an
additional event is received, matching the sibling test’s pattern. Preserve the
existing assertions for the two expected events and the sorting by
attempt_index.
In `@crates/aisix-proxy/src/messages.rs`:
- Around line 3766-3775: Update the typed pump’s Err arm to call the existing
anthropic_error_frame helper instead of constructing the Anthropic error frame
with an inline format!, preserving the current BridgeError value and response
behavior while eliminating the duplicated wire-format logic.
- Around line 959-961: Document in the PR description that
routing.stream_failure remains unwired for /v1/responses, state the reason it is
intentionally deferred, and link the specific follow-up issue; retain the
existing AISIX-Cloud#1222 reference for the implemented mid-stream failover
context.
In `@tests/e2e/src/cases/mid-stream-fallback-messages-e2e.test.ts`:
- Around line 101-107: Rename anthropicUpstream to a provider-neutral helper
such as trackedUpstream, preserving its start-and-register cleanup behavior.
Update both the Anthropic setup and the cross-protocol test call site to use
this helper, removing the duplicated direct startOpenAiUpstream call and manual
upstreams.push.
In `@tests/e2e/src/harness/upstream-openai.ts`:
- Around line 218-227: Update the stall branch in the upstream request handler
to wait on a promise that resolves when the response socket closes, and clean up
the associated listener/timer when either event occurs. Replace the fixed
10-minute sleep while preserving the intentional open-connection behavior until
the client disconnects.
🪄 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: fbf0219c-aad0-4af0-962a-4018049820d1
📒 Files selected for processing (10)
crates/aisix-provider-anthropic/src/lib.rscrates/aisix-provider-anthropic/src/wire.rscrates/aisix-proxy/src/chat.rscrates/aisix-proxy/src/lib.rscrates/aisix-proxy/src/messages.rscrates/aisix-proxy/src/stream_failover.rstests/e2e/src/cases/mid-stream-fallback-e2e.test.tstests/e2e/src/cases/mid-stream-fallback-messages-e2e.test.tstests/e2e/src/harness/index.tstests/e2e/src/harness/upstream-openai.ts
Comment on lines
+1213
to
+1258
| // Mid-stream failover (AISIX-Cloud#1222): built up front because | ||
| // it decides who owns the read timeout below. The continuation | ||
| // dispatch runs on the internal ChatFormat — a body the parser | ||
| // can't map disarms (terminate as before). The parse runs on the | ||
| // outbound body (post-redaction, model already rewritten), same | ||
| // source the cross-provider path translates; the client-facing | ||
| // model name is restored explicitly. | ||
| let mid_stream_ctx = mid_stream.and_then(|arm| { | ||
| let mut chat = aisix_provider_anthropic::parse_inbound_request(&body).ok()?; | ||
| chat.model = model_name.to_string(); | ||
| aisix_provider_anthropic::translate_extras_to_openai_shape(&mut chat.extra); | ||
| let serving = Arc::new(std::sync::Mutex::new( | ||
| crate::stream_failover::ServingAttempt { | ||
| target_id: model_id.to_string(), | ||
| target_model: attempt.model.clone(), | ||
| provider: provider_label.clone(), | ||
| provider_key_id: pk_id.to_string(), | ||
| upstream_model: upstream_model.clone(), | ||
| cooldown: model.cooldown.clone(), | ||
| attempt_index: arm.winner_attempt_index, | ||
| attempt_kind: arm.winner_attempt_kind, | ||
| attempt_started, | ||
| }, | ||
| )); | ||
| let shared = crate::stream_failover::MidStreamShared::new(); | ||
| let plan = crate::stream_failover::MidStreamPlan { | ||
| cfg: arm.cfg, | ||
| endpoint: crate::stream_failover::MidStreamEndpoint::Messages, | ||
| structured_output: crate::stream_failover::expects_structured_output(&chat), | ||
| remaining: arm.remaining, | ||
| state: state.clone(), | ||
| auth: arm.auth, | ||
| group: arm.group, | ||
| req: chat, | ||
| request_id: request_id.to_string(), | ||
| client: client_ctx.clone(), | ||
| retry_on_429: arm.retry_on_429, | ||
| fallback_on_statuses: arm.fallback_on_statuses, | ||
| requested_model: model_name.to_string(), | ||
| api_key_id: api_key_id.to_string(), | ||
| applied_guardrails: resolved_chain.applied().to_vec(), | ||
| serving: Arc::clone(&serving), | ||
| shared: shared.clone(), | ||
| }; | ||
| Some((plan, serving, shared)) | ||
| }); |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Compare the body source used for the continuation request on both /v1/messages sub-paths,
# and check whether the fallback dispatch re-applies the fallback PK's request overrides.
set -euo pipefail
# The passthrough path applies PK overrides before parsing; find both call sites.
rg -n -C 12 'apply_param_renames|apply_default_body_fields|apply_param_constraints' crates/aisix-proxy/src crates/aisix-provider-openai/src
# Does the mid-stream fallback dispatch apply the fallback PK's request overrides?
rg -n -C 15 'fn acquire_fallback_stream|continuation_request' crates/aisix-proxy/src/stream_failover.rsRepository: api7/aisix
Length of output: 49051
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the fallback dispatch path and the target-specific ProviderKey used by
# the continuation request.
sed -n '403,570p' crates/aisix-proxy/src/stream_failover.rs
# Inspect the plan fields and the messages passthrough caller to determine
# whether the continuation receives the fallback target's request overrides.
rg -n -C 12 'MidStreamPlan|acquire_fallback_stream|chat_stream|provider_key_id|attempt\.model|attempt\.provider' \
crates/aisix-proxy/src/stream_failover.rs crates/aisix-proxy/src/messages.rsRepository: api7/aisix
Length of output: 50367
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Show the exact ChatFormat serialization and override pipeline used by
# bridge.chat_stream().
rg -n -C 20 'async fn chat_stream|fn chat_stream|prepare_outbound_body|serialize.*ChatFormat|ChatFormat' \
crates/aisix-provider-openai/src crates/aisix-provider-anthropic/src crates/aisix-gateway/src
# Show the source body used for the other messages streaming path and the
# inbound parser's handling of unknown fields.
rg -n -C 18 'parse_inbound_request|translate_extras_to_openai_shape|let mut body|body\.clone' \
crates/aisix-proxy/src/messages.rs crates/aisix-provider-anthropic/srcRepository: api7/aisix
Length of output: 50367
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- OpenAI bridge chat_stream ---'
rg -n 'chat_stream|prepare_outbound_body|build_request' crates/aisix-provider-openai/src/bridge.rs
sed -n '600,760p' crates/aisix-provider-openai/src/bridge.rs
printf '%s\n' '--- Anthropic bridge chat_stream ---'
rg -n 'chat_stream|build_request' crates/aisix-provider-anthropic/src/bridge.rs
sed -n '280,430p' crates/aisix-provider-anthropic/src/bridge.rs
printf '%s\n' '--- Anthropic inbound extra handling ---'
sed -n '1347,1388p' crates/aisix-provider-anthropic/src/wire.rs
printf '%s\n' '--- ChatFormat and ChatMessage serialization ---'
rg -n -C 14 'pub struct ChatFormat|struct ChatFormat|pub struct ChatMessage|struct ChatMessage' cratesRepository: api7/aisix
Length of output: 23898
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
messages = Path("crates/aisix-proxy/src/messages.rs").read_text()
failover = Path("crates/aisix-proxy/src/stream_failover.rs").read_text()
anthropic_wire = Path("crates/aisix-provider-anthropic/src/wire.rs").read_text()
chat = Path("crates/aisix-gateway/src/chat.rs").read_text()
openai_bridge = Path("crates/aisix-provider-openai/src/bridge.rs").read_text()
anthropic_bridge = Path("crates/aisix-provider-anthropic/src/bridge.rs").read_text()
checks = {
"primary overrides precede continuation parse":
messages.index("apply_param_renames(&mut body") <
messages.index("parse_inbound_request(&body)"),
"inbound parser stores unknown top-level fields in ChatFormat.extra":
'chat.extra.insert(key.clone(), value.clone())' in anthropic_wire,
"ChatFormat flattens extra fields during serialization":
'#[serde(default, skip_serializing_if = "serde_json::Map::is_empty", flatten)]'
in chat,
"fallback continuation uses the shared cont_req":
"let cont_req = continuation_request(&plan.req, partial);" in failover and
"bridge.chat_stream(&cont_req, &ctx)" in failover,
"fallback context resolves the candidate ProviderKey":
"resolve_provider_key(&snapshot, model)" in failover and
"Arc::new(pk_entry.value.clone())" in failover,
"OpenAI bridge applies overrides after ChatFormat serialization":
"prepare_outbound_body(" in openai_bridge and
"apply_param_renames(&mut body, &r.param_renames)" in openai_bridge,
"Anthropic fallback bridge serializes the same ChatFormat request":
"build_request(req, upstream, system, messages, true)" in anthropic_bridge,
}
for name, ok in checks.items():
print(f"{'PASS' if ok else 'FAIL'}: {name}")
# Minimal behavioral witness for a source-wins rename:
# the primary rename changes the key before parsing; a different fallback
# rename cannot see the original source key and therefore cannot undo it.
body = {"primary_name": 1}
primary_rename = {"original_name": "primary_name"}
fallback_rename = {"original_name": "fallback_name"}
assert "primary_name" in body
assert "original_name" not in body
assert "fallback_name" not in body
print("PASS: primary rename survives as a fallback-visible field")
PYRepository: api7/aisix
Length of output: 641
Build fallback continuations before ProviderKey overrides
mid_stream_ctx parses body after applying the primary ProviderKey’s request.* overrides. These fields survive in ChatFormat.extra and reach bridge.chat_stream for the fallback target. Keep a post-redaction, pre-override request for continuation construction, and add a test with different primary and fallback param_renames.
🤖 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/messages.rs` around lines 1213 - 1258, The mid-stream
continuation currently parses the post-override body, causing primary
ProviderKey request overrides to leak into fallback requests. Preserve a
post-redaction, pre-override request body for continuation construction, then
have the mid_stream_ctx block parse that body while retaining the existing model
restoration and translation behavior. Add coverage using different primary and
fallback param_renames to verify fallback continuations use the fallback
mapping.
Comment on lines
+1479
to
+1499
| // Logical stream outcome (AISIX-Cloud#1222) — same | ||
| // derivation as the typed pumps; the terminal-failure | ||
| // signal comes from the byte combinator via the shared | ||
| // slot (the parser below it only sees clean bytes). | ||
| let stream_failed = terminal_failure.is_some(); | ||
| let stream_outcome = if usage.guardrail_blocked || !usage.reached_end { | ||
| "" | ||
| } else if stream_failed { | ||
| "partial_failed" | ||
| } else if mid_stream_fallbacks > 0 { | ||
| "partial_recovered" | ||
| } else { | ||
| "success" | ||
| }; | ||
| if mid_stream_fallbacks > 0 { | ||
| // Recovered = the fallback kept the stream alive; | ||
| // only a terminal upstream error counts as failed. | ||
| state_c | ||
| .metrics | ||
| .record_mid_stream_fallback(&model_name_c, !stream_failed); | ||
| } |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
An ineligible in-band upstream error is reported as stream_outcome: success.
terminal_failure is set in only two places: handle_passthrough_failure on fallback exhaustion (line 4096) and the ineligible mid-continuation branch (line 4033). The verbatim-phase ineligible branch at lines 3939-3947 sets neither. It yields the upstream's own error frame and returns.
Trace the resulting state:
- The combinator's stream ends with
Okitems only. build_anthropic_passthrough_streamexits its loop and setsreached_end = trueat line 4264.terminal_failureisNoneandattempt_seqis0.- The derivation here therefore yields
"success", and the event records status 200.
A non-retryable upstream in-band error (for example an in-band 400) is then indistinguishable from a fully delivered stream. The typed cross-provider pump handles the same case correctly at lines 2743-2755 by setting stream_failed.
Record the failure in the shared slot on the verbatim ineligible-in-band path so the derivation reports partial_failed.
🐛 Proposed fix in the verbatim ineligible branch
PassthroughFailure::No(err, exhausted) => {
if exhausted {
// Fallbacks attempted and exhausted: terminate
// with the synthesized frame carrying the last
// failure, like the typed pumps.
yield Ok(anthropic_error_frame(&err));
} else if let Some(frame) = withheld {
// Ineligible in-band error: release the
// upstream's own error frame byte-for-byte.
+ // The wire ends on an error — record it so the
+ // telemetry closure reports `partial_failed`
+ // instead of `success`.
+ *plan
+ .shared
+ .terminal_failure
+ .lock()
+ .expect("terminal failure lock") = Some((
+ crate::attempt::routing_error_class(&err).to_string(),
+ crate::attempt::attempt_error_message(&err),
+ ));
yield Ok(Bytes::from(frame));
} else if let Some(e) = original {🤖 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/messages.rs` around lines 1479 - 1499, Update the
verbatim-phase ineligible in-band error branch, near the existing upstream
error-frame return, to record the failure in the shared terminal_failure slot
before returning. Reuse the same failure representation used by
handle_passthrough_failure and the ineligible mid-continuation branch, so the
downstream stream_outcome derivation reports partial_failed and preserves the
existing error-frame behavior.
Comment on lines
+2743
to
+2755
| // Terminal upstream failure: the wire ends with an | ||
| // in-band error frame — distinct from a client | ||
| // abandon, so `reached_end` is still recorded and | ||
| // the telemetry closure reports `stream_outcome: | ||
| // partial_failed` (AISIX-Cloud#1222). | ||
| { | ||
| let comp = guard.comp(); | ||
| comp.stream_failed = true; | ||
| comp.stream_error_class = | ||
| crate::attempt::routing_error_class(&e).to_string(); | ||
| comp.stream_error_message = crate::attempt::attempt_error_message(&e); | ||
| comp.reached_end = true; | ||
| } |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Check whether the sibling endpoints apply the same reached_end/stream_failed
# treatment in their mid-stream error arms, and find consumers of the 499 status on streams.
set -euo pipefail
rg -n -C 10 'stream_failed' crates/aisix-proxy/src/chat.rs crates/aisix-proxy/src/responses.rs
rg -n -C 6 'CLIENT_CLOSED_REQUEST' crates/aisix-proxy/srcRepository: api7/aisix
Length of output: 20009
🏁 Script executed:
#!/bin/bash
set -euo pipefail
for f in crates/aisix-proxy/src/messages.rs crates/aisix-proxy/src/chat.rs crates/aisix-proxy/src/responses.rs; do
echo "===== $f: stream error arms ====="
rg -n -C 18 'Err\((e|err)\) =>|stream_failed = true|reached_end = true' "$f" | head -n 240
done
echo "===== telemetry status selection ====="
rg -n -C 8 'if (comp|usage)\.reached_end|status_code|stream_outcome' \
crates/aisix-proxy/src/messages.rs \
crates/aisix-proxy/src/chat.rs \
crates/aisix-proxy/src/responses.rs | head -n 300Repository: api7/aisix
Length of output: 25021
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "===== messages.rs changed arm ====="
sed -n '2680,2775p' crates/aisix-proxy/src/messages.rs
echo "===== chat.rs stream pump ====="
sed -n '4740,4840p' crates/aisix-proxy/src/chat.rs
echo "===== responses.rs stream pump candidates ====="
rg -n -C 24 'error frame|error event|upstream.*error|stream.*error|Some\(Err|Err\(err\)|Err\(e\)' \
crates/aisix-proxy/src/responses.rs | tail -n 320
echo "===== sibling reached_end assignments ====="
rg -n -C 5 'reached_end\s*=' \
crates/aisix-proxy/src/messages.rs \
crates/aisix-proxy/src/chat.rs \
crates/aisix-proxy/src/responses.rsRepository: api7/aisix
Length of output: 31425
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "===== responses.rs streaming loop ====="
sed -n '2440,2670p' crates/aisix-proxy/src/responses.rs
echo "===== stream_failed declarations and assignments ====="
rg -n -C 5 'stream_failed|stream_error_class|stream_error_message' \
crates/aisix-proxy/src/messages.rs \
crates/aisix-proxy/src/chat.rs \
crates/aisix-proxy/src/responses.rs
echo "===== deterministic sibling error-arm check ====="
python3 - <<'PY'
from pathlib import Path
import re
files = [
Path("crates/aisix-proxy/src/messages.rs"),
Path("crates/aisix-proxy/src/chat.rs"),
Path("crates/aisix-proxy/src/responses.rs"),
]
for path in files:
text = path.read_text()
print(f"== {path} ==")
for match in re.finditer(r'\bErr\((?:e|err)\)\s*=>', text):
start = match.start()
block = text[start:start + 1800]
if any(token in block for token in ("error frame", "error_frame_payload", "event(\"error\")")):
print("error-arm:", text.count("\n", 0, start) + 1)
print(" stream_failed:", "stream_failed" in block)
print(" reached_end:", bool(re.search(r'\breached_end\s*=', block)))
PYRepository: api7/aisix
Length of output: 26527
Align mid-stream error telemetry across streaming endpoints.
messages.rs now records reached_end = true, which changes status_code from 499 to 200 and emits stream_outcome: partial_failed. chat.rs does not set reached_end in its equivalent error arm. responses.rs also has no equivalent terminal-error state and only sets reached_end after a later upstream EOF. Apply one contract to all three endpoints and add coverage for mid-stream errors.
🤖 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/messages.rs` around lines 2743 - 2755, Align terminal
mid-stream error handling across the equivalent error arms in messages.rs,
chat.rs, and responses.rs: record the same reached_end, stream_failed, error
class, and error message state so all endpoints produce identical status and
stream_outcome telemetry. Update the shared endpoint-specific logic rather than
changing unrelated EOF handling, and add coverage verifying mid-stream errors
emit the unified contract.
Comment on lines
+227
to
+231
| const primary = await anthropicUpstream({ | ||
| rawStreamFrames: [MESSAGE_START, BLOCK_START, delta("Once upon")], | ||
| eventDelayMs: 200, | ||
| disconnectAfterEvents: 3, | ||
| }); |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
The drop test never triggers the disconnect branch.
rawStreamFrames has 3 entries and disconnectAfterEvents is 3. In the harness loop the guard runs at the top of iteration i, so the highest index reached is 2 and 2 >= 3 is false. The loop ends normally and res.end() closes the stream cleanly. The upstream therefore produces an EOF without message_stop, not an RST.
The test still passes because wrap_anthropic_passthrough maps that EOF to BridgeError::StreamAborted, which is also a qualifying trigger. But the case named in the test title and in the file comment at lines 71-76 (a real mid-body connection drop) is not exercised. Add a trailing frame so index 3 reaches the destroy branch, the same way the stall test at lines 278-287 does.
🐛 Proposed fix
const primary = await anthropicUpstream({
- rawStreamFrames: [MESSAGE_START, BLOCK_START, delta("Once upon")],
+ rawStreamFrames: [
+ MESSAGE_START,
+ BLOCK_START,
+ delta("Once upon"),
+ delta(" (never sent)"),
+ ],
eventDelayMs: 200,
disconnectAfterEvents: 3,
});📝 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.
Suggested change
| const primary = await anthropicUpstream({ | |
| rawStreamFrames: [MESSAGE_START, BLOCK_START, delta("Once upon")], | |
| eventDelayMs: 200, | |
| disconnectAfterEvents: 3, | |
| }); | |
| const primary = await anthropicUpstream({ | |
| rawStreamFrames: [ | |
| MESSAGE_START, | |
| BLOCK_START, | |
| delta("Once upon"), | |
| delta(" (never sent)"), | |
| ], | |
| eventDelayMs: 200, | |
| disconnectAfterEvents: 3, | |
| }); |
🤖 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 `@tests/e2e/src/cases/mid-stream-fallback-messages-e2e.test.ts` around lines
227 - 231, Update the primary upstream setup in the mid-stream fallback test to
append a fourth raw stream frame after the existing three, ensuring the harness
reaches index 3 and triggers disconnectAfterEvents. Keep the existing delay and
disconnect configuration unchanged so the test exercises an actual mid-body
connection drop.
jarvis9443
marked this pull request as draft
August 5, 2026 07:52
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.
Extends mid-stream fallback (
routing.stream_failure, AISIX-Cloud#1222) from/v1/chat/completionsto/v1/messages, on both of its dispatch legs. Phase 2 of the feature: the endpoint most Anthropic-SDK/Claude-Code traffic actually uses can now recover a committed (already-200) stream on the remaining fallback targets instead of truncating.What changes
Cross-provider leg (non-Anthropic targets,
ChatChunkStream+AnthropicSseEncoder): the existingstream_failover::wrapcombinator is inserted between the first-chunk peek and the SSE pump. Because the encoder lives outside the combinator, its message envelope and content-block state survive the switch — the spliced fallback chunks continue the same wire message with no client-visible seam.Anthropic passthrough leg (byte-verbatim forwarding): a new byte-level combinator (
wrap_anthropic_passthrough). Armed, it forwards complete SSE frames verbatim while side-channel tracking the wire state (message.id, open text block, next block index, delivered partial). On a qualifying failure it dispatches the remaining targets through the chat bridges and re-encodes their chunks onto the client's committed envelope via a resumedAnthropicSseEncoder(AnthropicSseEncoder::resume) — block indices continue, no secondmessage_start. The trailing assistant partial becomes native prefill on Anthropic-wire fallback targets (split_systemalready maps it). Unarmed streams keep today's byte path bit-for-bit.Failure classes on the passthrough leg (previously undetectable or silent):
errorframes are withheld (never partially forwarded — forwarding is frame-granular when armed) and either recovered or released verbatim;message_stopclassifies astransport_error;read_timeout(the silent-truncate byte wrapper stays on unarmed streams).Safety gates (same set as chat): tool_use / thinking blocks on the wire, a
message_deltaalready delivered (stop reason on the client's wire), partial text past the 1 MiB accumulation cap, or a structured-output request — all disarm continuation; the failure then terminates exactly as before.Telemetry parity with chat:
inbound_protocol: anthropic/ sink labelmessages(emit_usage_eventgrew a stamped inner variant for this);stream_outcome(success/partial_failed/partial_recovered) on this endpoint, are attributed to the serving target (id, provider, attempt kindmid_stream_fallback, attempt-scoped latency), and fallback targets' TPM is billed post-stream;aisix_mid_stream_fallbacks_totalrecords recoveries for this endpoint too;stream_outcome: partial_failedinstead of being conflated with client aborts (499) — the 499 label stays reserved for actual client abandons.Continuation prompt: switched to LiteLLM's current instruction verbatim ("The previous assistant response was interrupted mid-stream. Continue exactly where it stopped — do not repeat any of its content. Your response must read as a seamless continuation."). Upstream recently removed its old chat-completions continuation prompt; the Responses-API continuation is the one surviving baseline, so all our endpoints now share its text. Sent as a system turn; on Anthropic-wire targets the mid-conversation instruction becomes a user turn and the partial stays native prefill.
LiteLLM baseline
LiteLLM has no mid-stream fallback for its
/v1/messagessurface (anthropic_messagesroutes through_ageneric_api_call_with_fallbacks, which never wraps the stream) — a stream failure there emits one error frame and dies. So there is no behavior to align with on this endpoint; the design transplants our chat-completions semantics (which the #1222 review already settled) onto the Anthropic wire. Where LiteLLM's Responses continuation does something comparable, we match its mechanics (continuation input shape, usage merging into the terminal frame).Tests
Rust integration (wiremock,
crates/aisix-proxy/src/lib.rs):message_start, singlecontent_block_start, no error frame, native-prefill continuation body asserted);stream_failure): upstream error frame forwards verbatim, fallback never contacted;anthropicstamp) + terminalpartial_recoveredattributed to the fallback target.TS e2e (
tests/e2e, realaisixbinary + real sockets):read_timeout-classified recovery (armed-combinator-owned timeout);The mock upstream harness gained
rawStreamFrames(verbatim SSE frames incl.event:lines) to express Anthropic-wire streams./v1/responsesfollows in a separate PR on top of this one.Ref api7/AISIX-Cloud#1222
🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Bug Fixes
Tests