scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-04 07:27:51 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 195

perf(gateway): admit rate limits in one redis round trip and send scripts by hash (#47197)

Browse source 
* feat(gateway): hooks and observers around inference requests, with usage and cost

Adds the litellm-gateway-hooks crate: request facts built only when something
reads them, on-path hooks that declare what they read and what they may
change and run in config order, off-path observers on bounded queues, a
usage tap for SSE and JSON responses, and per-deployment prices from the
model cost map overlaid by config prices. The messages and chat handlers run
through it when anything is installed; with nothing installed requests take
the old path.

The key store now also returns the key alias, user, team alias and key and
team metadata, and router attempts carry the deployment id.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5

* fix(model-catalog): read token counts the cost map writes as whole floats

The published cost map now has prompt_cache_min_tokens: 4096.0, which the
typed u64 fields rejected, so the gateway dropped the whole map at startup
and priced every request at zero. Counts accept any whole non-negative
number, as Python reads any JSON number, and a test keeps the checked-in
map parsing.

* feat(gateway): write spend to the proxy's database from a native observer

Adds a spend writer observer that records every gateway request in the
proxy's own Postgres tables with the same statements the Python proxy
uses: spend logs, key/user/team/team member/end user spend, and the
daily user, team and end-user activity tables. Writes are batched on the
proxy's timers and are additive, so replicas and Python proxies can
write side by side. The database pool is now shared between virtual
keys and the writer.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5

* feat(gateway): enforce key, team, user, member and end-user budgets on shared spend counters

Adds a built-in budgets hook that runs the proxy's budget checks, in its
order and with its messages, on the proxy's Redis spend:* counters. Each
request reads and seeds the counters, checks them and reserves its
maximum cost in one Redis script, then settles the reservation to its
real cost off the request path, so replicas and Python proxies sharing
the Redis enforce one budget. Without Redis the counters are kept per
replica. The key store now reads user, team member and end-user budgets,
and auth no longer checks budgets against stored spend.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5

* feat(gateway): enforce key, user, team, member and end-user rate limits on shared counters

Adds a built-in rate limits hook with the proxy's v3 limiter semantics:
requests per window, tokens per window and in-flight requests per key,
on the proxy's own {scope:id}:window|requests|tokens|max_parallel_requests
Redis keys and with its scripts, so replicas and Python proxies sharing
the Redis enforce one limit. Tokens are reserved from the proxy's
estimate and settled to usage when the response ends. Refusals carry the
proxy's message, body and headers on each endpoint, and complete
responses carry its x-ratelimit-* headers. Without Redis the counters are
kept per replica. The key store now reads key, user, team and member
rpm/tpm limits and the key's max_parallel_requests, and end users' budget
limits are read with their budgets; keys with plain rate limits are no
longer refused.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5

* test(gateway-auth): expect the narrowed unsupported rate limit message

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5

* perf(gateway): admit rate limits in one redis round trip and send scripts by hash

A key with requests, in-flight and tokens limits made three sequential EVALs
before its request was routed, each carrying the whole Lua script. One script
now counts the request, takes the slot and reserves the tokens, stopping at the
first limit it is over (giving the slot back when tokens are over), with the
same keys and semantics as before. Every rate limit and budget script is sent
with EVALSHA and loaded on NOSCRIPT.

---------

Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Nate Armstrong 2026-10-19 14:24:02 -07:00 • committed by GitHub
parent ac33b81be5
commit 8fb0643e35
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 535 additions and 195 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm-rust/crates/gateway/Cargo.toml
View file

@ -31,7 +31,7 @@ litellm-db.workspace = true
litellm-http = { workspace = true, features = ["mcp"] }
litellm-llms.workspace = true
litellm-router.workspace = true
redis = { version = "1.7.0", default-features = false, features = ["connection-manager", "tokio-comp", "tokio-rustls-comp"] }
redis = { version = "1.7.0", default-features = false, features = ["connection-manager", "script", "tokio-comp", "tokio-rustls-comp"] }
litellm-secrets.workspace = true
litellm-tracing.workspace = true
serde_json.workspace = true

24
litellm-rust/crates/gateway/src/budgets.rs
View file

@ -11,7 +11,7 @@
use std::{
collections::HashMap,
sync::{Arc, Mutex, OnceLock},
sync::{Arc, LazyLock, Mutex, OnceLock},
time::{Duration, Instant},
};
@ -117,6 +117,10 @@ end
return #KEYS
"#;
/// Sent by hash, loaded once per Redis.
static ADMIT_SCRIPT: LazyLock<redis::Script> = LazyLock::new(|| redis::Script::new(ADMIT));
static SETTLE_SCRIPT: LazyLock<redis::Script> = LazyLock::new(|| redis::Script::new(SETTLE));
#[derive(Clone, Debug, PartialEq)]
pub struct BudgetSettings {
/// `litellm_settings.budget_exceeded_status_code`, 422 by default.
@ -459,19 +463,15 @@ impl Counters {
async fn settle(&self, keys: Vec<String>, changes: Vec<Change>, shared: bool) {
if shared && let Some((mut connection, redis)) = self.connection() {
let keys: Vec<String> = keys.iter().map(|key| redis.namespaced(key)).collect();
let mut command = redis::cmd("EVAL");
command
.arg(SETTLE)
.arg(keys.len())
.arg(&keys)
.arg(self.ttl.as_secs());
let mut command = SETTLE_SCRIPT.prepare_invoke();
command.key(&keys).arg(self.ttl.as_secs());
for change in &changes {
command
.arg(change.floor.to_string())
.arg(change.delta.to_string())
.arg(change.fresh.to_string());
}
if let Err(error) = command.query_async::<i64>(&mut connection).await {
if let Err(error) = command.invoke_async::<i64>(&mut connection).await {
// A counter left half-written would mislead every replica; drop it so the next
// read reseeds it from the database, as the proxy does.
tracing::warn!(%error, "updating spend counters failed; dropping them");
@ -518,11 +518,9 @@ async fn admit_redis(
estimate: f64,
ttl: Duration,
) -> redis::RedisResult<Admitted> {
let mut command = redis::cmd("EVAL");
let mut command = ADMIT_SCRIPT.prepare_invoke();
command
.arg(ADMIT)
.arg(keys.len())
.arg(keys)
.key(keys)
.arg(ttl.as_secs())
.arg(estimate.to_string())
.arg(
@ -553,7 +551,7 @@ async fn admit_redis(
"0"
});
}
let reply: Vec<redis::Value> = command.query_async(connection).await?;
let reply: Vec<redis::Value> = command.invoke_async(connection).await?;
let text = |value: &redis::Value| -> Option<String> {
match value {
redis::Value::BulkString(bytes) => String::from_utf8(bytes.clone()).ok(),

630
litellm-rust/crates/gateway/src/rate_limits.rs
View file

@ -12,7 +12,7 @@
use std::{
collections::HashMap,
sync::{Arc, Mutex},
sync::{Arc, LazyLock, Mutex},
time::{Duration, Instant, SystemTime, UNIX_EPOCH},
};
@ -38,6 +38,149 @@ const MAX_OUTPUT_FLOOR: i64 = 1024;
const MAX_LOCAL: usize = 100_000;
const FAIL_CLOSED: &str = "Rate limit enforcement unavailable: request counters could not be verified against Redis, and fail_closed_rate_limit_enforcement is enabled, so the request was rejected to avoid exceeding the configured rate limit. Retry shortly.";
/// Admits a request in one round trip, with the proxy's semantics for each counter: counts it on
/// every `(window, requests)` pair, starting a window that is missing or over; unless one is over
/// its limit, takes an in-flight slot unless the gauge is full; then adds the token estimate to
/// every `(window, tokens)` pair, or to none (returning the slot) when one would go over.
///
/// KEYS: each requests limit's window and counter, the in-flight set when the key caps it, then
/// each tokens limit's window and counter. ARGV: now, window size, the requests limits' count and
/// limits, the slot (`1` or `0`, limit, ttl, id), the token increment, then each tokens limit's
/// limit, ttl and window size.
///
/// Replies `{code, request counts..., in_flight, rest...}`: code 0 admitted with each tokens
/// count, 1 over a requests limit (no in-flight), 2 the gauge was full, 3 over the tokens limit
/// at `index` holding `current`.
const ADMIT: &str = r#"
local now = tonumber(ARGV[1])
local window_size = tonumber(ARGV[2])
local requests = tonumber(ARGV[3])
local arg = 4
local key = 1
local reply = { 0 }
local over = false
for i = 1, requests do
local limit = tonumber(ARGV[arg])
arg = arg + 1
local window_key = KEYS[key]
local counter_key = KEYS[key + 1]
key = key + 2
local window_start = redis.call('GET', window_key)
local counter
if not window_start or (now - tonumber(window_start)) >= window_size then
local prefix = string.sub(window_key, 1, -(#':window') - 1)
redis.call('DEL', prefix .. ':requests', prefix .. ':tokens')
redis.call('SET', window_key, tostring(now))
redis.call('SET', counter_key, 1)
redis.call('EXPIRE', window_key, window_size)
redis.call('EXPIRE', counter_key, window_size)
counter = 1
else
counter = redis.call('INCR', counter_key)
if redis.call('TTL', counter_key) == -1 then
redis.call('EXPIRE', counter_key, window_size)
end
end
if counter > limit then
over = true
end
table.insert(reply, counter)
end
if over then
reply[1] = 1
return reply
end
local slot_key
local slot_id = ARGV[arg + 3]
local in_flight = 0
if ARGV[arg] == '1' then
local limit = tonumber(ARGV[arg + 1])
local slot_ttl = tonumber(ARGV[arg + 2])
slot_key = KEYS[key]
key = key + 1
local time = tonumber(redis.call('TIME')[1])
redis.call('ZREMRANGEBYSCORE', slot_key, '-inf', time - slot_ttl)
in_flight = redis.call('ZCARD', slot_key)
if in_flight + 1 > limit then
reply[1] = 2
table.insert(reply, in_flight)
return reply
end
redis.call('ZADD', slot_key, time, slot_id)
redis.call('EXPIRE', slot_key, slot_ttl)
in_flight = redis.call('ZCARD', slot_key)
end
table.insert(reply, in_flight)
arg = arg + 4
local increment = tonumber(ARGV[arg])
arg = arg + 1
local first = key
local count = (#KEYS - first + 1) / 2
if count == 0 then
return reply
end
local time = tonumber(redis.call('TIME')[1])
local state = {}
for i = 1, count do
local window_key = KEYS[first + (i - 1) * 2]
local counter_key = KEYS[first + (i - 1) * 2 + 1]
local limit = tonumber(ARGV[arg + (i - 1) * 3])
local size = tonumber(ARGV[arg + (i - 1) * 3 + 2])
local window_start = redis.call('GET', window_key)
local expired = (not window_start) or ((time - tonumber(window_start)) >= size)
local current = 0
if not expired then
current = tonumber(redis.call('GET', counter_key) or 0)
end
local blocked
if increment > 0 then
blocked = current + increment > limit
else
blocked = current >= limit
end
if blocked then
if slot_key then
redis.call('ZREM', slot_key, slot_id)
end
reply[1] = 3
table.insert(reply, i)
table.insert(reply, current)
return reply
end
state[i] = expired
end
local reset = {}
for i = 1, count do
local window_key = KEYS[first + (i - 1) * 2]
local counter_key = KEYS[first + (i - 1) * 2 + 1]
local ttl = tonumber(ARGV[arg + (i - 1) * 3 + 1])
local size = tonumber(ARGV[arg + (i - 1) * 3 + 2])
if state[i] then
if not reset[window_key] then
local prefix = string.sub(window_key, 1, -(#':window') - 1)
redis.call('DEL', prefix .. ':requests', prefix .. ':tokens')
reset[window_key] = true
end
redis.call('SET', window_key, tostring(time))
redis.call('SET', counter_key, increment)
redis.call('EXPIRE', window_key, size)
if ttl > 0 then
redis.call('EXPIRE', counter_key, ttl)
end
table.insert(reply, increment)
else
local counter = redis.call('INCRBY', counter_key, increment)
if redis.call('TTL', counter_key) == -1 and ttl > 0 then
redis.call('EXPIRE', counter_key, ttl)
end
table.insert(reply, counter)
end
end
return reply
"#;
/// Counts a request on every `(window, requests)` pair, starting a window that is missing or
/// over, and returns each window's start and count. The caller compares; the count always rises.
const REQUESTS: &str = r#"
@ -134,21 +277,6 @@ end
return results
"#;
/// Takes an in-flight slot unless the gauge is full: `{0, in_flight}` or `{1, 1, in_flight, limit}`.
const ACQUIRE: &str = r#"
local now = tonumber(redis.call('TIME')[1])
local limit = tonumber(ARGV[1])
local slot_ttl = tonumber(ARGV[2])
redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', now - slot_ttl)
local in_flight = redis.call('ZCARD', KEYS[1])
if in_flight + 1 > limit then
return { 1, 1, in_flight, limit }
end
redis.call('ZADD', KEYS[1], now, ARGV[3])
redis.call('EXPIRE', KEYS[1], slot_ttl)
return { 0, redis.call('ZCARD', KEYS[1]) }
"#;
/// Adds each token delta, giving a counter without an expiry the window's.
const CHARGE: &str = r#"
local results = {}
@ -163,6 +291,12 @@ end
return results
"#;
/// Sent by hash, loaded once per Redis.
static ADMIT_SCRIPT: LazyLock<redis::Script> = LazyLock::new(|| redis::Script::new(ADMIT));
static REQUESTS_SCRIPT: LazyLock<redis::Script> = LazyLock::new(|| redis::Script::new(REQUESTS));
static RESERVE_SCRIPT: LazyLock<redis::Script> = LazyLock::new(|| redis::Script::new(RESERVE));
static CHARGE_SCRIPT: LazyLock<redis::Script> = LazyLock::new(|| redis::Script::new(CHARGE));
#[derive(Clone, Debug, PartialEq)]
pub struct RateLimitSettings {
/// `LITELLM_RATE_LIMIT_WINDOW_SIZE`, in seconds; 60 by default.
@ -448,15 +582,14 @@ impl Store {
if self.redis.is_some() {
let error = match self.connection() {
Some((mut connection, redis)) => {
let mut command = redis::cmd("EVAL");
command.arg(REQUESTS).arg(prefixes.len() * 2);
let mut invocation = REQUESTS_SCRIPT.prepare_invoke();
for prefix in prefixes {
command
.arg(redis.namespaced(&format!("{prefix}:window")))
.arg(redis.namespaced(&format!("{prefix}:requests")));
invocation
.key(redis.namespaced(&format!("{prefix}:window")))
.key(redis.namespaced(&format!("{prefix}:requests")));
}
command.arg(now).arg(self.window);
match integers(command.query_async(&mut connection).await) {
invocation.arg(now).arg(self.window);
match integers(invocation.invoke_async(&mut connection).await) {
Ok(reply) if reply.len() == prefixes.len() * 2 => {
return Ok(reply.iter().skip(1).step_by(2).copied().collect());
}
@ -471,32 +604,14 @@ impl Store {
Ok(self.local().requests(prefixes, now, self.window))
}
/// Takes an in-flight slot: the count with it, or the full count without. `true` when the
/// slot is in Redis.
async fn acquire(
&self,
prefix: &str,
limit: i64,
slot: &str,
) -> Result<(Result<i64, i64>, bool), Unavailable> {
/// Counts the request, takes its in-flight slot and reserves its tokens, stopping at the first
/// limit it is over: one script with Redis, else this replica's counters.
async fn admit(&self, plan: &Plan) -> Result<(Admission, bool), Unavailable> {
if self.redis.is_some() {
let error = match self.connection() {
Some((mut connection, redis)) => {
let reply = redis::cmd("EVAL")
.arg(ACQUIRE)
.arg(1)
.arg(redis.namespaced(&format!("{prefix}:max_parallel_requests")))
.arg(limit)
.arg(SLOT_TTL)
.arg(slot)
.query_async(&mut connection)
.await;
match integers(reply) {
Ok(reply) => match reply.as_slice() {
[0, in_flight] => return Ok((Ok(*in_flight), true)),
[1, _, in_flight, _] => return Ok((Err(*in_flight), true)),
_ => Some(malformed()),
},
match admit_redis(&mut connection, redis, plan, self.window).await {
Ok(admission) => return Ok((admission, true)),
Err(error) => Some(error),
}
}
@ -504,7 +619,60 @@ impl Store {
};
self.fall_back(error)?;
}
Ok((self.local().acquire(prefix, limit, slot), false))
Ok((self.admit_local(plan), false))
}
fn admit_local(&self, plan: &Plan) -> Admission {
let mut local = self.local();
let now = epoch_seconds();
let prefixes: Vec<String> = plan
.requests
.iter()
.map(|(prefix, _)| prefix.clone())
.collect();
let requests = local.requests(&prefixes, now, self.window);
if plan
.requests
.iter()
.zip(&requests)
.any(|((_, limit), count)| count > limit)
{
return Admission {
requests,
outcome: Outcome::RequestsOver,
};
}
let in_flight = match &plan.slot {
Some(slot) => match local.acquire(&slot.prefix, slot.limit, &slot.id) {
Ok(in_flight) => Some(in_flight),
Err(in_flight) => {
return Admission {
requests,
outcome: Outcome::SlotFull(in_flight),
};
}
},
None => None,
};
let tokens = match plan.tokens.is_empty() {
true => Vec::new(),
false => match local.check_and_add(&plan.token_pairs(), now, self.window) {
Reserved::Counts(counts) => counts,
Reserved::Over { index, current } => {
if let Some(slot) = &plan.slot {
local.release(&slot.prefix, &slot.id);
}
return Admission {
requests,
outcome: Outcome::TokensOver { index, current },
};
}
},
};
Admission {
requests,
outcome: Outcome::Admitted { in_flight, tokens },
}
}
async fn release(&self, prefix: &str, slot: &str, shared: bool) {
@ -523,25 +691,6 @@ impl Store {
self.local().release(prefix, slot);
}
/// Reserves `increment` tokens on every `(prefix, limit)`, or on none. `true` when the
/// reservation is in Redis.
async fn reserve(
&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(
@ -551,25 +700,24 @@ impl Store {
if self.redis.is_some() {
let error = match self.connection() {
Some((mut connection, redis)) => {
let mut command = redis::cmd("EVAL");
command.arg(RESERVE).arg(pairs.len() * 2);
let mut invocation = RESERVE_SCRIPT.prepare_invoke();
for pair in pairs {
command
.arg(redis.namespaced(&format!("{}:window", pair.prefix)))
.arg(redis.namespaced(&format!(
invocation
.key(redis.namespaced(&format!("{}:window", pair.prefix)))
.key(redis.namespaced(&format!(
"{}:{}",
pair.prefix,
pair.kind.name()
)));
}
for pair in pairs {
command
invocation
.arg(pair.limit)
.arg(pair.increment)
.arg(self.window)
.arg(self.window);
}
match integers(command.query_async(&mut connection).await) {
match integers(invocation.invoke_async(&mut connection).await) {
Ok(reply) if reply.first() == Some(&0) => {
return Ok((
Reserved::Counts(
@ -628,15 +776,17 @@ impl Store {
/// Adds each delta to its prefix's tokens.
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());
let mut invocation = CHARGE_SCRIPT.prepare_invoke();
for (prefix, _) in &deltas {
command.arg(redis.namespaced(&format!("{prefix}:tokens")));
invocation.key(redis.namespaced(&format!("{prefix}:tokens")));
}
for (_, delta) in &deltas {
command.arg(*delta).arg(self.window);
invocation.arg(*delta).arg(self.window);
}
if let Err(error) = command.query_async::<redis::Value>(&mut connection).await {
if let Err(error) = invocation
.invoke_async::<redis::Value>(&mut connection)
.await
{
tracing::warn!(%error, "settling reserved tokens failed");
}
return;
@ -645,6 +795,123 @@ impl Store {
}
}
/// What one request asks of the counters, in the proxy's order.
struct Plan {
/// Each requests limit: its prefix and limit.
requests: Vec<(String, i64)>,
slot: Option<SlotPlan>,
/// Each tokens limit: its prefix and limit.
tokens: Vec<(String, i64)>,
increment: i64,
}
impl Plan {
/// The tokens limits as the pairs to reserve the increment on.
fn token_pairs(&self) -> Vec<Pair> {
self.tokens
.iter()
.map(|(prefix, limit)| Pair {
prefix: prefix.clone(),
kind: Kind::Tokens,
limit: *limit,
increment: self.increment,
})
.collect()
}
}
struct SlotPlan {
prefix: String,
limit: i64,
id: String,
}
/// Each requests count, then how far the request got.
#[derive(Debug, PartialEq)]
struct Admission {
requests: Vec<i64>,
outcome: Outcome,
}
#[derive(Debug, PartialEq)]
enum Outcome {
RequestsOver,
SlotFull(i64),
TokensOver {
index: usize,
current: i64,
},
Admitted {
in_flight: Option<i64>,
tokens: Vec<i64>,
},
}
async fn admit_redis(
connection: &mut ConnectionManager,
redis: &Coordination,
plan: &Plan,
window: u64,
) -> redis::RedisResult<Admission> {
let mut invocation = ADMIT_SCRIPT.prepare_invoke();
invocation
.arg(epoch_seconds())
.arg(window)
.arg(plan.requests.len());
for (prefix, limit) in &plan.requests {
invocation
.key(redis.namespaced(&format!("{prefix}:window")))
.key(redis.namespaced(&format!("{prefix}:requests")))
.arg(*limit);
}
match &plan.slot {
Some(slot) => {
invocation
.key(redis.namespaced(&format!("{}:max_parallel_requests", slot.prefix)))
.arg(1)
.arg(slot.limit)
.arg(SLOT_TTL)
.arg(&slot.id);
}
None => {
invocation.arg(0).arg(0).arg(0).arg("");
}
}
invocation.arg(plan.increment);
for (prefix, limit) in &plan.tokens {
invocation
.key(redis.namespaced(&format!("{prefix}:window")))
.key(redis.namespaced(&format!("{prefix}:tokens")))
.arg(*limit)
.arg(window)
.arg(window);
}
let reply = integers(invocation.invoke_async(connection).await)?;
let counted = plan.requests.len();
let (code, rest) = reply.split_first().ok_or_else(malformed)?;
if rest.len() < counted {
return Err(malformed());
}
let (requests, rest) = rest.split_at(counted);
let requests = requests.to_vec();
let outcome = match (code, rest) {
(1, []) => Outcome::RequestsOver,
(2, [in_flight]) => Outcome::SlotFull(*in_flight),
(3, [_, index, current]) if (1..=plan.tokens.len() as i64).contains(index) => {
Outcome::TokensOver {
index: (*index - 1) as usize,
current: *current,
}
}
(0, [in_flight, tokens @ ..]) if tokens.len() == plan.tokens.len() => Outcome::Admitted {
in_flight: plan.slot.as_ref().map(|_| *in_flight),
tokens: tokens.to_vec(),
},
_ => return Err(malformed()),
};
Ok(Admission { requests, outcome })
}
fn malformed() -> redis::RedisError {
redis::RedisError::from((
redis::ErrorKind::UnexpectedReturnType,
@ -957,76 +1224,21 @@ impl Hook for RateLimiter {
if descriptors.is_empty() {
return Ok(None);
}
let unavailable = |Unavailable| self.unavailable(request);
let mut statuses = Vec::new();
// Requests per window, counted whether or not the request goes through.
// Requests per window (counted whether or not the request goes through), then
// requests in flight for the key only, then tokens per window, reserved from an
// estimate and settled when the response ends.
let windowed: Vec<usize> = (0..descriptors.len())
.filter(|index| descriptors[*index].limits.rpm.is_some())
.collect();
if !windowed.is_empty() {
let prefixes: Vec<String> = windowed
.iter()
.map(|index| descriptors[*index].prefix())
.collect();
let counts = self.store.requests(&prefixes).await.map_err(unavailable)?;
for (index, count) in windowed.iter().zip(counts) {
let limit = descriptors[*index].limits.rpm.unwrap_or_default();
statuses.push(Status {
descriptor: *index,
kind: Kind::Requests,
limit,
remaining: limit - count,
});
}
if let Some(over) = statuses.iter().find(|status| status.remaining < 0) {
return Err(self.refusal(request, &descriptors[over.descriptor], *over));
}
}
// Requests in flight, for the key only.
let mut slot = None;
if let Some(index) = descriptors
let parallel = descriptors
.iter()
.position(|descriptor| descriptor.max_parallel_requests.is_some())
{
let descriptor = &descriptors[index];
let limit = descriptor.max_parallel_requests.unwrap_or_default();
let prefix = descriptor.prefix();
let id = uuid::Uuid::new_v4().simple().to_string();
let (taken, shared) = self
.store
.acquire(&prefix, limit, &id)
.await
.map_err(unavailable)?;
let in_flight = match taken {
Ok(in_flight) => in_flight,
Err(in_flight) => {
let status = Status {
descriptor: index,
kind: Kind::Parallel,
limit,
remaining: limit - in_flight,
};
return Err(self.refusal(request, descriptor, status));
}
};
statuses.push(Status {
descriptor: index,
kind: Kind::Parallel,
limit,
remaining: (limit - in_flight).max(0),
});
slot = Some(Slot { prefix, id, shared });
}
// Tokens per window, reserved from an estimate and settled when the response ends.
.position(|descriptor| descriptor.max_parallel_requests.is_some());
let limited: Vec<(usize, i64)> = descriptors
.iter()
.enumerate()
.filter_map(|(index, descriptor)| Some((index, descriptor.limits.tpm?)))
.collect();
let mut tokens = None;
let mut increment = 0;
if let Some(min_tpm) = limited.iter().map(|(_, limit)| *limit).min() {
let body: Map<String, Value> =
serde_json::from_slice(&request.body()).unwrap_or_default();
@ -1039,54 +1251,110 @@ impl Hook for RateLimiter {
if let Some(max_tokens) = estimate.implicit_max_tokens {
edit.set_upstream_body_field("max_tokens", Value::from(max_tokens));
}
let scopes: Vec<(String, i64)> = limited
increment = estimate.tokens;
}
let plan = Plan {
requests: windowed
.iter()
.map(|index| {
let descriptor = &descriptors[*index];
(
descriptor.prefix(),
descriptor.limits.rpm.unwrap_or_default(),
)
})
.collect(),
slot: parallel.map(|index| SlotPlan {
prefix: descriptors[index].prefix(),
limit: descriptors[index].max_parallel_requests.unwrap_or_default(),
id: uuid::Uuid::new_v4().simple().to_string(),
}),
tokens: limited
.iter()
.map(|(index, limit)| (descriptors[*index].prefix(), *limit))
.collect();
let reserved = self.store.reserve(&scopes, estimate.tokens).await;
let release = |slot: Option<Slot>| async {
if let Some(slot) = slot {
self.store
.release(&slot.prefix, &slot.id, slot.shared)
.await;
}
};
let (reserved, shared) = match reserved {
Ok(reserved) => reserved,
Err(Unavailable) => {
release(slot).await;
return Err(self.unavailable(request));
}
};
match reserved {
Reserved::Counts(counts) => {
for ((index, limit), count) in limited.iter().zip(counts) {
statuses.push(Status {
descriptor: *index,
kind: Kind::Tokens,
limit: *limit,
remaining: (limit - count).max(0),
});
}
}
Reserved::Over { index, current } => {
release(slot).await;
let (descriptor, limit) = limited[index];
let status = Status {
descriptor,
kind: Kind::Tokens,
limit,
remaining: limit - current,
};
return Err(self.refusal(request, &descriptors[descriptor], status));
.collect(),
increment,
};
let (admission, shared) = self
.store
.admit(&plan)
.await
.map_err(|Unavailable| self.unavailable(request))?;
let mut statuses: Vec<Status> = windowed
.iter()
.zip(&admission.requests)
.map(|(index, count)| {
let limit = descriptors[*index].limits.rpm.unwrap_or_default();
Status {
descriptor: *index,
kind: Kind::Requests,
limit,
remaining: limit - count,
}
})
.collect();
let (in_flight, counts) = match admission.outcome {
Outcome::RequestsOver => {
let over = statuses
.iter()
.find(|status| status.remaining < 0)
.copied()
.unwrap_or(statuses[0]);
return Err(self.refusal(request, &descriptors[over.descriptor], over));
}
tokens = Some(Tokens {
prefixes: scopes.into_iter().map(|(prefix, _)| prefix).collect(),
reserved: estimate.tokens,
shared,
Outcome::SlotFull(in_flight) => {
let (index, slot) = (parallel.unwrap_or_default(), &plan.slot);
let limit = slot.as_ref().map_or(0, |slot| slot.limit);
let status = Status {
descriptor: index,
kind: Kind::Parallel,
limit,
remaining: limit - in_flight,
};
return Err(self.refusal(request, &descriptors[index], status));
}
Outcome::TokensOver { index, current } => {
let (descriptor, limit) = limited[index];
let status = Status {
descriptor,
kind: Kind::Tokens,
limit,
remaining: limit - current,
};
return Err(self.refusal(request, &descriptors[descriptor], status));
}
Outcome::Admitted { in_flight, tokens } => (in_flight, tokens),
};
let slot = match (parallel, plan.slot, in_flight) {
(Some(index), Some(slot), Some(in_flight)) => {
statuses.push(Status {
descriptor: index,
kind: Kind::Parallel,
limit: slot.limit,
remaining: (slot.limit - in_flight).max(0),
});
Some(Slot {
prefix: slot.prefix,
id: slot.id,
shared,
})
}
_ => None,
};
for ((index, limit), count) in limited.iter().zip(&counts) {
statuses.push(Status {
descriptor: *index,
kind: Kind::Tokens,
limit: *limit,
remaining: (limit - count).max(0),
});
}
let tokens = (!plan.tokens.is_empty()).then(|| Tokens {
prefixes: plan.tokens.into_iter().map(|(prefix, _)| prefix).collect(),
reserved: increment,
shared,
});
// The proxy sends where each limit stands on complete responses only.
if !request.stream {

74
litellm-rust/crates/gateway/tests/rate_limits.rs
View file

@ -466,6 +466,80 @@ async fn replicas_share_the_proxy_counters() {
assert_eq!(in_flight, 0);
}
#[rstest]
#[tokio::test]
async fn shared_refusals_count_and_hold_as_the_proxy() {
let Some(url) = redis_url() else { return };
let mut connection = redis::Client::open(url.as_str())
.unwrap()
.get_multiplexed_async_connection()
.await
.unwrap();
// Over the tokens limit: the request still counts, and its in-flight slot is given back.
let namespace = format!("rate-limits-{}", uuid::Uuid::new_v4());
let limits = KeyLimits {
key: Limits {
rpm: Some(10),
tpm: Some(1),
},
max_parallel_requests: Some(1),
..KeyLimits::default()
};
let setup = Setup {
redis: Some(redis(&url, &namespace).await),
..Setup::default()
};
let single = gateway(vec![("sk-key", key(None, limits))], setup).await;
let refused = messages(&single.app).await;
assert_eq!(refused.status, 429);
assert!(
refusal_message(&refused, &["error", "message"])
.contains("Limit type: tokens. Current limit: 1, Remaining: 1.")
);
let prefix = format!("{namespace}:{{api_key:{}}}", hash("sk-key").as_str());
let requests: Option<String> = redis::cmd("GET")
.arg(format!("{prefix}:requests"))
.query_async(&mut connection)
.await
.unwrap();
assert_eq!(requests.as_deref(), Some("1"));
let in_flight: i64 = redis::cmd("ZCARD")
.arg(format!("{prefix}:max_parallel_requests"))
.query_async(&mut connection)
.await
.unwrap();
assert_eq!(in_flight, 0);
// A full gauge on one replica refuses on another.
let namespace = format!("rate-limits-{}", uuid::Uuid::new_v4());
let limits = KeyLimits {
max_parallel_requests: Some(1),
..KeyLimits::default()
};
let mut replicas = Vec::new();
for _ in 0..2 {
let setup = Setup {
redis: Some(redis(&url, &namespace).await),
delay: Duration::from_millis(300),
..Setup::default()
};
replicas.push(gateway(vec![("sk-key", key(None, limits))], setup).await);
}
let (first, second) = tokio::join!(messages(&replicas[0].app), async {
tokio::time::sleep(Duration::from_millis(100)).await;
messages(&replicas[1].app).await
});
assert_eq!(first.status, 200);
assert_eq!(second.status, 429);
assert!(
refusal_message(&second, &["error", "message"])
.contains("Limit type: max_parallel_requests. Current limit: 1, Remaining: 0.")
);
settled().await;
assert_eq!(messages(&replicas[1].app).await.status, 200);
}
#[expect(
clippy::disallowed_methods,
reason = "test fixture rows; the lookups under test use the checked macro"