diff --git a/crates/aisix-provider-anthropic/src/lib.rs b/crates/aisix-provider-anthropic/src/lib.rs index 315186d1..fb12a3f6 100644 --- a/crates/aisix-provider-anthropic/src/lib.rs +++ b/crates/aisix-provider-anthropic/src/lib.rs @@ -29,8 +29,8 @@ pub use bridge::{AnthropicBridge, ANTHROPIC_DEFAULT_BASE, ANTHROPIC_VERSION}; /// - [`AnthropicSseEncoder`] re-encodes the bridge's `ChatChunk` /// stream as Anthropic typed SSE events. pub use wire::{ - chat_response_into_anthropic_json, parse_inbound_request, + chat_response_into_anthropic_json, parse_inbound_request, stream_error_into_bridge_error, translate_anthropic_tool_choice_to_openai, translate_anthropic_tools_to_openai, translate_extras_to_openai_shape, AnthropicInboundError, AnthropicSseEncoder, - AnthropicSseEvent, + AnthropicSseEvent, AnthropicStreamErrorBody, }; diff --git a/crates/aisix-provider-anthropic/src/wire.rs b/crates/aisix-provider-anthropic/src/wire.rs index 30ac0dce..538ff0e2 100644 --- a/crates/aisix-provider-anthropic/src/wire.rs +++ b/crates/aisix-provider-anthropic/src/wire.rs @@ -1744,6 +1744,13 @@ pub struct AnthropicSseEncoder { seen_input_tokens: u32, seen_output_tokens: u32, usage_seen: bool, + /// Usage folded ADDITIVELY into the closing pair on top of the + /// `seen_*` counters (mid-stream failover, AISIX-Cloud#1222): the + /// estimated spend of failed partial attempts. Anthropic wire usage + /// is cumulative per message, so the terminal counts must also + /// cover the partial text the client received before the switch. + extra_input_tokens: u32, + extra_output_tokens: u32, } impl AnthropicSseEncoder { @@ -1770,9 +1777,49 @@ impl AnthropicSseEncoder { seen_input_tokens: 0, seen_output_tokens: 0, usage_seen: false, + extra_input_tokens: 0, + extra_output_tokens: 0, } } + /// Rebuild an encoder positioned mid-message, for the passthrough + /// mid-stream failover (AISIX-Cloud#1222): the client already holds + /// the upstream's own `message_start` (and possibly an open text + /// block), so the resumed encoder must continue that envelope + /// instead of opening its own. `open_text_block` is the index of + /// the text content block the upstream left open (`None` → the + /// next text delta opens a fresh block at `next_block_index`). + pub fn resume( + message_id: impl Into, + model_display_name: impl Into, + open_text_block: Option, + next_block_index: usize, + ) -> Self { + let mut enc = Self::new(message_id, model_display_name, 0); + enc.sent_message_start = true; + enc.text_block_index = open_text_block; + enc.next_block_index = next_block_index; + enc + } + + /// Fold failed-attempt usage additively into the closing pair (see + /// `extra_input_tokens`). Refreshed by the pump whenever the + /// serving attempt changes. + pub fn set_extra_usage(&mut self, input_tokens: u32, output_tokens: u32) { + self.extra_input_tokens = input_tokens; + self.extra_output_tokens = output_tokens; + } + + /// Zero the cumulative usage counters when the serving attempt + /// changes (mid-stream failover): max-wins folding across attempts + /// would otherwise mix the failed attempt's counters into the + /// serving attempt's totals. + pub fn reset_usage_counters(&mut self) { + self.seen_input_tokens = 0; + self.seen_output_tokens = 0; + self.usage_seen = false; + } + /// Translate one chunk into the Anthropic SSE events to emit. /// Returns an empty Vec on no-op chunks. pub fn next_events(&mut self, chunk: &ChatChunk) -> Vec { @@ -1951,10 +1998,18 @@ impl AnthropicSseEncoder { /// place the client can learn the prompt token count. fn closing_pair(&mut self, stop_reason: &'static str) -> Vec { let mut usage = serde_json::Map::new(); - if self.seen_input_tokens > 0 { - usage.insert("input_tokens".into(), self.seen_input_tokens.into()); + let input_total = self + .seen_input_tokens + .saturating_add(self.extra_input_tokens); + if input_total > 0 { + usage.insert("input_tokens".into(), input_total.into()); } - usage.insert("output_tokens".into(), self.seen_output_tokens.into()); + usage.insert( + "output_tokens".into(), + self.seen_output_tokens + .saturating_add(self.extra_output_tokens) + .into(), + ); self.finished = true; vec![ AnthropicSseEvent { diff --git a/crates/aisix-proxy/src/chat.rs b/crates/aisix-proxy/src/chat.rs index 289e34eb..0afb3409 100644 --- a/crates/aisix-proxy/src/chat.rs +++ b/crates/aisix-proxy/src/chat.rs @@ -1665,6 +1665,8 @@ async fn dispatch( upstream, crate::stream_failover::MidStreamPlan { cfg, + endpoint: crate::stream_failover::MidStreamEndpoint::Chat, + structured_output: crate::stream_failover::expects_structured_output(req), remaining: attempt_models[winner_target_idx + 1..].to_vec(), state: state.clone(), auth: auth.clone(), @@ -3794,6 +3796,50 @@ fn emit_usage_event( guardrail_blocked: bool, client: &ClientContext, content: Option, +) { + emit_usage_event_stamped( + state, + "chat", + "openai", + request_id, + model_id, + requested_model, + api_key_id, + status_code, + elapsed, + prompt_tokens, + completion_tokens, + extras, + cost_usd, + guardrail_blocked, + client, + content, + ) +} + +/// `emit_usage_event` with the endpoint stamps explicit — the +/// mid-stream failover per-attempt events reuse chat's emit pipeline +/// (attribution-tag lookup, guardrail counters, OTLP fan-out) from the +/// `/v1/messages` and `/v1/responses` combinators, which stamp their +/// own `inbound_protocol` + usage-sink handler label. +#[allow(clippy::too_many_arguments)] +fn emit_usage_event_stamped( + state: &ProxyState, + sink_label: &'static str, + inbound_protocol: &'static str, + request_id: &str, + model_id: &str, + requested_model: &str, + api_key_id: &str, + status_code: u16, + elapsed: Duration, + prompt_tokens: u32, + completion_tokens: u32, + extras: UsageExtras, + cost_usd: f64, + guardrail_blocked: bool, + client: &ClientContext, + content: Option, ) { // Look up per-PK telemetry attribution tags from the live snapshot. // Empty `provider_key_id` (pre-dispatch error paths) → default @@ -3839,11 +3885,7 @@ fn emit_usage_event( cache_hit_saved_input_tokens: extras.cache_hit_saved_input_tokens, cache_hit_saved_output_tokens: extras.cache_hit_saved_output_tokens, upstream_ttft_ms: extras.upstream_ttft_ms, - // chat.rs is the OpenAI-shape /v1/chat/completions handler. - // /v1/responses / /v1/embeddings / /v1/audio* / /v1/images* / - // /v1/rerank don't emit UsageEvents today; when they do they - // also pass `"openai"` here. - inbound_protocol: "openai".to_string(), + inbound_protocol: inbound_protocol.to_string(), attempt_index: extras.attempt_index, attempt_kind: extras.attempt_kind, attempt_model: extras.attempt_model, @@ -3870,10 +3912,10 @@ fn emit_usage_event( // MCP attribution does not apply to the chat path. ..Default::default() }; - // Handler label "chat" matches the documented enumeration for + // Handler labels match the documented enumeration for // `aisix_usage_events_emitted_total` (#408). Keep `&'static str` // so prometheus cardinality stays bounded. - state.usage_sink.try_emit("chat", event.clone()); + state.usage_sink.try_emit(sink_label, event.clone()); // Guardrail outcome counters (#379). Recorded here — the one place every // chat path (success / error / streaming / cache-hit) funnels through — // from the same guardrail fields the UsageEvent carries. @@ -4077,6 +4119,7 @@ fn emit_failed_attempts( #[allow(clippy::too_many_arguments)] pub(crate) fn emit_mid_stream_failed_attempt( state: &ProxyState, + endpoint: crate::stream_failover::MidStreamEndpoint, request_id: &str, requested_model: &str, api_key_id: &str, @@ -4086,8 +4129,10 @@ pub(crate) fn emit_mid_stream_failed_attempt( prompt_tokens: u32, completion_tokens: u32, ) { - emit_usage_event( + emit_usage_event_stamped( state, + endpoint.sink_label(), + endpoint.inbound_protocol(), request_id, &rec.target_model_id, requested_model, diff --git a/crates/aisix-proxy/src/lib.rs b/crates/aisix-proxy/src/lib.rs index 3218c4d6..8847b378 100644 --- a/crates/aisix-proxy/src/lib.rs +++ b/crates/aisix-proxy/src/lib.rs @@ -4846,7 +4846,7 @@ data: [DONE]\n\n"; assert!(messages[1]["content"] .as_str() .unwrap() - .contains("Do not repeat the same content")); + .contains("do not repeat any of its content")); assert_eq!(messages[2]["role"], "assistant"); assert_eq!(messages[2]["content"], "Once upon"); } @@ -5034,6 +5034,441 @@ data: {\"error\":{\"message\":\"The server had an error\",\"type\":\"server_erro ); } + // ── Mid-stream failover on /v1/messages (AISIX-Cloud#1222 phase 2) ── + + fn anthropic_target_model(id: &str, name: &str, pk_id: &str) -> ResourceEntry { + let cfg = serde_json::json!({ + "display_name": name, + "provider": "anthropic", + "model_name": "claude-3-5-haiku-20241022", + "provider_key_id": pk_id, + }); + ResourceEntry::new(id, serde_json::from_value(cfg).unwrap(), 1) + } + + fn anthropic_pk_with( + pk_id: &'static str, + api_base: &str, + ) -> ResourceEntry { + matrix_pk_entry(pk_id, "sk-ant-test", api_base, "anthropic", "anthropic") + } + + /// Anthropic-wire stream head: committed envelope + one text delta, + /// then an in-band `error` frame. + const MSG_MID_STREAM_HEAD_SSE: &str = "\ +event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_up1\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"model\":\"claude-3-5-haiku-20241022\",\"stop_reason\":null,\"usage\":{\"input_tokens\":5,\"output_tokens\":1}}}\n\n\ +event: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n\ +event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"Once upon\"}}\n\n\ +event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",\"message\":\"Overloaded\"}}\n\n"; + + /// Anthropic-wire full recovery stream served by the fallback target. + const MSG_RECOVERY_SSE: &str = "\ +event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_up2\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"model\":\"claude-3-5-haiku-20241022\",\"stop_reason\":null,\"usage\":{\"input_tokens\":6,\"output_tokens\":0}}}\n\n\ +event: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n\ +event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\" a time\"}}\n\n\ +event: content_block_stop\ndata: {\"type\":\"content_block_stop\",\"index\":0}\n\n\ +event: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\",\"stop_sequence\":null},\"usage\":{\"output_tokens\":7}}\n\n\ +event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"; + + fn messages_mid_stream_snapshot( + primary_uri: &str, + secondary_uri: &str, + stream_failure: Option, + ) -> AisixSnapshot { + let snap = AisixSnapshot::new(); + snap.provider_keys + .insert(anthropic_pk_with("pk-ant-primary", primary_uri)); + snap.provider_keys + .insert(anthropic_pk_with("pk-ant-secondary", secondary_uri)); + snap.models.insert(anthropic_target_model( + "m-ant-primary", + "primary", + "pk-ant-primary", + )); + snap.models.insert(anthropic_target_model( + "m-ant-secondary", + "secondary", + "pk-ant-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_messages_wire(app: axum::Router, model: &str) -> String { + let body = serde_json::json!({ + "model": model, + "max_tokens": 128, + "messages": [{"role": "user", "content": "tell me a story"}], + "stream": true + }); + let req = Request::builder() + .method("POST") + .uri("/v1/messages") + .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 phase 2 core acceptance (passthrough leg): an + /// Anthropic upstream failing in-band inside its committed stream + /// moves the SAME client stream onto the fallback target; the + /// resumed encoder continues the client's message envelope — no + /// second `message_start`, no new `content_block_start`, text keeps + /// flowing in block 0 — and the fallback (an Anthropic-wire target + /// reached through the chat bridge) receives the continuation with + /// the partial as a native prefill assistant message. + #[tokio::test] + async fn messages_mid_stream_failover_passthrough_continues_envelope() { + use aisix_provider_anthropic::AnthropicBridge; + + let primary = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/messages")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(MSG_MID_STREAM_HEAD_SSE), + ) + .mount(&primary) + .await; + let secondary = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/messages")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(MSG_RECOVERY_SSE), + ) + .mount(&secondary) + .await; + + let hub = Arc::new(Hub::new()); + hub.register_specialized("anthropic", Arc::new(AnthropicBridge::new())); + let snap = messages_mid_stream_snapshot( + &primary.uri(), + &secondary.uri(), + Some(serde_json::json!({"mode": "continue"})), + ); + let app = build_router(build_state(snap, hub)); + + let wire = streaming_messages_wire(app, "smart").await; + assert_eq!( + wire.matches("event: message_start").count(), + 1, + "the client keeps its original message envelope:\n{wire}" + ); + assert_eq!( + wire.matches("event: content_block_start").count(), + 1, + "the continuation rides the already-open text block:\n{wire}" + ); + assert!(wire.contains("Once upon"), "primary partial:\n{wire}"); + assert!(wire.contains(" a time"), "fallback continuation:\n{wire}"); + assert!( + !wire.contains("event: error"), + "recovered stream carries no error frame:\n{wire}" + ); + assert_eq!( + wire.matches("event: message_stop").count(), + 1, + "exactly one clean terminal:\n{wire}" + ); + assert!( + wire.contains("\"output_tokens\""), + "closing message_delta carries usage:\n{wire}" + ); + + // The fallback got the continuation request on the Anthropic + // wire: the instruction (a mid-conversation system message + // becomes a user turn) and the partial as the trailing + // assistant message — native prefill. + 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 body_str = body.to_string(); + assert!( + body_str.contains("interrupted mid-stream"), + "continuation instruction present:\n{body_str}" + ); + let messages = body["messages"].as_array().unwrap(); + let last = messages.last().unwrap(); + assert_eq!(last["role"], "assistant", "trailing prefill:\n{body_str}"); + assert!( + last["content"].to_string().contains("Once upon"), + "prefill carries the partial:\n{body_str}" + ); + } + + /// Cross-provider leg: /v1/messages served by OpenAI-protocol + /// targets (chat bridge + AnthropicSseEncoder). The encoder lives + /// outside the combinator, so the spliced fallback chunks continue + /// the same synthesized envelope. + #[tokio::test] + async fn messages_mid_stream_failover_cross_provider_continues_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: {\"id\":\"up-2\",\"model\":\"gpt-4o\",\"choices\":[],\"usage\":{\"prompt_tokens\":9,\"completion_tokens\":7,\"total_tokens\":16}}\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_messages_wire(app, "smart").await; + assert_eq!( + wire.matches("event: message_start").count(), + 1, + "one synthesized envelope across the switch:\n{wire}" + ); + assert!(wire.contains("Once upon"), "primary partial:\n{wire}"); + assert!(wire.contains(" a time"), "fallback continuation:\n{wire}"); + assert!( + !wire.contains("event: error"), + "recovered stream carries no error frame:\n{wire}" + ); + assert_eq!(wire.matches("event: message_stop").count(), 1); + + // OpenAI-wire continuation body: original user message + the + // continuation instruction + the partial as assistant. + 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 + instruction + partial"); + assert_eq!(messages[1]["role"], "system"); + assert!(messages[1]["content"] + .as_str() + .unwrap() + .contains("do not repeat any of its content")); + assert_eq!(messages[2]["role"], "assistant"); + assert_eq!(messages[2]["content"], "Once upon"); + } + + /// Without `stream_failure` config the passthrough keeps its + /// historical behavior byte-for-byte: the upstream's own error + /// frame is forwarded and the fallback target is never contacted. + #[tokio::test] + async fn messages_mid_stream_default_terminate_forwards_upstream_error_frame() { + use aisix_provider_anthropic::AnthropicBridge; + + let primary = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/messages")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(MSG_MID_STREAM_HEAD_SSE), + ) + .mount(&primary) + .await; + let secondary = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/messages")) + .respond_with(ResponseTemplate::new(200)) + .expect(0) + .mount(&secondary) + .await; + + let hub = Arc::new(Hub::new()); + hub.register_specialized("anthropic", Arc::new(AnthropicBridge::new())); + let snap = messages_mid_stream_snapshot(&primary.uri(), &secondary.uri(), None); + let app = build_router(build_state(snap, hub)); + + let wire = streaming_messages_wire(app, "smart").await; + assert!(wire.contains("Once upon")); + assert!( + wire.contains("overloaded_error"), + "upstream error frame forwards verbatim:\n{wire}" + ); + let reqs = secondary.received_requests().await.unwrap(); + assert!(reqs.is_empty(), "fallback never contacted"); + } + + /// Output-shape safety gate on the passthrough leg: a thinking + /// block on the wire disarms continuation — the withheld upstream + /// error frame is released verbatim and no fallback is dispatched. + #[tokio::test] + async fn messages_mid_stream_fallback_skipped_after_thinking_block() { + use aisix_provider_anthropic::AnthropicBridge; + + let thinking_sse = "\ +event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_up1\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"model\":\"claude-3-5-haiku-20241022\",\"stop_reason\":null,\"usage\":{\"input_tokens\":5,\"output_tokens\":1}}}\n\n\ +event: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"thinking\",\"thinking\":\"\"}}\n\n\ +event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"thinking_delta\",\"thinking\":\"hmm\"}}\n\n\ +event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",\"message\":\"Overloaded\"}}\n\n"; + let primary = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/messages")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(thinking_sse), + ) + .mount(&primary) + .await; + let secondary = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/messages")) + .respond_with(ResponseTemplate::new(200)) + .expect(0) + .mount(&secondary) + .await; + + let hub = Arc::new(Hub::new()); + hub.register_specialized("anthropic", Arc::new(AnthropicBridge::new())); + let snap = messages_mid_stream_snapshot( + &primary.uri(), + &secondary.uri(), + Some(serde_json::json!({"mode": "continue"})), + ); + let app = build_router(build_state(snap, hub)); + + let wire = streaming_messages_wire(app, "smart").await; + assert!( + wire.contains("thinking"), + "thinking bytes forwarded:\n{wire}" + ); + assert!( + wire.contains("overloaded_error"), + "withheld upstream error frame released verbatim:\n{wire}" + ); + let reqs = secondary.received_requests().await.unwrap(); + assert!( + reqs.is_empty(), + "no continuation after a thinking block on the wire" + ); + } + + /// Telemetry contract on /v1/messages: the failed serving attempt + /// emits its own per-attempt event (estimated partial spend, + /// anthropic protocol stamp) and the terminal event is attributed + /// to the fallback target with `stream_outcome: partial_recovered`. + #[tokio::test] + async fn messages_mid_stream_failover_emits_attempt_events_and_outcome() { + use aisix_obs::UsageSink; + + 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\":{\"content\":\" a time\"},\"finish_reason\":\"stop\"}]}\n\n\ +data: {\"id\":\"up-2\",\"model\":\"gpt-4o\",\"choices\":[],\"usage\":{\"prompt_tokens\":9,\"completion_tokens\":7,\"total_tokens\":16}}\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 (tx, mut rx) = tokio::sync::mpsc::channel(8); + let app = build_router(build_state(snap, hub).with_usage_sink(UsageSink::new(tx))); + + let wire = streaming_messages_wire(app, "smart").await; + assert!(wire.contains(" a time"), "recovered:\n{wire}"); + + let mut events = Vec::new(); + for _ in 0..2 { + let ev = tokio::time::timeout(std::time::Duration::from_millis(500), rx.recv()) + .await + .expect("usage event was never emitted") + .expect("sender dropped"); + events.push(ev); + } + events.sort_by_key(|e| e.attempt_index); + + // Failed serving attempt: estimated partial spend, protocol + // stamped for this endpoint. + assert_eq!(events[0].attempt_index, 0); + assert_eq!(events[0].attempt_kind, "initial"); + assert_eq!(events[0].model_id, "m-primary"); + assert_eq!(events[0].error_class, "upstream_in_band"); + assert_eq!(events[0].inbound_protocol, "anthropic"); + assert_eq!(events[0].requested_model, "smart"); + assert!(events[0].usage_estimated); + assert!(events[0].completion_tokens > 0, "partial text billed"); + + // Terminal event: attributed to the serving fallback target. + assert_eq!(events[1].attempt_kind, "mid_stream_fallback"); + assert_eq!(events[1].model_id, "m-secondary"); + assert_eq!(events[1].attempt_model, "secondary"); + assert_eq!(events[1].status_code, 200); + assert_eq!(events[1].stream_outcome, "partial_recovered"); + assert_eq!(events[1].inbound_protocol, "anthropic"); + assert_eq!( + events[1].completion_tokens, 7, + "serving attempt's own count" + ); + } + #[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/messages.rs b/crates/aisix-proxy/src/messages.rs index 4fd8339f..1123498e 100644 --- a/crates/aisix-proxy/src/messages.rs +++ b/crates/aisix-proxy/src/messages.rs @@ -271,6 +271,7 @@ pub async fn messages( applied_guardrails.clone(), redaction_counts.clone(), monitor_hits.clone(), + "", captured_content, ); } @@ -394,6 +395,7 @@ pub async fn messages( // Input masking may have fired before the failure. redaction_counts.clone(), monitor_hits.clone(), + "", failure_content.take(), ); } @@ -466,11 +468,30 @@ fn emit_failed_attempts_anthropic( // terminal (winner / pre-dispatch) event does. crate::redact::RedactionCounts::new(), Vec::new(), + "", content, ); } } +/// Per-attempt mid-stream failover context (AISIX-Cloud#1222), built by +/// `dispatch` when `routing.stream_failure` is `mode: continue` and +/// fallback targets remain after the current one. Threaded into the +/// streaming sub-paths, which build the [`MidStreamPlan`] from it; +/// `None` keeps today's terminate behavior. +/// +/// [`MidStreamPlan`]: crate::stream_failover::MidStreamPlan +struct MidStreamArm { + cfg: aisix_core::StreamFailure, + remaining: Vec, + auth: AuthenticatedKey, + group: aisix_core::Model, + retry_on_429: bool, + fallback_on_statuses: Vec, + winner_attempt_index: u32, + winner_attempt_kind: &'static str, +} + #[allow(clippy::too_many_arguments)] async fn dispatch( state: &ProxyState, @@ -655,6 +676,16 @@ async fn dispatch( // pre-dispatch 4xx); if this endpoint ever grows semantic support, // widen this flag or the deferred policies are silently skipped. let is_routing_request = model_entry.value.routing.is_some(); + // Mid-stream failover config (AISIX-Cloud#1222): armed only for + // `mode: continue` on a routing request, mirroring chat.rs. + let stream_failure_cfg = model_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 mut routing = RoutingTelemetry::default(); // Walk targets, failing over to the next only on a retryable upstream @@ -741,6 +772,21 @@ async fn dispatch( } }; let attempt_started = Instant::now(); + // Arm mid-stream failover for this attempt when fallback + // targets remain after it (AISIX-Cloud#1222). + let mid_stream_arm = stream_failure_cfg.as_ref().and_then(|cfg| { + let remaining = &attempt_models[i + 1..]; + (!remaining.is_empty()).then(|| MidStreamArm { + cfg: cfg.clone(), + remaining: remaining.to_vec(), + auth: auth.clone(), + group: model_entry.value.clone(), + retry_on_429, + fallback_on_statuses: fallback_statuses.to_vec(), + winner_attempt_index: idx, + winner_attempt_kind: kind, + }) + }); match dispatch_to_target( state, &snapshot, @@ -767,6 +813,7 @@ async fn dispatch( &mut member_reservation, redactions_out.clone(), monitor_hits_out.clone(), + mid_stream_arm, ) .await { @@ -909,6 +956,9 @@ async fn dispatch_to_target( // Input-side monitor hits (AISIX-Cloud#562), same lifecycle as // `input_redactions`. input_monitor_hits: Vec, + // Mid-stream failover context (AISIX-Cloud#1222); consumed by the + // streaming branches of both sub-paths. + mid_stream: Option, ) -> Result { let model = &target.model; let pk_entry = crate::dispatch::resolve_provider_key(snapshot, model)?; @@ -937,6 +987,7 @@ async fn dispatch_to_target( member_reservation, input_redactions, input_monitor_hits, + mid_stream, ) .await; } @@ -964,6 +1015,7 @@ async fn dispatch_to_target( member_reservation, input_redactions, input_monitor_hits, + mid_stream, ) .await } @@ -1000,6 +1052,9 @@ async fn anthropic_passthrough_dispatch( member_reservation: &mut Option, input_redactions: crate::redact::RedactionCounts, input_monitor_hits: Vec, + // Mid-stream failover context (AISIX-Cloud#1222); `None` disarms — + // a mid-stream failure keeps today's truncate/forward behavior. + mid_stream: Option, ) -> Result { let mut body = body.clone(); let api_key = crate::dispatch::require_api_key(pk_value, model)?; @@ -1155,22 +1210,86 @@ async fn anthropic_passthrough_dispatch( // For SSE streaming: pass through the response body as a streaming // `text/event-stream` response. let headers = upstream_resp.headers().clone(); + // Mid-stream failover (AISIX-Cloud#1222): built up front because + // it decides who owns the read timeout below. The continuation + // dispatch runs on the internal ChatFormat — a body the parser + // can't map disarms (terminate as before). The parse runs on the + // outbound body (post-redaction, model already rewritten), same + // source the cross-provider path translates; the client-facing + // model name is restored explicitly. + let mid_stream_ctx = mid_stream.and_then(|arm| { + let mut chat = aisix_provider_anthropic::parse_inbound_request(&body).ok()?; + chat.model = model_name.to_string(); + aisix_provider_anthropic::translate_extras_to_openai_shape(&mut chat.extra); + let serving = Arc::new(std::sync::Mutex::new( + crate::stream_failover::ServingAttempt { + target_id: model_id.to_string(), + target_model: attempt.model.clone(), + provider: provider_label.clone(), + provider_key_id: pk_id.to_string(), + upstream_model: upstream_model.clone(), + cooldown: model.cooldown.clone(), + attempt_index: arm.winner_attempt_index, + attempt_kind: arm.winner_attempt_kind, + attempt_started, + }, + )); + let shared = crate::stream_failover::MidStreamShared::new(); + let plan = crate::stream_failover::MidStreamPlan { + cfg: arm.cfg, + endpoint: crate::stream_failover::MidStreamEndpoint::Messages, + structured_output: crate::stream_failover::expects_structured_output(&chat), + remaining: arm.remaining, + state: state.clone(), + auth: arm.auth, + group: arm.group, + req: chat, + request_id: request_id.to_string(), + client: client_ctx.clone(), + retry_on_429: arm.retry_on_429, + fallback_on_statuses: arm.fallback_on_statuses, + requested_model: model_name.to_string(), + api_key_id: api_key_id.to_string(), + applied_guardrails: resolved_chain.applied().to_vec(), + serving: Arc::clone(&serving), + shared: shared.clone(), + }; + Some((plan, serving, shared)) + }); // #554: enforce the per-chunk read timeout on the forwarded bytes. // When a `stream_timeout` is configured, peek the first byte so a // slow/erroring first token fails over (the caller loops to the next // target) before the 200 is committed; without one, forward directly // (pre-#554 behavior). A mid-stream stall truncates the forwarded - // stream — there is no in-band error frame for an opaque passthrough. + // stream — there is no in-band error frame for an opaque passthrough + // — unless mid-stream failover is armed, in which case the failover + // combinator owns the timeout (it must classify the stall as + // `read_timeout`, which the silent-truncate wrapper can't signal). let stream_budget = timeouts.stream; - let wrapped: std::pin::Pin> + Send>> = + let armed = mid_stream_ctx.is_some(); + let wrapped: std::pin::Pin> + Send>> = if armed + { + Box::pin(upstream_resp.bytes_stream()) + } else { Box::pin(crate::stream_timeout::with_read_timeout_bytes( upstream_resp.bytes_stream(), stream_budget, - )); + )) + }; let body_stream: std::pin::Pin> + Send>> = if timeouts.stream_configured { let mut wrapped = wrapped; - let first_bytes = match wrapped.next().await { + // The armed pipeline has no outer timeout wrapper — bound + // the first-byte peek explicitly so a stalled head still + // fails over pre-200 (identical outcome to the unarmed + // wrapper, which ends the stream on an elapsed timeout). + let first = match (armed, stream_budget) { + (true, Some(d)) => tokio::time::timeout(d, wrapped.next()) + .await + .unwrap_or_default(), + _ => wrapped.next().await, + }; + let first_bytes = match first { Some(Ok(b)) => b, Some(Err(e)) => { let err = crate::dispatch::reqwest_error_to_bridge(&e, send_started); @@ -1203,6 +1322,20 @@ async fn anthropic_passthrough_dispatch( } else { wrapped }; + let (mid_stream_plan, mid_stream_for_telem) = match mid_stream_ctx { + Some((plan, serving, shared)) => (Some(plan), Some((serving, shared))), + None => (None, None), + }; + let body_stream: std::pin::Pin> + Send>> = + match mid_stream_plan { + Some(plan) => Box::pin(wrap_anthropic_passthrough( + body_stream, + plan, + stream_budget, + model_name.to_string(), + )), + None => body_stream, + }; // Issue #245: parity with the OpenAI streaming fix (#225 / // #196). Pre-fix this path forwarded raw bytes and emitted a @@ -1279,6 +1412,9 @@ async fn anthropic_passthrough_dispatch( &upstream_model, crate::token_estimate::PromptInput::Anthropic(body.clone()), ); + let mid_stream_shared_for_parser = mid_stream_for_telem + .as_ref() + .map(|(_, shared)| shared.clone()); let parsed_stream = build_anthropic_passthrough_stream( body_stream, started, @@ -1287,16 +1423,87 @@ async fn anthropic_passthrough_dispatch( model_name.to_string(), content_cap, Some(estimator), + mid_stream_shared_for_parser, move |usage| { // Streaming responses that got this far are 200 — the // !status.is_success() guard above returned early on // upstream errors. // + // 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 + // (AISIX-Cloud#1222). + let ( + model_id_c, + provider_c, + provider_key_id_c, + upstream_model_c, + attempt_c, + attempt_elapsed, + mid_stream_fallbacks, + terminal_failure, + ) = match &mid_stream_for_telem { + Some((serving, shared)) => { + let s = serving.lock().expect("serving lock"); + ( + s.target_id.clone(), + s.provider.clone(), + s.provider_key_id.clone(), + s.upstream_model.clone(), + AttemptInfo { + index: s.attempt_index, + kind: s.attempt_kind.to_string(), + model: s.target_model.clone(), + ..Default::default() + }, + s.attempt_started.elapsed(), + shared.attempt_seq.load(Ordering::Relaxed), + shared + .terminal_failure + .lock() + .expect("terminal failure lock") + .clone(), + ) + } + None => ( + model_id_c, + provider_c, + provider_key_id_c, + upstream_model_c, + attempt_c, + attempt_started.elapsed(), + 0, + None, + ), + }; + // Logical stream outcome (AISIX-Cloud#1222) — same + // derivation as the typed pumps; the terminal-failure + // signal comes from the byte combinator via the shared + // slot (the parser below it only sees clean bytes). + let stream_failed = terminal_failure.is_some(); + let stream_outcome = if usage.guardrail_blocked || !usage.reached_end { + "" + } else if stream_failed { + "partial_failed" + } else if mid_stream_fallbacks > 0 { + "partial_recovered" + } else { + "success" + }; + if mid_stream_fallbacks > 0 { + // Recovered = the fallback kept the stream alive; + // only a terminal upstream error counts as failed. + state_c + .metrics + .record_mid_stream_fallback(&model_name_c, !stream_failed); + } // #688: apply the terminal token cost to TPM/TPD and release the // concurrency hold now the stream has ended. `add_tokens_post_stream` // is the sync analog of the reservation's async `commit_tokens` // (this end-of-stream closure can't await); dropping the hold frees // the concurrency slot(s) held for the stream's full lifetime. + // Mid-stream fallback targets that served part of this stream + // bill their TPM here too (AISIX-Cloud#1222). let streamed_tokens = total_tokens_with_cache( usage.prompt_tokens, usage.completion_tokens, @@ -1306,6 +1513,16 @@ async fn anthropic_passthrough_dispatch( for key in &post_stream_keys { limiter_c.add_tokens_post_stream(key, streamed_tokens); } + if let Some((_, shared)) = &mid_stream_for_telem { + for key in shared + .extra_post_stream_keys + .lock() + .expect("extra keys lock") + .iter() + { + limiter_c.add_tokens_post_stream(key, streamed_tokens); + } + } drop(stream_hold); // least_busy: stream over — this target is no longer // in-flight. @@ -1356,10 +1573,17 @@ async fn anthropic_passthrough_dispatch( }, // Attempt-scoped, unlike the e2e histogram above: any // failed attempt before this one emitted its own event. - attempt_started.elapsed(), + attempt_elapsed, metrics, &client_ctx_c, - attempt_c.clone(), + { + let mut attempt = attempt_c.clone(); + if let Some((class, message)) = &terminal_failure { + attempt.error_class = class.clone(); + attempt.error_message = message.clone(); + } + attempt + }, applied_guardrails_c.clone(), // #932: input-side mask counts captured before dispatch, // merged with the hold-back release's output-side counts. @@ -1373,6 +1597,7 @@ async fn anthropic_passthrough_dispatch( merged.extend(usage.monitor_hits); merged }, + stream_outcome, // Prompt captured up front; response assembled by the frame // parser into `usage.response_text`. Both gated on the cap. match (&captured_prompt_c, content_cap) { @@ -1741,6 +1966,9 @@ async fn cross_provider_dispatch( member_reservation: &mut Option, input_redactions: crate::redact::RedactionCounts, input_monitor_hits: Vec, + // Mid-stream failover context (AISIX-Cloud#1222); `None` disarms — + // a mid-stream failure keeps today's terminate behavior. + mid_stream: Option, ) -> Result { use aisix_gateway::Bridge; use aisix_provider_anthropic::{ @@ -1864,6 +2092,61 @@ async fn cross_provider_dispatch( state.health.record_success(model_name); state.runtime_status.mark_healthy(model_id); + // Mid-stream failover (AISIX-Cloud#1222): when armed, wrap the + // upstream so a qualifying mid-stream failure continues on the + // remaining targets inside the SAME client stream. The + // AnthropicSseEncoder lives OUTSIDE the combinator, so its + // message envelope + content-block state survive the switch and + // the spliced fallback chunks continue the same wire message. + let mid_stream_state = mid_stream.map(|arm| { + let serving = Arc::new(std::sync::Mutex::new( + crate::stream_failover::ServingAttempt { + target_id: model_id.to_string(), + target_model: attempt.model.clone(), + provider: provider_label.clone(), + provider_key_id: provider_key_id.to_string(), + upstream_model: upstream_model.clone(), + cooldown: model.cooldown.clone(), + attempt_index: arm.winner_attempt_index, + attempt_kind: arm.winner_attempt_kind, + attempt_started, + }, + )); + let shared = crate::stream_failover::MidStreamShared::new(); + (arm, serving, shared) + }); + let upstream = match &mid_stream_state { + Some((arm, serving, shared)) => crate::stream_failover::wrap( + upstream, + crate::stream_failover::MidStreamPlan { + cfg: arm.cfg.clone(), + endpoint: crate::stream_failover::MidStreamEndpoint::Messages, + structured_output: crate::stream_failover::expects_structured_output(&chat), + remaining: arm.remaining.clone(), + state: state.clone(), + auth: arm.auth.clone(), + group: arm.group.clone(), + req: chat.clone(), + request_id: request_id.to_string(), + client: client.clone(), + retry_on_429: arm.retry_on_429, + fallback_on_statuses: arm.fallback_on_statuses.clone(), + requested_model: model_name.to_string(), + api_key_id: api_key_id.to_string(), + applied_guardrails: resolved_chain.applied().to_vec(), + serving: Arc::clone(serving), + shared: shared.clone(), + }, + ), + None => upstream, + }; + let mid_stream_for_telem = mid_stream_state + .as_ref() + .map(|(_, serving, shared)| (Arc::clone(serving), shared.clone())); + let mid_stream_shared_for_pump = mid_stream_state + .as_ref() + .map(|(_, _, shared)| shared.clone()); + let message_id = format!("msg_{}", Uuid::new_v4().simple()); let encoder = AnthropicSseEncoder::new(message_id, model_name, 0); let state_for_telem = state.clone(); @@ -1937,11 +2220,76 @@ async fn cross_provider_dispatch( model_name.to_string(), content_cap, Some(estimator), + mid_stream_shared_for_pump, move |comp| { + // 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 + // (AISIX-Cloud#1222). + let ( + model_id_for_telem, + provider_for_telem, + provider_key_id_for_telem, + upstream_model_for_telem, + attempt_for_telem, + attempt_elapsed, + mid_stream_fallbacks, + ) = match &mid_stream_for_telem { + Some((serving, shared)) => { + let s = serving.lock().expect("serving lock"); + ( + s.target_id.clone(), + s.provider.clone(), + s.provider_key_id.clone(), + s.upstream_model.clone(), + AttemptInfo { + index: s.attempt_index, + kind: s.attempt_kind.to_string(), + model: s.target_model.clone(), + ..Default::default() + }, + s.attempt_started.elapsed(), + shared + .attempt_seq + .load(std::sync::atomic::Ordering::Relaxed), + ) + } + None => ( + model_id_for_telem, + provider_for_telem, + provider_key_id_for_telem, + upstream_model_for_telem, + attempt_for_telem, + attempt_started_for_telem.elapsed(), + 0, + ), + }; + // Logical stream outcome (AISIX-Cloud#1222) — the HTTP + // status froze at 200 when the head committed; this is + // the only signal separating "delivered in full" from + // "terminated mid-stream". Same derivation as chat.rs. + 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; only + // a terminal upstream error counts as failed (client + // aborts after a successful switch still recovered). + state_for_telem + .metrics + .record_mid_stream_fallback(&model_for_telem, !comp.stream_failed); + } // #688: apply the terminal token cost to TPM/TPD and release the // concurrency hold now the stream has ended (sync analog of the // reservation's async `commit_tokens`, which this closure can't - // await). + // await). Mid-stream fallback targets that served part of this + // stream bill their TPM here too (AISIX-Cloud#1222). let streamed_tokens = total_tokens_with_cache( comp.prompt_tokens, comp.completion_tokens, @@ -1951,6 +2299,16 @@ async fn cross_provider_dispatch( for key in &post_stream_keys { limiter_for_stream.add_tokens_post_stream(key, streamed_tokens); } + if let Some((_, shared)) = &mid_stream_for_telem { + for key in shared + .extra_post_stream_keys + .lock() + .expect("extra keys lock") + .iter() + { + limiter_for_stream.add_tokens_post_stream(key, streamed_tokens); + } + } drop(stream_hold); // least_busy: stream over — this target is no longer // in-flight. @@ -1998,10 +2356,17 @@ async fn cross_provider_dispatch( crate::CLIENT_CLOSED_REQUEST }, // Attempt-scoped — see the sibling passthrough path. - attempt_started_for_telem.elapsed(), + attempt_elapsed, metrics, &client_for_telem, - attempt_for_telem.clone(), + { + let mut attempt = attempt_for_telem.clone(); + if comp.stream_failed { + attempt.error_class = comp.stream_error_class.clone(); + attempt.error_message = comp.stream_error_message.clone(); + } + attempt + }, applied_guardrails_for_telem.clone(), // #932: input-side mask counts captured before dispatch, // merged with the hold-back release's output-side counts. @@ -2015,6 +2380,7 @@ async fn cross_provider_dispatch( merged.extend(comp.monitor_hits); merged }, + stream_outcome, // Prompt captured up front, response assembled across the // stream into `comp.response_text`; both gated on the cap. match (&captured_prompt_for_telem, content_cap) { @@ -2190,6 +2556,12 @@ fn build_anthropic_sse_stream( // Token-estimation fallback context (AISIX-Cloud#1074); see // `CompleteAnthropicStreamOnDrop::estimator`. estimator: Option, + // Mid-stream failover shared state (AISIX-Cloud#1222): the pump + // watches `attempt_seq` to reset its usage accumulators when the + // serving attempt changes, and folds `extra_usage` (failed + // partial attempts, estimated) into the encoder's closing usage. + // `None` on unarmed streams. + mid_stream: Option, on_complete: impl FnOnce(AnthropicStreamCompletion) + Send + 'static, ) -> axum::body::Body { use futures::StreamExt; @@ -2230,6 +2602,14 @@ fn build_anthropic_sse_stream( }; let mut upstream = upstream; let mut first_chunk_seen = false; + // Serving-attempt sequence snapshot (mid-stream failover, + // AISIX-Cloud#1222) — see chat.rs's pump for the why: max-wins + // folding across attempts would mix the failed attempt's usage + // into the serving attempt's totals. + let mut mid_stream_seq = mid_stream + .as_ref() + .map(|ms| ms.attempt_seq.load(std::sync::atomic::Ordering::Relaxed)) + .unwrap_or(0); // Accumulate assistant text for the end-of-stream output guardrail // (#448). Without a hold-back policy, bytes are forwarded live and // a blocked response is signalled with a terminal `error` event. @@ -2247,6 +2627,28 @@ fn build_anthropic_sse_stream( while let Some(item) = upstream.next().await { match item { Ok(chunk) => { + if let Some(ms) = mid_stream.as_ref() { + let seq = ms.attempt_seq.load(std::sync::atomic::Ordering::Relaxed); + if seq != mid_stream_seq { + mid_stream_seq = seq; + // Serving attempt changed: zero the + // per-attempt accumulators and fold the + // failed partials' estimated usage into the + // encoder's closing counts instead. + let comp = guard.comp(); + comp.prompt_tokens = 0; + comp.completion_tokens = 0; + comp.cache_creation_tokens = 0; + comp.cache_read_tokens = 0; + encoder.reset_usage_counters(); + let extra = ms + .extra_usage + .lock() + .expect("mid-stream extra lock") + .clone(); + encoder.set_extra_usage(extra.prompt_tokens, extra.completion_tokens); + } + } if !first_chunk_seen && chunk.delta.carries_generated_output() { first_chunk_seen = true; guard.comp().upstream_ttft_ms = @@ -2323,6 +2725,7 @@ fn build_anthropic_sse_stream( max_buffer_bytes = max_hold, "streaming /v1/messages response exceeded hold-back cap; failing closed", ); + guard.comp().guardrail_blocked = true; yield Ok(bytes::Bytes::from(guardrail_block_frame(None))); return; } @@ -2337,6 +2740,19 @@ fn build_anthropic_sse_stream( } } Err(e) => { + // Terminal upstream failure: the wire ends with an + // in-band error frame — distinct from a client + // abandon, so `reached_end` is still recorded and + // the telemetry closure reports `stream_outcome: + // partial_failed` (AISIX-Cloud#1222). + { + let comp = guard.comp(); + comp.stream_failed = true; + comp.stream_error_class = + crate::attempt::routing_error_class(&e).to_string(); + comp.stream_error_message = crate::attempt::attempt_error_message(&e); + comp.reached_end = true; + } // Hold-back: the held (unscanned) chunks are dropped — // fail closed; only the error frame reaches the client. let frame = format!( @@ -2428,6 +2844,7 @@ fn build_anthropic_sse_stream( reason = %reason, "guardrail blocked streaming /v1/messages response", ); + guard.comp().guardrail_blocked = true; // Hold-back: the held chunks are dropped — the matched // content never reached the wire. let frame = guardrail_block_frame(guardrail_name.as_deref()); @@ -2549,6 +2966,23 @@ struct AnthropicStreamCompletion { /// check (AISIX-Cloud#562). Merged with the input-side hits by the /// on_complete emit. monitor_hits: Vec, + /// The upstream errored mid-stream and the pump terminated the + /// response with an in-band error frame (AISIX-Cloud#1222). The + /// telemetry closure turns this into `stream_outcome: + /// partial_failed` — distinct from a client abandon, which leaves + /// `reached_end` false. + stream_failed: bool, + /// Attempt-taxonomy class of the terminal stream error; empty + /// unless `stream_failed`. + stream_error_class: String, + /// Client-safe message of the terminal stream error; empty unless + /// `stream_failed`. + stream_error_message: String, + /// The end-of-stream output guardrail blocked the response + /// (terminal error frame instead of clean completion). Keeps the + /// `stream_outcome` derivation from calling a blocked stream + /// `success` (AISIX-Cloud#1222). + guardrail_blocked: bool, } struct CompleteAnthropicStreamOnDrop { @@ -2705,6 +3139,11 @@ fn emit_anthropic_usage_event( // Monitor-mode guardrail observations (AISIX-Cloud#562), input + // output merged. guardrail_monitor_hits: Vec, + // Logical stream outcome (AISIX-Cloud#1222): `success` / + // `partial_failed` / `partial_recovered` on streaming terminal + // events; empty on non-streaming, aborted, and failed-attempt + // events. + stream_outcome: &str, content: Option, ) { // Per-PK telemetry attribution (#302 M17 / AISIX-Cloud#436). @@ -2760,6 +3199,7 @@ fn emit_anthropic_usage_event( applied_guardrails, redacted_entity_counts, guardrail_monitor_hits, + stream_outcome: stream_outcome.to_string(), ..Default::default() }; // Handler label "messages" — Anthropic /v1/messages inbound @@ -2909,6 +3349,10 @@ struct AnthropicStreamUsage { /// check (AISIX-Cloud#562). Merged with the input-side hits by the /// on_complete emit. monitor_hits: Vec, + /// The output guardrail blocked the response (terminal block frame + /// instead of clean bytes) — see the sibling field on + /// `AnthropicStreamCompletion`. + guardrail_blocked: bool, } /// Update the accumulator from one parsed SSE `data:` JSON object. @@ -3209,6 +3653,459 @@ impl Stream for AnthropicDeliveryCounter { } } +/// Side-channel wire state for the passthrough mid-stream failover +/// (AISIX-Cloud#1222): everything needed to decide whether the stream +/// is still continuation-safe and to seed a resumed +/// [`AnthropicSseEncoder`] that picks up the client's already-committed +/// message envelope. +/// +/// [`AnthropicSseEncoder`]: aisix_provider_anthropic::AnthropicSseEncoder +#[derive(Default)] +struct AnthropicWireTracker { + /// `message.id` from the forwarded `message_start`; `None` until + /// one arrives (a pre-envelope failure resumes with a fresh + /// encoder that emits its own `message_start`). + message_id: Option, + /// Index of the currently open text content block. + open_text_block: Option, + /// Next content-block index a resumed encoder may assign. + next_block_index: usize, + /// Delivered assistant text across the stream — the continuation + /// baseline. Bounded by `OUTPUT_ACCUMULATION_CAP`; + /// `partial_overflow` disarms past it (a faithful continuation + /// prompt can no longer be built). + partial: String, + partial_overflow: bool, + /// Wire shapes a fallback model cannot safely continue: tool_use / + /// thinking blocks, or the closing `message_delta` already told + /// the client the stop reason. Sticky. + unsafe_output: bool, + /// Clean terminal `message_stop` forwarded — EOF after this is a + /// normal end, not a truncation. + terminal_seen: bool, +} + +impl AnthropicWireTracker { + /// Update from one complete SSE frame's `data:` JSON. Returns the + /// in-band error when the frame is an `error` event — the caller + /// withholds that frame and runs the failover decision. + fn observe(&mut self, json: &Value) -> Option { + match json.get("type").and_then(Value::as_str) { + Some("message_start") => { + if let Some(id) = json + .get("message") + .and_then(|m| m.get("id")) + .and_then(Value::as_str) + { + self.message_id = Some(id.to_string()); + } + } + Some("content_block_start") => { + let index = json.get("index").and_then(Value::as_u64).unwrap_or(0) as usize; + self.next_block_index = self.next_block_index.max(index + 1); + match json + .get("content_block") + .and_then(|cb| cb.get("type")) + .and_then(Value::as_str) + { + Some("text") => self.open_text_block = Some(index), + // tool_use, thinking, redacted_thinking, and any + // future block type: not safely continuable. + _ => self.unsafe_output = true, + } + } + Some("content_block_delta") => { + if let Some(t) = json + .get("delta") + .and_then(|d| d.get("text")) + .and_then(Value::as_str) + { + if self.partial.len() + t.len() > crate::token_estimate::OUTPUT_ACCUMULATION_CAP + { + self.partial_overflow = true; + } else { + self.partial.push_str(t); + } + } + } + Some("content_block_stop") => { + let index = json.get("index").and_then(Value::as_u64).unwrap_or(0) as usize; + if self.open_text_block == Some(index) { + self.open_text_block = None; + } + } + Some("message_delta") => { + // The stop reason is on the client's wire — a + // continuation after it can't read as one message. + self.unsafe_output = true; + } + Some("message_stop") => self.terminal_seen = true, + Some("error") => { + let body = json + .get("error") + .and_then(|e| { + serde_json::from_value::< + aisix_provider_anthropic::AnthropicStreamErrorBody, + >(e.clone()) + .ok() + }) + .unwrap_or(aisix_provider_anthropic::AnthropicStreamErrorBody { + kind: None, + message: None, + }); + return Some(aisix_provider_anthropic::stream_error_into_bridge_error( + &body, + )); + } + _ => {} + } + None + } +} + +/// The terminal Anthropic-wire error frame the failover paths emit when +/// the stream cannot be continued — same shape as the cross-provider +/// pump's `Err` arm. +fn anthropic_error_frame(err: &aisix_gateway::BridgeError) -> Bytes { + Bytes::from(format!( + "event: error\ndata: {{\"type\":\"error\",\"error\":{{\"type\":\"{}\",\"message\":{}}}}}\n\n", + err.error_type(), + serde_json::to_string(&err.to_string()).unwrap_or_else(|_| "\"error\"".into()), + )) +} + +/// Wrap the passthrough byte stream with the mid-stream failover +/// combinator (AISIX-Cloud#1222). Forwarding is frame-granular: bytes +/// are held until a complete SSE frame is buffered, so an in-band +/// `error` frame can be withheld (never partially forwarded) while the +/// fallback decision runs. On a qualifying failure the remaining +/// targets are dispatched through the chat bridges and their chunks are +/// re-encoded onto the client's committed Anthropic envelope by a +/// resumed [`AnthropicSseEncoder`] — block indices and the message id +/// continue seamlessly. Ineligible failures reproduce today's behavior +/// (error item forwarded / silent truncation). +/// +/// The combinator owns the per-chunk read timeout when armed (the +/// silent-truncate byte wrapper upstream of it cannot signal +/// `read_timeout`). +/// +/// [`AnthropicSseEncoder`]: aisix_provider_anthropic::AnthropicSseEncoder +fn wrap_anthropic_passthrough( + upstream: S, + plan: crate::stream_failover::MidStreamPlan, + per_chunk: Option, + model_display_name: String, +) -> impl Stream> + Send +where + S: Stream> + Send + 'static, +{ + use aisix_provider_anthropic::AnthropicSseEncoder; + + /// One poll of the verbatim upstream, timeout folded in. + enum Polled { + Item(reqwest::Result), + Eof, + TimedOut, + } + + async_stream::stream! { + let mut upstream = std::pin::pin!(upstream); + let mut tracker = AnthropicWireTracker::default(); + let mut buf: Vec = Vec::new(); + let mut cursor = 0usize; + let mut used: u32 = 0; + // Frame parsing disarmed (cap overflow on a non-conformant + // stream): flush + forward raw items; no failover possible. + let mut parse_disarmed = false; + let send_started = Instant::now(); + // The serving fallback target's concurrency hold; replaced on + // every switch, dropped at generator end (#450 semantics). + let mut _fallback_hold: Option = None; + // Set when a qualifying failure acquired a fallback stream — + // the verbatim phase ends and the continuation phase drives the + // resumed encoder below. + let mut continuation: Option = None; + + // ---- Verbatim phase: forward frames, watch for failure ---- + while continuation.is_none() { + let polled = match per_chunk { + Some(d) => match tokio::time::timeout(d, upstream.next()).await { + Ok(Some(item)) => Polled::Item(item), + Ok(None) => Polled::Eof, + Err(_) => Polled::TimedOut, + }, + None => match upstream.next().await { + Some(item) => Polled::Item(item), + None => Polled::Eof, + }, + }; + // Resolve this poll into either forwarded bytes (continue) + // or a failure to run the fallback decision on. `original` + // keeps the raw reqwest error and `withheld` the raw + // in-band error frame, so an ineligible failure surfaces + // the exact bytes/item the unarmed path would. + let mut original: Option = None; + let mut withheld: Option> = None; + let failure: aisix_gateway::BridgeError = match polled { + Polled::Item(Ok(bytes)) => { + if parse_disarmed { + yield Ok(bytes); + continue; + } + buf.extend_from_slice(&bytes); + if buf.len() > MAX_SSE_FRAME_BUF_BYTES { + // Non-conformant stream (no frame terminator + // in 1 MiB): stop parsing, flush verbatim, + // disarm failover for the rest of the stream. + tracing::warn!( + buffered = buf.len(), + "anthropic passthrough failover: frame buffer exceeded cap; \ + disarming mid-stream failover for this stream" + ); + parse_disarmed = true; + yield Ok(Bytes::from(std::mem::take(&mut buf))); + continue; + } + // Forward complete frames; withhold an in-band + // `error` frame and run the failover decision. + let mut in_band: Option = None; + let mut forward = Vec::new(); + while let Some(end) = find_frame_end(&buf) { + let frame: Vec = buf.drain(..end).collect(); + if let Some(data) = extract_sse_data_line(&frame) { + if let Ok(json) = serde_json::from_slice::(data) { + if let Some(err) = tracker.observe(&json) { + in_band = Some(err); + withheld = Some(frame); + break; + } + } + } + forward.extend_from_slice(&frame); + } + if !forward.is_empty() { + yield Ok(Bytes::from(forward)); + } + match in_band { + Some(err) => err, + None => continue, + } + } + Polled::Item(Err(e)) => { + if parse_disarmed { + yield Err(e); + return; + } + let err = crate::dispatch::reqwest_error_to_bridge(&e, send_started); + original = Some(e); + err + } + Polled::Eof => { + if parse_disarmed || tracker.terminal_seen { + return; + } + // EOF without `message_stop`: the upstream broke + // the stream — same class the typed bridges map + // to `StreamAborted`. + aisix_gateway::BridgeError::StreamAborted + } + Polled::TimedOut => { + if parse_disarmed { + return; + } + aisix_gateway::BridgeError::Timeout { + elapsed_ms: per_chunk.map(|d| d.as_millis() as u64).unwrap_or_default(), + cause: String::new(), + } + } + }; + match handle_passthrough_failure(&plan, &mut cursor, &mut used, &tracker, failure) + .await + { + PassthroughFailure::Continue(stream, hold) => { + _fallback_hold = hold; + // Discard any incomplete trailing frame of the + // failed upstream — never forwarded, replaced by + // the continuation. + buf.clear(); + continuation = Some(stream); + } + PassthroughFailure::No(err, exhausted) => { + if exhausted { + // Fallbacks attempted and exhausted: terminate + // with the synthesized frame carrying the last + // failure, like the typed pumps. + yield Ok(anthropic_error_frame(&err)); + } else if let Some(frame) = withheld { + // Ineligible in-band error: release the + // upstream's own error frame byte-for-byte. + yield Ok(Bytes::from(frame)); + } else if let Some(e) = original { + // Ineligible transport error: forward the + // original item verbatim, like the unarmed path. + yield Err(e); + } + // Ineligible timeout / EOF: silent truncation, + // today's behavior. + return; + } + } + } + + // ---- Continuation phase: re-encode fallback chunks ---- + // Resume the client's committed envelope; a pre-envelope + // failure (no message_start forwarded yet) starts a fresh one + // instead — the encoder then emits its own message_start. + let mut encoder = match tracker.message_id.clone() { + Some(id) => AnthropicSseEncoder::resume( + id, + model_display_name.clone(), + tracker.open_text_block, + tracker.next_block_index, + ), + None => AnthropicSseEncoder::new( + format!("msg_{}", Uuid::new_v4().simple()), + model_display_name.clone(), + 0, + ), + }; + { + let extra = plan + .shared + .extra_usage + .lock() + .expect("mid-stream extra lock") + .clone(); + encoder.set_extra_usage(extra.prompt_tokens, extra.completion_tokens); + } + let mut stream = continuation.expect("continuation stream set above"); + loop { + match stream.next().await { + Some(Ok(chunk)) => { + // The same output-shape gates as the typed + // combinator, for any FURTHER fallback. + if chunk.delta.tool_calls.is_some() || chunk.delta.reasoning_content.is_some() + { + tracker.unsafe_output = true; + } + if let Some(t) = chunk.delta.content.as_deref() { + if tracker.partial.len() + t.len() + > crate::token_estimate::OUTPUT_ACCUMULATION_CAP + { + tracker.partial_overflow = true; + } else { + tracker.partial.push_str(t); + } + } + for ev in encoder.next_events(&chunk) { + yield Ok(Bytes::from(ev.to_sse_string())); + } + if encoder.is_finished() { + return; + } + } + Some(Err(err)) => { + match handle_passthrough_failure( + &plan, &mut cursor, &mut used, &tracker, err, + ) + .await + { + PassthroughFailure::Continue(next, hold) => { + _fallback_hold = hold; + stream = next; + // New serving attempt: zero the encoder's + // per-attempt counters and re-fold the + // grown failed-partial estimate. + encoder.reset_usage_counters(); + let extra = plan + .shared + .extra_usage + .lock() + .expect("mid-stream extra lock") + .clone(); + encoder.set_extra_usage(extra.prompt_tokens, extra.completion_tokens); + } + PassthroughFailure::No(err, exhausted) => { + if !exhausted { + // Ineligible mid-continuation failure — + // the wire is gateway-synthesized now, + // so record it like the typed pumps do. + *plan + .shared + .terminal_failure + .lock() + .expect("terminal failure lock") = Some(( + crate::attempt::routing_error_class(&err).to_string(), + crate::attempt::attempt_error_message(&err), + )); + } + yield Ok(anthropic_error_frame(&err)); + return; + } + } + } + None => { + for ev in encoder.force_finish() { + yield Ok(Bytes::from(ev.to_sse_string())); + } + return; + } + } + } + } +} + +/// Outcome of one passthrough failure decision. +enum PassthroughFailure { + /// A fallback target is streaming — splice it in. + Continue( + aisix_gateway::ChatChunkStream, + Option, + ), + /// No continuation. Carries the failure back (the last fallback + /// error on exhaustion, else the original) so the caller renders + /// phase-appropriate termination; `bool` = fallbacks were attempted + /// and exhausted (terminal failure already recorded). + No(aisix_gateway::BridgeError, bool), +} + +/// Run the eligibility gate + fallback acquisition for a passthrough +/// failure. Mirrors `stream_failover::wrap`'s decision order. +async fn handle_passthrough_failure( + plan: &crate::stream_failover::MidStreamPlan, + cursor: &mut usize, + used: &mut u32, + tracker: &AnthropicWireTracker, + err: aisix_gateway::BridgeError, +) -> PassthroughFailure { + let eligible = crate::stream_failover::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) + && !tracker.unsafe_output + && !tracker.partial_overflow + && !plan.structured_output + && *used < plan.cfg.max_fallbacks_or_default(); + if !eligible { + return PassthroughFailure::No(err, false); + } + match crate::stream_failover::acquire_fallback_stream(plan, cursor, used, err, &tracker.partial) + .await + { + Ok((stream, hold)) => PassthroughFailure::Continue(stream, hold), + Err(last) => { + *plan + .shared + .terminal_failure + .lock() + .expect("terminal failure lock") = Some(( + crate::attempt::routing_error_class(&last).to_string(), + crate::attempt::attempt_error_message(&last), + )); + PassthroughFailure::No(last, true) + } + } +} + /// Wrap an Anthropic upstream byte stream so token usage is parsed /// in-flight and `on_complete` fires once at end-of-stream (or /// client-disconnect) with the accumulated counts. Bytes are forwarded @@ -3230,6 +4127,11 @@ fn build_anthropic_passthrough_stream( // Token-estimation fallback context (AISIX-Cloud#1074); see // `AnthropicStreamGuard::estimator`. estimator: Option, + // Mid-stream failover shared state (AISIX-Cloud#1222): the parser + // resets its usage accumulators when the serving attempt changes, + // so the failed attempt's counters don't max-wins-mix into the + // serving attempt's totals. `None` on unarmed streams. + mid_stream: Option, on_complete: F, ) -> AnthropicDeliveryCounter> where @@ -3262,10 +4164,34 @@ where futures::pin_mut!(upstream); let mut buf: Vec = Vec::new(); let mut first_token_seen = false; + // Serving-attempt sequence snapshot (mid-stream failover, + // AISIX-Cloud#1222) — see chat.rs's pump for the why. + let mut mid_stream_seq = mid_stream + .as_ref() + .map(|ms| ms.attempt_seq.load(Ordering::Relaxed)) + .unwrap_or(0); // Whole-response hold-back buffer (BufferFull policies only). let mut held: Vec = Vec::new(); while let Some(item) = upstream.next().await { if let Ok(bytes) = &item { + 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; + // Serving attempt changed: zero the per-attempt + // usage accumulators; the resumed encoder's + // closing frame carries the serving attempt's + // counts plus the folded failed-partial + // estimate, and gets harvested below like any + // upstream frame. + let usage = guard.usage(); + usage.prompt_tokens = 0; + usage.completion_tokens = 0; + usage.cache_creation_tokens = 0; + usage.cache_read_tokens = 0; + usage.output_tokens_from_delta = false; + } + } // Side-channel parse: copy into the frame buffer (the // original `bytes` is yielded unchanged below) and drain // any complete SSE frames into the accumulator. @@ -3309,6 +4235,7 @@ where max_buffer_bytes = max_hold, "streaming /v1/messages passthrough exceeded hold-back cap; failing closed", ); + guard.usage().guardrail_blocked = true; yield Ok(Bytes::from(guardrail_block_frame(None))); return; } @@ -3420,6 +4347,7 @@ where "guardrail blocked streaming /v1/messages passthrough response", ); blocked = true; + guard.usage().guardrail_blocked = true; let frame = guardrail_block_frame(guardrail_name.as_deref()); yield Ok(Bytes::from(frame)); } diff --git a/crates/aisix-proxy/src/stream_failover.rs b/crates/aisix-proxy/src/stream_failover.rs index 1351e831..dd137a2b 100644 --- a/crates/aisix-proxy/src/stream_failover.rs +++ b/crates/aisix-proxy/src/stream_failover.rs @@ -1,5 +1,6 @@ -//! Mid-stream failover for `/v1/chat/completions` streaming -//! (AISIX-Cloud#1222, `routing.stream_failure: continue`). +//! Mid-stream failover for streaming responses (AISIX-Cloud#1222, +//! `routing.stream_failure: continue`) — `/v1/chat/completions`, +//! `/v1/messages`, and `/v1/responses`. //! //! 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 @@ -13,6 +14,15 @@ //! partial; Anthropic-wire targets consume the trailing assistant //! message as native prefill). //! +//! The endpoints' native-protocol passthrough legs (`/v1/messages` on +//! an Anthropic target, `/v1/responses` on an OpenAI target) forward +//! opaque bytes, not [`ChatChunk`]s — their byte-level combinators live +//! in their endpoint modules and reuse [`acquire_fallback_stream`] + +//! [`classify_trigger`] from here, re-encoding the fallback's chunks +//! onto the client's already-committed wire envelope. +//! +//! [`ChatChunk`]: aisix_gateway::ChatChunk +//! //! 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 @@ -33,11 +43,41 @@ 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: "; +/// `_build_responses_continuation_input` — the one continuation surface +/// LiteLLM still ships, after upstream removed the chat-completions +/// one) — kept byte-identical so the two gateways' fallback models +/// receive the same steering. LiteLLM sends it as a `developer` turn on the +/// Responses input; our internal shape carries it as a system message, +/// which each provider bridge maps to its own instruction tier. The +/// partial text is NOT interpolated here; it rides the assistant +/// message that follows. +const CONTINUATION_SYSTEM_PROMPT: &str = "The previous assistant response was interrupted mid-stream. Continue exactly where it stopped — do not repeat any of its content. Your response must read as a seamless continuation."; + +/// Which endpoint's stream the combinator is serving. Selects the +/// UsageEvent stamps (`inbound_protocol` + the usage-sink handler +/// label) for the per-attempt events emitted on a mid-stream switch, so +/// they land in telemetry alongside the endpoint's own events. +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +pub(crate) enum MidStreamEndpoint { + Chat, + Messages, +} + +impl MidStreamEndpoint { + pub(crate) fn sink_label(self) -> &'static str { + match self { + Self::Chat => "chat", + Self::Messages => "messages", + } + } + + pub(crate) fn inbound_protocol(self) -> &'static str { + match self { + Self::Chat => "openai", + Self::Messages => "anthropic", + } + } +} /// The serving attempt behind the live client stream. Starts as the /// pre-stream winner; rewritten by the combinator on every mid-stream @@ -65,6 +105,15 @@ pub(crate) struct ServingAttempt { /// keep the request's telemetry coherent while doing so. pub(crate) struct MidStreamPlan { pub cfg: StreamFailure, + /// Endpoint stamps for the per-attempt events (see + /// [`MidStreamEndpoint`]). + pub endpoint: MidStreamEndpoint, + /// Whether the request pins the output to a structured shape the + /// continuation cannot safely extend. Computed by the caller from + /// its own inbound request model (`response_format` on the chat + /// shape, `text.format` on the Responses shape; the Anthropic shape + /// has no equivalent). + pub structured_output: bool, /// Targets after the pre-stream winner, in strategy order. pub remaining: Vec, pub state: ProxyState, @@ -112,6 +161,13 @@ pub(crate) struct MidStreamShared { /// it bills the pre-stream reservation's keys (#450 / #1087 /// family). pub extra_post_stream_keys: Arc>>, + /// Terminal stream failure `(error_class, error_message)` recorded + /// by a byte-level (passthrough) combinator, which ends the wire + /// with a synthesized protocol error frame the downstream builder + /// cannot distinguish from clean bytes — unlike the typed pumps, + /// which observe the `Err` item directly. `None` = no terminal + /// failure. + pub terminal_failure: Arc>>, } impl MidStreamShared { @@ -120,6 +176,7 @@ impl MidStreamShared { 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())), + terminal_failure: Arc::new(Mutex::new(None)), } } } @@ -189,7 +246,7 @@ pub(crate) fn wrap( // 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 unsafe_output = plan.structured_output; let mut used: u32 = 0; let mut cursor = 0usize; let max = plan.cfg.max_fallbacks_or_default(); @@ -256,7 +313,7 @@ pub(crate) fn wrap( /// 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) { +pub(crate) 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"); @@ -315,6 +372,7 @@ fn finalize_failed_attempt(plan: &MidStreamPlan, err: &BridgeError, partial: &st } crate::chat::emit_mid_stream_failed_attempt( &plan.state, + plan.endpoint, &plan.request_id, &plan.requested_model, &plan.api_key_id, @@ -342,7 +400,7 @@ fn finalize_failed_attempt(plan: &MidStreamPlan, err: &BridgeError, partial: &st /// 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( +pub(crate) async fn acquire_fallback_stream( plan: &MidStreamPlan, cursor: &mut usize, used: &mut u32, @@ -436,6 +494,7 @@ async fn acquire_fallback_stream( }; crate::chat::emit_mid_stream_failed_attempt( &plan.state, + plan.endpoint, &plan.request_id, &plan.requested_model, &plan.api_key_id, @@ -538,6 +597,7 @@ async fn acquire_fallback_stream( } crate::chat::emit_mid_stream_failed_attempt( &plan.state, + plan.endpoint, &plan.request_id, &plan.requested_model, &plan.api_key_id, 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 6bcfb273..da9311cd 100644 --- a/tests/e2e/src/cases/mid-stream-fallback-e2e.test.ts +++ b/tests/e2e/src/cases/mid-stream-fallback-e2e.test.ts @@ -188,7 +188,7 @@ describe("mid-stream fallback e2e", () => { 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", + "do not repeat any of its content", ); expect(body.messages[2].role).toBe("assistant"); expect(body.messages[2].content).toBe("Once upon "); diff --git a/tests/e2e/src/cases/mid-stream-fallback-messages-e2e.test.ts b/tests/e2e/src/cases/mid-stream-fallback-messages-e2e.test.ts new file mode 100644 index 00000000..b94090cc --- /dev/null +++ b/tests/e2e/src/cases/mid-stream-fallback-messages-e2e.test.ts @@ -0,0 +1,421 @@ +import { createHash } from "node:crypto"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; +import { + EtcdClient, + SeedClient, + spawnApp, + startOpenAiUpstream, + waitConfigPropagation, + type OpenAiUpstream, + type OpenAiUpstreamOptions, + type SpawnedApp, +} from "../harness/index.js"; + +const CALLER_PLAINTEXT = "sk-mid-stream-messages-caller"; +const CALLER_KEY_HASH = createHash("sha256") + .update(CALLER_PLAINTEXT) + .digest("hex"); + +// Anthropic-wire SSE frames (verbatim, incl. framing) for the +// passthrough upstream mocks. +const frame = (event: string, data: Record) => + `event: ${event}\ndata: ${JSON.stringify(data)}\n\n`; + +const MESSAGE_START = frame("message_start", { + type: "message_start", + message: { + id: "msg_up1", + type: "message", + role: "assistant", + content: [], + model: "claude-3-5-haiku-20241022", + stop_reason: null, + usage: { input_tokens: 5, output_tokens: 1 }, + }, +}); +const BLOCK_START = frame("content_block_start", { + type: "content_block_start", + index: 0, + content_block: { type: "text", text: "" }, +}); +const delta = (text: string) => + frame("content_block_delta", { + type: "content_block_delta", + index: 0, + delta: { type: "text_delta", text }, + }); +const RECOVERY_FRAMES = [ + frame("message_start", { + type: "message_start", + message: { + id: "msg_up2", + type: "message", + role: "assistant", + content: [], + model: "claude-3-5-haiku-20241022", + stop_reason: null, + usage: { input_tokens: 6, output_tokens: 0 }, + }, + }), + BLOCK_START, + delta(" a time."), + frame("content_block_stop", { type: "content_block_stop", index: 0 }), + frame("message_delta", { + type: "message_delta", + delta: { stop_reason: "end_turn", stop_sequence: null }, + usage: { output_tokens: 7 }, + }), + frame("message_stop", { type: "message_stop" }), +]; + +// AISIX-Cloud#1222 phase 2: mid-stream fallback on /v1/messages. +// Covers the transports the Rust wiremock suite cannot simulate — a +// real mid-body connection drop and a real inter-frame stall on the +// Anthropic PASSTHROUGH leg (byte-verbatim, resumed-encoder splice), +// plus the cross-protocol recovery (Anthropic-wire head, OpenAI-wire +// fallback) and the client-cancel non-trigger. +describe("mid-stream fallback /v1/messages 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 anthropicUpstream( + opts: OpenAiUpstreamOptions, + ): Promise { + const upstream = await startOpenAiUpstream(opts); + upstreams.push(upstream); + return upstream; + } + + async function createAnthropicTarget( + 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-ant-mock", + api_base: upstream.baseUrl, + provider: "anthropic", + adapter: "anthropic", + }); + await seed.createModel({ + display_name: displayName, + provider: "anthropic", + model_name: "claude-3-5-haiku-20241022", + provider_key_id: providerKey.id, + }); + } + + async function createOpenAiTarget( + 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; + }); + } + + async function streamMessages( + model: string, + signal?: AbortSignal, + ): Promise { + const res = await fetch(`${app!.proxyUrl}/v1/messages`, { + method: "POST", + headers: { + authorization: `Bearer ${CALLER_PLAINTEXT}`, + "content-type": "application/json", + }, + body: JSON.stringify({ + model, + max_tokens: 128, + stream: true, + messages: [{ role: "user", content: "tell me a story" }], + }), + signal, + }); + expect(res.status).toBe(200); + return res; + } + + async function readAll(res: Response): Promise { + const reader = res.body!.getReader(); + const decoder = new TextDecoder(); + let wire = ""; + for (;;) { + const { done, value } = await reader.read(); + if (done) break; + wire += decoder.decode(value, { stream: true }); + } + return wire; + } + + test( + "passthrough: connection drop mid-stream continues the client envelope on the fallback target", + { timeout: 30_000 }, + async (ctx) => { + if (!etcdReachable || !app || !seed) { + ctx.skip(); + return; + } + // eventDelayMs keeps a real inter-frame gap so the delivered + // frames are flushed (and ACKed) before the RST — without it the + // kernel discards receiver-buffered data on destroy() and the + // failure degrades into a pre-first-frame abort. + const primary = await anthropicUpstream({ + rawStreamFrames: [MESSAGE_START, BLOCK_START, delta("Once upon")], + eventDelayMs: 200, + disconnectAfterEvents: 3, + }); + const secondary = await anthropicUpstream({ + rawStreamFrames: RECOVERY_FRAMES, + }); + await createAnthropicTarget("msf-msg-drop-a", primary); + await createAnthropicTarget("msf-msg-drop-b", secondary); + await createGroup( + "msf-msg-drop", + ["msf-msg-drop-a", "msf-msg-drop-b"], + { mode: "continue" }, + ); + await waitSeedApplied("msf-msg-drop"); + + const wire = await readAll(await streamMessages("msf-msg-drop")); + // One committed envelope across the switch, text block continues. + expect(wire.match(/event: message_start/g)).toHaveLength(1); + expect(wire.match(/event: content_block_start/g)).toHaveLength(1); + expect(wire).toContain("Once upon"); + expect(wire).toContain(" a time."); + expect(wire).not.toContain("event: error"); + expect(wire.match(/event: message_stop/g)).toHaveLength(1); + + // The fallback saw the continuation: instruction + the partial as + // the trailing (prefill) assistant message on the Anthropic wire. + const cont = secondary.receivedRequests.find((r) => + r.path.includes("/messages"), + ); + expect(cont).toBeDefined(); + expect(cont!.body).toContain("interrupted mid-stream"); + const body = JSON.parse(cont!.body) as { + messages: Array<{ role: string }>; + }; + expect(body.messages.at(-1)?.role).toBe("assistant"); + }, + ); + + test( + "passthrough: inter-frame stall (read timeout) fails over inside the stream", + { timeout: 30_000 }, + async (ctx) => { + if (!etcdReachable || !app || !seed) { + ctx.skip(); + return; + } + // Head flows quickly, then the upstream hangs without closing — + // the armed combinator owns the read timeout and classifies the + // stall as `read_timeout`. + const primary = await anthropicUpstream({ + rawStreamFrames: [ + MESSAGE_START, + BLOCK_START, + delta("Once upon"), + delta(" (never sent)"), + ], + eventDelayMs: 100, + stallAfterEvents: 3, + }); + const secondary = await anthropicUpstream({ + rawStreamFrames: RECOVERY_FRAMES, + }); + await createAnthropicTarget("msf-msg-stall-a", primary); + await createAnthropicTarget("msf-msg-stall-b", secondary); + await createGroup( + "msf-msg-stall", + ["msf-msg-stall-a", "msf-msg-stall-b"], + { mode: "continue", on: ["read_timeout"] }, + { stream_timeout: 1500 }, + ); + await waitSeedApplied("msf-msg-stall"); + + const wire = await readAll(await streamMessages("msf-msg-stall")); + expect(wire.match(/event: message_start/g)).toHaveLength(1); + expect(wire).toContain("Once upon"); + expect(wire).toContain(" a time."); + expect(wire).not.toContain("event: error"); + expect(wire.match(/event: message_stop/g)).toHaveLength(1); + }, + ); + + test( + "cross-protocol: anthropic-wire head resumes on an OpenAI-protocol fallback", + { timeout: 30_000 }, + async (ctx) => { + if (!etcdReachable || !app || !seed) { + ctx.skip(); + return; + } + const primary = await anthropicUpstream({ + rawStreamFrames: [ + MESSAGE_START, + BLOCK_START, + delta("Once upon"), + frame("error", { + type: "error", + error: { type: "overloaded_error", message: "Overloaded" }, + }), + ], + eventDelayMs: 100, + }); + const chunk = (content: string, finish: string | null = null) => + JSON.stringify({ + id: "up-2", + object: "chat.completion.chunk", + created: 1, + model: "gpt-4o-mini", + choices: [{ index: 0, delta: { content }, finish_reason: finish }], + }); + const secondary = await startOpenAiUpstream({ + streamEvents: [chunk(" a time."), chunk("", "stop"), "[DONE]"], + }); + upstreams.push(secondary); + await createAnthropicTarget("msf-msg-xp-a", primary); + await createOpenAiTarget("msf-msg-xp-b", secondary); + await createGroup("msf-msg-xp", ["msf-msg-xp-a", "msf-msg-xp-b"], { + mode: "continue", + }); + await waitSeedApplied("msf-msg-xp"); + + const wire = await readAll(await streamMessages("msf-msg-xp")); + expect(wire.match(/event: message_start/g)).toHaveLength(1); + expect(wire).toContain("Once upon"); + expect(wire).toContain(" a time."); + expect(wire).not.toContain("event: error"); + expect(wire.match(/event: message_stop/g)).toHaveLength(1); + + // OpenAI-wire continuation body: instruction + partial assistant. + const cont = secondary.receivedRequests.find((r) => + r.path.includes("/chat/completions"), + ); + expect(cont).toBeDefined(); + const body = JSON.parse(cont!.body) as { + messages: Array<{ role: string; content: string }>; + }; + const last = body.messages.at(-1); + expect(last?.role).toBe("assistant"); + expect(last?.content).toContain("Once upon"); + expect( + body.messages.some( + (m) => + m.role === "system" && + m.content.includes("do not repeat any of its content"), + ), + ).toBe(true); + }, + ); + + test( + "client cancel mid-stream never dispatches the fallback target", + { timeout: 30_000 }, + async (ctx) => { + if (!etcdReachable || !app || !seed) { + ctx.skip(); + return; + } + const primary = await anthropicUpstream({ + rawStreamFrames: [ + MESSAGE_START, + BLOCK_START, + delta("Once upon"), + delta(" a slow"), + delta(" stream"), + ], + eventDelayMs: 500, + }); + const secondary = await anthropicUpstream({ + rawStreamFrames: RECOVERY_FRAMES, + }); + await createAnthropicTarget("msf-msg-cancel-a", primary); + await createAnthropicTarget("msf-msg-cancel-b", secondary); + await createGroup( + "msf-msg-cancel", + ["msf-msg-cancel-a", "msf-msg-cancel-b"], + { mode: "continue" }, + ); + await waitSeedApplied("msf-msg-cancel"); + + const abort = new AbortController(); + const res = await streamMessages("msf-msg-cancel", abort.signal); + const reader = res.body!.getReader(); + // Read a couple of chunks, then walk away mid-stream. + await reader.read(); + await reader.read(); + abort.abort(); + // Give any (incorrect) fallback dispatch time to happen. + await new Promise((r) => setTimeout(r, 2_000)); + expect( + secondary.receivedRequests.filter((r) => r.path.includes("/messages")), + ).toHaveLength(0); + }, + ); +}); diff --git a/tests/e2e/src/harness/index.ts b/tests/e2e/src/harness/index.ts index f5bd5ce6..8fd21072 100644 --- a/tests/e2e/src/harness/index.ts +++ b/tests/e2e/src/harness/index.ts @@ -3,7 +3,7 @@ export { AdminClient, waitConfigPropagation, awaitWindowHeadroom } from "./admin export { ProxyClient } from "./proxy.js"; export { EtcdClient } from "./etcd.js"; export { SeedClient } from "./seed.js"; -export { startOpenAiUpstream, type OpenAiUpstream, type ReceivedRequest } from "./upstream-openai.js"; +export { startOpenAiUpstream, type OpenAiUpstream, type OpenAiUpstreamOptions, type ReceivedRequest } from "./upstream-openai.js"; export { startMcpUpstream, type McpUpstream } from "./upstream-mcp.js"; export { startRestUpstream, type RestUpstream } from "./upstream-rest.js"; export { pickFreePort, pickFreePorts } from "./ports.js"; diff --git a/tests/e2e/src/harness/upstream-openai.ts b/tests/e2e/src/harness/upstream-openai.ts index c32a2098..f2504757 100644 --- a/tests/e2e/src/harness/upstream-openai.ts +++ b/tests/e2e/src/harness/upstream-openai.ts @@ -31,6 +31,22 @@ export interface OpenAiUpstreamOptions { errorContentType?: string; /** Drop the connection after writing this many SSE events. */ disconnectAfterEvents?: number; + /** + * Stall (indefinitely, until the peer goes away) after writing this + * many SSE events, WITHOUT closing the socket — models a hung + * upstream mid-stream, i.e. the read-timeout scenario, as opposed to + * `disconnectAfterEvents` (connection drop). + */ + stallAfterEvents?: number; + /** + * Pre-formed SSE frames written verbatim — each entry must carry its + * own `event:`/`data:` lines and trailing blank line. For + * Anthropic-wire mocks (`event: message_start` etc.), which the + * `data:`-only `streamEvents` shape can't express. Same + * `eventDelayMs`/`disconnectAfterEvents` semantics, counted per + * frame. Takes precedence over `streamEvents`. + */ + rawStreamFrames?: string[]; /** * Raw (non-JSON) 200 response body — e.g. MP4 bytes for the `/v1/videos` * content-proxy path. When set (and `status` < 400), the reply is these @@ -71,6 +87,10 @@ export interface OpenAiUpstreamStep { /** Content-Type for the error body (default `application/json`). See #543. */ errorContentType?: string; disconnectAfterEvents?: number; + /** See `OpenAiUpstreamOptions.stallAfterEvents`. */ + stallAfterEvents?: number; + /** See `OpenAiUpstreamOptions.rawStreamFrames`. */ + rawStreamFrames?: string[]; /** Extra response headers, same semantics as on the top-level options. */ responseHeaders?: Record; /** Raw (non-JSON) 200 body — see `OpenAiUpstreamOptions.rawBody`. */ @@ -168,7 +188,7 @@ export async function startOpenAiUpstream( return; } - const isStream = !!step.streamEvents; + const isStream = !!step.streamEvents || !!step.rawStreamFrames; if (isStream) { res.statusCode = 200; res.setHeader("content-type", "text/event-stream"); @@ -178,8 +198,12 @@ export async function startOpenAiUpstream( // first token (TTFT timeout) independently of the headers (#554). res.flushHeaders(); if (step.firstEventDelayMs) await sleep(step.firstEventDelayMs); - const events = step.streamEvents ?? []; - for (let i = 0; i < events.length; i++) { + // Raw frames are written verbatim (they carry their own SSE + // framing); `streamEvents` entries get the `data:` framing here. + const frames = + step.rawStreamFrames ?? + (step.streamEvents ?? []).map((e) => `data: ${e}\n\n`); + for (let i = 0; i < frames.length; i++) { // The gateway may have abandoned a stalled stream (#554 read // timeout) and closed the connection; stop writing rather than // throwing on a dead socket. @@ -191,7 +215,16 @@ export async function startOpenAiUpstream( res.destroy(); return; } - res.write(`data: ${events[i]}\n\n`); + if ( + step.stallAfterEvents !== undefined && + i >= step.stallAfterEvents + ) { + // Hang without closing: the socket stays open until the + // gateway abandons it (read timeout) or the test ends. + await sleep(600_000); + return; + } + res.write(frames[i]); if (step.eventDelayMs) await sleep(step.eventDelayMs); } if (!res.writableEnded && !res.destroyed) res.end();