From b39ad41f13be5a75259204fbe50bc2b562edb66d Mon Sep 17 00:00:00 2001 From: Jarvis Date: Wed, 5 Aug 2026 15:58:28 +0800 Subject: [PATCH] =?UTF-8?q?Revert=20"feat(routing):=20mid-stream=20fallbac?= =?UTF-8?q?k=20=E2=80=94=20resume=20a=20committed=20stream=20on=20fallback?= =?UTF-8?q?=20targets=20(#882)"?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This reverts commit 7d519b7022cf4d81cf7928883520828b11a9c625. --- 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 | 287 +------- crates/aisix-proxy/src/lib.rs | 353 ---------- crates/aisix-proxy/src/routing.rs | 1 - crates/aisix-proxy/src/stream_failover.rs | 635 ------------------ schemas/resources/model.schema.json | 88 --- schemas/resources/routing.schema.json | 100 --- .../src/cases/mid-stream-fallback-e2e.test.ts | 304 --------- 13 files changed, 25 insertions(+), 1888 deletions(-) delete mode 100644 crates/aisix-proxy/src/stream_failover.rs delete 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 f30c9c53..092379a4 100644 --- a/crates/aisix-admin/src/openapi.rs +++ b/crates/aisix-admin/src/openapi.rs @@ -4343,19 +4343,6 @@ 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 7118a9fe..bc3fda26 100644 --- a/crates/aisix-core/src/lib.rs +++ b/crates/aisix-core/src/lib.rs @@ -49,9 +49,8 @@ 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, StreamFailure, - StreamFailureMode, StreamFailureTrigger, TelemetryKind, TelemetryTags, - WhenAllUnavailablePolicy, DEFAULT_COOLDOWN_TRIGGER_STATUSES, + Routing, RoutingStrategy, RoutingTarget, SchemaError, StreamDoneMarker, 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 c0fb057c..dff6e8f1 100644 --- a/crates/aisix-core/src/models/mod.rs +++ b/crates/aisix-core/src/models/mod.rs @@ -69,10 +69,7 @@ pub use provider_key::{ }; pub use rate_limit::{McpRateLimit, RateLimit}; pub use rate_limit_policy::{PolicyScope, PolicyWindow, RateLimitPolicy}; -pub use routing::{ - Routing, RoutingStrategy, RoutingTarget, StreamFailure, StreamFailureMode, - StreamFailureTrigger, WhenAllUnavailablePolicy, -}; +pub use routing::{Routing, RoutingStrategy, RoutingTarget, 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 14b88d87..2c9f579b 100644 --- a/crates/aisix-core/src/models/routing.rs +++ b/crates/aisix-core/src/models/routing.rs @@ -184,90 +184,6 @@ 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. 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) - } -} - -/// 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, -)] -#[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 { @@ -355,7 +271,6 @@ mod tests { fallback_on_statuses: None, when_all_unavailable: None, sticky: None, - stream_failure: None, }; assert_eq!(r.max_fallbacks_or_default(), 0); } @@ -371,7 +286,6 @@ 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 19ba97da..9da440cf 100644 --- a/crates/aisix-obs/src/metrics.rs +++ b/crates/aisix-obs/src/metrics.rs @@ -79,12 +79,6 @@ 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"; @@ -1216,20 +1210,6 @@ 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 8089cb9f..78200cc7 100644 --- a/crates/aisix-obs/src/usage.rs +++ b/crates/aisix-obs/src/usage.rs @@ -195,22 +195,6 @@ 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 289e34eb..4835eab2 100644 --- a/crates/aisix-proxy/src/chat.rs +++ b/crates/aisix-proxy/src/chat.rs @@ -242,7 +242,6 @@ 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(), @@ -514,7 +513,6 @@ 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(), @@ -1217,9 +1215,6 @@ 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 @@ -1423,7 +1418,6 @@ async fn dispatch( upstream, idx, kind, - target_idx, attempt_started, }); won_member_reservation = member_reservation; @@ -1510,7 +1504,6 @@ async fn dispatch( upstream, idx: winner_idx, kind: winner_kind, - target_idx: winner_target_idx, attempt_started: winner_attempt_started, } = won; let model = &model; @@ -1558,15 +1551,26 @@ 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(); - // #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 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) + }; 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 @@ -1631,61 +1635,6 @@ 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_shared = crate::stream_failover::MidStreamShared::new(); - 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(), - auth: auth.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), - shared: mid_stream_shared.clone(), - }, - ) - } - _ => upstream, - }; - let serving_for_telem = Arc::clone(&serving); - let mid_stream_shared_for_telem = mid_stream_shared.clone(); let sse_stream = build_sse_stream( upstream, now, @@ -1697,80 +1646,12 @@ 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_shared), 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_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 - // "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 { - // 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, - // including any mid-stream fallback targets that - // served part of this stream. + // Rate-limit accounting (TPM cap) for all layers. 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- @@ -1834,20 +1715,8 @@ async fn dispatch( attempt_index: winner_idx, attempt_kind: winner_kind.to_string(), attempt_model: attempt_model_for_telem.clone(), - // 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(), + error_class: String::new(), + error_message: String::new(), applied_guardrails: applied_guardrails_for_telem.clone(), provider_key_id: provider_key_id_for_telem.clone(), redacted_entity_counts: { @@ -1893,7 +1762,10 @@ async fn dispatch( upstream_model: &upstream_model_for_metrics, provider_key_id: &provider_key_id_for_metrics, stream: true, - is_fallback: mid_stream_fallbacks > 0, + // The serving target is fixed once the stream + // commits, so fallback attribution is the + // pre-stream winner's attempt kind. + is_fallback: winner_kind == "fallback", }, crate::request_metrics::Tokens { input: comp.prompt_tokens, @@ -3215,8 +3087,6 @@ 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) @@ -3849,7 +3719,6 @@ 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 @@ -3972,10 +3841,6 @@ 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 @@ -4065,55 +3930,6 @@ 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, - ); -} - #[allow(clippy::too_many_arguments)] fn emit_access_log( method: &str, @@ -4290,16 +4106,6 @@ 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 @@ -4452,13 +4258,6 @@ 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, - // 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, @@ -4546,16 +4345,6 @@ 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. @@ -4587,20 +4376,6 @@ 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`. @@ -4743,29 +4518,11 @@ 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(ms) = mid_stream.as_ref() { - let extra = ms - .extra_usage - .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 3218c4d6..8295c9f3 100644 --- a/crates/aisix-proxy/src/lib.rs +++ b/crates/aisix-proxy/src/lib.rs @@ -66,7 +66,6 @@ mod routing; mod semantic; pub mod sse_keepalive; mod state; -mod stream_failover; mod stream_timeout; mod token_estimate; mod usage_attr; @@ -4682,358 +4681,6 @@ 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 b8879f1e..550328b9 100644 --- a/crates/aisix-proxy/src/routing.rs +++ b/crates/aisix-proxy/src/routing.rs @@ -898,7 +898,6 @@ 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 deleted file mode 100644 index 1351e831..00000000 --- a/crates/aisix-proxy/src/stream_failover.rs +++ /dev/null @@ -1,635 +0,0 @@ -//! 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, - /// Caller identity — the per-target quota gate needs the identity - /// dimensions for conditional policy rows (AISIX-Cloud#892). - pub auth: crate::auth::AuthenticatedKey, - /// 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>, - /// 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 -/// 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(); - // 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)) => { - 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, hold)) => { - current = next; - _fallback_hold = hold; - } - 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.shared.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< - ( - 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; - 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; - }; - // 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, - &plan.auth, - 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.shared.attempt_seq.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); - // 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; - 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, hold)); - } - Err(err) => { - // 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"); - 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 9ac96a8f..219de767 100644 --- a/schemas/resources/model.schema.json +++ b/schemas/resources/model.schema.json @@ -430,14 +430,6 @@ "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": { @@ -666,86 +658,6 @@ ], "type": "object" }, - "StreamFailure": { - "additionalProperties": false, - "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`.", - "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": "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).", - "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 237a5952..3f194c48 100644 --- a/schemas/resources/routing.schema.json +++ b/schemas/resources/routing.schema.json @@ -60,17 +60,6 @@ } ] }, - "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", @@ -174,95 +163,6 @@ }, "additionalProperties": false }, - "StreamFailure": { - "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": { - "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": "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).", - "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 deleted file mode 100644 index 6bcfb273..00000000 --- a/tests/e2e/src/cases/mid-stream-fallback-e2e.test.ts +++ /dev/null @@ -1,304 +0,0 @@ -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 "); - }, 30_000); - - 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); -});