diff --git a/crates/aisix-proxy/AGENTS.md b/crates/aisix-proxy/AGENTS.md index 489241ad..fcf104b1 100644 --- a/crates/aisix-proxy/AGENTS.md +++ b/crates/aisix-proxy/AGENTS.md @@ -20,3 +20,26 @@ silently uncorrelated, which reads exactly like working code: Do not hold a span guard across an await to work around this — it leaks the span onto whatever the executor runs next on that thread. + +## A per-model gate must say whether it binds the requested entry or each target + +`resolve_attempt_models` expands a routing model into targets, so `model_entry` / +`virtual_entry` is the **group**, which carries none of a member's config. A gate +written against it silently never runs for group traffic, and nothing errors — +requests keep succeeding on a target that should have been excluded. + +Decide, and encode the decision at the call site: + +- **Binds each target** (anything protecting the upstream behind it — rate limits, + cooldown, health, timeouts): resolve it from the attempt model *inside* the + dispatch loop, in all four group-capable endpoints (chat, messages, + count_tokens, responses) and in both the streaming and non-streaming branches. + A limit-shaped gate should skip the target and let dispatch continue rather than + failing the whole request — see `quota::reserve_routing_target`. +- **Binds the requested entry** (anything scoped to the alias the caller named — + `allowed_cidrs`, guardrail attachment): keep it pre-dispatch, and say so in the + user-facing docs, because the group/member split is otherwise invisible. + +A reservation-shaped gate additionally must not double-charge: `reserve_routing_target` +returns `None` for non-routing dispatch, whose model layers the pre-dispatch +`quota::enforce*` already reserved. diff --git a/crates/aisix-proxy/src/chat.rs b/crates/aisix-proxy/src/chat.rs index 9adae0c8..98c3d183 100644 --- a/crates/aisix-proxy/src/chat.rs +++ b/crates/aisix-proxy/src/chat.rs @@ -1058,12 +1058,14 @@ async fn dispatch( } // Multi-layer rate-limit reservation (api_key inline + model inline + policies). + // `mut` so a routing dispatch can fold the winning target's model-layer + // reservation into it once the winner is known (AISIX-Cloud#1087). let model_rl = crate::quota::ModelRateLimit::from_model( &req.model, &virtual_entry.id, &virtual_entry.value, ); - let reservation = crate::quota::enforce_rate_limit(state, auth, Some(&model_rl)) + let mut reservation = crate::quota::enforce_rate_limit(state, auth, Some(&model_rl)) .await .map_err(&with_model)?; @@ -1138,6 +1140,11 @@ async fn dispatch( kind: &'static str, } let mut won: Option = None; + // The winning target's own model-layer reservation (routing dispatch + // only) — folded into `reservation` after the loop so the stream hold + // and post-stream token accounting cover the member's limits too + // (AISIX-Cloud#1087). + let mut won_member_reservation: Option = None; 'targets: for attempt in &attempt_models { let model = &attempt.model; @@ -1163,6 +1170,48 @@ async fn dispatch( } else { String::new() }; + // Reserve THIS target's own model rate-limit layers before + // dispatching to it (AISIX-Cloud#1087). Over-limit → record a + // 429 attempt and move on to the remaining targets in strategy + // order (same-target retries can't help — the window won't + // reset mid-loop). + let member_reservation = match crate::quota::reserve_routing_target( + state, + is_routing_request, + &model.display_name, + &attempt.id, + model, + ) + .await + { + Ok(r) => r, + Err(e) => { + stream_routing.attempts.push(AttemptRecord { + index: idx, + kind, + target_model, + target_model_id: attempt.id.clone(), + provider_key_id: pk_entry.id.clone(), + status: 429, + success: false, + error_class: "rate_limit_exceeded".to_string(), + error_message: e.to_string(), + latency_ms: 0, + }); + // Keep the limiter's own Retry-After hint on the wire: + // when every target is exhausted this error becomes the + // client's 429, and SDKs back off on that header. + last_err = Some(BridgeError::upstream_status_with_retry_after( + 429, + format!( + "routing target {:?} is over its model rate limit: {e}", + model.display_name + ), + crate::quota::retry_after_of(&e).map(Duration::from_secs), + )); + continue 'targets; + } + }; let model_arc = Arc::new(model.clone()); let pk_arc = Arc::new(pk_entry.value.clone()); // Streaming deadline (#554): bound the connect by the effective @@ -1242,6 +1291,7 @@ async fn dispatch( idx, kind, }); + won_member_reservation = member_reservation; break 'targets; } Err(err) => { @@ -1306,6 +1356,13 @@ async fn dispatch( // which the CompleteOnDrop guard fires on both paths). Pre-fix the // permit was released here, letting a key capped at N run far more // than N simultaneous streams (#450). + // + // Fold the winning target's model-layer reservation in first, so the + // stream hold keeps its concurrency slot(s) and `post_stream_keys` + // bills its TPM/TPD at stream end too (AISIX-Cloud#1087). + if let Some(member) = won_member_reservation.take() { + reservation.merge(member); + } let post_stream_keys = reservation.keys(); let stream_concurrency_hold = reservation.into_stream_hold(); // least_busy: keep this target counted as in-flight for the stream's @@ -1864,8 +1921,13 @@ async fn dispatch( let is_routing_request = virtual_entry.value.routing.is_some() || virtual_entry.value.is_semantic(); let mut routing = RoutingTelemetry::default(); + // The winning target's own model-layer reservation (routing dispatch + // only) — folded into `reservation` at the commit point below so the + // member's TPM/TPD bills with the request-level layers + // (AISIX-Cloud#1087). + let mut won_member_reservation: Option = None; - for attempt in &attempt_models { + 'targets: for attempt in &attempt_models { let model = &attempt.model; let Some(provider) = model.provider.as_deref() else { last_err = Some(BridgeError::Config("model has no provider".into())); @@ -1915,6 +1977,49 @@ async fn dispatch( String::new() }; + // Reserve THIS target's own model rate-limit layers before + // dispatching to it (AISIX-Cloud#1087). Over-limit → record a + // 429 attempt and move on to the remaining targets in strategy + // order (same-target retries can't help — the window won't + // reset mid-loop). + let member_reservation = match crate::quota::reserve_routing_target( + state, + is_routing_request, + &model.display_name, + &attempt.id, + model, + ) + .await + { + Ok(r) => r, + Err(e) => { + routing.attempts.push(AttemptRecord { + index: attempt_index, + kind, + target_model, + target_model_id: attempt.id.clone(), + provider_key_id: pk_entry.id.clone(), + status: 429, + success: false, + error_class: "rate_limit_exceeded".to_string(), + error_message: e.to_string(), + latency_ms: 0, + }); + // Keep the limiter's own Retry-After hint on the wire: + // when every target is exhausted this error becomes the + // client's 429, and SDKs back off on that header. + last_err = Some(BridgeError::upstream_status_with_retry_after( + 429, + format!( + "routing target {:?} is over its model rate limit: {e}", + model.display_name + ), + crate::quota::retry_after_of(&e).map(Duration::from_secs), + )); + continue 'targets; + } + }; + let attempt_started = Instant::now(); // least_busy: count this target as in-flight for the upstream // call. The response is fully buffered, so the target is done @@ -1950,6 +2055,7 @@ async fn dispatch( error_message: String::new(), latency_ms: attempt_latency_ms, }); + won_member_reservation = member_reservation; upstream = Some(resp); break; } @@ -2046,6 +2152,12 @@ async fn dispatch( let provider_request_id = upstream.id.clone(); let provider_model_version = upstream.model.clone(); let finish_reason = finish_reason_label(&upstream.finish_reason); + // Fold the winning target's model-layer reservation in so one commit + // bills the member's TPM/TPD alongside the request-level layers + // (AISIX-Cloud#1087). + if let Some(member) = won_member_reservation.take() { + reservation.merge(member); + } reservation.commit_tokens(total).await; // cp-api recomputes cost server-side from its pricing catalog when diff --git a/crates/aisix-proxy/src/count_tokens.rs b/crates/aisix-proxy/src/count_tokens.rs index 576c58e2..a6f19d1f 100644 --- a/crates/aisix-proxy/src/count_tokens.rs +++ b/crates/aisix-proxy/src/count_tokens.rs @@ -200,6 +200,7 @@ async fn dispatch( .map(|r| r.fallback_on_statuses_or_default()) .unwrap_or(&[]); + let is_routing_request = model_entry.value.routing.is_some(); let mut last_err: Option = None; let mut any_anthropic = false; for target in &attempt_models { @@ -210,6 +211,26 @@ async fn dispatch( continue; } any_anthropic = true; + // Reserve THIS target's own model rate-limit layers before + // dispatching to it (AISIX-Cloud#1087); over-limit → skip it and + // try the remaining targets. Like the handler-level `_reservation` it is + // never token-committed — count_tokens burns no generation tokens; + // the drop at scope end releases the concurrency slot. + let _member_reservation = match crate::quota::reserve_routing_target( + state, + is_routing_request, + &target.model.display_name, + &target.id, + &target.model, + ) + .await + { + Ok(r) => r, + Err(e) => { + last_err = Some(e); + continue; + } + }; match count_tokens_to_target( state, &snapshot, diff --git a/crates/aisix-proxy/src/messages.rs b/crates/aisix-proxy/src/messages.rs index 2c177bf3..2a5f7da8 100644 --- a/crates/aisix-proxy/src/messages.rs +++ b/crates/aisix-proxy/src/messages.rs @@ -697,6 +697,38 @@ async fn dispatch( } else { String::new() }; + // Reserve THIS target's own model rate-limit layers before + // dispatching to it (AISIX-Cloud#1087). Over-limit → record a + // 429 attempt and move on to the remaining targets in strategy + // order (same-target retries can't help — the window won't + // reset mid-loop). + let mut member_reservation = match crate::quota::reserve_routing_target( + state, + is_routing_request, + &target.model.display_name, + &target.id, + &target.model, + ) + .await + { + Ok(r) => r, + Err(e) => { + routing.attempts.push(AttemptRecord { + index: idx, + kind, + target_model, + target_model_id: target.id.clone(), + provider_key_id: pk_id.clone(), + status: 429, + success: false, + error_class: "rate_limit_exceeded".to_string(), + error_message: e.to_string(), + latency_ms: 0, + }); + last_err = Some(e); + continue 'targets; + } + }; let attempt_started = Instant::now(); match dispatch_to_target( state, @@ -719,6 +751,7 @@ async fn dispatch( ..Default::default() }, &mut reservation, + &mut member_reservation, redactions_out.clone(), monitor_hits_out.clone(), ) @@ -749,7 +782,14 @@ async fn dispatch( // end-of-stream guard (#688), so `reservation` is `None` and // this is skipped. if !outcome.usage_handled_by_stream { - if let Some(r) = reservation.take() { + if let Some(mut r) = reservation.take() { + // Fold this target's model-layer reservation in + // (AISIX-Cloud#1087) so one commit bills the + // member's TPM/TPD too. Already `None` when the + // streaming path folded it into the guard. + if let Some(member) = member_reservation.take() { + r.merge(member); + } let total = total_tokens_with_cache( outcome.metrics.prompt_tokens, outcome.metrics.completion_tokens, @@ -829,6 +869,12 @@ async fn dispatch_to_target( // post-stream token accounting into the end-of-stream guard. Left in place // on the non-streaming / error paths for the handler to commit or retry. reservation: &mut Option, + // This target's own model-layer reservation (routing dispatch only, + // AISIX-Cloud#1087). The streaming path folds it into `reservation` + // before the take above so the end-of-stream guard covers the member's + // limits; the non-streaming path leaves it for the handler to commit + // alongside `reservation`. + member_reservation: &mut Option, // Input-side PII mask counts (#932) — the streaming paths merge these // into their end-of-stream telemetry emit (the non-streaming emit // happens in `messages()`, which already holds them). @@ -858,6 +904,7 @@ async fn dispatch_to_target( client, attempt, reservation, + member_reservation, input_redactions, input_monitor_hits, ) @@ -882,6 +929,7 @@ async fn dispatch_to_target( client, attempt, reservation, + member_reservation, input_redactions, input_monitor_hits, ) @@ -911,6 +959,10 @@ async fn anthropic_passthrough_dispatch( client_ctx: &ClientContext, attempt: AttemptInfo, reservation: &mut Option, + // This target's own model-layer reservation (AISIX-Cloud#1087); folded + // into `reservation` before the streaming take so the end-of-stream + // guard covers the member's limits. + member_reservation: &mut Option, input_redactions: crate::redact::RedactionCounts, input_monitor_hits: Vec, ) -> Result { @@ -1173,6 +1225,16 @@ async fn anthropic_passthrough_dispatch( // TPM/TPD accounting and `into_stream_hold` keeps the concurrency slot(s) // until the stream ends (mirrors chat.rs). `take()` leaves the handler's // `reservation` as `None`, so it won't also `commit_tokens`. + // + // Fold this target's model-layer reservation in first (AISIX-Cloud#1087) + // so the guard covers the member's limits too; `take()` leaves it `None` + // for the same reason. + if let Some(member) = member_reservation.take() { + match reservation.as_mut() { + Some(main) => main.merge(member), + None => *reservation = Some(member), + } + } let post_stream_keys = reservation.as_ref().map(|r| r.keys()).unwrap_or_default(); let stream_hold = reservation.take().map(|r| r.into_stream_hold()); let limiter_c = std::sync::Arc::clone(&state.limiter); @@ -1535,6 +1597,10 @@ async fn cross_provider_dispatch( client: &ClientContext, attempt: AttemptInfo, reservation: &mut Option, + // This target's own model-layer reservation (AISIX-Cloud#1087); folded + // into `reservation` before the streaming take so the end-of-stream + // guard covers the member's limits. + member_reservation: &mut Option, input_redactions: crate::redact::RedactionCounts, input_monitor_hits: Vec, ) -> Result { @@ -1697,6 +1763,16 @@ async fn cross_provider_dispatch( // keys drive post-stream TPM/TPD accounting, the hold keeps the // concurrency slot(s) until the stream ends. `take()` leaves the // handler's `reservation` as `None` so it won't also `commit_tokens`. + // + // Fold this target's model-layer reservation in first (AISIX-Cloud#1087) + // so the guard covers the member's limits too; `take()` leaves it `None` + // so the handler won't also commit it. + if let Some(member) = member_reservation.take() { + match reservation.as_mut() { + Some(main) => main.merge(member), + None => *reservation = Some(member), + } + } let post_stream_keys = reservation.as_ref().map(|r| r.keys()).unwrap_or_default(); let stream_hold = reservation.take().map(|r| r.into_stream_hold()); let limiter_for_stream = std::sync::Arc::clone(&state.limiter); diff --git a/crates/aisix-proxy/src/quota.rs b/crates/aisix-proxy/src/quota.rs index 5570fdb9..49bea4e7 100644 --- a/crates/aisix-proxy/src/quota.rs +++ b/crates/aisix-proxy/src/quota.rs @@ -293,6 +293,42 @@ pub(crate) async fn reserve_model_only( Ok(MultiReservation::new(reservations)) } +/// Reserve the model-scoped layers for one routing-dispatch target (Model +/// Group / semantic-router member), mirroring the ensemble per-sub-call +/// reservation (#620). Returns `Ok(None)` for a direct (non-routing) +/// dispatch: there the target IS the requested entry, whose model layers +/// were already reserved pre-dispatch by [`enforce`]/[`enforce_rate_limit`], +/// so reserving again would double-count the request (AISIX-Cloud#1087). +/// +/// An `Err` means this target is over one of its own limits right now — +/// the dispatch loops treat that as a failed 429 attempt and continue with +/// the remaining targets (matching LiteLLM, which filters rate-limited +/// deployments out of the candidate set). +pub(crate) async fn reserve_routing_target( + state: &ProxyState, + is_routing_request: bool, + target_name: &str, + target_entry_id: &str, + target: &aisix_core::Model, +) -> Result, ProxyError> { + if !is_routing_request { + return Ok(None); + } + reserve_model_only(state, target_name, target_entry_id, target) + .await + .map(Some) +} + +/// Seconds until the offending window reopens, for a +/// [`reserve_routing_target`] rejection. `chat.rs` funnels its rejection +/// through a `BridgeError`, which would otherwise drop the hint the +/// `/v1/messages` and `/v1/responses` loops keep by carrying the +/// `ProxyError::RateLimit` itself — so every endpoint's all-targets-exhausted +/// 429 lands with the same `Retry-After`. +pub(crate) fn retry_after_of(err: &ProxyError) -> Option { + err.retry_after_secs() +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/aisix-proxy/src/responses.rs b/crates/aisix-proxy/src/responses.rs index c089202d..03856a34 100644 --- a/crates/aisix-proxy/src/responses.rs +++ b/crates/aisix-proxy/src/responses.rs @@ -577,6 +577,38 @@ async fn dispatch( model: target_model.clone(), ..Default::default() }; + // Reserve THIS target's own model rate-limit layers before + // dispatching to it (AISIX-Cloud#1087). Over-limit → record a + // 429 attempt and move on to the remaining targets in strategy + // order (same-target retries can't help — the window won't + // reset mid-loop). + let mut member_reservation = match crate::quota::reserve_routing_target( + state, + is_routing_request, + &target.model.display_name, + &target.id, + &target.model, + ) + .await + { + Ok(r) => r, + Err(e) => { + routing.attempts.push(AttemptRecord { + index: idx, + kind, + target_model, + target_model_id: target.id.clone(), + provider_key_id: pk_id.clone(), + status: 429, + success: false, + error_class: "rate_limit_exceeded".to_string(), + error_message: e.to_string(), + latency_ms: 0, + }); + last_err = Some(e); + continue 'targets; + } + }; let result = if target.model.provider.as_deref() == Some("openai") { responses_to_target( state, @@ -592,6 +624,7 @@ async fn dispatch( client, attempt, &mut reservation, + &mut member_reservation, redactions_out.clone(), monitor_hits_out.clone(), ) @@ -611,6 +644,7 @@ async fn dispatch( client, attempt, &mut reservation, + &mut member_reservation, redactions_out.clone(), monitor_hits_out.clone(), ) @@ -642,7 +676,14 @@ async fn dispatch( // guard (#688), so `reservation` is `None` and this is // skipped. if !success.usage_handled_by_stream { - if let Some(r) = reservation.take() { + if let Some(mut r) = reservation.take() { + // Fold this target's model-layer reservation in + // (AISIX-Cloud#1087) so one commit bills the + // member's TPM/TPD too. Already `None` when the + // streaming path folded it into the guard. + if let Some(member) = member_reservation.take() { + r.merge(member); + } let total = success .usage .as_ref() @@ -820,6 +861,11 @@ async fn responses_to_target( client_ctx: &ClientContext, attempt: AttemptInfo, reservation: &mut Option, + // This target's own model-layer reservation (routing dispatch only, + // AISIX-Cloud#1087). The streaming path folds it into `reservation` + // before the take below; the non-streaming path leaves it for the + // handler to commit alongside `reservation`. + member_reservation: &mut Option, // Input-side PII mask counts (#932) for the verbatim streaming path's // end-of-stream emit; the non-streaming/buffered emits happen in the // handler, which already holds them. @@ -1226,6 +1272,15 @@ async fn responses_to_target( // post-stream TPM/TPD accounting, the hold keeps the concurrency slot(s) // until the stream ends. `take()` leaves the handler's `reservation` as // `None` so it won't also `commit_tokens`. + // Fold this target's model-layer reservation in first (AISIX-Cloud#1087) + // so the guard covers the member's limits too; `take()` leaves it `None` + // so the handler won't also commit it. + if let Some(member) = member_reservation.take() { + match reservation.as_mut() { + Some(main) => main.merge(member), + None => *reservation = Some(member), + } + } let post_stream_keys = reservation.as_ref().map(|r| r.keys()).unwrap_or_default(); let stream_hold = reservation.take().map(|r| r.into_stream_hold()); let limiter_c = std::sync::Arc::clone(&state.limiter); @@ -1468,6 +1523,11 @@ async fn responses_cross_provider_to_target( client_ctx: &ClientContext, attempt: AttemptInfo, reservation: &mut Option, + // This target's own model-layer reservation (routing dispatch only, + // AISIX-Cloud#1087). The streaming path folds it into `reservation` + // before the take below; the non-streaming path leaves it for the + // handler to commit alongside `reservation`. + member_reservation: &mut Option, // Input-side PII mask counts (#932), merged into the streamed judge // path's end-of-stream emit; non-streaming emits happen in the handler. input_redactions: crate::redact::RedactionCounts, @@ -1619,6 +1679,15 @@ async fn responses_cross_provider_to_target( // post-stream TPM/TPD accounting, the hold keeps the concurrency slot(s) // until the stream ends. `take()` leaves the handler's `reservation` as // `None` so it won't also `commit_tokens`. + // Fold this target's model-layer reservation in first (AISIX-Cloud#1087) + // so the guard covers the member's limits too; `take()` leaves it `None` + // so the handler won't also commit it. + if let Some(member) = member_reservation.take() { + match reservation.as_mut() { + Some(main) => main.merge(member), + None => *reservation = Some(member), + } + } let post_stream_keys = reservation.as_ref().map(|r| r.keys()).unwrap_or_default(); let stream_hold = reservation.take().map(|r| r.into_stream_hold()); let limiter_c = std::sync::Arc::clone(&state.limiter); diff --git a/crates/aisix-ratelimit/src/limiter.rs b/crates/aisix-ratelimit/src/limiter.rs index d1dd0837..5369622b 100644 --- a/crates/aisix-ratelimit/src/limiter.rs +++ b/crates/aisix-ratelimit/src/limiter.rs @@ -195,6 +195,15 @@ impl MultiReservation { self.reservations.iter().map(|r| r.key.clone()).collect() } + /// Absorb another reservation's layers into this one, so a single + /// `commit_tokens` / `into_stream_hold` finalises both. Used by the + /// routing dispatch to fold the winning target's model-layer + /// reservation into the request-level reservation once the winner + /// is known. + pub fn merge(&mut self, other: MultiReservation) { + self.reservations.extend(other.reservations); + } + /// Convert into an owned [`StreamConcurrencyGuard`] for the streaming /// path. The per-layer concurrency slots stay held — they are NOT /// released here — and are released only when the returned guard drops, @@ -670,6 +679,27 @@ mod tests { assert_eq!(keys, vec!["api_key:k1", "model:m1", "team:t1"]); } + #[tokio::test] + async fn multi_reservation_merge_commits_and_releases_absorbed_layers() { + let clock = TestClock::new(100); + let limiter = Limiter::local_with_clock(clock.clone()); + let l = limits(None, Some(1000), Some(1)); + + let main = limiter.pre_commit("api_key:k1", &l).await.unwrap(); + let member = limiter.pre_commit("model:target", &l).await.unwrap(); + + let mut multi = MultiReservation::new(vec![main]); + multi.merge(MultiReservation::new(vec![member])); + assert_eq!(multi.keys(), vec!["api_key:k1", "model:target"]); + + // One commit finalises both layers: tokens land on each and the + // absorbed layer's concurrency slot is released. + multi.commit_tokens(300).await; + let s = limiter.peek("model:target", &l).await.unwrap(); + assert_eq!(s.tpm_used, 300); + assert!(limiter.pre_commit("model:target", &l).await.is_ok()); + } + #[tokio::test] async fn multi_reservation_partial_failure_releases_acquired_layers() { let clock = TestClock::new(100); diff --git a/tests/e2e/src/cases/model-group-member-ratelimit-e2e.test.ts b/tests/e2e/src/cases/model-group-member-ratelimit-e2e.test.ts new file mode 100644 index 00000000..72072437 --- /dev/null +++ b/tests/e2e/src/cases/model-group-member-ratelimit-e2e.test.ts @@ -0,0 +1,447 @@ +import { createHash } from "node:crypto"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; +import { + EtcdClient, + SeedClient, + ProxyClient, + spawnApp, + startOpenAiUpstream, + waitConfigPropagation, + type OpenAiUpstream, + type SpawnedApp, +} from "../harness/index.js"; + +// E2E for AISIX-Cloud#1087: a Model Group dispatch must honor each +// TARGET's own `rate_limit`, not just the group entry's. Pre-fix the +// pre-dispatch reservation covered only the requested entry (the group, +// which carries no limits), so a member's RPM/TPM was silently ignored +// and every request kept landing on the over-limit first target. +// +// Post-fix each routing attempt reserves the target's model-scoped +// layers first (mirroring the ensemble per-sub-call reservation, #620): +// an over-limit member becomes a failed 429 attempt that fails over to +// the next target — LiteLLM's semantics, where rate-limited deployments +// are filtered from the candidate set — and a request served by a +// member commits its token cost to that member's TPM bucket. + +const CALLER_PLAINTEXT = "sk-1087-caller"; +const CALLER_KEY_HASH = createHash("sha256") + .update(CALLER_PLAINTEXT) + .digest("hex"); +// Readiness-probe caller allowed to access NOTHING: a chat call with it +// returns 404 while a name is absent from the snapshot and 403 once it +// propagated. The 403 fires at the ACL gate, before any rate-limit +// reservation, so probing never consumes the member quotas under test +// (and routing models never appear in /v1/models, so listing can't be +// the probe). +const PROBE_PLAINTEXT = "sk-1087-probe"; +const PROBE_KEY_HASH = createHash("sha256") + .update(PROBE_PLAINTEXT) + .digest("hex"); + +function chatBody(content: string) { + return { + id: "cmpl-1087", + object: "chat.completion", + created: 0, + model: "gpt-4o-mini", + choices: [ + { + index: 0, + message: { role: "assistant", content }, + finish_reason: "stop", + }, + ], + usage: { prompt_tokens: 5, completion_tokens: 3, total_tokens: 8 }, + }; +} + +function chunk(json: Record): string { + return JSON.stringify({ + id: "chatcmpl-1087", + object: "chat.completion.chunk", + created: 0, + model: "gpt-4o-mini", + ...json, + }); +} + +function streamEvents(content: string, totalTokens: number): string[] { + return [ + chunk({ choices: [{ index: 0, delta: { role: "assistant" }, finish_reason: null }] }), + chunk({ choices: [{ index: 0, delta: { content }, finish_reason: null }] }), + chunk({ choices: [{ index: 0, delta: {}, finish_reason: "stop" }] }), + chunk({ + choices: [], + usage: { + prompt_tokens: totalTokens - 2, + completion_tokens: 2, + total_tokens: totalTokens, + }, + }), + "[DONE]", + ]; +} + +function anthropicMessageBody(text: string) { + return { + id: `msg_${text}`, + type: "message", + role: "assistant", + content: [{ type: "text", text }], + model: "claude-3-5-haiku-20241022", + stop_reason: "end_turn", + usage: { input_tokens: 5, output_tokens: 4 }, + }; +} + +type ChatResult = { + status: number; + body: { + choices?: Array<{ message?: { content?: string } }>; + error?: { message?: string }; + }; +}; + +describe("model group member rate limit e2e (AISIX-Cloud#1087)", () => { + let app: SpawnedApp | undefined; + let seed: SeedClient | undefined; + let proxy: ProxyClient | 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); + proxy = new ProxyClient(app.proxyUrl, CALLER_PLAINTEXT); + await seed.createApiKey({ + key_hash: CALLER_KEY_HASH, + allowed_models: ["*"], + }); + await seed.createApiKey({ + key_hash: PROBE_KEY_HASH, + allowed_models: ["__probe-none__"], + }); + }); + + afterAll(async () => { + await app?.exit(); + await Promise.all(upstreams.map((u) => u.close())); + }); + + async function newUpstream(opts: Parameters[0]): Promise { + const u = await startOpenAiUpstream(opts); + upstreams.push(u); + return u; + } + + async function createOpenAiModel( + displayName: string, + upstream: OpenAiUpstream, + extra: Record = {}, + ): Promise { + if (!seed) throw new Error("seed client not initialized"); + const pk = await seed.createProviderKey({ + display_name: `${displayName}-pk`, + secret: "sk-openai-mock", + api_base: `${upstream.baseUrl}/v1`, + }); + await seed.createModel({ + display_name: displayName, + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: pk.id, + ...extra, + }); + } + + // Readiness: probe every name with the no-access key until each + // returns 403 (in snapshot, ACL-rejected) instead of 404 (not yet + // propagated). Deliberately NOT a real chat call with the main caller + // — that would burn the member's own rate-limit quota under test. + async function waitModelsListed(names: string[]): Promise { + if (!app) throw new Error("app not initialized"); + const probe = new ProxyClient(app.proxyUrl, PROBE_PLAINTEXT); + await waitConfigPropagation(async () => { + for (const n of names) { + const res = await probe.chat({ + model: n, + messages: [{ role: "user", content: "probe" }], + }); + if (res.status !== 403) return false; + } + return true; + }); + } + + async function callGroup(model: string): Promise { + if (!proxy) throw new Error("proxy client not initialized"); + return (await proxy.chat({ + model, + messages: [{ role: "user", content: "hello" }], + })) as ChatResult; + } + + function servedContent(r: ChatResult): string { + expect(r.status).toBe(200); + return r.body.choices?.[0]?.message?.content ?? ""; + } + + /** + * Sleep until the current wall-clock minute has at least `headroomSecs` + * left. The limiter buckets on fixed wall-clock minutes + * (`window_start = now - now % 60`), so a burst that straddles a boundary + * silently gets a fresh quota and the failover/429 assertions would flap. + * Waiting for headroom keeps each burst inside one window. + */ + async function awaitWindowHeadroom(headroomSecs: number): Promise { + const secondsLeft = 60 - (Math.floor(Date.now() / 1000) % 60); + if (secondsLeft >= headroomSecs) return; + await new Promise((r) => setTimeout(r, secondsLeft * 1000 + 100)); + } + + test("over-limit member fails over to the next target (RPM)", async (ctx) => { + if (!etcdReachable || !app) { + ctx.skip(); + return; + } + const limited = await newUpstream({ nonStreamBody: chatBody("served-by-limited") }); + const backup = await newUpstream({ nonStreamBody: chatBody("served-by-backup") }); + await createOpenAiModel("mgrl-limited", limited, { rate_limit: { rpm: 1 } }); + await createOpenAiModel("mgrl-backup", backup); + await seed!.createModel({ + display_name: "mgrl-group", + routing: { + strategy: "failover", + targets: [{ model: "mgrl-limited" }, { model: "mgrl-backup" }], + }, + }); + await waitModelsListed(["mgrl-limited", "mgrl-backup", "mgrl-group"]); + await awaitWindowHeadroom(5); + + // First call through the group lands on the first target. + expect(servedContent(await callGroup("mgrl-group"))).toBe("served-by-limited"); + + // Second call in the same minute: the first target is over its own + // RPM=1, so dispatch must fail over to the backup. Pre-fix the + // member's limit was never consulted and this still returned + // "served-by-limited". + expect(servedContent(await callGroup("mgrl-group"))).toBe("served-by-backup"); + }); + + test("all members over limit surfaces as 429", async (ctx) => { + if (!etcdReachable || !app) { + ctx.skip(); + return; + } + const c1 = await newUpstream({ nonStreamBody: chatBody("served-by-c1") }); + const c2 = await newUpstream({ nonStreamBody: chatBody("served-by-c2") }); + await createOpenAiModel("mgrl-c1", c1, { rate_limit: { rpm: 1 } }); + await createOpenAiModel("mgrl-c2", c2, { rate_limit: { rpm: 1 } }); + await seed!.createModel({ + display_name: "mgrl-both-limited", + routing: { + strategy: "failover", + targets: [{ model: "mgrl-c1" }, { model: "mgrl-c2" }], + }, + }); + await waitModelsListed(["mgrl-c1", "mgrl-c2", "mgrl-both-limited"]); + await awaitWindowHeadroom(5); + + expect(servedContent(await callGroup("mgrl-both-limited"))).toBe("served-by-c1"); + expect(servedContent(await callGroup("mgrl-both-limited"))).toBe("served-by-c2"); + + // Both members exhausted → the request fails with 429, not a 5xx, and + // carries the limiter's `Retry-After` so SDK back-off still works. The + // hint is the load-bearing half: a bare 429 makes clients retry + // immediately against a window that has not reopened. + const third = await fetch(`${app!.proxyUrl}/v1/chat/completions`, { + method: "POST", + headers: { + authorization: `Bearer ${CALLER_PLAINTEXT}`, + "content-type": "application/json", + }, + body: JSON.stringify({ + model: "mgrl-both-limited", + messages: [{ role: "user", content: "hello" }], + }), + }); + expect(third.status).toBe(429); + const body = (await third.json()) as { error?: { message?: string } }; + expect(body.error?.message ?? "").toContain("rate limit"); + const retryAfter = Number.parseInt(third.headers.get("retry-after") ?? "", 10); + expect(retryAfter).toBeGreaterThan(0); + expect(retryAfter).toBeLessThanOrEqual(60); + }); + + test("winning member's TPM bucket is committed, throttling the next call", async (ctx) => { + if (!etcdReachable || !app) { + ctx.skip(); + return; + } + // Upstream reports total_tokens=8 per call (chatBody default); the + // member's TPM=5 admits the first call (window empty at pre-commit) + // and must reject the second (8 committed ≥ 5). + const tpmUp = await newUpstream({ nonStreamBody: chatBody("served-by-tpm") }); + const backup = await newUpstream({ nonStreamBody: chatBody("served-by-tpm-backup") }); + await createOpenAiModel("mgrl-tpm", tpmUp, { rate_limit: { tpm: 5 } }); + await createOpenAiModel("mgrl-tpm-backup", backup); + await seed!.createModel({ + display_name: "mgrl-tpm-group", + routing: { + strategy: "failover", + targets: [{ model: "mgrl-tpm" }, { model: "mgrl-tpm-backup" }], + }, + }); + await waitModelsListed(["mgrl-tpm", "mgrl-tpm-backup", "mgrl-tpm-group"]); + await awaitWindowHeadroom(5); + + // First call is served by the TPM-capped member and commits 8 tokens + // to ITS bucket (the reservation-merge under test — pre-fix the + // tokens landed nowhere and the second call stayed on the member). + expect(servedContent(await callGroup("mgrl-tpm-group"))).toBe("served-by-tpm"); + expect(servedContent(await callGroup("mgrl-tpm-group"))).toBe("served-by-tpm-backup"); + }); + + test("streaming: member TPM commits at stream end and the next stream fails over", async (ctx) => { + if (!etcdReachable || !app) { + ctx.skip(); + return; + } + const streamLimited = await newUpstream({ + streamEvents: streamEvents("stream-from-limited", 16), + }); + const streamBackup = await newUpstream({ + streamEvents: streamEvents("stream-from-backup", 16), + }); + await createOpenAiModel("mgrl-stream-tpm", streamLimited, { + rate_limit: { tpm: 10 }, + }); + await createOpenAiModel("mgrl-stream-backup", streamBackup); + await seed!.createModel({ + display_name: "mgrl-stream-group", + routing: { + strategy: "failover", + targets: [{ model: "mgrl-stream-tpm" }, { model: "mgrl-stream-backup" }], + }, + }); + await waitModelsListed([ + "mgrl-stream-tpm", + "mgrl-stream-backup", + "mgrl-stream-group", + ]); + // The follow-up poll below runs for up to 5s after the first stream, + // so this burst needs more headroom than the non-streaming cases. + await awaitWindowHeadroom(15); + + const streamCall = async (): Promise<{ status: number; text: string }> => { + const res = await fetch(`${app!.proxyUrl}/v1/chat/completions`, { + method: "POST", + headers: { + authorization: `Bearer ${CALLER_PLAINTEXT}`, + "content-type": "application/json", + }, + body: JSON.stringify({ + model: "mgrl-stream-group", + messages: [{ role: "user", content: "hello" }], + stream: true, + }), + }); + return { status: res.status, text: await res.text() }; + }; + + // First stream is served by the capped member; its terminal usage + // frame (16 tokens > TPM=10) must be committed to the MEMBER's + // bucket via the merged post-stream keys. + const first = await streamCall(); + expect(first.status).toBe(200); + expect(first.text).toContain("stream-from-limited"); + + // The post-stream commit fires when the server drops the response + // body — a tick after the client finishes reading — so poll until + // the follow-up stream lands on the backup. + const deadline = Date.now() + 5_000; + let servedBy = ""; + while (Date.now() < deadline) { + const next = await streamCall(); + expect(next.status).toBe(200); + if (next.text.includes("stream-from-backup")) { + servedBy = "backup"; + break; + } + await new Promise((r) => setTimeout(r, 100)); + } + expect(servedBy).toBe("backup"); + }); + + test("/v1/messages: over-limit member fails over to the next target", async (ctx) => { + if (!etcdReachable || !app) { + ctx.skip(); + return; + } + const anLimited = await newUpstream({ + nonStreamBody: anthropicMessageBody("an-served-by-limited"), + }); + const anBackup = await newUpstream({ + nonStreamBody: anthropicMessageBody("an-served-by-backup"), + }); + const mkAnthropic = async (name: string, up: OpenAiUpstream, extra: Record = {}) => { + const pk = await seed!.createProviderKey({ + display_name: `${name}-pk`, + provider: "anthropic", + adapter: "anthropic", + secret: "sk-ant-mock", + // Anthropic bridge appends /v1/messages: point at the bare host. + api_base: up.baseUrl, + }); + await seed!.createModel({ + display_name: name, + provider: "anthropic", + model_name: "claude-3-5-haiku-20241022", + provider_key_id: pk.id, + ...extra, + }); + }; + await mkAnthropic("mgrl-an-limited", anLimited, { rate_limit: { rpm: 1 } }); + await mkAnthropic("mgrl-an-backup", anBackup); + await seed!.createModel({ + display_name: "mgrl-an-group", + routing: { + strategy: "failover", + targets: [{ model: "mgrl-an-limited" }, { model: "mgrl-an-backup" }], + }, + }); + await waitModelsListed(["mgrl-an-limited", "mgrl-an-backup", "mgrl-an-group"]); + await awaitWindowHeadroom(5); + + const callMessages = async (): Promise<{ status: number; text: string }> => { + const res = await fetch(`${app!.proxyUrl}/v1/messages`, { + method: "POST", + headers: { + authorization: `Bearer ${CALLER_PLAINTEXT}`, + "content-type": "application/json", + }, + body: JSON.stringify({ + model: "mgrl-an-group", + max_tokens: 32, + messages: [{ role: "user", content: "hello" }], + }), + }); + return { status: res.status, text: await res.text() }; + }; + + const first = await callMessages(); + expect(first.status).toBe(200); + expect(first.text).toContain("an-served-by-limited"); + + // Same minute, member over its RPM=1 → the /v1/messages dispatch + // loop must fail over exactly like /v1/chat/completions. + const second = await callMessages(); + expect(second.status).toBe(200); + expect(second.text).toContain("an-served-by-backup"); + }); +});