feat(proxy): mid-stream fallback for /v1/messages — both dispatch legs

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,
 };

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<String>,
+        model_display_name: impl Into<String>,
+        open_text_block: Option<usize>,
+        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<AnthropicSseEvent> {
@@ -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<AnthropicSseEvent> {
         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 {

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<CapturedContent>,
+) {
+    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<CapturedContent>,
 ) {
     // 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,