fix(proxy): enforce each routing target's own model rate limit

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.

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<StreamWin> = 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<aisix_ratelimit::MultiReservation> = 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<aisix_ratelimit::MultiReservation> = 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

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<ProxyError> = 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,

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<aisix_ratelimit::MultiReservation>,
+    // 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<aisix_ratelimit::MultiReservation>,
     // 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<aisix_ratelimit::MultiReservation>,
+    // 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<aisix_ratelimit::MultiReservation>,
     input_redactions: crate::redact::RedactionCounts,
     input_monitor_hits: Vec<aisix_core::GuardrailMonitorHit>,
 ) -> Result<DispatchOutcome, ProxyError> {
@@ -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<aisix_ratelimit::MultiReservation>,
+    // 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<aisix_ratelimit::MultiReservation>,
     input_redactions: crate::redact::RedactionCounts,
     input_monitor_hits: Vec<aisix_core::GuardrailMonitorHit>,
 ) -> Result<DispatchOutcome, ProxyError> {
@@ -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);