feat(gateway): model capacity hook, the proxy's dynamic rate limiter (#47055)

Adds the model_capacity hook, turned on by listing dynamic_rate_limiter_v3
(or model_capacity) in litellm_settings.callbacks. It caps each model
group at the sum of its deployments' rpm and tpm across every replica,
and with priority_reservation holds each priority's pool to its share
once the group's saturation reaches the threshold, counting it below.
It uses the proxy's model_saturation_check and priority_model keys and
scripts, its refusal messages and headers, and adds a successful
response's tokens to the group and the pool admission chose. The rate
limit counters gain a general check-and-add over request and token
pairs, and the throttling refusal shape is shared.


Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5

Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

litellm-rust/crates/gateway/src/lib.rs

@@ -9,6 +9,7 @@ pub mod hooks;
 mod key_status;
 mod keys;
 pub mod metrics;
+pub mod model_capacity;
 pub mod rate_limits;
 pub mod request_hooks;
 mod secrets;

litellm-rust/crates/gateway/src/main.rs

@@ -128,6 +128,11 @@ async fn main() -> Result<(), Box<dyn Error>> {
             ),
         ));
     }
+    if let Some(capacity) =
+        litellm_gateway::model_capacity::ModelCapacity::for_gateway(&config, coordination.clone())
+    {
+        hooks.push(Arc::new(capacity));
+    }
     hooks.extend(litellm_gateway::request_hooks::from_config(&config)?);
     let mut observers: Vec<Arc<dyn litellm_gateway_hooks::Observer>> = Vec::new();
     let spend = database.zip(litellm_gateway::spend::SpendSettings::from_config(&config));

litellm-rust/crates/gateway/src/model_capacity.rs

@@ -0,0 +1,654 @@
+//! The model capacity hook, the proxy's `dynamic_rate_limiter_v3`: caps each model group at the
+//! sum of its deployments' `rpm` and `tpm` across every replica, and with `priority_reservation`
+//! splits that capacity between priorities once the group is busy.
+//!
+//! Each request adds one to `{model_saturation_check:<group>}:requests` and checks the group's
+//! tokens, all or nothing, with the proxy's keys and scripts. With priorities, the request's pool
+//! `{priority_model:<group>:<priority>}` (or `:default_pool`) is held to its share of capacity
+//! once the group's saturation reaches `saturation_threshold`, and only counted below it. Tokens
+//! are added to both counters when the response ends.
+
+use std::{
+    collections::HashMap,
+    sync::{Arc, Mutex},
+    time::{Duration, Instant},
+};
+
+use futures_util::future::BoxFuture;
+use http::{HeaderName, HeaderValue};
+use litellm_config::Config;
+use litellm_gateway_hooks::{
+    Facts, Hook, KeyFacts, Record, Refusal, RequestEdit, RequestFacts, Settle, Targets,
+};
+use serde_json::Value;
+
+use crate::{
+    coordination::Coordination,
+    rate_limits::{Kind, Pair, RateLimitSettings, Reserved, Store, Unavailable},
+};
+
+/// The names a config's `callbacks` list turns this hook on with.
+const NAMES: [&str; 2] = ["dynamic_rate_limiter_v3", "model_capacity"];
+
+/// A priority's share of a group's capacity, as `priority_reservation` writes it.
+#[derive(Clone, Copy, Debug, PartialEq)]
+pub enum Share {
+    Percent(f64),
+    Rpm(f64),
+    Tpm(f64),
+}
+
+impl Share {
+    /// The share as a fraction of the group, as the proxy's `convert_priority_to_percent`.
+    fn fraction(self, capacity: &Capacity) -> f64 {
+        match self {
+            Self::Percent(value) => value,
+            Self::Rpm(value) => match capacity.rpm {
+                Some(rpm) if rpm > 0 => value / rpm as f64,
+                _ => value,
+            },
+            Self::Tpm(value) => match capacity.tpm {
+                Some(tpm) if tpm > 0 => value / tpm as f64,
+                _ => value,
+            },
+        }
+    }
+}
+
+/// Which tokens count: `general_settings.token_rate_limit_type`.
+#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
+pub enum TokenCount {
+    #[default]
+    Total,
+    Input,
+    Output,
+}
+
+#[derive(Clone, Debug, PartialEq)]
+pub struct CapacitySettings {
+    pub window: u64,
+    pub fail_closed: bool,
+    /// `litellm_settings.priority_reservation`, in config order.
+    pub reservation: Option<Vec<(String, Share)>>,
+    /// `priority_reservation_settings.default_priority`: the share of keys without a listed
+    /// priority, together.
+    pub default_priority: f64,
+    /// `priority_reservation_settings.saturation_threshold`.
+    pub saturation_threshold: f64,
+    /// `priority_reservation_settings.saturation_check_cache_ttl`: how long a group's
+    /// saturation is reused before it is read again.
+    pub saturation_ttl: Duration,
+    pub tokens: TokenCount,
+}
+
+impl Default for CapacitySettings {
+    fn default() -> Self {
+        let rate_limits = RateLimitSettings::default();
+        Self {
+            window: rate_limits.window,
+            fail_closed: rate_limits.fail_closed,
+            reservation: None,
+            default_priority: 0.25,
+            saturation_threshold: 0.5,
+            saturation_ttl: Duration::from_secs(60),
+            tokens: TokenCount::Total,
+        }
+    }
+}
+
+impl CapacitySettings {
+    /// `None` when the config's `callbacks` do not list the hook.
+    pub fn from_config(config: &Config) -> Option<Self> {
+        let litellm = &config.litellm_settings;
+        let listed = litellm.callbacks.iter().flatten().any(|callback| {
+            callback
+                .as_str()
+                .is_some_and(|name| NAMES.contains(&name.trim()))
+        });
+        if !listed {
+            return None;
+        }
+        let rate_limits = RateLimitSettings::from_config(config);
+        let number = |value: &litellm_config::Value| match value {
+            litellm_config::Value::Number(number) => number.as_f64(),
+            litellm_config::Value::String(text) => text.trim().parse().ok(),
+            _ => None,
+        };
+        let reservation = match litellm.additional_fields.get("priority_reservation") {
+            Some(litellm_config::Value::Mapping(mapping)) => Some(
+                mapping
+                    .iter()
+                    .filter_map(|(priority, share)| {
+                        let priority = match priority {
+                            litellm_config::Value::String(text) => text.clone(),
+                            litellm_config::Value::Number(number) => number.to_string(),
+                            _ => return None,
+                        };
+                        let share = match share {
+                            litellm_config::Value::Mapping(share) => {
+                                let value = share.get("value").and_then(number).unwrap_or(1.0);
+                                match share.get("type").and_then(litellm_config::Value::as_str) {
+                                    Some("rpm") => Share::Rpm(value),
+                                    Some("tpm") => Share::Tpm(value),
+                                    _ => Share::Percent(value),
+                                }
+                            }
+                            other => Share::Percent(number(other)?),
+                        };
+                        Some((priority, share))
+                    })
+                    .collect(),
+            ),
+            _ => None,
+        };
+        let settings = litellm
+            .additional_fields
+            .get("priority_reservation_settings");
+        let setting = |name: &str| {
+            settings
+                .and_then(|settings| settings.get(name))
+                .and_then(number)
+        };
+        let defaults = Self::default();
+        Some(Self {
+            window: rate_limits.window,
+            fail_closed: rate_limits.fail_closed,
+            reservation,
+            default_priority: setting("default_priority").unwrap_or(defaults.default_priority),
+            saturation_threshold: setting("saturation_threshold")
+                .unwrap_or(defaults.saturation_threshold),
+            saturation_ttl: setting("saturation_check_cache_ttl")
+                .filter(|seconds| *seconds >= 0.0)
+                .map_or(defaults.saturation_ttl, Duration::from_secs_f64),
+            tokens: match config
+                .general_settings
+                .additional_fields
+                .get("token_rate_limit_type")
+                .and_then(litellm_config::Value::as_str)
+            {
+                Some("input") => TokenCount::Input,
+                Some("output") => TokenCount::Output,
+                _ => TokenCount::Total,
+            },
+        })
+    }
+}
+
+/// A model group's requests and tokens per window: its deployments' `rpm` and `tpm`, summed.
+#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
+pub struct Capacity {
+    pub rpm: Option<i64>,
+    pub tpm: Option<i64>,
+}
+
+/// Each configured group's capacity, from `rpm`/`tpm` on the deployment, its `litellm_params`
+/// or its `model_info`, in that order, as the proxy's router sums them.
+pub fn capacities(config: &Config) -> HashMap<String, Capacity> {
+    let mut groups: HashMap<String, Capacity> = HashMap::new();
+    for model in &config.model_list {
+        let configured = |field: &str| -> Option<f64> {
+            let from_value = |value: &litellm_config::Value| match value {
+                litellm_config::Value::Number(number) => number.as_f64(),
+                litellm_config::Value::String(text) => text.trim().parse().ok(),
+                _ => None,
+            };
+            let spelled = |value: &litellm_config::Spelled<f64>| match value {
+                litellm_config::Spelled::Value(value) => Some(*value),
+                litellm_config::Spelled::Text(text) => text.trim().parse().ok(),
+            };
+            let params = match field {
+                "tpm" => model.litellm_params.tpm.as_ref(),
+                _ => model.litellm_params.rpm.as_ref(),
+            };
+            model
+                .additional_fields
+                .get(field)
+                .and_then(from_value)
+                .or_else(|| params.and_then(spelled))
+                .or_else(|| model.model_info.get(field).and_then(from_value))
+        };
+        let group = groups.entry(model.model_name.clone()).or_default();
+        let add = |total: &mut Option<i64>, value: Option<f64>| {
+            if let Some(value) = value {
+                *total = Some(total.unwrap_or(0) + value as i64);
+            }
+        };
+        add(&mut group.tpm, configured("tpm"));
+        add(&mut group.rpm, configured("rpm"));
+    }
+    groups
+}
+
+/// The pool a request draws on and its share of the group.
+#[derive(Clone, Debug, PartialEq)]
+struct Pool {
+    name: String,
+    share: f64,
+}
+
+/// The `model_capacity` hook.
+pub struct ModelCapacity {
+    settings: CapacitySettings,
+    groups: HashMap<String, Capacity>,
+    store: Arc<Store>,
+    saturation: Mutex<HashMap<String, (Option<i64>, Instant)>>,
+}
+
+impl ModelCapacity {
+    pub fn new(
+        settings: CapacitySettings,
+        groups: HashMap<String, Capacity>,
+        redis: Option<Arc<Coordination>>,
+    ) -> Self {
+        Self {
+            store: Arc::new(Store::new(redis, settings.window, settings.fail_closed)),
+            settings,
+            groups,
+            saturation: Mutex::default(),
+        }
+    }
+
+    /// The hook for a gateway whose config lists it, else `None`.
+    pub fn for_gateway(config: &Config, redis: Option<Arc<Coordination>>) -> Option<Self> {
+        Some(Self::new(
+            CapacitySettings::from_config(config)?,
+            capacities(config),
+            redis,
+        ))
+    }
+
+    /// The pool the proxy's `_get_priority_allocation` picks, with its normalized share.
+    fn pool(&self, group: &str, capacity: &Capacity, priority: Option<&str>) -> Pool {
+        let explicit = self.settings.reservation.as_ref().and_then(|reservation| {
+            let priority = priority?;
+            reservation
+                .iter()
+                .any(|(name, _)| name == priority)
+                .then(|| {
+                    let fractions: Vec<(&str, f64)> = reservation
+                        .iter()
+                        .map(|(name, share)| (name.as_str(), share.fraction(capacity)))
+                        .collect();
+                    let total: f64 = fractions.iter().map(|(_, fraction)| fraction).sum();
+                    let fraction = fractions
+                        .iter()
+                        .find(|(name, _)| *name == priority)
+                        .map_or(0.0, |(_, fraction)| *fraction);
+                    (
+                        priority,
+                        if total > 1.0 {
+                            fraction / total
+                        } else {
+                            fraction
+                        },
+                    )
+                })
+        });
+        match explicit {
+            Some((priority, share)) => Pool {
+                name: format!("{group}:{priority}"),
+                share,
+            },
+            None => Pool {
+                name: format!("{group}:default_pool"),
+                share: self.settings.default_priority,
+            },
+        }
+    }
+
+    /// A counter's value, reused for `saturation_check_cache_ttl`.
+    async fn cached(&self, prefix: &str, kind: Kind) -> Option<i64> {
+        let key = format!("{prefix}:{}", kind.name());
+        if let Some((value, read_at)) = self
+            .saturation
+            .lock()
+            .unwrap_or_else(|e| e.into_inner())
+            .get(&key)
+            && read_at.elapsed() < self.settings.saturation_ttl
+        {
+            return *value;
+        }
+        let value = self.store.read(prefix, kind).await;
+        let mut cache = self.saturation.lock().unwrap_or_else(|e| e.into_inner());
+        if cache.len() >= 10_000 {
+            cache.clear();
+        }
+        cache.insert(key, (value, Instant::now()));
+        value
+    }
+
+    /// How full the group is: its busiest of requests and tokens against capacity.
+    async fn saturation(&self, prefix: &str, capacity: &Capacity) -> f64 {
+        let mut saturation = 0.0_f64;
+        for (kind, limit) in [(Kind::Requests, capacity.rpm), (Kind::Tokens, capacity.tpm)] {
+            if let Some(limit) = limit.filter(|limit| *limit > 0)
+                && let Some(count) = self.cached(prefix, kind).await
+            {
+                saturation = saturation.max(count as f64 / limit as f64);
+            }
+        }
+        saturation
+    }
+
+    fn unavailable(&self, request: &RequestFacts) -> Refusal {
+        crate::rate_limits::unavailable(request.route)
+    }
+}
+
+/// `team_metadata.priority`, else the key's `metadata.priority`.
+fn priority(key: Option<&KeyFacts>) -> Option<String> {
+    let key = key?;
+    let read = |value: &Value| match value {
+        Value::String(text) => Some(text.clone()),
+        Value::Number(number) => Some(number.to_string()),
+        _ => None,
+    };
+    key.team_metadata
+        .get("priority")
+        .and_then(read)
+        .or_else(|| key.metadata.get("priority").and_then(read))
+}
+
+fn limit_text(limit: Option<i64>) -> String {
+    limit.map_or_else(|| "not configured".to_owned(), |limit| limit.to_string())
+}
+
+/// `{:.<places>%}` as Python formats a fraction.
+fn percent(fraction: f64, places: usize) -> String {
+    format!("{:.*}%", places, fraction * 100.0)
+}
+
+impl Hook for ModelCapacity {
+    fn name(&self) -> &'static str {
+        "model_capacity"
+    }
+
+    fn reads(&self) -> Facts {
+        Facts::KEY | Facts::USAGE
+    }
+
+    fn writes(&self) -> Targets {
+        Targets::ADMISSION | Targets::RESPONSE_HEADERS
+    }
+
+    fn run<'a>(
+        &'a self,
+        request: &'a Arc<RequestFacts>,
+        edit: &'a mut RequestEdit,
+    ) -> BoxFuture<'a, Result<Option<Box<dyn Settle>>, Refusal>> {
+        Box::pin(async move {
+            let priority = priority(request.key.as_deref());
+            let priority_header = priority.as_deref().unwrap_or("default");
+            if !request.stream {
+                if priority_header.chars().all(|char| (char as u32) < 256)
+                    && let Ok(value) = HeaderValue::from_str(priority_header)
+                {
+                    edit.add_response_header(HeaderName::from_static("x-litellm-priority"), value);
+                }
+                edit.add_response_header(
+                    HeaderName::from_static("x-litellm-rate-limiter-version"),
+                    HeaderValue::from_static("v3"),
+                );
+            }
+            let group = &*request.requested_model;
+            let Some(capacity) = self.groups.get(group).copied() else {
+                return Ok(None);
+            };
+            if capacity.rpm.is_none() && capacity.tpm.is_none() {
+                return Ok(None);
+            }
+            let model_prefix = format!("{{model_saturation_check:{group}}}");
+            let saturation = self.saturation(&model_prefix, &capacity).await;
+            let pool = self.pool(group, &capacity, priority.as_deref());
+            let pool_prefix = format!("{{priority_model:{}}}", pool.name);
+
+            // The group counts every request against its whole capacity; a zero limit is off.
+            let mut pairs = Vec::new();
+            let push = |pairs: &mut Vec<(&'static str, Pair)>,
+                        key: &'static str,
+                        prefix: &str,
+                        limits: [Option<i64>; 2]| {
+                for (kind, limit, increment) in
+                    [(Kind::Requests, limits[0], 1), (Kind::Tokens, limits[1], 0)]
+                {
+                    if let Some(limit) = limit {
+                        pairs.push((
+                            key,
+                            Pair {
+                                prefix: prefix.to_owned(),
+                                kind,
+                                limit,
+                                increment,
+                            },
+                        ));
+                    }
+                }
+            };
+            push(
+                &mut pairs,
+                "model_saturation_check",
+                &model_prefix,
+                [
+                    capacity.rpm.filter(|rpm| *rpm != 0),
+                    capacity.tpm.filter(|tpm| *tpm != 0),
+                ],
+            );
+            let share = |limit: Option<i64>| limit.map(|limit| (limit as f64 * pool.share) as i64);
+            let pool_limits = [share(capacity.rpm), share(capacity.tpm)];
+            let reserving =
+                self.settings.reservation.is_some() && pool_limits.iter().any(Option::is_some);
+            let enforced = reserving && saturation >= self.settings.saturation_threshold;
+            if enforced {
+                push(&mut pairs, "priority_model", &pool_prefix, pool_limits);
+            }
+            let checks: Vec<Pair> = pairs.iter().map(|(_, pair)| pair.clone()).collect();
+            let (reserved, shared) = self
+                .store
+                .check_and_add(&checks)
+                .await
+                .map_err(|Unavailable| self.unavailable(request))?;
+            let counts = match reserved {
+                Reserved::Counts(counts) => counts,
+                Reserved::Over { index, current } => {
+                    let (key, pair) = &pairs[index];
+                    let remaining = pair.limit - current;
+                    let kind = pair.kind.name();
+                    let priority_text = priority.as_deref().unwrap_or("None");
+                    let (message, saturation_header) = match *key {
+                        "model_saturation_check" => (
+                            format!(
+                                "Model capacity reached for {group}. Priority: {priority_text}, Rate limit type: {kind}, Model TPM: {}, Model RPM: {}, Remaining: {remaining}",
+                                limit_text(capacity.tpm),
+                                limit_text(capacity.rpm),
+                            ),
+                            None,
+                        ),
+                        _ => (
+                            format!(
+                                "Priority-based rate limit exceeded. Model: {group}, Priority: {priority_text}, Rate limit type: {kind}, Model TPM: {}, Model RPM: {}, Remaining: {remaining}, Model saturation: {}",
+                                limit_text(capacity.tpm),
+                                limit_text(capacity.rpm),
+                                percent(saturation, 1),
+                            ),
+                            Some(percent(saturation, 2)),
+                        ),
+                    };
+                    let mut headers = vec![
+                        ("retry-after", self.settings.window.to_string()),
+                        ("rate_limit_type", kind.to_owned()),
+                        ("x-litellm-priority", priority_header.to_owned()),
+                    ];
+                    if let Some(saturation) = saturation_header {
+                        headers.push(("x-litellm-saturation", saturation));
+                    }
+                    return Err(crate::rate_limits::throttled(
+                        request.route,
+                        message,
+                        headers,
+                    ));
+                }
+            };
+            let mut statuses: Vec<(&str, Kind, i64, i64)> = pairs
+                .iter()
+                .zip(&counts)
+                .map(|((key, pair), count)| (*key, pair.kind, pair.limit, pair.limit - count))
+                .collect();
+            // Below the threshold the pool's requests are only counted, so its share is known
+            // once it is enforced.
+            if reserving
+                && !enforced
+                && let Some(rpm) = pool_limits[0]
+                && let Ok(counts) = self
+                    .store
+                    .requests(std::slice::from_ref(&pool_prefix))
+                    .await
+            {
+                statuses.push(("priority_model", Kind::Requests, rpm, rpm - counts[0]));
+            }
+            if !request.stream {
+                for (key, kind, limit, remaining) in statuses {
+                    let kind = kind.name();
+                    for (name, value) in [
+                        (format!("x-ratelimit-{key}-remaining-{kind}"), remaining),
+                        (format!("x-ratelimit-{key}-limit-{kind}"), limit),
+                    ] {
+                        if let Ok(name) = HeaderName::from_bytes(name.as_bytes()) {
+                            edit.add_response_header(name, HeaderValue::from(value));
+                        }
+                    }
+                }
+            }
+            Ok(Some(Box::new(Charge {
+                store: self.store.clone(),
+                prefixes: [model_prefix, pool_prefix],
+                tokens: self.settings.tokens,
+                shared,
+            }) as Box<dyn Settle>))
+        })
+    }
+}
+
+/// Adds a successful response's tokens to the group and its pool.
+struct Charge {
+    store: Arc<Store>,
+    prefixes: [String; 2],
+    tokens: TokenCount,
+    shared: bool,
+}
+
+impl Settle for Charge {
+    fn settle(self: Box<Self>, record: &Record) {
+        if !(200..300).contains(&record.outcome.status) {
+            return;
+        }
+        let Some(usage) = record.outcome.usage else {
+            return;
+        };
+        let tokens = match self.tokens {
+            TokenCount::Total => usage.total_tokens(),
+            TokenCount::Input => usage.prompt_tokens(),
+            TokenCount::Output => usage.output_tokens,
+        };
+        let Ok(tokens) = i64::try_from(tokens) else {
+            return;
+        };
+        if tokens == 0 {
+            return;
+        }
+        let Charge {
+            store,
+            prefixes,
+            shared,
+            ..
+        } = *self;
+        tokio::spawn(async move {
+            let deltas = prefixes
+                .into_iter()
+                .map(|prefix| (prefix, tokens))
+                .collect();
+            store.charge(deltas, shared).await;
+        });
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use rstest::rstest;
+
+    use super::*;
+
+    fn capacity_hook(reservation: Option<Vec<(String, Share)>>) -> ModelCapacity {
+        ModelCapacity::new(
+            CapacitySettings {
+                reservation,
+                ..CapacitySettings::default()
+            },
+            HashMap::new(),
+            None,
+        )
+    }
+
+    #[rstest]
+    fn pools_follow_the_proxy_allocation() {
+        let capacity = Capacity {
+            rpm: Some(100),
+            tpm: Some(1_000),
+        };
+        let hook = capacity_hook(Some(vec![
+            ("prod".into(), Share::Percent(0.6)),
+            ("dev".into(), Share::Rpm(80.0)),
+        ]));
+        // 0.6 + 0.8 is over-allocated, so both shrink to their proportion.
+        let prod = hook.pool("m", &capacity, Some("prod"));
+        assert_eq!(prod.name, "m:prod");
+        assert!((prod.share - 0.6 / 1.4).abs() < 1e-12);
+        assert_eq!(
+            hook.pool("m", &capacity, Some("other")),
+            Pool {
+                name: "m:default_pool".into(),
+                share: 0.25
+            }
+        );
+        assert_eq!(
+            capacity_hook(None).pool("m", &capacity, Some("prod")).name,
+            "m:default_pool"
+        );
+    }
+
+    #[rstest]
+    #[case(0.625, 1, "62.5%")]
+    #[case(0.5, 2, "50.00%")]
+    fn formats_percentages_like_python(
+        #[case] fraction: f64,
+        #[case] places: usize,
+        #[case] expected: &str,
+    ) {
+        assert_eq!(percent(fraction, places), expected);
+    }
+
+    #[rstest]
+    fn capacities_sum_deployments() {
+        let config = Config::from_yaml(
+            "model_list:
+  - model_name: a
+    litellm_params: {model: openai/x, rpm: 10, tpm: 1000}
+  - model_name: a
+    rpm: 5
+    litellm_params: {model: openai/y, rpm: 99}
+    model_info: {tpm: 500}
+  - model_name: b
+    litellm_params: {model: openai/z}
+",
+        )
+        .unwrap();
+        let groups = capacities(&config);
+        assert_eq!(
+            groups["a"],
+            Capacity {
+                rpm: Some(15),
+                tpm: Some(1_500)
+            }
+        );
+        assert_eq!(groups["b"], Capacity::default());
+    }
+}

litellm-rust/crates/gateway/src/rate_limits.rs

@@ -225,14 +225,14 @@ impl Descriptor {
 }
 
 #[derive(Clone, Copy, Debug, Eq, PartialEq)]
-enum Kind {
+pub(crate) enum Kind {
     Requests,
     Tokens,
     Parallel,
 }
 
 impl Kind {
-    fn name(self) -> &'static str {
+    pub(crate) fn name(self) -> &'static str {
         match self {
             Self::Requests => "requests",
             Self::Tokens => "tokens",
@@ -250,15 +250,25 @@ struct Status {
     remaining: i64,
 }
 
-/// The token reservation went over on the descriptor at `index`, which held `current`.
+/// A counter to check and add to: `{prefix}:requests` or `{prefix}:tokens` in its window.
+#[derive(Clone, Debug, PartialEq)]
+pub(crate) struct Pair {
+    pub(crate) prefix: String,
+    pub(crate) kind: Kind,
+    pub(crate) limit: i64,
+    pub(crate) increment: i64,
+}
+
+/// Every pair's count after adding, or the pair at `index` that would go over, which held
+/// `current`.
 #[derive(Debug, PartialEq)]
-enum Reserved {
+pub(crate) enum Reserved {
     Counts(Vec<i64>),
     Over { index: usize, current: i64 },
 }
 
 /// Redis could not be used and the settings say to refuse rather than fall back.
-struct Unavailable;
+pub(crate) struct Unavailable;
 
 #[derive(Debug)]
 struct Window {
@@ -307,27 +317,24 @@ impl Local {
             .collect()
     }
 
-    fn reserve(
-        &mut self,
-        scopes: &[(String, i64)],
-        increment: i64,
-        now: u64,
-        size: u64,
-    ) -> Reserved {
-        for (index, (prefix, limit)) in scopes.iter().enumerate() {
-            let current = match self.windows.get(prefix) {
-                Some(window) if now.saturating_sub(window.start) < size => window.tokens,
+    fn check_and_add(&mut self, pairs: &[Pair], now: u64, size: u64) -> Reserved {
+        for (index, pair) in pairs.iter().enumerate() {
+            let current = match self.windows.get(&pair.prefix) {
+                Some(window) if now.saturating_sub(window.start) < size => match pair.kind {
+                    Kind::Requests => window.requests,
+                    _ => window.tokens,
+                },
                 _ => 0,
             };
-            if blocked(current, increment, *limit) {
+            if blocked(current, pair.increment, pair.limit) {
                 return Reserved::Over { index, current };
             }
         }
         Reserved::Counts(
-            scopes
+            pairs
                 .iter()
-                .map(|(prefix, _)| {
-                    let (window, expired) = self.window(prefix, now, size);
+                .map(|pair| {
+                    let (window, expired) = self.window(&pair.prefix, now, size);
                     if expired {
                         *window = Window {
                             start: now,
@@ -335,13 +342,26 @@ impl Local {
                             tokens: 0,
                         };
                     }
-                    window.tokens += increment;
-                    window.tokens
+                    let count = match pair.kind {
+                        Kind::Requests => &mut window.requests,
+                        _ => &mut window.tokens,
+                    };
+                    *count += pair.increment;
+                    *count
                 })
                 .collect(),
         )
     }
 
+    /// A counter's value in its current window.
+    fn read(&self, prefix: &str, kind: Kind, now: u64, size: u64) -> Option<i64> {
+        let window = self.windows.get(prefix)?;
+        (now.saturating_sub(window.start) < size).then_some(match kind {
+            Kind::Requests => window.requests,
+            _ => window.tokens,
+        })
+    }
+
     fn charge(&mut self, deltas: &[(String, i64)]) {
         for (prefix, delta) in deltas {
             if let Some(window) = self.windows.get_mut(prefix) {
@@ -380,7 +400,7 @@ fn blocked(current: i64, increment: i64, limit: i64) -> bool {
 }
 
 /// Where the counters live.
-struct Store {
+pub(crate) struct Store {
     redis: Option<Arc<Coordination>>,
     local: Mutex<Local>,
     window: u64,
@@ -388,6 +408,15 @@ struct Store {
 }
 
 impl Store {
+    pub(crate) fn new(redis: Option<Arc<Coordination>>, window: u64, fail_closed: bool) -> Self {
+        Self {
+            redis,
+            local: Mutex::default(),
+            window,
+            fail_closed,
+        }
+    }
+
     fn connection(&self) -> Option<(ConnectionManager, &Coordination)> {
         let redis = self.redis.as_deref()?;
         Some((redis.connection()?, redis))
@@ -414,7 +443,7 @@ impl Store {
     }
 
     /// Counts the request on each prefix's window and returns the counts.
-    async fn requests(&self, prefixes: &[String]) -> Result<Vec<i64>, Unavailable> {
+    pub(crate) async fn requests(&self, prefixes: &[String]) -> Result<Vec<i64>, Unavailable> {
         let now = epoch_seconds();
         if self.redis.is_some() {
             let error = match self.connection() {
@@ -500,21 +529,43 @@ impl Store {
         &self,
         scopes: &[(String, i64)],
         increment: i64,
+    ) -> Result<(Reserved, bool), Unavailable> {
+        let pairs: Vec<Pair> = scopes
+            .iter()
+            .map(|(prefix, limit)| Pair {
+                prefix: prefix.clone(),
+                kind: Kind::Tokens,
+                limit: *limit,
+                increment,
+            })
+            .collect();
+        self.check_and_add(&pairs).await
+    }
+
+    /// Adds every pair's increment, or none when one would go over its limit. `true` when the
+    /// counters are in Redis.
+    pub(crate) async fn check_and_add(
+        &self,
+        pairs: &[Pair],
     ) -> Result<(Reserved, bool), Unavailable> {
         if self.redis.is_some() {
             let error = match self.connection() {
                 Some((mut connection, redis)) => {
                     let mut command = redis::cmd("EVAL");
-                    command.arg(RESERVE).arg(scopes.len() * 2);
-                    for (prefix, _) in scopes {
+                    command.arg(RESERVE).arg(pairs.len() * 2);
+                    for pair in pairs {
                         command
-                            .arg(redis.namespaced(&format!("{prefix}:window")))
-                            .arg(redis.namespaced(&format!("{prefix}:tokens")));
+                            .arg(redis.namespaced(&format!("{}:window", pair.prefix)))
+                            .arg(redis.namespaced(&format!(
+                                "{}:{}",
+                                pair.prefix,
+                                pair.kind.name()
+                            )));
                     }
-                    for (_, limit) in scopes {
+                    for pair in pairs {
                         command
-                            .arg(*limit)
-                            .arg(increment)
+                            .arg(pair.limit)
+                            .arg(pair.increment)
                             .arg(self.window)
                             .arg(self.window);
                     }
@@ -528,7 +579,7 @@ impl Store {
                             ));
                         }
                         Ok(reply) => match reply.as_slice() {
-                            [1, index, current, _] if (1..=scopes.len() as i64).contains(index) => {
+                            [1, index, current, _] if (1..=pairs.len() as i64).contains(index) => {
                                 return Ok((
                                     Reserved::Over {
                                         index: (*index - 1) as usize,
@@ -548,13 +599,34 @@ impl Store {
         }
         Ok((
             self.local()
-                .reserve(scopes, increment, epoch_seconds(), self.window),
+                .check_and_add(pairs, epoch_seconds(), self.window),
             false,
         ))
     }
 
+    /// A counter's current value, `None` when it is not set.
+    pub(crate) async fn read(&self, prefix: &str, kind: Kind) -> Option<i64> {
+        if let Some((mut connection, redis)) = self.connection() {
+            let key = redis.namespaced(&format!("{prefix}:{}", kind.name()));
+            match redis::cmd("GET")
+                .arg(key)
+                .query_async::<Option<String>>(&mut connection)
+                .await
+            {
+                Ok(value) => {
+                    return value
+                        .and_then(|value| value.parse::<f64>().ok())
+                        .map(|value| value as i64);
+                }
+                Err(error) => tracing::warn!(%error, "reading a rate limit counter failed"),
+            }
+        }
+        self.local()
+            .read(prefix, kind, epoch_seconds(), self.window)
+    }
+
     /// Adds each delta to its prefix's tokens.
-    async fn charge(&self, deltas: Vec<(String, i64)>, shared: bool) {
+    pub(crate) async fn charge(&self, deltas: Vec<(String, i64)>, shared: bool) {
         if shared && let Some((mut connection, redis)) = self.connection() {
             let mut command = redis::cmd("EVAL");
             command.arg(CHARGE).arg(deltas.len());
@@ -711,12 +783,7 @@ impl RateLimiter {
         database: Option<PgPool>,
     ) -> Self {
         Self {
-            store: Arc::new(Store {
-                redis,
-                local: Mutex::default(),
-                window: settings.window,
-                fail_closed: settings.fail_closed,
-            }),
+            store: Arc::new(Store::new(redis, settings.window, settings.fail_closed)),
             settings,
             end_users: database.map(EndUsers::new),
         }
@@ -801,22 +868,15 @@ impl RateLimiter {
             status.limit,
             status.remaining.max(0),
         );
-        let mut refusal = match request.route {
-            Route::Messages => Refusal::new(429, "rate_limit_error", message).anthropic(),
-            Route::ChatCompletions => {
-                let mut refusal = Refusal::new(429, "throttling_error", message).without_param();
-                for (name, value) in [
-                    ("retry-after", self.settings.window.to_string()),
-                    ("rate_limit_type", status.kind.name().to_owned()),
-                    ("reset_at", reset_at),
-                ] {
-                    if let Ok(value) = HeaderValue::from_str(&value) {
-                        refusal = refusal.with_header(HeaderName::from_static(name), value);
-                    }
-                }
-                refusal
-            }
-        };
+        let mut refusal = throttled(
+            request.route,
+            message,
+            vec![
+                ("retry-after", self.settings.window.to_string()),
+                ("rate_limit_type", status.kind.name().to_owned()),
+                ("reset_at", reset_at),
+            ],
+        );
         for (name, value) in key_limit_headers(request.key.as_deref()) {
             refusal = refusal.with_header(name, value);
         }
@@ -824,11 +884,37 @@ impl RateLimiter {
     }
 
     fn unavailable(&self, request: &RequestFacts) -> Refusal {
-        match request.route {
-            Route::Messages => Refusal::new(503, "api_error", FAIL_CLOSED).anthropic(),
-            Route::ChatCompletions => {
-                Refusal::new(503, "internal_server_error", FAIL_CLOSED).without_param()
+        unavailable(request.route)
+    }
+}
+
+/// The proxy's 429 from one of its own limiters: on Chat Completions its error body with the
+/// limiter's headers, on `/v1/messages` Anthropic's envelope, where the proxy drops them.
+pub(crate) fn throttled(
+    route: Route,
+    message: String,
+    headers: Vec<(&'static str, String)>,
+) -> Refusal {
+    match route {
+        Route::Messages => Refusal::new(429, "rate_limit_error", message).anthropic(),
+        Route::ChatCompletions => {
+            let mut refusal = Refusal::new(429, "throttling_error", message).without_param();
+            for (name, value) in headers {
+                if let Ok(value) = HeaderValue::from_str(&value) {
+                    refusal = refusal.with_header(HeaderName::from_static(name), value);
+                }
             }
+            refusal
+        }
+    }
+}
+
+/// The proxy's 503 when counters cannot be verified and enforcement fails closed.
+pub(crate) fn unavailable(route: Route) -> Refusal {
+    match route {
+        Route::Messages => Refusal::new(503, "api_error", FAIL_CLOSED).anthropic(),
+        Route::ChatCompletions => {
+            Refusal::new(503, "internal_server_error", FAIL_CLOSED).without_param()
         }
     }
 }
@@ -1150,35 +1236,35 @@ mod tests {
         assert_eq!(estimate(&request, Some(&key), "other", 100_000).tokens, 10);
     }
 
+    fn reserve(local: &mut Local, increment: i64, now: u64) -> Reserved {
+        let pair = Pair {
+            prefix: "{api_key:h}".into(),
+            kind: Kind::Tokens,
+            limit: 10,
+            increment,
+        };
+        local.check_and_add(&[pair], now, 60)
+    }
+
     #[rstest]
     fn local_windows_count_reserve_and_reset() {
         let mut local = Local::default();
         let prefixes = vec!["{api_key:h}".to_owned()];
         assert_eq!(local.requests(&prefixes, 100, 60), [1]);
         assert_eq!(local.requests(&prefixes, 159, 60), [2]);
-        let scopes = vec![("{api_key:h}".to_owned(), 10)];
+        assert_eq!(reserve(&mut local, 6, 159), Reserved::Counts(vec![6]));
         assert_eq!(
-            local.reserve(&scopes, 6, 159, 60),
-            Reserved::Counts(vec![6])
-        );
-        assert_eq!(
-            local.reserve(&scopes, 6, 159, 60),
+            reserve(&mut local, 6, 159),
             Reserved::Over {
                 index: 0,
                 current: 6
             }
         );
         local.charge(&[("{api_key:h}".to_owned(), -4)]);
-        assert_eq!(
-            local.reserve(&scopes, 8, 159, 60),
-            Reserved::Counts(vec![10])
-        );
+        assert_eq!(reserve(&mut local, 8, 159), Reserved::Counts(vec![10]));
         // The window started at 100 is over at 160: requests and tokens start again.
         assert_eq!(local.requests(&prefixes, 160, 60), [1]);
-        assert_eq!(
-            local.reserve(&scopes, 10, 160, 60),
-            Reserved::Counts(vec![10])
-        );
+        assert_eq!(reserve(&mut local, 10, 160), Reserved::Counts(vec![10]));
     }
 
     #[rstest]

litellm-rust/crates/gateway/tests/model_capacity.rs

@@ -0,0 +1,306 @@
+//! The model capacity hook on real requests: a group held to its deployments' summed capacity,
+//! priorities held to their share, refusals and headers as the proxy sends them, and counters
+//! shared through Redis. Set `LITELLM_TEST_REDIS_URL` to also run the shared case.
+
+use std::{collections::HashMap, sync::Arc, time::Duration};
+
+use axum::{
+    Router,
+    body::{Body, to_bytes},
+    http::{HeaderMap, Request, header::CONTENT_TYPE},
+    response::IntoResponse,
+};
+use litellm_auth_types::SecretValue;
+use litellm_config::Config;
+use litellm_gateway::{
+    coordination::{Coordination, RedisSettings},
+    model_capacity::{CapacitySettings, ModelCapacity},
+};
+use litellm_gateway_auth::keys::{
+    AllowedRoutes, KeyHash, KeyLookup, KeyModels, KeyProfile, KeyStatus, LookupFuture,
+};
+use rstest::rstest;
+use serde_json::{Value, json};
+use tokio::net::TcpListener;
+use tower::ServiceExt;
+
+/// 3 input and 4 output tokens: 7 count against the group.
+const MESSAGE: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-x","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":3,"output_tokens":4}}"#;
+
+async fn upstream() -> String {
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let address = listener.local_addr().unwrap();
+    let app = Router::new().fallback(|_: Request<Body>| async {
+        ([(CONTENT_TYPE, "application/json")], MESSAGE).into_response()
+    });
+    tokio::spawn(async move { axum::serve(listener, app).await });
+    format!("http://{address}")
+}
+
+struct Keys(HashMap<KeyHash, KeyStatus>);
+
+impl KeyLookup for Keys {
+    fn lookup<'a>(&'a self, hash: &'a KeyHash) -> LookupFuture<'a> {
+        Box::pin(async move { Ok(self.0.get(hash).cloned()) })
+    }
+}
+
+/// Two deployments of `claude`, each with `limits` (such as `rpm: 1`).
+async fn gateway(limits: &str, settings: &str, redis: Option<Arc<Coordination>>) -> Router {
+    let address = upstream().await;
+    let config = Config::from_yaml(&format!(
+        "model_list:
+  - model_name: claude
+    litellm_params: {{model: anthropic/claude-x, api_base: {address}, api_key: sk-a, {limits}}}
+  - model_name: claude
+    litellm_params: {{model: anthropic/claude-y, api_base: {address}, api_key: sk-b, {limits}}}
+router_settings:
+  num_retries: 0
+general_settings:
+  master_key: sk-master
+litellm_settings:
+  callbacks: [dynamic_rate_limiter_v3]
+{settings}"
+    ))
+    .unwrap();
+    let inference = litellm_gateway::build_inference(&config).unwrap();
+    let hook = ModelCapacity::for_gateway(&config, redis).unwrap();
+    litellm_gateway::hooks::install(&inference, &config, vec![Arc::new(hook)], Vec::new())
+        .await
+        .unwrap();
+    let mut metadata = serde_json::Map::new();
+    metadata.insert("priority".into(), json!("prod"));
+    let key = KeyStatus::Active {
+        expires_at: None,
+        models: KeyModels::default(),
+        allowed_routes: AllowedRoutes::default(),
+        team_id: None,
+        profile: Arc::new(KeyProfile {
+            metadata: Arc::new(metadata),
+            ..KeyProfile::default()
+        }),
+    };
+    let keys: Arc<dyn KeyLookup> = Arc::new(Keys(HashMap::from([(
+        KeyHash::from_token(&SecretValue::new("sk-prod")),
+        key,
+    )])));
+    litellm_gateway::router_with_keys(inference, &config, None, None, Some(keys))
+}
+
+struct Reply {
+    status: u16,
+    headers: HeaderMap,
+    body: Value,
+}
+
+async fn send(app: &Router, path: &str, token: &str) -> Reply {
+    let body = json!({
+        "model": "claude",
+        "max_tokens": 8,
+        "messages": [{"role": "user", "content": "hi"}],
+    });
+    let response = app
+        .clone()
+        .oneshot(
+            Request::post(path)
+                .header("content-type", "application/json")
+                .header("authorization", format!("Bearer {token}"))
+                .body(Body::from(body.to_string()))
+                .unwrap(),
+        )
+        .await
+        .unwrap();
+    let status = response.status().as_u16();
+    let headers = response.headers().clone();
+    let body = to_bytes(response.into_body(), 1 << 20).await.unwrap();
+    Reply {
+        status,
+        headers,
+        body: serde_json::from_slice(&body).unwrap_or(Value::Null),
+    }
+}
+
+fn header<'a>(reply: &'a Reply, name: &str) -> Option<&'a str> {
+    reply.headers.get(name).map(|value| value.to_str().unwrap())
+}
+
+async fn settled() {
+    tokio::time::sleep(Duration::from_millis(100)).await;
+}
+
+#[rstest]
+#[tokio::test]
+async fn a_group_is_held_to_its_deployments_capacity() {
+    let app = gateway("rpm: 1", "", None).await;
+    let first = send(&app, "/v1/messages", "sk-master").await;
+    assert_eq!(first.status, 200);
+    assert_eq!(
+        header(
+            &first,
+            "x-ratelimit-model_saturation_check-remaining-requests"
+        ),
+        Some("1")
+    );
+    assert_eq!(
+        header(&first, "x-ratelimit-model_saturation_check-limit-requests"),
+        Some("2")
+    );
+    assert_eq!(header(&first, "x-litellm-rate-limiter-version"), Some("v3"));
+    assert_eq!(header(&first, "x-litellm-priority"), Some("default"));
+    assert_eq!(send(&app, "/v1/messages", "sk-master").await.status, 200);
+
+    let message = "Model capacity reached for claude. Priority: None, Rate limit type: requests, Model TPM: not configured, Model RPM: 2, Remaining: 0";
+    let refused = send(&app, "/v1/chat/completions", "sk-master").await;
+    assert_eq!(refused.status, 429);
+    assert_eq!(
+        refused.body,
+        json!({"error": {"message": message, "type": "throttling_error", "param": null, "code": "429"}})
+    );
+    assert_eq!(header(&refused, "retry-after"), Some("60"));
+    assert_eq!(header(&refused, "rate_limit_type"), Some("requests"));
+    assert_eq!(header(&refused, "x-litellm-priority"), Some("default"));
+    let refused = send(&app, "/v1/messages", "sk-master").await;
+    assert_eq!(
+        refused.body,
+        json!({"type": "error", "error": {"type": "rate_limit_error", "message": message}})
+    );
+    assert_eq!(header(&refused, "retry-after"), None);
+}
+
+#[rstest]
+#[tokio::test]
+async fn tokens_count_after_the_response() {
+    // 2 x 5 tokens of capacity; each response uses 7.
+    let app = gateway("tpm: 5", "", None).await;
+    assert_eq!(send(&app, "/v1/messages", "sk-master").await.status, 200);
+    settled().await;
+    assert_eq!(send(&app, "/v1/messages", "sk-master").await.status, 200);
+    settled().await;
+    let refused = send(&app, "/v1/messages", "sk-master").await;
+    assert_eq!(refused.status, 429);
+    assert_eq!(
+        refused.body["error"]["message"],
+        "Model capacity reached for claude. Priority: None, Rate limit type: tokens, Model TPM: 10, Model RPM: not configured, Remaining: -4"
+    );
+}
+
+#[rstest]
+#[tokio::test]
+async fn busy_groups_hold_priorities_to_their_share() {
+    // 20 requests a window; `prod` gets a fifth once the group is busy, here from the start.
+    let app = gateway(
+        "rpm: 10",
+        "  priority_reservation: {prod: 0.1}\n  priority_reservation_settings: {saturation_threshold: 0.0}\n",
+        None,
+    )
+    .await;
+    assert_eq!(
+        send(&app, "/v1/chat/completions", "sk-prod").await.status,
+        200
+    );
+    let second = send(&app, "/v1/chat/completions", "sk-prod").await;
+    assert_eq!(second.status, 200);
+    assert_eq!(header(&second, "x-litellm-priority"), Some("prod"));
+    assert_eq!(
+        header(&second, "x-ratelimit-priority_model-remaining-requests"),
+        Some("0")
+    );
+    let refused = send(&app, "/v1/chat/completions", "sk-prod").await;
+    assert_eq!(refused.status, 429);
+    assert_eq!(
+        refused.body["error"]["message"],
+        "Priority-based rate limit exceeded. Model: claude, Priority: prod, Rate limit type: requests, Model TPM: not configured, Model RPM: 20, Remaining: 0, Model saturation: 0.0%"
+    );
+    assert_eq!(header(&refused, "x-litellm-saturation"), Some("0.00%"));
+    // Other callers share the default pool and the rest of the group.
+    assert_eq!(
+        send(&app, "/v1/chat/completions", "sk-master").await.status,
+        200
+    );
+}
+
+#[rstest]
+fn the_hook_is_on_only_when_listed() {
+    let config = Config::from_yaml("model_list: []").unwrap();
+    assert_eq!(CapacitySettings::from_config(&config), None);
+    let config =
+        Config::from_yaml("model_list: []\nlitellm_settings:\n  callbacks: model_capacity\n")
+            .unwrap();
+    assert_eq!(
+        CapacitySettings::from_config(&config),
+        Some(CapacitySettings::default())
+    );
+}
+
+async fn redis(url: &str, namespace: &str) -> Arc<Coordination> {
+    let coordination =
+        Coordination::new(RedisSettings::new(url, Some(namespace.to_owned()))).unwrap();
+    let router: Arc<dyn litellm_gateway::coordination::RouterHandle> = Arc::new(NoRouter(
+        litellm_gateway_inference::ModelRouter::from_config(
+            &Config::from_yaml("model_list: []").unwrap(),
+        )
+        .unwrap(),
+    ));
+    coordination.start(router, None, tokio_util::sync::CancellationToken::new());
+    for _ in 0..100 {
+        if coordination.connected() {
+            return coordination;
+        }
+        tokio::time::sleep(Duration::from_millis(20)).await;
+    }
+    panic!("Redis did not connect");
+}
+
+struct NoRouter(litellm_gateway_inference::ModelRouter);
+
+impl litellm_gateway::coordination::RouterHandle for NoRouter {
+    fn router(&self) -> &litellm_gateway_inference::ModelRouter {
+        &self.0
+    }
+}
+
+#[rstest]
+#[tokio::test]
+async fn replicas_share_the_group_counters() {
+    let Ok(url) = std::env::var("LITELLM_TEST_REDIS_URL") else {
+        return;
+    };
+    let namespace = format!("model-capacity-{}", uuid::Uuid::new_v4());
+    let limits = "rpm: 2, tpm: 1000";
+    let first = gateway(limits, "", Some(redis(&url, &namespace).await)).await;
+    let second = gateway(limits, "", Some(redis(&url, &namespace).await)).await;
+    for app in [&first, &second, &first, &second] {
+        assert_eq!(send(app, "/v1/messages", "sk-master").await.status, 200);
+    }
+    settled().await;
+    assert_eq!(send(&second, "/v1/messages", "sk-master").await.status, 429);
+
+    let mut connection = redis::Client::open(url.as_str())
+        .unwrap()
+        .get_multiplexed_async_connection()
+        .await
+        .unwrap();
+    let get = async |connection: &mut redis::aio::MultiplexedConnection, key: String| {
+        redis::cmd("GET")
+            .arg(key)
+            .query_async::<Option<String>>(connection)
+            .await
+            .unwrap()
+    };
+    let group = format!("{namespace}:{{model_saturation_check:claude}}");
+    // The refused request is not counted; four responses used 7 tokens each.
+    assert_eq!(
+        get(&mut connection, format!("{group}:requests"))
+            .await
+            .as_deref(),
+        Some("4")
+    );
+    assert_eq!(
+        get(&mut connection, format!("{group}:tokens"))
+            .await
+            .as_deref(),
+        Some("28")
+    );
+    let pool = format!("{namespace}:{{priority_model:claude:default_pool}}:tokens");
+    assert_eq!(get(&mut connection, pool).await.as_deref(), Some("28"));
+}

litellm-rust/docs/customer-proxy/model-capacity.svg

@@ -0,0 +1,52 @@
+<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 1000 560" font-family="-apple-system, Segoe UI, Helvetica, Arial, sans-serif" font-size="14">
+  <defs>
+    <marker id="a" viewBox="0 0 10 10" refX="9" refY="5" markerWidth="7" markerHeight="7" orient="auto"><path d="M0 0L10 5L0 10z" fill="#475569"/></marker>
+    <marker id="r" viewBox="0 0 10 10" refX="9" refY="5" markerWidth="7" markerHeight="7" orient="auto"><path d="M0 0L10 5L0 10z" fill="#dc2626"/></marker>
+  </defs>
+  <rect width="1000" height="560" fill="#ffffff"/>
+  <text x="24" y="34" font-size="18" font-weight="600" fill="#0f172a">Model capacity hook (dynamic_rate_limiter_v3): fleet-wide group capacity and priority shares</text>
+
+  <rect x="24" y="60" width="300" height="160" rx="8" fill="#eff6ff" stroke="#3b82f6"/>
+  <text x="38" y="84" font-weight="600" fill="#1e3a8a">At startup</text>
+  <text x="38" y="106" fill="#1e293b" font-size="13">group capacity = sum of deployments'</text>
+  <text x="38" y="124" fill="#1e293b" font-size="13">   rpm / tpm (deployment, litellm_params,</text>
+  <text x="38" y="142" fill="#1e293b" font-size="13">   model_info)</text>
+  <text x="38" y="164" fill="#1e293b" font-size="13">priority_reservation shares: percent,</text>
+  <text x="38" y="182" fill="#1e293b" font-size="13">   rpm or tpm; normalized when over 1</text>
+  <text x="38" y="204" fill="#64748b" font-size="13">on when callbacks list the hook</text>
+
+  <rect x="360" y="60" width="616" height="160" rx="8" fill="#fffbeb" stroke="#f59e0b"/>
+  <text x="374" y="84" font-weight="600" fill="#92400e">Per request (one Redis script, all or nothing)</text>
+  <text x="374" y="106" fill="#1e293b" font-size="13">priority: team metadata, else key metadata; pool &lt;group&gt;:&lt;priority&gt; or &lt;group&gt;:default_pool</text>
+  <text x="374" y="128" fill="#1e293b" font-size="13">saturation = max(requests / rpm, tokens / tpm), read and reused for saturation_check_cache_ttl</text>
+  <text x="374" y="150" fill="#1e293b" font-size="13">{model_saturation_check:&lt;group&gt;}: +1 request, tokens checked; always enforced</text>
+  <text x="374" y="172" fill="#1e293b" font-size="13">{priority_model:&lt;pool&gt;}: share x capacity, enforced at saturation &gt;= threshold,</text>
+  <text x="374" y="190" fill="#1e293b" font-size="13">   only counted below it</text>
+  <line x1="324" y1="140" x2="356" y2="140" stroke="#475569" stroke-width="1.5" marker-end="url(#a)"/>
+
+  <rect x="24" y="256" width="460" height="130" rx="8" fill="#fef2f2" stroke="#dc2626"/>
+  <text x="38" y="280" font-weight="600" fill="#991b1b">Refusal: 429, the proxy's messages</text>
+  <text x="38" y="302" fill="#1e293b" font-size="13">"Model capacity reached for &lt;group&gt;. Priority: ..., Rate limit</text>
+  <text x="38" y="320" fill="#1e293b" font-size="13">  type: ..., Model TPM: ..., Model RPM: ..., Remaining: ..."</text>
+  <text x="38" y="340" fill="#1e293b" font-size="13">"Priority-based rate limit exceeded. ... Model saturation: 62.5%"</text>
+  <text x="38" y="362" fill="#1e293b" font-size="13">chat: retry-after, rate_limit_type, x-litellm-priority,</text>
+  <text x="38" y="378" fill="#1e293b" font-size="13">   x-litellm-saturation; messages: Anthropic envelope only</text>
+  <path d="M500 220 L500 240 L254 240 L254 252" fill="none" stroke="#dc2626" stroke-width="1.5" marker-end="url(#r)"/>
+
+  <rect x="516" y="256" width="460" height="130" rx="8" fill="#f8fafc" stroke="#64748b"/>
+  <text x="530" y="280" font-weight="600" fill="#0f172a">Admitted: request runs</text>
+  <text x="530" y="302" fill="#334155" font-size="13">complete responses carry x-litellm-priority,</text>
+  <text x="530" y="320" fill="#334155" font-size="13">x-litellm-rate-limiter-version: v3 and</text>
+  <text x="530" y="338" fill="#334155" font-size="13">x-ratelimit-model_saturation_check-* /</text>
+  <text x="530" y="356" fill="#334155" font-size="13">x-ratelimit-priority_model-*</text>
+  <path d="M760 220 L760 252" fill="none" stroke="#475569" stroke-width="1.5" marker-end="url(#a)"/>
+
+  <rect x="24" y="420" width="952" height="56" rx="8" fill="#f5f3ff" stroke="#8b5cf6"/>
+  <text x="38" y="444" font-weight="600" fill="#5b21b6">Settle: a successful response's tokens (token_rate_limit_type: total, input or output)</text>
+  <text x="38" y="464" fill="#1e293b" font-size="13">added to the group's and the admitted pool's :tokens, with the window TTL</text>
+  <line x1="746" y1="386" x2="746" y2="416" stroke="#475569" stroke-width="1.5" marker-end="url(#a)"/>
+
+  <rect x="24" y="496" width="952" height="52" rx="8" fill="#ecfdf5" stroke="#10b981"/>
+  <text x="38" y="518" font-weight="600" fill="#065f46">Same keys and scripts as the Python proxy, under its namespace: Rust and Python replicas share one count</text>
+  <text x="38" y="538" fill="#64748b" font-size="13">No Redis, or Redis unreachable: this replica's counters (fail_closed_rate_limit_enforcement: 503 instead)</text>
+</svg>