From ad0253dc0704aa55d5212737fc99e5f59c25f130 Mon Sep 17 00:00:00 2001 From: Ming Wen Date: Mon, 15 Jun 2026 12:41:51 +0800 Subject: [PATCH] fix(proxy): port ensemble streaming to the #607 MultiReservation API #606 (ensemble streaming) was written against the pre-#607 reservation API and merged onto a main that already had #607's cluster-Redis change, which dropped MultiReservation's generics/lifetime, made commit_tokens async, and made into_stream_hold arg-less. The semantic merge conflict broke main's build (no textual conflict, so it merged CLEAN). Port the streaming-ensemble reservation code to the new API: - MultiReservation<'a, SystemClock> -> MultiReservation. - into_stream_hold(arg) -> into_stream_hold(). - The error-exit billing helper was a sync closure calling the now-async commit_tokens with no .await, so the commit was a dropped future and panel tokens were never billed on streaming error exits. Split into a sync survivor_total + emit_panel_then_fail helper and .await commit_tokens inline at each of the 6 exits (mirroring the non-streaming path). One commit per request, every survivor billed once -- now actually executed. Verified: cargo check + clippy -D warnings clean, 263 core + 440 proxy tests green (incl. both streaming e2e), schema-drift gate clean. Refs api7/ai-gateway#601. --- crates/aisix-proxy/src/chat.rs | 72 +++++++++++++++++----------------- 1 file changed, 36 insertions(+), 36 deletions(-) diff --git a/crates/aisix-proxy/src/chat.rs b/crates/aisix-proxy/src/chat.rs index 4545246e..a05e808d 100644 --- a/crates/aisix-proxy/src/chat.rs +++ b/crates/aisix-proxy/src/chat.rs @@ -1977,19 +1977,25 @@ async fn dispatch_ensemble( // path, which likewise holds the socket bytes-free across panel + judge; // warming it would need an early SSE handshake before the panel completes. if req.is_streaming() { - // Commit + emit the already-billed panel survivors, then fail. Shared - // by both streaming error exits (panel exhausted, or judge connect - // failed): in either case every surviving panel member round-tripped - // an upstream and must be billed, exactly like the non-streaming error - // paths. `charge: None` — the per-member events already carry the bill. - let bill_panel_and_fail = - |panel: Vec, - reservation: aisix_ratelimit::MultiReservation<'a, aisix_ratelimit::SystemClock>, - proxy_err: ProxyError| - -> DispatchFailure { - let survivor_total: u64 = - panel.iter().map(|p| u64::from(p.usage.total_tokens)).sum(); - reservation.commit_tokens(survivor_total); + // Emit the already-billed panel survivors, then build the failure. + // Shared by both streaming error exits (panel exhausted, or judge + // connect failed): in either case every surviving panel member + // round-tripped an upstream and must be billed, exactly like the + // non-streaming error paths. `charge: None` — the per-member events + // already carry the bill. + // + // The token commit is NOT done here: `MultiReservation::commit_tokens` + // is async (#607's cluster-Redis path), and this helper is sync so + // `emit_panel_member` stays callable across every exit. Each call site + // therefore `.await`s `reservation.commit_tokens(survivor_total)` + // inline FIRST (mirroring the non-streaming ensemble path), then calls + // this for the emit + `DispatchFailure`. `panel` is borrowed so the + // call site still owns it to compute `survivor_total`. + let survivor_total = |panel: &[crate::ensemble::PanelOutcome]| -> u64 { + panel.iter().map(|p| u64::from(p.usage.total_tokens)).sum() + }; + let emit_panel_then_fail = + |panel: &[crate::ensemble::PanelOutcome], proxy_err: ProxyError| -> DispatchFailure { for (index, member) in panel.iter().enumerate() { emit_panel_member( member, index, /* blocked */ false, /* bypass */ "", @@ -2005,9 +2011,9 @@ async fn dispatch_ensemble( match crate::ensemble::run_ensemble_panel(req, ensemble_cfg, &caller).await { Ok(triple) => triple, Err(crate::ensemble::EnsembleError::InsufficientPanel { panel, .. }) => { - return Err(bill_panel_and_fail( - panel, - reservation, + reservation.commit_tokens(survivor_total(&panel)).await; + return Err(emit_panel_then_fail( + &panel, ProxyError::Bridge(BridgeError::upstream_status( 502, "ensemble panel did not reach the required number of responses", @@ -2017,11 +2023,8 @@ async fn dispatch_ensemble( // `run_ensemble_panel` never returns `Judge` (it stops before // the judge call), but the enum is non-exhaustive to us here. Err(crate::ensemble::EnsembleError::Judge { panel, source }) => { - return Err(bill_panel_and_fail( - panel, - reservation, - ProxyError::Bridge(source), - )); + reservation.commit_tokens(survivor_total(&panel)).await; + return Err(emit_panel_then_fail(&panel, ProxyError::Bridge(source))); } }; @@ -2035,9 +2038,9 @@ async fn dispatch_ensemble( // snapshot, and the judge is required config), but stay total: bill the // panel + fail rather than panic if the entry vanished mid-request. let Some(judge_entry) = snapshot.models.get_by_name(&ensemble_cfg.judge.model) else { - return Err(bill_panel_and_fail( - panel, - reservation, + reservation.commit_tokens(survivor_total(&panel)).await; + return Err(emit_panel_then_fail( + &panel, ProxyError::Bridge(BridgeError::InvalidUpstreamConfig( "ensemble judge references an unknown model".into(), )), @@ -2048,9 +2051,9 @@ async fn dispatch_ensemble( Ok(pk) => pk, Err(e) => { tracing::warn!(error = %e, "ensemble judge provider key unresolved"); - return Err(bill_panel_and_fail( - panel, - reservation, + reservation.commit_tokens(survivor_total(&panel)).await; + return Err(emit_panel_then_fail( + &panel, ProxyError::Bridge(BridgeError::InvalidUpstreamConfig( "ensemble judge has an unresolved provider key".into(), )), @@ -2059,9 +2062,9 @@ async fn dispatch_ensemble( }; let Some(judge_bridge) = crate::dispatch::resolve_bridge(&state.hub, &judge_pk.value) else { - return Err(bill_panel_and_fail( - panel, - reservation, + reservation.commit_tokens(survivor_total(&panel)).await; + return Err(emit_panel_then_fail( + &panel, ProxyError::Bridge(BridgeError::Config( "ensemble judge has no registered bridge".into(), )), @@ -2082,11 +2085,8 @@ async fn dispatch_ensemble( // panel (same invariant as the non-streaming judge-failure path), // then surface the bridge error (5xx → 502, 4xx preserved). Err(be) => { - return Err(bill_panel_and_fail( - panel, - reservation, - ProxyError::Bridge(be), - )); + reservation.commit_tokens(survivor_total(&panel)).await; + return Err(emit_panel_then_fail(&panel, ProxyError::Bridge(be))); } }; @@ -2156,7 +2156,7 @@ async fn dispatch_ensemble( // Hold concurrency for the stream's full lifetime (#450). Snapshot the // keys BEFORE consuming the reservation into the owned guard. let post_stream_keys = reservation.keys(); - let stream_concurrency_hold = reservation.into_stream_hold(Arc::clone(&state.limiter)); + let stream_concurrency_hold = reservation.into_stream_hold(); let limiter = Arc::clone(&state.limiter); let sse_stream = build_sse_stream(