From 64e2bf2ac935a31adba48c4c59e9acab915b9628 Mon Sep 17 00:00:00 2001 From: Jarvis Date: Tue, 4 Aug 2026 18:04:12 +0800 Subject: [PATCH 1/2] =?UTF-8?q?feat(routing):=20mid-stream=20fallback=20?= =?UTF-8?q?=E2=80=94=20resume=20a=20committed=20stream=20on=20fallback=20t?= =?UTF-8?q?argets?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Once a streaming response's first chunk is delivered the 200 is committed and the pre-stream retry/failover loop is out of reach: a mid-generation transport break, inter-chunk stall, decode failure, or provider in-band error could only terminate the stream (in-band error frame, no [DONE]). Long generations died with no recovery even when healthy fallback targets existed (AISIX-Cloud#1222). Add routing.stream_failure — explicit opt-in, default terminate: stream_failure: mode: continue # terminate (default) | continue on: [transport_error, read_timeout, upstream_decode_error, upstream_in_band_error] # default: all max_fallbacks: 1 With continue, a qualifying mid-stream error hands the SAME client stream to the remaining targets (strategy order, runtime state re-checked, cooldown/health recorded). The fallback request carries the original messages plus LiteLLM's verbatim continuation instruction and the partial text as an assistant message — native prefill on Anthropic-wire targets. Retryability reuses is_retryable (non-429 in-band 4xx never triggers; retry_on_429 and fallback_on_statuses apply). Safety boundaries (deliberate divergence from LiteLLM, which falls back with no output-shape guards): streams that emitted tool-call or reasoning deltas, structured-output requests (response_format json_object/json_schema), and partials past the 1MiB continuation cap keep the terminate behavior. Telemetry: the failed attempt emits its own per-attempt UsageEvent with estimated partial spend (prompt + delivered partial, #655/#1074 contracts); the terminal event is attributed to the SERVING target (attempt_kind mid_stream_fallback) with the terminal error class when the stream still dies; new request-level stream_outcome (success | partial_failed | partial_recovered) separates 'HTTP 200' from 'stream completed' — partial_failed also lands on plain terminate-mode failures. Client-facing usage frames fold the partial estimate in (LiteLLM's merge semantics). New counter aisix_mid_stream_fallbacks_total{model,outcome}. Client-cancel safety is structural: the failover combinator advances only when the client pulls, so an abandoned stream can never dispatch a ghost fallback request (#1094 family). First phase covers /v1/chat/completions; /v1/responses and /v1/messages keep terminate semantics pending protocol-specific continuation design. --- crates/aisix-admin/src/openapi.rs | 13 + crates/aisix-core/src/lib.rs | 5 +- crates/aisix-core/src/models/mod.rs | 5 +- crates/aisix-core/src/models/routing.rs | 86 +++ crates/aisix-obs/src/metrics.rs | 20 + crates/aisix-obs/src/usage.rs | 16 + crates/aisix-proxy/src/chat.rs | 238 +++++++- crates/aisix-proxy/src/lib.rs | 353 ++++++++++++ crates/aisix-proxy/src/routing.rs | 1 + crates/aisix-proxy/src/stream_failover.rs | 529 ++++++++++++++++++ schemas/resources/model.schema.json | 88 +++ schemas/resources/routing.schema.json | 100 ++++ .../src/cases/mid-stream-fallback-e2e.test.ts | 304 ++++++++++ 13 files changed, 1738 insertions(+), 20 deletions(-) create mode 100644 crates/aisix-proxy/src/stream_failover.rs create mode 100644 tests/e2e/src/cases/mid-stream-fallback-e2e.test.ts diff --git a/crates/aisix-admin/src/openapi.rs b/crates/aisix-admin/src/openapi.rs index 092379a4..f30c9c53 100644 --- a/crates/aisix-admin/src/openapi.rs +++ b/crates/aisix-admin/src/openapi.rs @@ -4343,6 +4343,19 @@ fn add_variant_titles(doc: &mut Value) { "/components/schemas/StreamDoneMarker/oneOf", &["Required", "Optional", "None"], ), + ( + "/components/schemas/StreamFailureMode/oneOf", + &["Terminate", "Continue"], + ), + ( + "/components/schemas/StreamFailureTrigger/oneOf", + &[ + "Transport error", + "Read timeout", + "Upstream decode error", + "Upstream in-band error", + ], + ), ]; for (pointer, titles) in variant_titles { diff --git a/crates/aisix-core/src/lib.rs b/crates/aisix-core/src/lib.rs index a379c6b8..fed08dad 100644 --- a/crates/aisix-core/src/lib.rs +++ b/crates/aisix-core/src/lib.rs @@ -49,8 +49,9 @@ pub use models::{ GuardrailMonitorHit, KeywordConfig, KeywordPattern, McpAuthType, McpRateLimit, McpServer, McpServerType, McpTransport, Model, ObservabilityExporter, ParamConstraints, PolicyScope, PolicyWindow, ProviderKey, RateLimit, RateLimitPolicy, RequestOverrides, ResponseOverrides, - Routing, RoutingStrategy, RoutingTarget, SchemaError, StreamDoneMarker, TelemetryKind, - TelemetryTags, WhenAllUnavailablePolicy, DEFAULT_COOLDOWN_TRIGGER_STATUSES, + Routing, RoutingStrategy, RoutingTarget, SchemaError, StreamDoneMarker, StreamFailure, + StreamFailureMode, StreamFailureTrigger, TelemetryKind, TelemetryTags, + WhenAllUnavailablePolicy, DEFAULT_COOLDOWN_TRIGGER_STATUSES, }; pub use resource::{Resource, ResourceEntry}; pub use snapshot::{ResourceTable, SnapshotHandle}; diff --git a/crates/aisix-core/src/models/mod.rs b/crates/aisix-core/src/models/mod.rs index c324ba5f..0a1ce6fe 100644 --- a/crates/aisix-core/src/models/mod.rs +++ b/crates/aisix-core/src/models/mod.rs @@ -63,7 +63,10 @@ pub use provider_key::{ }; pub use rate_limit::{McpRateLimit, RateLimit}; pub use rate_limit_policy::{PolicyScope, PolicyWindow, RateLimitPolicy}; -pub use routing::{Routing, RoutingStrategy, RoutingTarget, WhenAllUnavailablePolicy}; +pub use routing::{ + Routing, RoutingStrategy, RoutingTarget, StreamFailure, StreamFailureMode, + StreamFailureTrigger, WhenAllUnavailablePolicy, +}; pub use schema::{ validate_a2a_agent, validate_a2a_agent_lenient, validate_apikey, validate_apikey_lenient, validate_cache_policy, validate_cache_policy_lenient, validate_guardrail, diff --git a/crates/aisix-core/src/models/routing.rs b/crates/aisix-core/src/models/routing.rs index 2c9f579b..73103612 100644 --- a/crates/aisix-core/src/models/routing.rs +++ b/crates/aisix-core/src/models/routing.rs @@ -184,6 +184,90 @@ pub struct Routing { /// default). Ignored by non-`weighted` strategies. #[serde(default, skip_serializing_if = "Option::is_none")] pub sticky: Option, + /// What to do when a streaming response fails AFTER its first chunk was + /// already delivered to the client (the HTTP 200 is committed and cannot + /// be revised). Omitted keeps the historical behavior: terminate the + /// stream with an in-band error frame and no `[DONE]`. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub stream_failure: Option, +} + +/// Mid-stream failure policy for streaming responses (AISIX-Cloud#1222). +/// +/// Applies only to failures that occur after the response head (and +/// possibly some chunks) reached the client; failures before the first +/// chunk keep using the regular retry/failover loop. +#[derive(Debug, Clone, Serialize, Deserialize, schemars::JsonSchema)] +pub struct StreamFailure { + /// `terminate` (default) keeps the current behavior. `continue` lets + /// the router call the remaining fallback targets and resume the SAME + /// client stream with a best-effort continuation of the partial text. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub mode: Option, + /// Which mid-stream error classes trigger the fallback. Omitted = + /// all of them. Non-retryable errors (an in-band 4xx other than 429, + /// unless listed in `fallback_on_statuses`) never trigger regardless. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub on: Option>, + /// Max fallback targets tried for one mid-stream failure. Defaults + /// to 1 — mid-stream recovery burns client-visible latency per + /// attempt, so the default is deliberately tighter than the + /// pre-stream `max_fallbacks`. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub max_fallbacks: Option, +} + +impl StreamFailure { + pub fn mode_or_default(&self) -> StreamFailureMode { + self.mode.unwrap_or_default() + } + + pub fn max_fallbacks_or_default(&self) -> u32 { + self.max_fallbacks.unwrap_or(1) + } + + /// Configured trigger classes; all classes when unset. `continue` + /// is itself the explicit opt-in, so the default set is the full + /// one rather than a conservative subset. + pub fn on_or_default(&self) -> &[StreamFailureTrigger] { + const ALL: &[StreamFailureTrigger] = &[ + StreamFailureTrigger::TransportError, + StreamFailureTrigger::ReadTimeout, + StreamFailureTrigger::UpstreamDecodeError, + StreamFailureTrigger::UpstreamInBandError, + ]; + self.on.as_deref().unwrap_or(ALL) + } +} + +/// See [`StreamFailure::mode`]. +#[derive( + Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize, schemars::JsonSchema, +)] +#[serde(rename_all = "snake_case")] +pub enum StreamFailureMode { + /// Terminate the stream: in-band error frame, no `[DONE]` (the + /// historical behavior). + #[default] + Terminate, + /// Continue on a fallback target inside the same client stream. + Continue, +} + +/// Mid-stream error classes eligible for [`StreamFailureMode::Continue`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)] +#[serde(rename_all = "snake_case")] +pub enum StreamFailureTrigger { + /// The upstream connection broke mid-stream (reset, premature close). + TransportError, + /// The gap between chunks exceeded the effective `stream_timeout`. + ReadTimeout, + /// A frame failed to parse as a chunk (and was not a recognizable + /// in-band error envelope). + UpstreamDecodeError, + /// The provider reported an error inside the committed 200 stream + /// (an SSE error frame / event-stream modeled exception). + UpstreamInBandError, } impl Routing { @@ -271,6 +355,7 @@ mod tests { fallback_on_statuses: None, when_all_unavailable: None, sticky: None, + stream_failure: None, }; assert_eq!(r.max_fallbacks_or_default(), 0); } @@ -286,6 +371,7 @@ mod tests { fallback_on_statuses: None, when_all_unavailable: None, sticky: None, + stream_failure: None, }; assert_eq!(r.max_fallbacks_or_default(), 0); } diff --git a/crates/aisix-obs/src/metrics.rs b/crates/aisix-obs/src/metrics.rs index 2b9592e0..ee39ce95 100644 --- a/crates/aisix-obs/src/metrics.rs +++ b/crates/aisix-obs/src/metrics.rs @@ -79,6 +79,12 @@ pub const M_DEPLOYMENT_STATE: &str = "aisix_deployment_state"; pub const M_DEPLOYMENT_COOLED_DOWN_TOTAL: &str = "aisix_deployment_cooled_down_total"; pub const M_ROUTING_SUCCESSFUL_FALLBACKS_TOTAL: &str = "aisix_routing_successful_fallbacks_total"; pub const M_ROUTING_FAILED_FALLBACKS_TOTAL: &str = "aisix_routing_failed_fallbacks_total"; +/// Mid-stream fallback outcomes (`routing.stream_failure: continue`, +/// AISIX-Cloud#1222). `outcome` is `recovered` (the client stream +/// completed on a fallback target) or `failed` (every eligible fallback +/// target also failed and the stream terminated). Labelled by the +/// requested (routing) model. +pub const M_MID_STREAM_FALLBACKS_TOTAL: &str = "aisix_mid_stream_fallbacks_total"; pub const M_RATELIMIT_REMAINING_REQUESTS: &str = "aisix_ratelimit_remaining_requests"; pub const M_RATELIMIT_REMAINING_TOKENS: &str = "aisix_ratelimit_remaining_tokens"; pub const M_BUDGET_LIMIT_USD: &str = "aisix_budget_limit_usd"; @@ -1181,6 +1187,20 @@ impl Metrics { }); } + /// One mid-stream fallback episode resolved (AISIX-Cloud#1222): + /// `recovered` when the client stream completed on a fallback + /// target, `failed` when the fallback chain was exhausted. + pub fn record_mid_stream_fallback(&self, model: &str, recovered: bool) { + metrics::with_local_recorder(&self.inner.recorder, || { + metrics::counter!( + M_MID_STREAM_FALLBACKS_TOTAL, + "model" => model.to_string(), + "outcome" => if recovered { "recovered" } else { "failed" }, + ) + .increment(1); + }); + } + pub fn set_rate_limit_remaining( &self, api_key_id: &str, diff --git a/crates/aisix-obs/src/usage.rs b/crates/aisix-obs/src/usage.rs index 78200cc7..8089cb9f 100644 --- a/crates/aisix-obs/src/usage.rs +++ b/crates/aisix-obs/src/usage.rs @@ -195,6 +195,22 @@ pub struct UsageEvent { #[serde(default, skip_serializing_if = "String::is_empty")] pub finish_reason: String, + /// Logical outcome of a STREAMING response, set on the serving + /// attempt's event only (AISIX-Cloud#1222). The HTTP status is + /// committed at 200 before the stream runs, so it cannot express a + /// mid-stream failure; this field can: + /// - `success` — the stream completed normally; + /// - `partial_failed` — the stream terminated after the 200 with an + /// in-band error frame (`[DONE]` withheld); + /// - `partial_recovered` — a mid-stream failure was recovered by + /// `routing.stream_failure: continue` on a fallback target and the + /// stream then completed normally. + /// + /// Empty on non-streaming events, failed-attempt events, and + /// client-abandoned streams (those keep `status_code: 499`). + #[serde(default, skip_serializing_if = "String::is_empty")] + pub stream_outcome: String, + /// Cost the DP computed for this request in US dollars. Zero when /// the request never reached cost calculation (e.g. blocked by a /// guardrail before dispatch). cp-api recomputes this server-side diff --git a/crates/aisix-proxy/src/chat.rs b/crates/aisix-proxy/src/chat.rs index 1b2da878..bdd76817 100644 --- a/crates/aisix-proxy/src/chat.rs +++ b/crates/aisix-proxy/src/chat.rs @@ -251,6 +251,7 @@ pub async fn chat_completions( attempt_model: winner.map(|w| w.target_model.clone()).unwrap_or_default(), error_class: String::new(), error_message: String::new(), + stream_outcome: String::new(), applied_guardrails: applied_guardrails.clone(), provider_key_id: success.provider_key_id.clone(), redacted_entity_counts: redaction_counts.clone(), @@ -530,6 +531,7 @@ pub async fn chat_completions( .unwrap_or_default(), error_class: String::new(), error_message: String::new(), + stream_outcome: String::new(), // The chain governed the request even though it // ultimately blocked on the output filter. applied_guardrails: applied_guardrails.clone(), @@ -1226,6 +1228,9 @@ async fn dispatch( upstream: aisix_gateway::ChatChunkStream, idx: u32, kind: &'static str, + /// Position of the winning target in `attempt_models` — the + /// mid-stream failover plan takes the targets after it. + target_idx: usize, /// When this attempt began. The end-of-stream UsageEvent /// reports `attempt_started.elapsed()` so `latency_ms` covers /// this attempt alone, matching the failed-attempt events and @@ -1416,6 +1421,7 @@ async fn dispatch( upstream, idx, kind, + target_idx, attempt_started, }); won_member_reservation = member_reservation; @@ -1494,6 +1500,7 @@ async fn dispatch( upstream, idx: winner_idx, kind: winner_kind, + target_idx: winner_target_idx, attempt_started: winner_attempt_started, } = won; let model = &model; @@ -1541,26 +1548,15 @@ async fn dispatch( let user_id_for_metrics = auth.key().user_id.clone(); let provider_for_metrics = provider.to_ascii_lowercase(); let model_for_metrics = req.model.clone(); - let provider_key_id_for_metrics = pk_id.clone(); - // #890 req-3/req-4: readable provider-key name + normalised inbound - // client type, captured for the streaming on_complete metric emission - // (mirrors the non-streaming `record_success` path). - let provider_key_name_for_metrics = { - let snap = state.snapshot.load(); - crate::usage_attr::provider_key_metric_name(&snap, &pk_id) - }; + // #890 req-3/req-4: the readable provider-key name is resolved + // inside the on_complete closure from the SERVING attempt's key + // (a mid-stream fallback may have changed it); the normalised + // inbound client type is request-scoped and captured here. let user_name_for_metrics = auth.key().user_name.clone(); let client_type_for_metrics = state .client_classifier .classify(&client.user_agent) .to_string(); - // Captured for the stream-end telemetry closure so - // emit_usage_event can look up `telemetry_tags` for per-PK - // attribution (#302 M17 / AISIX-Cloud#436). The metrics - // variant above is `&str`-scoped to inner scopes that consume - // it as a borrow; the telem variant is owned for the move - // into the on_complete closure. - let provider_key_id_for_telem = pk_id.clone(); let upstream_model_for_metrics = model.upstream_model().unwrap_or("unknown").to_string(); let bypass_reason_for_telem = bypass_reason.clone().unwrap_or_default(); // Applied guardrail set (#379), owned for the move into on_complete so @@ -1625,6 +1621,64 @@ async fn dispatch( &upstream_model_for_metrics, crate::token_estimate::PromptInput::Chat(Box::new(req.clone())), ); + // Mid-stream failover (AISIX-Cloud#1222): the serving-attempt + // handle starts as the pre-stream winner and is rewritten by the + // combinator on every switch, so the completion closure below + // attributes the terminal event to whichever target actually + // finished the stream. + let serving = Arc::new(std::sync::Mutex::new( + crate::stream_failover::ServingAttempt { + target_id: model_id_for_telem.clone(), + target_model: attempt_model_for_telem.clone(), + provider: provider_for_metrics.clone(), + provider_key_id: pk_id.clone(), + upstream_model: upstream_model_for_metrics.clone(), + cooldown: model.cooldown.clone(), + attempt_index: winner_idx, + attempt_kind: winner_kind, + attempt_started: winner_attempt_started, + }, + )); + let mid_stream_fallbacks = Arc::new(AtomicU32::new(0)); + let mid_stream_extra = Arc::new(std::sync::Mutex::new( + aisix_gateway::chat::UsageStats::default(), + )); + let stream_failure_cfg = virtual_entry + .value + .routing + .as_ref() + .and_then(|r| r.stream_failure.clone()) + .filter(|cfg| { + cfg.mode_or_default() == aisix_core::StreamFailureMode::Continue + && is_routing_request + }); + let upstream = match stream_failure_cfg { + Some(cfg) if winner_target_idx + 1 < attempt_models.len() => { + crate::stream_failover::wrap( + upstream, + crate::stream_failover::MidStreamPlan { + cfg, + remaining: attempt_models[winner_target_idx + 1..].to_vec(), + state: state.clone(), + group: virtual_entry.value.clone(), + req: req.clone(), + request_id: request_id.to_string(), + client: client.clone(), + retry_on_429, + fallback_on_statuses: fallback_statuses.to_vec(), + requested_model: req.model.clone(), + api_key_id: auth.entry.id.clone(), + applied_guardrails: applied_guardrails.clone(), + serving: Arc::clone(&serving), + fallbacks_attempted: Arc::clone(&mid_stream_fallbacks), + extra_usage: Arc::clone(&mid_stream_extra), + }, + ) + } + _ => upstream, + }; + let serving_for_telem = Arc::clone(&serving); + let mid_stream_fallbacks_for_telem = Arc::clone(&mid_stream_fallbacks); let sse_stream = build_sse_stream( upstream, now, @@ -1636,8 +1690,60 @@ async fn dispatch( client_requested_usage, // Single upstream: nothing pre-incurred, so no usage to fold in. aisix_gateway::chat::UsageStats::default(), + Some(mid_stream_extra), Some(estimator), move |comp: StreamCompletion| { + // Mid-stream failover may have moved the stream onto a + // fallback target — read the SERVING attempt (not the + // pre-stream winner) for everything target-scoped. + let ( + model_id_for_telem, + provider_for_metrics, + provider_key_id_for_telem, + upstream_model_for_metrics, + winner_idx, + winner_kind, + attempt_model_for_telem, + winner_attempt_started, + ) = { + let s = serving_for_telem.lock().expect("serving lock"); + ( + s.target_id.clone(), + s.provider.clone(), + s.provider_key_id.clone(), + s.upstream_model.clone(), + s.attempt_index, + s.attempt_kind, + s.target_model.clone(), + s.attempt_started, + ) + }; + let provider_key_id_for_metrics = provider_key_id_for_telem.clone(); + let provider_key_name_for_metrics = { + let snap = state_for_telem.snapshot.load(); + crate::usage_attr::provider_key_metric_name(&snap, &provider_key_id_for_metrics) + }; + let mid_stream_fallbacks = mid_stream_fallbacks_for_telem.load(Ordering::Relaxed); + // Logical stream outcome (AISIX-Cloud#1222): the HTTP + // status froze at 200 when the head committed, so this is + // the only signal that separates "delivered in full" from + // "terminated mid-stream". Guardrail blocks and client + // aborts keep their own dedicated signals. + let stream_outcome = if comp.guardrail_blocked || !comp.reached_end { + "" + } else if comp.stream_failed { + "partial_failed" + } else if mid_stream_fallbacks > 0 { + "partial_recovered" + } else { + "success" + }; + if mid_stream_fallbacks > 0 { + metrics_for_stream.record_mid_stream_fallback( + &model_for_metrics, + stream_outcome == "partial_recovered", + ); + } // Rate-limit accounting (TPM cap) for all layers. for key in &post_stream_keys { limiter.add_tokens_post_stream(key, comp.total_tokens); @@ -1705,8 +1811,20 @@ async fn dispatch( attempt_index: winner_idx, attempt_kind: winner_kind.to_string(), attempt_model: attempt_model_for_telem.clone(), - error_class: String::new(), - error_message: String::new(), + // A stream that terminated mid-flight carries the + // terminal error on its serving attempt — the HTTP + // status can no longer say so (AISIX-Cloud#1222). + error_class: if comp.stream_failed { + comp.stream_error_class.clone() + } else { + String::new() + }, + error_message: if comp.stream_failed { + comp.stream_error_message.clone() + } else { + String::new() + }, + stream_outcome: stream_outcome.to_string(), applied_guardrails: applied_guardrails_for_telem.clone(), provider_key_id: provider_key_id_for_telem.clone(), redacted_entity_counts: { @@ -3055,6 +3173,8 @@ async fn dispatch_ensemble( content_cap, client_requested_usage, panel_usage_sum, + // Ensembles don't participate in mid-stream failover. + None, Some(judge_estimator), move |comp: StreamCompletion| { // Rate-limit accounting: the panel tokens (already round-tripped) @@ -3706,6 +3826,7 @@ fn emit_usage_event( attempt_model: extras.attempt_model, error_class: extras.error_class, error_message: extras.error_message, + stream_outcome: extras.stream_outcome, // Per-PK telemetry attribution (#302 M17 / AISIX-Cloud#436). // Source struct is `aisix_core::TelemetryTags`; the wire // shape is flat strings + a bool, with skip_serializing_if @@ -3828,6 +3949,10 @@ struct UsageExtras { error_class: String, /// Short error message for a failed attempt; empty on success. error_message: String, + /// Logical streaming outcome (`success` / `partial_failed` / + /// `partial_recovered`); empty on non-streaming events + /// (AISIX-Cloud#1222). + stream_outcome: String, /// The `{kind, hook}` set of guardrails that governed this request, /// captured at chain-resolve time. Lands on /// `dpmgr_usage_events.applied_guardrails` so the dashboard can show @@ -3917,6 +4042,55 @@ fn emit_failed_attempts( } } +/// Per-attempt UsageEvent for a serving attempt that failed +/// MID-STREAM and is being handed off to a fallback target +/// (AISIX-Cloud#1222). Unlike [`emit_failed_attempts`] (whose records +/// are read back at handler return), this fires from inside the +/// stream-failover combinator at switch time — the routing telemetry +/// was already finalized when the 200 committed. The partial spend is +/// estimated (prompt from the original request, completion from the +/// delivered partial text), matching the "prompts always billed + +/// generated partial billed" contract. +#[allow(clippy::too_many_arguments)] +pub(crate) fn emit_mid_stream_failed_attempt( + state: &ProxyState, + request_id: &str, + requested_model: &str, + api_key_id: &str, + client: &ClientContext, + applied_guardrails: &[AppliedGuardrail], + rec: &AttemptRecord, + prompt_tokens: u32, + completion_tokens: u32, +) { + emit_usage_event( + state, + request_id, + &rec.target_model_id, + requested_model, + api_key_id, + rec.status, + Duration::from_millis(u64::from(rec.latency_ms)), + prompt_tokens, + completion_tokens, + UsageExtras { + usage_estimated: prompt_tokens > 0 || completion_tokens > 0, + attempt_index: rec.index, + attempt_kind: rec.kind.to_string(), + attempt_model: rec.target_model.clone(), + error_class: rec.error_class.clone(), + error_message: rec.error_message.clone(), + applied_guardrails: applied_guardrails.to_vec(), + provider_key_id: rec.provider_key_id.clone(), + ..UsageExtras::default() + }, + /* cost_usd */ 0.0, + /* guardrail_blocked */ false, + client, + None, + ); +} + fn record_error(metrics: &Metrics, err: &ProxyError, model: &str, status: u16, elapsed: Duration) { let outcome = RequestOutcome::from_status(status); // Provider is unknown for pre-dispatch errors (auth, 404, etc.). @@ -4102,6 +4276,16 @@ struct StreamCompletion { /// still be abandoned midway, and a zero-chunk stream can still /// legitimately reach its end (an immediate error frame). reached_end: bool, + /// `true` when the upstream stream terminated with an error frame + /// (post-fallback-exhaustion if mid-stream failover was armed). The + /// HTTP status is already committed at 200, so this is what the + /// telemetry closure turns into `stream_outcome: partial_failed` + /// (AISIX-Cloud#1222). + stream_failed: bool, + /// Attempt-taxonomy class/message of the terminal stream error; + /// empty unless `stream_failed`. + stream_error_class: String, + stream_error_message: String, } /// Parameters needed to run output-guardrail evaluation at @@ -4254,6 +4438,12 @@ fn build_sse_stream( // `on_complete` (`comp`) counts stay stream-only. Zero for single-upstream // callers, where the fold is a no-op. base_usage: aisix_gateway::chat::UsageStats, + // Estimated spend of mid-stream-failed partial attempts, written by the + // failover combinator at switch time and folded into the client-facing + // usage frames alongside `base_usage` (LiteLLM merges partial + fallback + // usage the same way; AISIX-Cloud#1222). Shared because the value is + // only known mid-stream; `None` on paths without mid-stream failover. + mid_stream_extra: Option>>, // Token-estimation fallback context (AISIX-Cloud#1074); see // `CompleteOnDrop::estimator`. estimator: Option, @@ -4514,11 +4704,25 @@ where // from `comp + base_usage`). if let Some(u) = chunk.usage.as_mut() { *u = u.saturating_add(&base_usage); + // Mid-stream failover: fold the failed partial + // attempts' estimated spend into the client-facing + // usage frame, AFTER `comp` captured the + // serving-attempt-only counts (AISIX-Cloud#1222). + if let Some(extra) = mid_stream_extra.as_ref() { + let extra = extra.lock().expect("mid-stream extra lock").clone(); + *u = u.saturating_add(&extra); + } } Some(chunk) } Err(err) => { errored = true; + { + let comp = guard.comp(); + comp.stream_failed = true; + comp.stream_error_class = routing_error_class(&err).to_string(); + comp.stream_error_message = attempt_error_message(&err); + } let etype = err.error_type(); yield Ok::<_, Infallible>( Event::default() diff --git a/crates/aisix-proxy/src/lib.rs b/crates/aisix-proxy/src/lib.rs index ddfe83a7..03cb6be1 100644 --- a/crates/aisix-proxy/src/lib.rs +++ b/crates/aisix-proxy/src/lib.rs @@ -64,6 +64,7 @@ mod routing; mod semantic; pub mod sse_keepalive; mod state; +mod stream_failover; mod stream_timeout; mod token_estimate; mod usage_attr; @@ -4635,6 +4636,358 @@ data: [DONE]\n\n"; ResourceEntry::new(format!("router-{name}"), model, 1) } + /// Like [`routing_entry`] but with a `stream_failure` block — + /// the AISIX-Cloud#1222 mid-stream failover knob. + fn routing_entry_with_stream_failure( + name: &str, + targets: &[&str], + stream_failure: serde_json::Value, + ) -> ResourceEntry { + let target_objs: Vec = targets + .iter() + .map(|t| serde_json::json!({"model": t})) + .collect(); + let cfg = serde_json::json!({ + "display_name": name, + "routing": { + "strategy": "failover", + "targets": target_objs, + "stream_failure": stream_failure, + } + }); + let model: Model = serde_json::from_value(cfg).unwrap(); + ResourceEntry::new(format!("router-{name}"), model, 1) + } + + /// SSE body: role preamble + one content delta + an in-band error + /// frame — a provider failing inside its committed 200 stream. + const MID_STREAM_FAILING_SSE: &str = "\ +data: {\"id\":\"up-1\",\"model\":\"gpt-4o\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\"},\"finish_reason\":null}]}\n\n\ +data: {\"id\":\"up-1\",\"model\":\"gpt-4o\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"Once upon\"},\"finish_reason\":null}]}\n\n\ +data: {\"error\":{\"message\":\"The server had an error\",\"type\":\"server_error\"}}\n\n"; + + fn mid_stream_snapshot( + primary_uri: &str, + secondary_uri: &str, + stream_failure: Option, + ) -> AisixSnapshot { + let snap = AisixSnapshot::new(); + snap.provider_keys + .insert(pk_entry_with_id("pk-primary", primary_uri)); + snap.provider_keys + .insert(pk_entry_with_id("pk-secondary", secondary_uri)); + snap.models + .insert(model_entry_with_id("m-primary", "primary", "pk-primary")); + snap.models.insert(model_entry_with_id( + "m-secondary", + "secondary", + "pk-secondary", + )); + match stream_failure { + Some(sf) => snap.models.insert(routing_entry_with_stream_failure( + "smart", + &["primary", "secondary"], + sf, + )), + None => snap.models.insert(routing_entry( + "smart", + "failover", + &["primary", "secondary"], + None, + None, + None, + )), + } + snap.apikeys.insert(apikey_entry("sk-caller", &["smart"])); + snap + } + + async fn streaming_chat_wire(app: axum::Router, model: &str) -> String { + let body = serde_json::json!({ + "model": model, + "messages": [{"role": "user", "content": "tell me a story"}], + "stream": true + }); + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("authorization", "Bearer sk-caller") + .header("content-type", "application/json") + .body(Body::from(body.to_string())) + .unwrap(); + let resp = run(app, req).await; + assert_eq!(resp.status(), StatusCode::OK); + let mut body_stream = resp.into_body().into_data_stream(); + let mut wire = Vec::new(); + while let Some(chunk) = body_stream.next().await { + wire.extend_from_slice(chunk.unwrap().as_ref()); + } + String::from_utf8(wire).expect("SSE bytes are utf8") + } + + /// AISIX-Cloud#1222 core acceptance: with `stream_failure: continue`, + /// a provider error inside the committed 200 stream moves the SAME + /// client stream onto the fallback target; the fallback gets the + /// original messages + the continuation instruction + the partial + /// text as an assistant message; the client sees primary content, + /// then fallback content, then exactly one `[DONE]` and no error + /// frame. + #[tokio::test] + async fn mid_stream_failure_continues_on_fallback_target_in_same_stream() { + let primary = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(MID_STREAM_FAILING_SSE), + ) + .mount(&primary) + .await; + let secondary = MockServer::start().await; + let recovery_sse = "\ +data: {\"id\":\"up-2\",\"model\":\"gpt-4o\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\"},\"finish_reason\":null}]}\n\n\ +data: {\"id\":\"up-2\",\"model\":\"gpt-4o\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\" a time\"},\"finish_reason\":\"stop\"}]}\n\n\ +data: [DONE]\n\n"; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(recovery_sse), + ) + .mount(&secondary) + .await; + + let hub = Arc::new(Hub::new()); + hub.register_specialized("openai", Arc::new(openai_test_bridge())); + let snap = mid_stream_snapshot( + &primary.uri(), + &secondary.uri(), + Some(serde_json::json!({"mode": "continue"})), + ); + let app = build_router(build_state(snap, hub)); + + let wire = streaming_chat_wire(app, "smart").await; + assert!( + wire.contains("Once upon"), + "primary partial must reach the client:\n{wire}" + ); + assert!( + wire.contains(" a time"), + "fallback continuation must reach the client:\n{wire}" + ); + assert!( + !wire.contains("event: error"), + "recovered stream must not carry an error frame:\n{wire}" + ); + assert_eq!( + wire.matches("data: [DONE]").count(), + 1, + "exactly one [DONE] on recovery:\n{wire}" + ); + + // The fallback target received the continuation request: the + // original user message, then the continuation instruction, + // then the partial as an assistant message (LiteLLM shape). + let reqs = secondary.received_requests().await.unwrap(); + assert_eq!(reqs.len(), 1, "secondary called exactly once"); + let body: serde_json::Value = serde_json::from_slice(&reqs[0].body).unwrap(); + let messages = body["messages"].as_array().unwrap(); + assert_eq!(messages.len(), 3, "user + continuation system + partial"); + assert_eq!(messages[0]["role"], "user"); + assert_eq!(messages[1]["role"], "system"); + assert!(messages[1]["content"] + .as_str() + .unwrap() + .contains("Do not repeat the same content")); + assert_eq!(messages[2]["role"], "assistant"); + assert_eq!(messages[2]["content"], "Once upon"); + } + + /// Default config (no `stream_failure`) keeps the historical + /// terminate behavior: in-band error frame, no `[DONE]`, and the + /// fallback target is never contacted. + #[tokio::test] + async fn mid_stream_failure_without_config_terminates_and_never_calls_fallback() { + let primary = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(MID_STREAM_FAILING_SSE), + ) + .mount(&primary) + .await; + let secondary = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with(ResponseTemplate::new(200)) + .mount(&secondary) + .await; + + let hub = Arc::new(Hub::new()); + hub.register_specialized("openai", Arc::new(openai_test_bridge())); + let snap = mid_stream_snapshot(&primary.uri(), &secondary.uri(), None); + let app = build_router(build_state(snap, hub)); + + let wire = streaming_chat_wire(app, "smart").await; + assert!(wire.contains("Once upon")); + assert!( + wire.contains("event: error"), + "terminate mode keeps the in-band error frame:\n{wire}" + ); + assert!( + wire.contains("upstream_in_band_error"), + "error frame carries the in-band error type:\n{wire}" + ); + assert!( + !wire.contains("data: [DONE]"), + "no [DONE] after abnormal termination:\n{wire}" + ); + let reqs = secondary.received_requests().await.unwrap(); + assert!(reqs.is_empty(), "fallback target must not be contacted"); + } + + /// Fallback exhaustion: the fallback target fails mid-stream too + /// (and `max_fallbacks` defaults to 1) — the client gets the + /// in-band error frame and no `[DONE]`, never a fabricated clean + /// completion. + #[tokio::test] + async fn mid_stream_fallback_exhaustion_surfaces_error_without_done() { + let primary = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(MID_STREAM_FAILING_SSE), + ) + .mount(&primary) + .await; + let secondary = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(MID_STREAM_FAILING_SSE), + ) + .mount(&secondary) + .await; + + let hub = Arc::new(Hub::new()); + hub.register_specialized("openai", Arc::new(openai_test_bridge())); + let snap = mid_stream_snapshot( + &primary.uri(), + &secondary.uri(), + Some(serde_json::json!({"mode": "continue"})), + ); + let app = build_router(build_state(snap, hub)); + + let wire = streaming_chat_wire(app, "smart").await; + assert!( + wire.contains("event: error"), + "exhausted fallback surfaces the error:\n{wire}" + ); + assert!( + !wire.contains("data: [DONE]"), + "no fabricated [DONE] when the fallback also fails:\n{wire}" + ); + let reqs = secondary.received_requests().await.unwrap(); + assert_eq!(reqs.len(), 1, "fallback was attempted once"); + } + + /// Safety boundary: a stream that already emitted tool-call deltas + /// must terminate even with `continue` configured — a fallback + /// model cannot safely continue half-emitted tool-call arguments + /// (the LiteLLM gap this design deliberately closes). + #[tokio::test] + async fn mid_stream_fallback_skipped_after_tool_call_delta() { + let primary = MockServer::start().await; + let tool_call_sse = "\ +data: {\"id\":\"up-1\",\"model\":\"gpt-4o\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\"},\"finish_reason\":null}]}\n\n\ +data: {\"id\":\"up-1\",\"model\":\"gpt-4o\",\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_1\",\"type\":\"function\",\"function\":{\"name\":\"get_weather\",\"arguments\":\"{\\\"ci\"}}]},\"finish_reason\":null}]}\n\n\ +data: {\"error\":{\"message\":\"The server had an error\",\"type\":\"server_error\"}}\n\n"; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(tool_call_sse), + ) + .mount(&primary) + .await; + let secondary = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with(ResponseTemplate::new(200)) + .mount(&secondary) + .await; + + let hub = Arc::new(Hub::new()); + hub.register_specialized("openai", Arc::new(openai_test_bridge())); + let snap = mid_stream_snapshot( + &primary.uri(), + &secondary.uri(), + Some(serde_json::json!({"mode": "continue"})), + ); + let app = build_router(build_state(snap, hub)); + + let wire = streaming_chat_wire(app, "smart").await; + assert!( + wire.contains("event: error"), + "tool-call streams terminate, not continue:\n{wire}" + ); + assert!(!wire.contains("data: [DONE]")); + let reqs = secondary.received_requests().await.unwrap(); + assert!( + reqs.is_empty(), + "no fallback dispatch after a tool-call delta" + ); + } + + /// `on` narrows the trigger set: a config listing only + /// `read_timeout` must NOT fall back on an in-band error. + #[tokio::test] + async fn mid_stream_fallback_respects_on_trigger_list() { + let primary = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(MID_STREAM_FAILING_SSE), + ) + .mount(&primary) + .await; + let secondary = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with(ResponseTemplate::new(200)) + .mount(&secondary) + .await; + + let hub = Arc::new(Hub::new()); + hub.register_specialized("openai", Arc::new(openai_test_bridge())); + let snap = mid_stream_snapshot( + &primary.uri(), + &secondary.uri(), + Some(serde_json::json!({"mode": "continue", "on": ["read_timeout"]})), + ); + let app = build_router(build_state(snap, hub)); + + let wire = streaming_chat_wire(app, "smart").await; + assert!(wire.contains("event: error")); + assert!(!wire.contains("data: [DONE]")); + let reqs = secondary.received_requests().await.unwrap(); + assert!( + reqs.is_empty(), + "in-band error is not in the configured trigger set" + ); + } + #[tokio::test] async fn routing_failover_retries_to_second_target_when_first_5xxs() { let bad_upstream = MockServer::start().await; diff --git a/crates/aisix-proxy/src/routing.rs b/crates/aisix-proxy/src/routing.rs index 550328b9..b8879f1e 100644 --- a/crates/aisix-proxy/src/routing.rs +++ b/crates/aisix-proxy/src/routing.rs @@ -898,6 +898,7 @@ mod tests { fallback_on_statuses: None, when_all_unavailable: None, sticky: None, + stream_failure: None, } } diff --git a/crates/aisix-proxy/src/stream_failover.rs b/crates/aisix-proxy/src/stream_failover.rs new file mode 100644 index 00000000..f615c693 --- /dev/null +++ b/crates/aisix-proxy/src/stream_failover.rs @@ -0,0 +1,529 @@ +//! Mid-stream failover for `/v1/chat/completions` streaming +//! (AISIX-Cloud#1222, `routing.stream_failure: continue`). +//! +//! Once the first chunk of a streaming response has been committed the +//! HTTP 200 can no longer be revised, so the pre-stream retry/failover +//! loop is out of reach. This module wraps the winning upstream +//! [`ChatChunkStream`] in a combinator that, when a qualifying error +//! arrives mid-stream, dispatches the remaining fallback targets and +//! splices their chunks into the SAME client stream — asking the +//! fallback model to continue the already-delivered partial text +//! (LiteLLM's mid-stream fallback semantics: original messages + a +//! continuation system instruction + an assistant message carrying the +//! partial; Anthropic-wire targets consume the trailing assistant +//! message as native prefill). +//! +//! Client-cancel safety is structural: the combinator only makes +//! progress when the pump polls it, and the pump only advances when +//! the client connection pulls — a disconnected client stops the +//! generator at its suspension point, so no fallback dispatch can +//! fire for an abandoned stream. + +use std::sync::atomic::{AtomicU32, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::Instant; + +use aisix_core::{Model, StreamFailure, StreamFailureTrigger}; +use aisix_gateway::{BridgeError, ChatFormat, ChatMessage}; +use futures::StreamExt; + +use crate::attempt::{attempt_error_message, routing_error_class, AttemptRecord}; +use crate::client_ip::ClientContext; +use crate::routing::AttemptModel; +use crate::ProxyState; + +/// Verbatim LiteLLM continuation instruction (`litellm/router.py`, +/// `_acompletion_streaming_iterator`) — kept byte-identical so the two +/// gateways' fallback models receive the same steering. The partial +/// text is NOT interpolated here; it rides the assistant message that +/// follows. +const CONTINUATION_SYSTEM_PROMPT: &str = "You are a helpful assistant. You are given a message and you need to respond to it. You are also given a generated content. You need to respond to the message in continuation of the generated content. Do not repeat the same content. Your response should be in continuation of this text: "; + +/// The serving attempt behind the live client stream. Starts as the +/// pre-stream winner; rewritten by the combinator on every mid-stream +/// switch. The pump's `on_complete` closure reads it at stream end so +/// the terminal UsageEvent attributes tokens/latency to the target +/// that actually finished the response. +pub(crate) struct ServingAttempt { + pub target_id: String, + /// Routing-target display name for the event's `attempt_model` + /// (empty for direct models, same convention as the dispatch loop). + pub target_model: String, + pub provider: String, + pub provider_key_id: String, + pub upstream_model: String, + /// The serving target's cooldown config, carried here so a + /// mid-stream failure can run the cooldown decision without a + /// snapshot lookup. + pub cooldown: Option, + pub attempt_index: u32, + pub attempt_kind: &'static str, + pub attempt_started: Instant, +} + +/// Everything the combinator needs to dispatch fallback targets and +/// keep the request's telemetry coherent while doing so. +pub(crate) struct MidStreamPlan { + pub cfg: StreamFailure, + /// Targets after the pre-stream winner, in strategy order. + pub remaining: Vec, + pub state: ProxyState, + /// The routing (group) model — resolves group-level timeout + /// defaults for each fallback target. + pub group: Model, + /// The original client request (pre-continuation). + pub req: ChatFormat, + pub request_id: String, + pub client: ClientContext, + pub retry_on_429: bool, + pub fallback_on_statuses: Vec, + /// Client-facing model name (`req.model`) for the failed-attempt + /// events. + pub requested_model: String, + pub api_key_id: String, + pub applied_guardrails: Vec, + /// Shared with the pump's completion closure. + pub serving: Arc>, + /// Count of mid-stream fallback targets attempted; the completion + /// closure derives `partial_recovered` and the fallback metric + /// from it. + pub fallbacks_attempted: Arc, + /// Estimated usage of failed partial attempts, folded into the + /// final stream's client-facing usage frames by the pump (LiteLLM + /// merges partial + fallback usage the same way). + pub extra_usage: Arc>, +} + +/// Classify a mid-stream [`BridgeError`] into the configurable trigger +/// taxonomy. `UpstreamStatus` cannot occur after the 200 is committed; +/// config/credential errors are pre-dispatch by construction. Both map +/// to `None` (never fall back) defensively. +pub(crate) fn classify_trigger(err: &BridgeError) -> Option { + match err { + BridgeError::Transport(_) | BridgeError::StreamAborted => { + Some(StreamFailureTrigger::TransportError) + } + BridgeError::Timeout { .. } => Some(StreamFailureTrigger::ReadTimeout), + BridgeError::UpstreamDecode(_) => Some(StreamFailureTrigger::UpstreamDecodeError), + BridgeError::UpstreamInBand { .. } => Some(StreamFailureTrigger::UpstreamInBandError), + BridgeError::UpstreamStatus { .. } + | BridgeError::Config(_) + | BridgeError::InvalidUpstreamConfig(_) + | BridgeError::InvalidUpstreamCredentials(_) => None, + } +} + +/// Whether the request pins the output to a structured shape +/// (`response_format: json_object` / `json_schema`). A fallback model +/// cannot safely continue a half-emitted JSON document, so these +/// requests keep the terminate behavior regardless of config. +pub(crate) fn expects_structured_output(req: &ChatFormat) -> bool { + req.extra + .get("response_format") + .and_then(|rf| rf.get("type")) + .and_then(|t| t.as_str()) + .is_some_and(|t| t == "json_object" || t == "json_schema") +} + +/// Build the continuation request for a fallback target: the original +/// messages, then the continuation instruction, then an assistant +/// message carrying the partial text. An empty partial (the failure +/// beat the first content delta) retries with the untouched messages — +/// LiteLLM's `is_pre_first_chunk` branch: a continuation prompt there +/// would only waste tokens and confuse the model. +pub(crate) fn continuation_request(orig: &ChatFormat, partial: &str) -> ChatFormat { + let mut req = orig.clone(); + if !partial.is_empty() { + req.messages + .push(ChatMessage::system(CONTINUATION_SYSTEM_PROMPT)); + req.messages.push(ChatMessage::assistant(partial)); + } + req +} + +/// Wrap the winning upstream stream with the mid-stream failover +/// combinator. The caller has already checked `mode: continue` and +/// that `plan.remaining` is non-empty. +pub(crate) fn wrap( + upstream: aisix_gateway::ChatChunkStream, + plan: MidStreamPlan, +) -> aisix_gateway::ChatChunkStream { + Box::pin(async_stream::stream! { + let mut current = upstream; + // Generated content accumulated across every attempt — the + // continuation baseline. Capped at the same bound as the + // pump's estimation buffer; past it a faithful continuation + // prompt can no longer be built, so fallback disarms. + let mut partial = String::new(); + let mut partial_overflow = false; + // Output shapes a fallback model cannot safely continue: + // half-emitted tool calls, provider-signed reasoning streams, + // structured output. Sticky once observed. + let mut unsafe_output = expects_structured_output(&plan.req); + let mut used: u32 = 0; + let mut cursor = 0usize; + let max = plan.cfg.max_fallbacks_or_default(); + loop { + match current.next().await { + Some(Ok(chunk)) => { + if chunk.delta.tool_calls.is_some() + || chunk.delta.reasoning_content.is_some() + { + unsafe_output = true; + } + if let Some(text) = chunk.delta.content.as_deref() { + if partial.len() + text.len() + > crate::token_estimate::OUTPUT_ACCUMULATION_CAP + { + partial_overflow = true; + } else { + partial.push_str(text); + } + } + yield Ok(chunk); + } + Some(Err(err)) => { + let eligible = classify_trigger(&err) + .is_some_and(|t| plan.cfg.on_or_default().contains(&t)) + && crate::routing::is_retryable( + &err, + plan.retry_on_429, + &plan.fallback_on_statuses, + ) + && !unsafe_output + && !partial_overflow + && used < max; + if !eligible { + yield Err(err); + return; + } + match acquire_fallback_stream(&plan, &mut cursor, &mut used, err, &partial) + .await + { + Ok(next) => current = next, + Err(last) => { + yield Err(last); + return; + } + } + } + None => return, + } + } + }) +} + +/// Record the outgoing (failed) serving attempt: per-attempt UsageEvent +/// with the estimated partial spend, cooldown + health bookkeeping. +/// Mirrors what the pre-stream loop does for a failed attempt, minus +/// the pieces that only exist before the 200 (routing telemetry is +/// already finalized; the access log already went out). +fn finalize_failed_attempt(plan: &MidStreamPlan, err: &BridgeError, partial: &str) { + let (rec, failed_cooldown, failed_target_id, failed_upstream_model, failed_display); + { + let serving = plan.serving.lock().expect("serving lock"); + rec = AttemptRecord { + index: serving.attempt_index, + kind: serving.attempt_kind, + target_model: serving.target_model.clone(), + target_model_id: serving.target_id.clone(), + provider_key_id: serving.provider_key_id.clone(), + status: err.http_status(), + success: false, + error_class: routing_error_class(err).to_string(), + error_message: attempt_error_message(err), + latency_ms: serving + .attempt_started + .elapsed() + .as_millis() + .min(u32::MAX as u128) as u32, + }; + failed_cooldown = serving.cooldown.clone(); + failed_target_id = serving.target_id.clone(); + failed_upstream_model = serving.upstream_model.clone(); + failed_display = if serving.target_model.is_empty() { + plan.requested_model.clone() + } else { + serving.target_model.clone() + }; + } + if let Some((ttl, reason)) = crate::cooldown::decide_cooldown(err, failed_cooldown.as_ref()) { + plan.state + .runtime_status + .mark_cooldown(&failed_target_id, ttl, reason); + } + plan.state.health.record_failure(&failed_display); + + // Bill the failed attempt's real partial spend: prompt from the + // original request, completion from the delivered partial text + // (the same estimator the pump uses when an upstream reports no + // usage — AISIX-Cloud#1074). + let est = crate::token_estimate::Estimator::new( + &failed_upstream_model, + crate::token_estimate::PromptInput::Chat(Box::new(plan.req.clone())), + ); + let prompt_tokens = est.count_prompt(); + let completion_tokens = if partial.is_empty() { + 0 + } else { + est.count_output(partial) + }; + { + let mut extra = plan.extra_usage.lock().expect("extra_usage lock"); + *extra = extra.saturating_add(&aisix_gateway::UsageStats::new( + prompt_tokens, + completion_tokens, + )); + } + crate::chat::emit_mid_stream_failed_attempt( + &plan.state, + &plan.request_id, + &plan.requested_model, + &plan.api_key_id, + &plan.client, + &plan.applied_guardrails, + &rec, + prompt_tokens, + completion_tokens, + ); + tracing::warn!( + request_id = %plan.request_id, + failed_target = %failed_display, + error = %err, + partial_bytes = partial.len(), + "mid-stream failure; attempting fallback targets", + ); +} + +/// Try the remaining targets (from `cursor`, bounded by the episode's +/// `max_fallbacks`) until one produces a live stream. Targets in +/// cooldown / unhealthy state are skipped without burning fallback +/// budget; a dispatched target that fails to connect burns one. On +/// success the serving handle is rewritten and the caller splices the +/// returned stream into the client response. On exhaustion the most +/// recent error is returned — the pump then terminates the stream with +/// it (in-band error frame, no `[DONE]`), same as LiteLLM surfacing +/// the fallback's own failure. +async fn acquire_fallback_stream( + plan: &MidStreamPlan, + cursor: &mut usize, + used: &mut u32, + original_err: BridgeError, + partial: &str, +) -> Result { + finalize_failed_attempt(plan, &original_err, partial); + let max = plan.cfg.max_fallbacks_or_default(); + let mut last_err = original_err; + let cont_req = continuation_request(&plan.req, partial); + + while *cursor < plan.remaining.len() && *used < max { + let attempt = &plan.remaining[*cursor]; + *cursor += 1; + // Re-check runtime state at switch time — the pre-stream filter + // ran before this stream started and the world has moved (the + // failed target itself may just have been cooled down). + let stale_after = attempt + .model + .background_model_check + .as_ref() + .map(|cfg| std::time::Duration::from_secs(cfg.stale_after_seconds)); + let status = plan + .state + .runtime_status + .status_with_stale(&attempt.id, stale_after) + .status; + if matches!( + status, + crate::RuntimeStatus::Unhealthy | crate::RuntimeStatus::Cooldown + ) { + tracing::debug!( + target = %attempt.model.display_name, + ?status, + "skipping mid-stream fallback candidate (runtime state)", + ); + continue; + } + let model = &attempt.model; + let Ok(provider) = crate::dispatch::require_provider(model) else { + continue; + }; + let provider = provider.to_ascii_lowercase(); + let snapshot = plan.state.snapshot.load(); + let Ok(pk_entry) = crate::dispatch::resolve_provider_key(&snapshot, model) else { + continue; + }; + let Some(bridge) = crate::dispatch::resolve_bridge(&plan.state.hub, &pk_entry.value) else { + continue; + }; + *used += 1; + plan.fallbacks_attempted.fetch_add(1, Ordering::Relaxed); + let attempt_started = Instant::now(); + let mut ctx = crate::dispatch::bridge_ctx( + &plan.request_id, + &attempt.id, + Arc::new(model.clone()), + &pk_entry.id, + Arc::new(pk_entry.value.clone()), + Some(&plan.client), + ); + let timeouts = crate::routing::effective_timeouts( + model, + Some(&plan.group), + plan.state.default_timeouts, + ); + if let Some(d) = timeouts.stream { + ctx = ctx.with_deadline(d); + } + match bridge.chat_stream(&cont_req, &ctx).await { + Ok(up) => { + let up = crate::stream_timeout::with_read_timeout(up, timeouts.stream); + plan.state.health.record_success(&model.display_name); + plan.state.runtime_status.mark_healthy(&attempt.id); + { + let mut serving = plan.serving.lock().expect("serving lock"); + serving.attempt_index += 1; + serving.target_id = attempt.id.clone(); + serving.target_model = model.display_name.clone(); + serving.provider = provider; + serving.provider_key_id = pk_entry.id.clone(); + serving.upstream_model = + model.upstream_model().unwrap_or("unknown").to_string(); + serving.cooldown = model.cooldown.clone(); + serving.attempt_kind = "mid_stream_fallback"; + serving.attempt_started = attempt_started; + } + tracing::info!( + request_id = %plan.request_id, + fallback_target = %model.display_name, + continuation_bytes = partial.len(), + "mid-stream fallback target streaming; continuing client response", + ); + return Ok(up); + } + Err(err) => { + // The candidate never produced a stream — record it as + // its own failed attempt (zero tokens) and move on. + let rec = AttemptRecord { + index: { + let mut serving = plan.serving.lock().expect("serving lock"); + serving.attempt_index += 1; + serving.attempt_index + }, + kind: "mid_stream_fallback", + target_model: model.display_name.clone(), + target_model_id: attempt.id.clone(), + provider_key_id: pk_entry.id.clone(), + status: err.http_status(), + success: false, + error_class: routing_error_class(&err).to_string(), + error_message: attempt_error_message(&err), + latency_ms: attempt_started.elapsed().as_millis().min(u32::MAX as u128) as u32, + }; + if let Some((ttl, reason)) = + crate::cooldown::decide_cooldown(&err, model.cooldown.as_ref()) + { + plan.state + .runtime_status + .mark_cooldown(&attempt.id, ttl, reason); + } + if crate::routing::is_retryable(&err, plan.retry_on_429, &plan.fallback_on_statuses) + { + plan.state.health.record_failure(&model.display_name); + } + crate::chat::emit_mid_stream_failed_attempt( + &plan.state, + &plan.request_id, + &plan.requested_model, + &plan.api_key_id, + &plan.client, + &plan.applied_guardrails, + &rec, + 0, + 0, + ); + last_err = err; + } + } + } + Err(last_err) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn req_with_extra(extra: serde_json::Value) -> ChatFormat { + let mut req = ChatFormat::new("m", vec![ChatMessage::user("hi")]); + if let serde_json::Value::Object(map) = extra { + req.extra = map; + } + req + } + + #[test] + fn continuation_appends_instruction_and_partial() { + let orig = ChatFormat::new("m", vec![ChatMessage::user("write a story")]); + let cont = continuation_request(&orig, "Once upon a time"); + assert_eq!(cont.messages.len(), 3); + assert_eq!(cont.messages[1].content_str(), CONTINUATION_SYSTEM_PROMPT); + assert_eq!(cont.messages[2].content_str(), "Once upon a time"); + // Empty partial → untouched messages (LiteLLM pre-first-chunk + // branch). + let plain = continuation_request(&orig, ""); + assert_eq!(plain.messages.len(), 1); + } + + #[test] + fn structured_output_detection() { + assert!(!expects_structured_output(&req_with_extra( + serde_json::json!({}) + ))); + assert!(expects_structured_output(&req_with_extra( + serde_json::json!({"response_format": {"type": "json_object"}}) + ))); + assert!(expects_structured_output(&req_with_extra( + serde_json::json!({"response_format": {"type": "json_schema", "json_schema": {}}}) + ))); + assert!(!expects_structured_output(&req_with_extra( + serde_json::json!({"response_format": {"type": "text"}}) + ))); + } + + #[test] + fn trigger_classification_covers_the_mid_stream_taxonomy() { + use StreamFailureTrigger as T; + assert_eq!( + classify_trigger(&BridgeError::Transport("reset".into())), + Some(T::TransportError) + ); + assert_eq!( + classify_trigger(&BridgeError::StreamAborted), + Some(T::TransportError) + ); + assert_eq!( + classify_trigger(&BridgeError::Timeout { + elapsed_ms: 1, + cause: String::new() + }), + Some(T::ReadTimeout) + ); + assert_eq!( + classify_trigger(&BridgeError::UpstreamDecode("x".into())), + Some(T::UpstreamDecodeError) + ); + assert_eq!( + classify_trigger(&BridgeError::UpstreamInBand { + status: Some(529), + message: "overloaded".into(), + parsed: None, + wire: aisix_gateway::UpstreamWire::Anthropic, + }), + Some(T::UpstreamInBandError) + ); + assert_eq!( + classify_trigger(&BridgeError::upstream_status(500, "http")), + None + ); + assert_eq!(classify_trigger(&BridgeError::Config("c".into())), None); + } +} diff --git a/schemas/resources/model.schema.json b/schemas/resources/model.schema.json index 219de767..5c24247c 100644 --- a/schemas/resources/model.schema.json +++ b/schemas/resources/model.schema.json @@ -430,6 +430,14 @@ "default": "failover", "description": "Strategy used to select a target for each request." }, + "stream_failure": { + "allOf": [ + { + "$ref": "#/definitions/StreamFailure" + } + ], + "description": "What to do when a streaming response fails AFTER its first chunk was already delivered to the client (the HTTP 200 is committed and cannot be revised). Omitted keeps the historical behavior: terminate the stream with an in-band error frame and no `[DONE]`." + }, "targets": { "description": "Ordered set of direct models available to this routing model.", "items": { @@ -658,6 +666,86 @@ ], "type": "object" }, + "StreamFailure": { + "additionalProperties": false, + "description": "Mid-stream failure policy for streaming responses (AISIX-Cloud#1222).\n\nApplies only to failures that occur after the response head (and possibly some chunks) reached the client; failures before the first chunk keep using the regular retry/failover loop.", + "properties": { + "max_fallbacks": { + "description": "Max fallback targets tried for one mid-stream failure. Defaults to 1 — mid-stream recovery burns client-visible latency per attempt, so the default is deliberately tighter than the pre-stream `max_fallbacks`.", + "format": "uint32", + "minimum": 0.0, + "type": "integer" + }, + "mode": { + "allOf": [ + { + "$ref": "#/definitions/StreamFailureMode" + } + ], + "description": "`terminate` (default) keeps the current behavior. `continue` lets the router call the remaining fallback targets and resume the SAME client stream with a best-effort continuation of the partial text." + }, + "on": { + "description": "Which mid-stream error classes trigger the fallback. Omitted = all of them. Non-retryable errors (an in-band 4xx other than 429, unless listed in `fallback_on_statuses`) never trigger regardless.", + "items": { + "$ref": "#/definitions/StreamFailureTrigger" + }, + "type": "array" + } + }, + "type": "object" + }, + "StreamFailureMode": { + "description": "See [`StreamFailure::mode`].", + "oneOf": [ + { + "description": "Terminate the stream: in-band error frame, no `[DONE]` (the historical behavior).", + "enum": [ + "terminate" + ], + "type": "string" + }, + { + "description": "Continue on a fallback target inside the same client stream.", + "enum": [ + "continue" + ], + "type": "string" + } + ] + }, + "StreamFailureTrigger": { + "description": "Mid-stream error classes eligible for [`StreamFailureMode::Continue`].", + "oneOf": [ + { + "description": "The upstream connection broke mid-stream (reset, premature close).", + "enum": [ + "transport_error" + ], + "type": "string" + }, + { + "description": "The gap between chunks exceeded the effective `stream_timeout`.", + "enum": [ + "read_timeout" + ], + "type": "string" + }, + { + "description": "A frame failed to parse as a chunk (and was not a recognizable in-band error envelope).", + "enum": [ + "upstream_decode_error" + ], + "type": "string" + }, + { + "description": "The provider reported an error inside the committed 200 stream (an SSE error frame / event-stream modeled exception).", + "enum": [ + "upstream_in_band_error" + ], + "type": "string" + } + ] + }, "WhenAllUnavailablePolicy": { "description": "Behavior when every routing target is unavailable because of runtime health or cooldown state.", "oneOf": [ diff --git a/schemas/resources/routing.schema.json b/schemas/resources/routing.schema.json index 3f194c48..f393288b 100644 --- a/schemas/resources/routing.schema.json +++ b/schemas/resources/routing.schema.json @@ -60,6 +60,17 @@ } ] }, + "stream_failure": { + "description": "What to do when a streaming response fails AFTER its first chunk was already delivered to the client (the HTTP 200 is committed and cannot be revised). Omitted keeps the historical behavior: terminate the stream with an in-band error frame and no `[DONE]`.", + "anyOf": [ + { + "$ref": "#/definitions/StreamFailure" + }, + { + "type": "null" + } + ] + }, "targets": { "description": "Ordered set of direct models available to this routing model.", "type": "array", @@ -163,6 +174,95 @@ }, "additionalProperties": false }, + "StreamFailure": { + "description": "Mid-stream failure policy for streaming responses (AISIX-Cloud#1222).\n\nApplies only to failures that occur after the response head (and possibly some chunks) reached the client; failures before the first chunk keep using the regular retry/failover loop.", + "type": "object", + "properties": { + "max_fallbacks": { + "description": "Max fallback targets tried for one mid-stream failure. Defaults to 1 — mid-stream recovery burns client-visible latency per attempt, so the default is deliberately tighter than the pre-stream `max_fallbacks`.", + "type": [ + "integer", + "null" + ], + "format": "uint32", + "minimum": 0.0 + }, + "mode": { + "description": "`terminate` (default) keeps the current behavior. `continue` lets the router call the remaining fallback targets and resume the SAME client stream with a best-effort continuation of the partial text.", + "anyOf": [ + { + "$ref": "#/definitions/StreamFailureMode" + }, + { + "type": "null" + } + ] + }, + "on": { + "description": "Which mid-stream error classes trigger the fallback. Omitted = all of them. Non-retryable errors (an in-band 4xx other than 429, unless listed in `fallback_on_statuses`) never trigger regardless.", + "type": [ + "array", + "null" + ], + "items": { + "$ref": "#/definitions/StreamFailureTrigger" + } + } + }, + "additionalProperties": false + }, + "StreamFailureMode": { + "description": "See [`StreamFailure::mode`].", + "oneOf": [ + { + "description": "Terminate the stream: in-band error frame, no `[DONE]` (the historical behavior).", + "type": "string", + "enum": [ + "terminate" + ] + }, + { + "description": "Continue on a fallback target inside the same client stream.", + "type": "string", + "enum": [ + "continue" + ] + } + ] + }, + "StreamFailureTrigger": { + "description": "Mid-stream error classes eligible for [`StreamFailureMode::Continue`].", + "oneOf": [ + { + "description": "The upstream connection broke mid-stream (reset, premature close).", + "type": "string", + "enum": [ + "transport_error" + ] + }, + { + "description": "The gap between chunks exceeded the effective `stream_timeout`.", + "type": "string", + "enum": [ + "read_timeout" + ] + }, + { + "description": "A frame failed to parse as a chunk (and was not a recognizable in-band error envelope).", + "type": "string", + "enum": [ + "upstream_decode_error" + ] + }, + { + "description": "The provider reported an error inside the committed 200 stream (an SSE error frame / event-stream modeled exception).", + "type": "string", + "enum": [ + "upstream_in_band_error" + ] + } + ] + }, "WhenAllUnavailablePolicy": { "description": "Behavior when every routing target is unavailable because of runtime health or cooldown state.", "oneOf": [ diff --git a/tests/e2e/src/cases/mid-stream-fallback-e2e.test.ts b/tests/e2e/src/cases/mid-stream-fallback-e2e.test.ts new file mode 100644 index 00000000..6d20bd85 --- /dev/null +++ b/tests/e2e/src/cases/mid-stream-fallback-e2e.test.ts @@ -0,0 +1,304 @@ +import { createHash } from "node:crypto"; +import OpenAI from "openai"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; +import { + EtcdClient, + SeedClient, + spawnApp, + startOpenAiUpstream, + waitConfigPropagation, + type OpenAiUpstream, + type SpawnedApp, +} from "../harness/index.js"; + +const CALLER_PLAINTEXT = "sk-mid-stream-fallback-caller"; +const CALLER_KEY_HASH = createHash("sha256") + .update(CALLER_PLAINTEXT) + .digest("hex"); + +const chunk = (content: string, finish: string | null = null) => + JSON.stringify({ + id: "up-1", + object: "chat.completion.chunk", + created: 1, + model: "gpt-4o-mini", + choices: [ + { index: 0, delta: { content }, finish_reason: finish }, + ], + }); + +// AISIX-Cloud#1222: `routing.stream_failure: continue` — recover a +// streaming response INSIDE the committed 200 after the upstream fails +// mid-generation. Covers the transports the Rust integration tests +// cannot simulate with wiremock: a real mid-body connection drop and a +// real inter-chunk stall, plus the client-cancel non-trigger. +describe("mid-stream fallback e2e", () => { + let app: SpawnedApp | undefined; + let seed: SeedClient | undefined; + let etcdReachable = false; + const upstreams: OpenAiUpstream[] = []; + + beforeAll(async () => { + const etcd = new EtcdClient(); + etcdReachable = await etcd.ping(); + if (!etcdReachable) return; + + app = await spawnApp(); + seed = new SeedClient(etcd, app.etcdPrefix); + await seed.createApiKey({ + key_hash: CALLER_KEY_HASH, + allowed_models: ["*"], + }); + }); + + afterAll(async () => { + await app?.exit(); + await Promise.all(upstreams.map((u) => u.close())); + }); + + async function createTarget( + displayName: string, + upstream: OpenAiUpstream, + ): Promise { + if (!seed) throw new Error("seed client not initialized"); + const providerKey = await seed.createProviderKey({ + display_name: `${displayName}-pk`, + secret: "sk-mock", + api_base: `${upstream.baseUrl}/v1`, + }); + await seed.createModel({ + display_name: displayName, + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: providerKey.id, + }); + } + + async function createGroup( + name: string, + targets: string[], + streamFailure: Record, + extra: Record = {}, + ): Promise { + if (!seed) throw new Error("seed client not initialized"); + await seed.createModel({ + display_name: name, + routing: { + strategy: "failover", + targets: targets.map((t) => ({ model: t })), + stream_failure: streamFailure, + }, + ...extra, + }); + } + + // Watch events apply in revision order: once a canary key written + // AFTER the resources authenticates, everything before it is live. + async function waitSeedApplied(label: string): Promise { + const canary = `sk-canary-${label}-${Date.now()}`; + await seed!.createApiKey({ + key_hash: createHash("sha256").update(canary).digest("hex"), + allowed_models: ["*"], + }); + await waitConfigPropagation(async () => { + const res = await fetch(`${app!.proxyUrl}/v1/models`, { + headers: { authorization: `Bearer ${canary}` }, + }); + return res.status === 200; + }); + } + + function sdk(): OpenAI { + return new OpenAI({ + apiKey: CALLER_PLAINTEXT, + baseURL: `${app?.proxyUrl}/v1`, + maxRetries: 0, + }); + } + + test("connection drop mid-stream continues on the fallback target in the same stream", async (ctx) => { + if (!etcdReachable || !app || !seed) { + ctx.skip(); + return; + } + // Primary streams two content chunks then destroys the socket — + // a real transport break, no error frame at all. The inter-event + // delay lets each write reach the gateway before the RST (a reset + // discards any data still sitting in the receiver's buffer, which + // would turn this into a pre-first-chunk failure instead). + const primary = await startOpenAiUpstream({ + streamEvents: [ + chunk("Once "), + chunk("upon "), + chunk("NEVER-SENT"), + ], + eventDelayMs: 200, + disconnectAfterEvents: 2, + }); + upstreams.push(primary); + const secondary = await startOpenAiUpstream({ + streamEvents: [chunk("a time."), chunk("", "stop"), "[DONE]"], + }); + upstreams.push(secondary); + + await createTarget("msf-drop-primary", primary); + await createTarget("msf-drop-secondary", secondary); + await createGroup( + "msf-drop-group", + ["msf-drop-primary", "msf-drop-secondary"], + { mode: "continue" }, + ); + await waitSeedApplied("msf-drop"); + + const collected: string[] = []; + let sawFinish = false; + let surfacedError = false; + const stream = await sdk().chat.completions.create({ + model: "msf-drop-group", + messages: [{ role: "user", content: "tell me a story" }], + stream: true, + }); + try { + for await (const c of stream) { + const delta = c.choices[0]?.delta; + if (delta?.content) collected.push(delta.content); + if (c.choices[0]?.finish_reason) sawFinish = true; + } + } catch { + surfacedError = true; + } + + // The client saw primary content, then the fallback's + // continuation, then a clean completion — one logical answer. + expect(collected.join("")).toBe("Once upon a time."); + expect(sawFinish).toBe(true); + expect(surfacedError).toBe(false); + + // The fallback received the original messages plus the + // continuation instruction and the partial as an assistant + // message (LiteLLM's mid-stream fallback shape). + const calls = secondary.receivedRequests.filter((r) => + r.path.endsWith("/chat/completions"), + ); + expect(calls.length).toBe(1); + const body = JSON.parse(calls[0].body) as { + messages: Array<{ role: string; content: string }>; + }; + expect(body.messages.length).toBe(3); + expect(body.messages[0].role).toBe("user"); + expect(body.messages[1].role).toBe("system"); + expect(body.messages[1].content).toContain( + "Do not repeat the same content", + ); + expect(body.messages[2].role).toBe("assistant"); + expect(body.messages[2].content).toBe("Once upon "); + }); + + test("inter-chunk stall past stream_timeout falls back when read_timeout is a trigger", async (ctx) => { + if (!etcdReachable || !app || !seed) { + ctx.skip(); + return; + } + // Primary sends one chunk fast, then stalls far past the group's + // stream_timeout before the next one. + const primary = await startOpenAiUpstream({ + streamEvents: [chunk("The answer "), chunk("NEVER-ARRIVES")], + eventDelayMs: 5_000, + }); + upstreams.push(primary); + const secondary = await startOpenAiUpstream({ + streamEvents: [chunk("is 42."), chunk("", "stop"), "[DONE]"], + }); + upstreams.push(secondary); + + await createTarget("msf-stall-primary", primary); + await createTarget("msf-stall-secondary", secondary); + await createGroup( + "msf-stall-group", + ["msf-stall-primary", "msf-stall-secondary"], + { mode: "continue", on: ["read_timeout", "transport_error"] }, + // Group-level per-chunk budget: 1.5s gaps time out (#809 — + // stream_timeout is per chunk, not whole-response). + { stream_timeout: 1_500 }, + ); + await waitSeedApplied("msf-stall"); + + const collected: string[] = []; + let sawFinish = false; + const stream = await sdk().chat.completions.create({ + model: "msf-stall-group", + messages: [{ role: "user", content: "what is the answer" }], + stream: true, + }); + for await (const c of stream) { + const delta = c.choices[0]?.delta; + if (delta?.content) collected.push(delta.content); + if (c.choices[0]?.finish_reason) sawFinish = true; + } + + expect(collected.join("")).toBe("The answer is 42."); + expect(sawFinish).toBe(true); + const calls = secondary.receivedRequests.filter((r) => + r.path.endsWith("/chat/completions"), + ); + expect(calls.length).toBe(1); + const body = JSON.parse(calls[0].body) as { + messages: Array<{ role: string; content: string }>; + }; + expect(body.messages[2]?.content).toBe("The answer "); + }, 30_000); + + test("client cancel mid-stream never dispatches the fallback target", async (ctx) => { + if (!etcdReachable || !app || !seed) { + ctx.skip(); + return; + } + // Primary drips chunks slowly enough for the client to abort + // between them; the eventual disconnect after the abort must NOT + // start a fallback request (no ghost upstream traffic, #1094). + const primary = await startOpenAiUpstream({ + streamEvents: [ + chunk("drip "), + chunk("drip "), + chunk("drip "), + chunk("", "stop"), + "[DONE]", + ], + eventDelayMs: 500, + }); + upstreams.push(primary); + const secondary = await startOpenAiUpstream({ + streamEvents: [chunk("ghost"), chunk("", "stop"), "[DONE]"], + }); + upstreams.push(secondary); + + await createTarget("msf-cancel-primary", primary); + await createTarget("msf-cancel-secondary", secondary); + await createGroup( + "msf-cancel-group", + ["msf-cancel-primary", "msf-cancel-secondary"], + { mode: "continue" }, + { stream_timeout: 1_000 }, + ); + await waitSeedApplied("msf-cancel"); + + const stream = await sdk().chat.completions.create({ + model: "msf-cancel-group", + messages: [{ role: "user", content: "drip feed" }], + stream: true, + }); + // Take the first content chunk, then abandon the stream. + for await (const c of stream) { + if (c.choices[0]?.delta?.content) break; + } + stream.controller.abort(); + + // Give the gateway ample time to (wrongly) fire a fallback if the + // cancel path were broken — including the 1s read-timeout window. + await new Promise((r) => setTimeout(r, 3_000)); + const calls = secondary.receivedRequests.filter((r) => + r.path.endsWith("/chat/completions"), + ); + expect(calls.length).toBe(0); + }, 30_000); +}); From 93722d618920d9cf1e4233ca16866a8248c049c9 Mon Sep 17 00:00:00 2001 From: Jarvis Date: Tue, 4 Aug 2026 18:32:45 +0800 Subject: [PATCH 2/2] =?UTF-8?q?fix(stream-failover):=20review=20round=20?= =?UTF-8?q?=E2=80=94=20reserve=20fallback=20quota,=20per-attempt=20usage?= =?UTF-8?q?=20isolation,=20honest=20metric?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Reserve the fallback target's own rate-limit layers before the continuation dispatch (reserve_routing_target, AISIX-Cloud#1087 parity): refused reservation skips the candidate as a recorded 429 attempt without burning fallback budget; a granted one converts to a stream-lifetime concurrency hold (released on switch/stream end) and registers its keys for post-stream TPM accounting in the completion closure. - Reset the pump's usage accumulators when the serving attempt switches (shared attempt_seq): a per-chunk-usage provider (Gemini) would otherwise max-wins-mix the failed attempt's partial counters into the serving attempt's terminal event and double-fold with the estimated partial. - The mid-stream fallback metric now reports recovered = no terminal stream error, so a client that disconnects after a successful switch (or a guardrail block on recovered content) no longer counts as a failed failover. - Schema polish: drop the issue-tracker shorthand and the rustdoc intra-doc link from the public StreamFailure/StreamFailureMode descriptions (regenerated). - e2e: explicit 30s timeout on the connection-drop test, matching its siblings. --- crates/aisix-core/src/models/routing.rs | 12 +- crates/aisix-proxy/src/chat.rs | 85 ++++++++---- crates/aisix-proxy/src/stream_failover.rs | 124 ++++++++++++++++-- schemas/resources/model.schema.json | 4 +- schemas/resources/routing.schema.json | 4 +- .../src/cases/mid-stream-fallback-e2e.test.ts | 2 +- 6 files changed, 187 insertions(+), 44 deletions(-) diff --git a/crates/aisix-core/src/models/routing.rs b/crates/aisix-core/src/models/routing.rs index 73103612..14b88d87 100644 --- a/crates/aisix-core/src/models/routing.rs +++ b/crates/aisix-core/src/models/routing.rs @@ -192,11 +192,10 @@ pub struct Routing { pub stream_failure: Option, } -/// Mid-stream failure policy for streaming responses (AISIX-Cloud#1222). -/// -/// Applies only to failures that occur after the response head (and -/// possibly some chunks) reached the client; failures before the first -/// chunk keep using the regular retry/failover loop. +/// Mid-stream failure policy for streaming responses. Applies only to +/// failures that occur after the response head (and possibly some +/// chunks) reached the client; failures before the first chunk keep +/// using the regular retry/failover loop. #[derive(Debug, Clone, Serialize, Deserialize, schemars::JsonSchema)] pub struct StreamFailure { /// `terminate` (default) keeps the current behavior. `continue` lets @@ -240,7 +239,8 @@ impl StreamFailure { } } -/// See [`StreamFailure::mode`]. +/// How a mid-stream failure is handled once the response is already +/// streaming to the client. #[derive( Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize, schemars::JsonSchema, )] diff --git a/crates/aisix-proxy/src/chat.rs b/crates/aisix-proxy/src/chat.rs index bdd76817..8d138f2f 100644 --- a/crates/aisix-proxy/src/chat.rs +++ b/crates/aisix-proxy/src/chat.rs @@ -1639,10 +1639,7 @@ async fn dispatch( attempt_started: winner_attempt_started, }, )); - let mid_stream_fallbacks = Arc::new(AtomicU32::new(0)); - let mid_stream_extra = Arc::new(std::sync::Mutex::new( - aisix_gateway::chat::UsageStats::default(), - )); + let mid_stream_shared = crate::stream_failover::MidStreamShared::new(); let stream_failure_cfg = virtual_entry .value .routing @@ -1670,15 +1667,14 @@ async fn dispatch( api_key_id: auth.entry.id.clone(), applied_guardrails: applied_guardrails.clone(), serving: Arc::clone(&serving), - fallbacks_attempted: Arc::clone(&mid_stream_fallbacks), - extra_usage: Arc::clone(&mid_stream_extra), + shared: mid_stream_shared.clone(), }, ) } _ => upstream, }; let serving_for_telem = Arc::clone(&serving); - let mid_stream_fallbacks_for_telem = Arc::clone(&mid_stream_fallbacks); + let mid_stream_shared_for_telem = mid_stream_shared.clone(); let sse_stream = build_sse_stream( upstream, now, @@ -1690,7 +1686,7 @@ async fn dispatch( client_requested_usage, // Single upstream: nothing pre-incurred, so no usage to fold in. aisix_gateway::chat::UsageStats::default(), - Some(mid_stream_extra), + Some(mid_stream_shared), Some(estimator), move |comp: StreamCompletion| { // Mid-stream failover may have moved the stream onto a @@ -1723,7 +1719,9 @@ async fn dispatch( let snap = state_for_telem.snapshot.load(); crate::usage_attr::provider_key_metric_name(&snap, &provider_key_id_for_metrics) }; - let mid_stream_fallbacks = mid_stream_fallbacks_for_telem.load(Ordering::Relaxed); + let mid_stream_fallbacks = mid_stream_shared_for_telem + .attempt_seq + .load(Ordering::Relaxed); // Logical stream outcome (AISIX-Cloud#1222): the HTTP // status froze at 200 when the head committed, so this is // the only signal that separates "delivered in full" from @@ -1739,15 +1737,29 @@ async fn dispatch( "success" }; if mid_stream_fallbacks > 0 { - metrics_for_stream.record_mid_stream_fallback( - &model_for_metrics, - stream_outcome == "partial_recovered", - ); + // Recovered = the fallback kept the stream alive + // (no terminal stream error). A client that + // disconnects after a successful switch, or a + // guardrail block on the recovered content, is + // still a working failover — only a terminal + // upstream error counts as failed. + metrics_for_stream + .record_mid_stream_fallback(&model_for_metrics, !comp.stream_failed); } - // Rate-limit accounting (TPM cap) for all layers. + // Rate-limit accounting (TPM cap) for all layers, + // including any mid-stream fallback targets that + // served part of this stream. for key in &post_stream_keys { limiter.add_tokens_post_stream(key, comp.total_tokens); } + for key in mid_stream_shared_for_telem + .extra_post_stream_keys + .lock() + .expect("extra keys lock") + .iter() + { + limiter.add_tokens_post_stream(key, comp.total_tokens); + } // Telemetry: emit with the actual upstream-reported counts. // cost_usd stays 0.0; cp-api recomputes server-side from // its model_pricing catalog (same pattern as the non- @@ -4438,12 +4450,13 @@ fn build_sse_stream( // `on_complete` (`comp`) counts stay stream-only. Zero for single-upstream // callers, where the fold is a no-op. base_usage: aisix_gateway::chat::UsageStats, - // Estimated spend of mid-stream-failed partial attempts, written by the - // failover combinator at switch time and folded into the client-facing - // usage frames alongside `base_usage` (LiteLLM merges partial + fallback - // usage the same way; AISIX-Cloud#1222). Shared because the value is - // only known mid-stream; `None` on paths without mid-stream failover. - mid_stream_extra: Option>>, + // Mid-stream failover shared state (AISIX-Cloud#1222): the failed + // partial attempts' estimated usage (folded into client-facing usage + // frames alongside `base_usage` — LiteLLM merges partial + fallback + // usage the same way) and the attempt sequence the pump watches to + // reset its accumulators on a serving-attempt switch. `None` on paths + // without mid-stream failover. + mid_stream: Option, // Token-estimation fallback context (AISIX-Cloud#1074); see // `CompleteOnDrop::estimator`. estimator: Option, @@ -4531,6 +4544,16 @@ where // the truncated response as a successful one. let mut errored = false; let mut first_chunk_seen = false; + // Serving-attempt sequence snapshot (mid-stream failover). When + // the combinator switches targets, the usage accumulated from + // the failed attempt must not max-wins-mix into the serving + // attempt's counters (a per-chunk-usage provider like Gemini + // would otherwise leak partial counters into the terminal + // event and double-fold with the estimated partial). + let mut mid_stream_seq = mid_stream + .as_ref() + .map(|ms| ms.attempt_seq.load(Ordering::Relaxed)) + .unwrap_or(0); // Render + serialise one held/live chunk into an SSE Event. // Serialisation of these plain structs can't realistically fail; // the Err arm mirrors the pre-hold-back defensive error frame. @@ -4562,6 +4585,20 @@ where while let Some(item) = upstream.next().await { let maybe_chunk = match item { Ok(mut chunk) => { + if let Some(ms) = mid_stream.as_ref() { + let seq = ms.attempt_seq.load(Ordering::Relaxed); + if seq != mid_stream_seq { + mid_stream_seq = seq; + let comp = guard.comp(); + comp.prompt_tokens = 0; + comp.completion_tokens = 0; + comp.total_tokens = 0; + comp.cached_prompt_tokens = 0; + comp.reasoning_tokens = 0; + comp.cache_creation_tokens = 0; + comp.cache_read_tokens = 0; + } + } // Record TTFT on the first chunk carrying generated // output — reasoning text included, role-only frames // excluded. See `ChatDelta::carries_generated_output`. @@ -4708,8 +4745,12 @@ where // attempts' estimated spend into the client-facing // usage frame, AFTER `comp` captured the // serving-attempt-only counts (AISIX-Cloud#1222). - if let Some(extra) = mid_stream_extra.as_ref() { - let extra = extra.lock().expect("mid-stream extra lock").clone(); + if let Some(ms) = mid_stream.as_ref() { + let extra = ms + .extra_usage + .lock() + .expect("mid-stream extra lock") + .clone(); *u = u.saturating_add(&extra); } } diff --git a/crates/aisix-proxy/src/stream_failover.rs b/crates/aisix-proxy/src/stream_failover.rs index f615c693..a10984a1 100644 --- a/crates/aisix-proxy/src/stream_failover.rs +++ b/crates/aisix-proxy/src/stream_failover.rs @@ -84,14 +84,41 @@ pub(crate) struct MidStreamPlan { pub applied_guardrails: Vec, /// Shared with the pump's completion closure. pub serving: Arc>, - /// Count of mid-stream fallback targets attempted; the completion - /// closure derives `partial_recovered` and the fallback metric - /// from it. - pub fallbacks_attempted: Arc, + /// Cross-task state shared with the SSE pump and its completion + /// closure. + pub shared: MidStreamShared, +} + +/// State the failover combinator shares with `build_sse_stream` and the +/// completion closure. Cheap to clone (all `Arc`s). +#[derive(Clone)] +pub(crate) struct MidStreamShared { /// Estimated usage of failed partial attempts, folded into the /// final stream's client-facing usage frames by the pump (LiteLLM /// merges partial + fallback usage the same way). pub extra_usage: Arc>, + /// Bumped on every fallback dispatch. The pump watches it to reset + /// its usage accumulators when the serving attempt changes — + /// max-wins folding across attempts would otherwise mix a + /// per-chunk-usage provider's (e.g. Gemini) partial counters into + /// the serving attempt's totals. The completion closure reads it + /// as the fallbacks-attempted count. + pub attempt_seq: Arc, + /// Rate-limit keys of fallback targets that served this stream — + /// the completion closure bills their TPM post-stream the same way + /// it bills the pre-stream reservation's keys (#450 / #1087 + /// family). + pub extra_post_stream_keys: Arc>>, +} + +impl MidStreamShared { + pub fn new() -> Self { + Self { + extra_usage: Arc::new(Mutex::new(aisix_gateway::UsageStats::default())), + attempt_seq: Arc::new(AtomicU32::new(0)), + extra_post_stream_keys: Arc::new(Mutex::new(Vec::new())), + } + } } /// Classify a mid-stream [`BridgeError`] into the configurable trigger @@ -163,6 +190,11 @@ pub(crate) fn wrap( let mut used: u32 = 0; let mut cursor = 0usize; let max = plan.cfg.max_fallbacks_or_default(); + // The serving fallback target's rate-limit hold (concurrency + // slot). Replaced on every switch — releasing the previous + // fallback's slot — and released when the generator drops at + // stream end or client cancellation (#450 semantics). + let mut _fallback_hold: Option = None; loop { match current.next().await { Some(Ok(chunk)) => { @@ -200,7 +232,10 @@ pub(crate) fn wrap( match acquire_fallback_stream(&plan, &mut cursor, &mut used, err, &partial) .await { - Ok(next) => current = next, + Ok((next, hold)) => { + current = next; + _fallback_hold = hold; + } Err(last) => { yield Err(last); return; @@ -269,7 +304,7 @@ fn finalize_failed_attempt(plan: &MidStreamPlan, err: &BridgeError, partial: &st est.count_output(partial) }; { - let mut extra = plan.extra_usage.lock().expect("extra_usage lock"); + let mut extra = plan.shared.extra_usage.lock().expect("extra_usage lock"); *extra = extra.saturating_add(&aisix_gateway::UsageStats::new( prompt_tokens, completion_tokens, @@ -310,7 +345,13 @@ async fn acquire_fallback_stream( used: &mut u32, original_err: BridgeError, partial: &str, -) -> Result { +) -> Result< + ( + aisix_gateway::ChatChunkStream, + Option, + ), + BridgeError, +> { finalize_failed_attempt(plan, &original_err, partial); let max = plan.cfg.max_fallbacks_or_default(); let mut last_err = original_err; @@ -355,8 +396,56 @@ async fn acquire_fallback_stream( let Some(bridge) = crate::dispatch::resolve_bridge(&plan.state.hub, &pk_entry.value) else { continue; }; + // Reserve the fallback target's own rate-limit layers before + // dispatching to it, exactly like the pre-stream loop + // (AISIX-Cloud#1087) — without this a mid-stream continuation + // would be invisible to the target model's rpm/concurrency + // caps. A refused reservation skips the candidate (recorded as + // a 429 attempt) without burning fallback budget: nothing was + // dispatched upstream. + let member_reservation = match crate::quota::reserve_routing_target( + &plan.state, + true, + &model.display_name, + &attempt.id, + model, + ) + .await + { + Ok(r) => r, + Err(e) => { + let rec = AttemptRecord { + index: { + let mut serving = plan.serving.lock().expect("serving lock"); + serving.attempt_index += 1; + serving.attempt_index + }, + kind: "mid_stream_fallback", + target_model: model.display_name.clone(), + target_model_id: attempt.id.clone(), + provider_key_id: pk_entry.id.clone(), + status: 429, + success: false, + error_class: "rate_limit_exceeded".to_string(), + error_message: e.to_string(), + latency_ms: 0, + }; + crate::chat::emit_mid_stream_failed_attempt( + &plan.state, + &plan.request_id, + &plan.requested_model, + &plan.api_key_id, + &plan.client, + &plan.applied_guardrails, + &rec, + 0, + 0, + ); + continue; + } + }; *used += 1; - plan.fallbacks_attempted.fetch_add(1, Ordering::Relaxed); + plan.shared.attempt_seq.fetch_add(1, Ordering::Relaxed); let attempt_started = Instant::now(); let mut ctx = crate::dispatch::bridge_ctx( &plan.request_id, @@ -379,6 +468,17 @@ async fn acquire_fallback_stream( let up = crate::stream_timeout::with_read_timeout(up, timeouts.stream); plan.state.health.record_success(&model.display_name); plan.state.runtime_status.mark_healthy(&attempt.id); + // Convert the reservation into a stream-lifetime hold + // and register its keys so the completion closure bills + // this target's TPM post-stream too. + let hold = member_reservation.map(|r| { + plan.shared + .extra_post_stream_keys + .lock() + .expect("extra keys lock") + .extend(r.keys()); + r.into_stream_hold() + }); { let mut serving = plan.serving.lock().expect("serving lock"); serving.attempt_index += 1; @@ -398,11 +498,13 @@ async fn acquire_fallback_stream( continuation_bytes = partial.len(), "mid-stream fallback target streaming; continuing client response", ); - return Ok(up); + return Ok((up, hold)); } Err(err) => { - // The candidate never produced a stream — record it as - // its own failed attempt (zero tokens) and move on. + // The candidate never produced a stream — the refused + // reservation drops here, rolling its counters back. + // Record it as a failed attempt (zero tokens) and move + // on. let rec = AttemptRecord { index: { let mut serving = plan.serving.lock().expect("serving lock"); diff --git a/schemas/resources/model.schema.json b/schemas/resources/model.schema.json index 5c24247c..9ac96a8f 100644 --- a/schemas/resources/model.schema.json +++ b/schemas/resources/model.schema.json @@ -668,7 +668,7 @@ }, "StreamFailure": { "additionalProperties": false, - "description": "Mid-stream failure policy for streaming responses (AISIX-Cloud#1222).\n\nApplies only to failures that occur after the response head (and possibly some chunks) reached the client; failures before the first chunk keep using the regular retry/failover loop.", + "description": "Mid-stream failure policy for streaming responses. Applies only to failures that occur after the response head (and possibly some chunks) reached the client; failures before the first chunk keep using the regular retry/failover loop.", "properties": { "max_fallbacks": { "description": "Max fallback targets tried for one mid-stream failure. Defaults to 1 — mid-stream recovery burns client-visible latency per attempt, so the default is deliberately tighter than the pre-stream `max_fallbacks`.", @@ -695,7 +695,7 @@ "type": "object" }, "StreamFailureMode": { - "description": "See [`StreamFailure::mode`].", + "description": "How a mid-stream failure is handled once the response is already streaming to the client.", "oneOf": [ { "description": "Terminate the stream: in-band error frame, no `[DONE]` (the historical behavior).", diff --git a/schemas/resources/routing.schema.json b/schemas/resources/routing.schema.json index f393288b..237a5952 100644 --- a/schemas/resources/routing.schema.json +++ b/schemas/resources/routing.schema.json @@ -175,7 +175,7 @@ "additionalProperties": false }, "StreamFailure": { - "description": "Mid-stream failure policy for streaming responses (AISIX-Cloud#1222).\n\nApplies only to failures that occur after the response head (and possibly some chunks) reached the client; failures before the first chunk keep using the regular retry/failover loop.", + "description": "Mid-stream failure policy for streaming responses. Applies only to failures that occur after the response head (and possibly some chunks) reached the client; failures before the first chunk keep using the regular retry/failover loop.", "type": "object", "properties": { "max_fallbacks": { @@ -212,7 +212,7 @@ "additionalProperties": false }, "StreamFailureMode": { - "description": "See [`StreamFailure::mode`].", + "description": "How a mid-stream failure is handled once the response is already streaming to the client.", "oneOf": [ { "description": "Terminate the stream: in-band error frame, no `[DONE]` (the historical behavior).", diff --git a/tests/e2e/src/cases/mid-stream-fallback-e2e.test.ts b/tests/e2e/src/cases/mid-stream-fallback-e2e.test.ts index 6d20bd85..6bcfb273 100644 --- a/tests/e2e/src/cases/mid-stream-fallback-e2e.test.ts +++ b/tests/e2e/src/cases/mid-stream-fallback-e2e.test.ts @@ -192,7 +192,7 @@ describe("mid-stream fallback e2e", () => { ); expect(body.messages[2].role).toBe("assistant"); expect(body.messages[2].content).toBe("Once upon "); - }); + }, 30_000); test("inter-chunk stall past stream_timeout falls back when read_timeout is a trigger", async (ctx) => { if (!etcdReachable || !app || !seed) {