feat(redis): move the spend counter keyspace onto the Rust store
Some checks are pending
LiteLLM Rust DB / db-queries (push) Waiting to run
LiteLLM Rust / rust-lint (push) Waiting to run
LiteLLM Rust / rust-test (push) Waiting to run
LiteLLM Rust / rust-wheel (push) Waiting to run

Every spend:* reader and writer (auth budget checks, reservations and lease
renewal, post-call increments, reseed, resets, key/team/user/tag endpoints,
shadow eval) now goes through litellm/redis/stores/spend_counters.py, backed by
litellm-redis-store's spend_counters module. Keys, value bytes (json.dumps for
SET, str() for the set-max script), TTLs and the set-max Lua script match what
main writes, pinned by matching inline cases in the Rust and Python tests.

The auth budget read rides the request's native pre-call batch through a
native flush hook, so it leaves in the same pipeline as the armed cooldown
read, and post-call increments ride the native post-call batch. The
auth-cache invalidation subscriber forwards spend:* broadcasts to the native
memory tier.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Yujong Lee 2026-11-03 15:55:33 +09:00
parent 70cad58f5e
commit fc4682727e
47 changed files with 3569 additions and 2502 deletions

View file

@ -47,6 +47,10 @@ mod _native {
#[pymodule_export]
use crate::redis::store::router_usage::NativeRouterUsage;
#[pymodule_export]
use crate::redis::store::spend_counters::{
NativeDeclaredSpendCounters, NativeSpendCounterError, NativeSpendCounters,
};
#[pymodule_export]
use crate::routes::audio_transcription::{atranscription, transcription};
#[pymodule_export]
use crate::routes::chat_completions::{acompletion, completion};
@ -128,6 +132,9 @@ mod tests {
"NativeDeclaredCooldowns",
"NativeRouterCooldowns",
"NativeRouterUsage",
"NativeDeclaredSpendCounters",
"NativeSpendCounterError",
"NativeSpendCounters",
"native_redis_breaker_states",
"NativeClickHouseSpendConfig",
"NativeClickHouseSpendStorage",

View file

@ -2,6 +2,7 @@ mod policy;
pub(crate) mod request;
pub(crate) mod router_cooldowns;
pub(crate) mod router_usage;
pub(crate) mod spend_counters;
use litellm_host_python::PythonContext;
use litellm_redis_store::Number;

View file

@ -0,0 +1,329 @@
use std::{
sync::{Arc, Mutex, PoisonError},
time::Duration,
};
use litellm_redis_store::spend_counters::{
DeclaredReads, Error, LostSeed, Reading, SpendCounter, SpendCounters,
};
use pyo3::{
exceptions::{PyRuntimeError, PyValueError},
prelude::*,
types::PyDict,
};
use super::{PyNumber, call, request::NativeRequestRedis, store_error};
use crate::execution::run_async_value;
use crate::redis::client::NativeRedisClient;
pyo3::create_exception!(
_native,
NativeSpendCounterError,
pyo3::exceptions::PyException,
"A spend counter operation failed. Args are (kind, message), kind one of \"breaker_open\", \"timeout\", \"transport\", \"rejected\", \"decode\" or \"memory\"."
);
fn spend_error(error: Error) -> PyErr {
use litellm_redis_client::Error as RedisError;
let kind = match &error {
Error::Redis(RedisError::BreakerOpen) => "breaker_open",
Error::Redis(RedisError::Transport(error)) if error.is_timeout() => "timeout",
Error::Redis(RedisError::Transport(_)) => "transport",
Error::Redis(RedisError::Server(_)) => "rejected",
Error::Redis(RedisError::Decode(_)) | Error::Decode(_) => "decode",
Error::Memory(_) => "memory",
};
NativeSpendCounterError::new_err((kind, litellm_error::Report::new(&error).to_string()))
}
/// A `spend:*` key, refused when it names anything outside the spend counter keyspace.
pub(crate) struct PySpendCounter(SpendCounter);
impl<'a, 'py> FromPyObject<'a, 'py> for PySpendCounter {
type Error = PyErr;
fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
let key: String = value.extract()?;
SpendCounter::parse(&key)
.map(Self)
.map_err(|error| PyValueError::new_err(error.to_string()))
}
}
fn counters(counters: Vec<PySpendCounter>) -> Vec<SpendCounter> {
counters.into_iter().map(|counter| counter.0).collect()
}
fn increments(
increments: Vec<(PySpendCounter, PyNumber)>,
) -> Vec<(SpendCounter, litellm_redis_store::Number)> {
increments
.into_iter()
.map(|(counter, amount)| (counter.0, amount.0))
.collect()
}
fn ttl(seconds: Option<f64>) -> Option<Duration> {
seconds.map(|seconds| Duration::from_secs_f64(seconds.max(0.0)))
}
fn numbers(values: Vec<Option<litellm_redis_store::Number>>) -> Vec<Option<PyNumber>> {
values
.into_iter()
.map(|value| value.map(PyNumber))
.collect()
}
/// The proxy's spend counters, Redis first with this worker's memory copy behind it.
#[pyclass(frozen, name = "NativeSpendCounters")]
pub(crate) struct NativeSpendCounters(Arc<SpendCounters>);
#[pymethods]
impl NativeSpendCounters {
#[new]
#[pyo3(signature = (client, redis_ttl_seconds = 60))]
fn new(client: Option<PyRef<'_, NativeRedisClient>>, redis_ttl_seconds: u64) -> Self {
Self(Arc::new(SpendCounters::new(
client.map(|client| client.0.clone()),
Duration::from_secs(redis_ttl_seconds),
)))
}
/// `("redis" | "memory", value)`: where the answer came from, a Redis miss included.
fn read<'py>(&self, py: Python<'py>, counter: PySpendCounter) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let scope = call(py, "async_get_cache")?;
run_async_value(py, async move {
let reading = scope
.scope(store.read(&counter.0))
.await
.map_err(store_error)?;
Ok(match reading {
Reading::Redis(value) => ("redis", value.map(PyNumber)),
Reading::Memory(value) => ("memory", value.map(PyNumber)),
})
})
}
fn read_many<'py>(
&self,
py: Python<'py>,
counters: Vec<PySpendCounter>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let counters = self::counters(counters);
let scope = call(py, "async_batch_get_cache")?;
run_async_value(py, async move {
Ok(scope.scope(store.read_many(&counters)).await.map(numbers))
})
}
fn declare_read_many(
&self,
request: &NativeRequestRedis,
counters: Vec<PySpendCounter>,
) -> Option<NativeDeclaredSpendCounters> {
self.0
.declare_read_many_pre_call(&request.0, &self::counters(counters))
.map(|declared| NativeDeclaredSpendCounters(Mutex::new(Some(declared))))
}
fn increment<'py>(
&self,
py: Python<'py>,
counter: PySpendCounter,
amount: PyNumber,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let scope = call(py, "async_increment")?;
run_async_value(py, async move {
scope
.scope(store.increment(&counter.0, amount.0))
.await
.map(PyNumber)
.map_err(spend_error)
})
}
fn increment_many<'py>(
&self,
py: Python<'py>,
increments: Vec<(PySpendCounter, PyNumber)>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let increments = self::increments(increments);
let scope = call(py, "async_increment_pipeline")?;
run_async_value(py, async move {
let totals = scope
.scope(store.increment_many(&increments))
.await
.map_err(spend_error)?;
Ok(totals.into_iter().map(PyNumber).collect::<Vec<_>>())
})
}
/// Declares the increments on the request's post-call batch; `False` when there is no Redis.
fn increment_many_post_call(
&self,
request: &NativeRequestRedis,
increments: Vec<(PySpendCounter, PyNumber)>,
) -> bool {
self.0
.increment_many_post_call(&request.0, &self::increments(increments))
}
fn set_max<'py>(
&self,
py: Python<'py>,
counter: PySpendCounter,
spend: PyNumber,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let scope = call(py, "async_set_max")?;
run_async_value(py, async move {
scope
.scope(store.set_max(&counter.0, spend.0))
.await
.map(|stored| stored.map(PyNumber))
.map_err(spend_error)
})
}
#[pyo3(signature = (counter, spend, *, increment_if_lost))]
fn seed<'py>(
&self,
py: Python<'py>,
counter: PySpendCounter,
spend: PyNumber,
increment_if_lost: bool,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let lost = if increment_if_lost {
LostSeed::Increment
} else {
LostSeed::Answer
};
let scope = call(py, "async_set_cache")?;
run_async_value(py, async move {
scope
.scope(store.seed(&counter.0, spend.0, lost))
.await
.map(PyNumber)
.map_err(spend_error)
})
}
#[pyo3(signature = (counter, spend, ttl_seconds = None))]
fn set<'py>(
&self,
py: Python<'py>,
counter: PySpendCounter,
spend: PyNumber,
ttl_seconds: Option<f64>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let scope = call(py, "async_set_cache")?;
run_async_value(py, async move {
scope
.scope(store.set(&counter.0, spend.0, ttl(ttl_seconds)))
.await
.map_err(spend_error)
})
}
fn delete<'py>(
&self,
py: Python<'py>,
counters: Vec<PySpendCounter>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let counters = self::counters(counters);
let scope = call(py, "async_delete_cache")?;
run_async_value(py, async move {
scope
.scope(store.delete(&counters))
.await
.map_err(spend_error)
})
}
fn refresh_ttl<'py>(
&self,
py: Python<'py>,
counter: PySpendCounter,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let scope = call(py, "async_refresh_ttl")?;
run_async_value(py, async move {
scope
.scope(store.refresh_ttl(&counter.0))
.await
.map_err(spend_error)
})
}
fn get_local(&self, counter: PySpendCounter) -> PyResult<Option<PyNumber>> {
Ok(self
.0
.get_local(&counter.0)
.map_err(store_error)?
.map(PyNumber))
}
#[pyo3(signature = (counter, spend, ttl_seconds = None))]
fn set_local(
&self,
counter: PySpendCounter,
spend: PyNumber,
ttl_seconds: Option<f64>,
) -> PyResult<()> {
self.0
.set_local(&counter.0, spend.0, ttl(ttl_seconds))
.map_err(store_error)
}
fn delete_local(&self, counter: PySpendCounter) -> PyResult<()> {
self.0.delete_local(&counter.0).map_err(store_error)
}
fn increment_local(&self, counter: PySpendCounter, amount: PyNumber) -> PyResult<PyNumber> {
self.0
.increment_local(&counter.0, amount.0)
.map(PyNumber)
.map_err(store_error)
}
fn flush_local(&self) -> PyResult<()> {
self.0.flush_local().map_err(store_error)
}
fn snapshot<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyDict>> {
let snapshot = PyDict::new(py);
for (key, spend) in self.0.snapshot().map_err(store_error)? {
snapshot.set_item(key, PyNumber(spend))?;
}
Ok(snapshot)
}
}
/// A `read_many` declared on a request's pre-call batch. `resolve` answers it once.
#[pyclass(frozen, name = "NativeDeclaredSpendCounters")]
pub(crate) struct NativeDeclaredSpendCounters(Mutex<Option<DeclaredReads>>);
#[pymethods]
impl NativeDeclaredSpendCounters {
fn resolve<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
let declared = self
.0
.lock()
.unwrap_or_else(PoisonError::into_inner)
.take()
.ok_or_else(|| {
PyRuntimeError::new_err("this spend counter read was already resolved")
})?;
let scope = call(py, "async_batch_get_cache")?;
run_async_value(py, async move {
Ok(scope.scope(declared.resolve()).await.map(numbers))
})
}
}

View file

@ -3,6 +3,7 @@ mod number;
mod request;
pub mod router_cooldowns;
pub mod router_usage;
pub mod spend_counters;
mod tier;
pub use codec::DecodeError;

View file

@ -0,0 +1,548 @@
use std::{
sync::{Arc, LazyLock, Mutex, MutexGuard, PoisonError},
time::Duration,
};
use futures_util::future::join_all;
use litellm_cache::Error as MemoryError;
use litellm_cache_memory::InMemoryCache;
use litellm_python_compat::{json, repr::to_str};
use litellm_redis_client::{Declared, Expiry, RedisBatch, RedisClient, Script};
use crate::tier::{Clock, system_clock};
use crate::{DecodeError, Number, RequestRedis, codec::python_value};
/// `InMemoryCache`'s defaults, which `spend_counter_cache` runs its memory tier with.
const MEMORY_ENTRIES: usize = 200;
pub const MEMORY_TTL: Duration = Duration::from_secs(600);
/// `spend_counter_cache`'s `default_in_memory_ttl`, which only its memory-only increment applies.
pub const LOCAL_INCREMENT_TTL: Duration = Duration::from_secs(60);
/// `RedisCache`'s TTL when `litellm.default_redis_ttl` is unset.
pub const REDIS_TTL: Duration = Duration::from_secs(60);
/// `RedisCache.async_set_max`, byte for byte: a counter only ever rises to the authoritative spend.
static SET_MAX: LazyLock<Script> = LazyLock::new(|| {
Script::new(concat!(
"local cur = redis.call('GET', KEYS[1]) ",
"if cur == false or tonumber(cur) < tonumber(ARGV[1]) then ",
"redis.call('SET', KEYS[1], ARGV[1]) ",
"if tonumber(ARGV[2]) > 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]) end ",
"return ARGV[1] end ",
"return cur"
))
});
/// One budget's spend counter. `parse` and `key` round-trip `main`'s `spend:*` keys exactly.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub enum SpendCounter {
Key {
token: String,
window: Option<String>,
},
Team {
team_id: String,
window: Option<String>,
},
TeamMember {
user_id: String,
team_id: String,
},
User(String),
EndUser(String),
Org(String),
Tag(String),
Project(String),
ModelAccessGroup(String),
ShadowEval(String),
}
#[derive(Debug, thiserror::Error)]
#[error("{0:?} is not a spend counter")]
pub struct NotASpendCounter(pub String);
impl SpendCounter {
pub fn key(&self) -> String {
let windowed = |scope: &str, id: &str, window: &Option<String>| match window {
Some(window) => format!("spend:{scope}:{id}:window:{window}"),
None => format!("spend:{scope}:{id}"),
};
match self {
Self::Key { token, window } => windowed("key", token, window),
Self::Team { team_id, window } => windowed("team", team_id, window),
Self::TeamMember { user_id, team_id } => {
format!("spend:team_member:{user_id}:{team_id}")
}
Self::User(id) => format!("spend:user:{id}"),
Self::EndUser(id) => format!("spend:end_user:{id}"),
Self::Org(id) => format!("spend:org:{id}"),
Self::Tag(name) => format!("spend:tag:{name}"),
Self::Project(id) => format!("spend:project:{id}"),
Self::ModelAccessGroup(name) => format!("spend:model_access_group:{name}"),
Self::ShadowEval(job_id) => format!("spend:shadow_eval:{job_id}"),
}
}
pub fn parse(key: &str) -> Result<Self, NotASpendCounter> {
let invalid = || NotASpendCounter(key.to_owned());
let rest = key.strip_prefix("spend:").ok_or_else(invalid)?;
let (scope, id) = rest.split_once(':').ok_or_else(invalid)?;
let windowed = || match id.rsplit_once(":window:") {
Some((id, window)) if !window.is_empty() => (id.to_owned(), Some(window.to_owned())),
_ => (id.to_owned(), None),
};
Ok(match scope {
"key" => {
let (token, window) = windowed();
Self::Key { token, window }
}
"team" => {
let (team_id, window) = windowed();
Self::Team { team_id, window }
}
"team_member" => {
let (user_id, team_id) = id.rsplit_once(':').ok_or_else(invalid)?;
Self::TeamMember {
user_id: user_id.to_owned(),
team_id: team_id.to_owned(),
}
}
"user" => Self::User(id.to_owned()),
"end_user" => Self::EndUser(id.to_owned()),
"org" => Self::Org(id.to_owned()),
"tag" => Self::Tag(id.to_owned()),
"project" => Self::Project(id.to_owned()),
"model_access_group" => Self::ModelAccessGroup(id.to_owned()),
"shadow_eval" => Self::ShadowEval(id.to_owned()),
_ => return Err(invalid()),
})
}
}
#[derive(Debug, thiserror::Error)]
pub enum Error {
#[error(transparent)]
Redis(#[from] litellm_redis_client::Error),
#[error(transparent)]
Memory(#[from] MemoryError),
#[error(transparent)]
Decode(#[from] DecodeError),
}
/// Where a read was answered. Redis is cross-pod truth, including a clean miss; memory only holds
/// this pod's writes and answers only when there is no Redis or it failed.
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum Reading {
Redis(Option<Number>),
Memory(Option<Number>),
}
/// Who wins when a seed's `SET NX` lost and the winner's value is gone by the time it is read.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum LostSeed {
/// `SpendCounterReseed.coalesced`: answer with the seed without writing it.
Answer,
/// `SpendCounterReseed.coalesced_window`: `INCRBYFLOAT` the seed, keeping any TTL.
Increment,
}
fn decode(raw: &[u8]) -> Result<Number, DecodeError> {
Number::from_value(&python_value(raw)?).ok_or(DecodeError::Shape("a spend counter is a number"))
}
fn decode_reply(raw: Option<Vec<u8>>) -> Result<Option<Number>, DecodeError> {
raw.as_deref().map(decode).transpose()
}
/// `json.dumps` of the value, as `RedisCache.async_set_cache` writes it.
fn json_bytes(value: Number) -> String {
json::dumps(&value.to_value()).unwrap_or_else(|_| to_str(&value.to_value()))
}
fn whole_seconds(ttl: Duration) -> u64 {
ttl.as_secs()
}
/// The proxy's spend counters: Redis shared by every pod, and this pod's memory copy that answers
/// only when Redis cannot.
pub struct SpendCounters {
memory: InMemoryCache<Number>,
remote: Option<RedisClient>,
redis_ttl: Duration,
updates: Mutex<()>,
}
impl SpendCounters {
pub fn new(client: Option<RedisClient>, redis_ttl: Duration) -> Self {
Self::with_clock(client, redis_ttl, system_clock())
}
pub fn with_clock(client: Option<RedisClient>, redis_ttl: Duration, now: Clock) -> Self {
Self {
memory: InMemoryCache::with_clock(Some(MEMORY_ENTRIES), Some(MEMORY_TTL), move || {
now()
}),
remote: client,
redis_ttl,
updates: Mutex::new(()),
}
}
pub fn has_redis(&self) -> bool {
self.remote.is_some()
}
pub fn redis_ttl(&self) -> Duration {
self.redis_ttl
}
/// `read_spend_counter_cache_value`: Redis first, memory when Redis is absent or failed.
pub async fn read(&self, counter: &SpendCounter) -> Result<Reading, MemoryError> {
if let Some(remote) = &self.remote {
match remote
.get::<Vec<u8>>(&counter.key())
.await
.map_err(Error::from)
.and_then(|raw| decode_reply(raw).map_err(Error::from))
{
Ok(value) => return Ok(Reading::Redis(value)),
Err(error) => {
tracing::debug!(%error, "spend counter Redis read failed, falling back to memory");
}
}
}
Ok(Reading::Memory(self.get_local(counter)?))
}
/// One round trip for many counters; `None` when there is no Redis or the read failed, so each
/// caller reads on its own.
pub async fn read_many(&self, counters: &[SpendCounter]) -> Option<Vec<Option<Number>>> {
let declared = self.declare_read_many(RedisClient::batch, counters)?;
declared.resolve().await
}
/// `read_many` declared on the request's pre-call batch, so it goes out with the request's other
/// pre-call reads.
pub fn declare_read_many_pre_call(
&self,
request: &RequestRedis,
counters: &[SpendCounter],
) -> Option<DeclaredReads> {
self.declare_read_many(|client| request.pre(client), counters)
}
fn declare_read_many(
&self,
batch_for: impl FnOnce(&RedisClient) -> RedisBatch,
counters: &[SpendCounter],
) -> Option<DeclaredReads> {
let batch = batch_for(self.remote.as_ref()?);
Some(DeclaredReads(
counters
.iter()
.map(|counter| batch.get::<Vec<u8>>(&counter.key()))
.collect(),
))
}
/// `increment_spend_counter_cache`: `INCRBYFLOAT` refreshing the TTL, mirrored into memory. A
/// failed increment may or may not have landed, so the counter is dropped everywhere before the
/// error returns. Without Redis it is a memory increment.
pub async fn increment(&self, counter: &SpendCounter, amount: Number) -> Result<Number, Error> {
let Some(remote) = &self.remote else {
return Ok(self.increment_local(counter, amount)?);
};
let key = counter.key();
match remote
.incr_by_float(&key, amount.as_f64(), Expiry::Refresh(self.redis_ttl))
.await
{
Ok(total) => {
let total = Number::Float(total);
self.memory.set_cache(key, total, None)?;
Ok(total)
}
Err(error) => {
self.invalidate(std::slice::from_ref(counter)).await;
Err(error.into())
}
}
}
/// `run_spend_counter_pipeline`: every increment in one round trip, totals in order. The caller
/// decides what a failure means for counters whose increment may or may not have landed.
pub async fn increment_many(
&self,
increments: &[(SpendCounter, Number)],
) -> Result<Vec<Number>, Error> {
let Some(remote) = &self.remote else {
return increments
.iter()
.map(|(counter, amount)| Ok(self.increment_local(counter, *amount)?))
.collect();
};
let declared = self.declare_increments(&remote.batch(), increments);
let mut totals = Vec::with_capacity(declared.len());
for (key, total) in declared {
let total = Number::Float(total.resolve().await?);
self.memory.set_cache(key, total, None)?;
totals.push(total);
}
Ok(totals)
}
/// `increment_many` on the request's post-call batch. Each total reaches memory once the batch
/// goes out, and a counter whose increment failed is dropped everywhere. `false` without Redis.
pub fn increment_many_post_call(
self: &Arc<Self>,
request: &RequestRedis,
increments: &[(SpendCounter, Number)],
) -> bool {
let Some(remote) = &self.remote else {
return false;
};
let declared = self.declare_increments(&request.post(remote), increments);
let counters = increments
.iter()
.map(|(counter, _)| counter.clone())
.collect::<Vec<_>>();
let store = Arc::clone(self);
request.spawn(async move {
for ((key, total), counter) in declared.into_iter().zip(counters) {
match total.settled().await {
Ok(total) => {
if let Err(error) = store.memory.set_cache(key, Number::Float(total), None)
{
tracing::warn!(%error, "mirroring a spend counter total into memory failed");
}
}
Err(error) => {
tracing::warn!(%error, %key, "spend counter increment did not land in the post-call pipeline; invalidating it");
store.invalidate(std::slice::from_ref(&counter)).await;
}
}
}
});
true
}
/// `_repair_stale_spend_counter`: raise memory and Redis to `spend`, never lowering either.
/// `None` without Redis, else what Redis holds afterwards.
pub async fn set_max(
&self,
counter: &SpendCounter,
spend: Number,
) -> Result<Option<Number>, Error> {
let key = counter.key();
{
let _guard = lock(&self.updates);
let raise = self
.memory
.get_cache(&key)?
.is_none_or(|cached| cached.as_f64() < spend.as_f64());
if raise {
self.memory.set_cache(key.clone(), spend, None)?;
}
}
let Some(remote) = &self.remote else {
return Ok(None);
};
let stored: Option<Vec<u8>> = remote
.eval(
&SET_MAX,
&[&key],
(
to_str(&spend.to_value()),
whole_seconds(self.redis_ttl).to_string(),
),
)
.await?;
Ok(stored
.as_deref()
.map(|raw| {
std::str::from_utf8(raw)
.ok()
.and_then(|text| text.trim().parse::<f64>().ok())
.map(Number::Float)
.ok_or(DecodeError::Shape("set_max returned a non-number"))
})
.transpose()?)
}
/// `SpendCounterReseed.coalesced(_window)`: seed a cold counter with `SET NX` so one pod wins,
/// answer with whatever the counter holds, and mirror it into memory. Without Redis, memory
/// keeps the larger of the seed and what it held.
pub async fn seed(
&self,
counter: &SpendCounter,
spend: Number,
lost: LostSeed,
) -> Result<Number, Error> {
let key = counter.key();
let Some(remote) = &self.remote else {
let _guard = lock(&self.updates);
let seeded = match self.memory.get_cache(&key)? {
Some(cached) if cached.as_f64() > spend.as_f64() => Number::Float(cached.as_f64()),
_ => spend,
};
self.memory.set_cache(key, seeded, None)?;
return Ok(seeded);
};
let current = if remote
.set_nx(&key, json_bytes(spend), self.redis_ttl)
.await?
{
spend
} else {
match decode_reply(remote.get::<Vec<u8>>(&key).await?)? {
Some(winner) => Number::Float(winner.as_f64()),
None => match lost {
LostSeed::Answer => spend,
LostSeed::Increment => Number::Float(
remote
.incr_by_float(
&key,
spend.as_f64(),
Expiry::KeepExisting(self.redis_ttl),
)
.await?,
),
},
}
};
self.memory.set_cache(key, current, None)?;
Ok(current)
}
/// A spend reset: memory, then Redis, each for `ttl` or its own default.
pub async fn set(
&self,
counter: &SpendCounter,
spend: Number,
ttl: Option<Duration>,
) -> Result<(), Error> {
let key = counter.key();
self.memory.set_cache(key.clone(), spend, ttl)?;
let Some(remote) = &self.remote else {
return Ok(());
};
remote
.set(&key, json_bytes(spend), ttl.unwrap_or(self.redis_ttl))
.await?;
Ok(())
}
/// Drop counters from memory, then from Redis in one round trip.
pub async fn delete(&self, counters: &[SpendCounter]) -> Result<(), Error> {
for counter in counters {
self.memory.delete_cache(&counter.key())?;
}
let Some(remote) = &self.remote else {
return Ok(());
};
let batch = remote.batch();
let deleted = counters
.iter()
.map(|counter| batch.del(&counter.key()))
.collect::<Vec<_>>();
for reply in deleted {
reply.resolve().await?;
}
Ok(())
}
/// `refresh_spend_counter_ttl`: `false` without Redis or for a counter Redis no longer holds.
pub async fn refresh_ttl(&self, counter: &SpendCounter) -> Result<bool, Error> {
let Some(remote) = &self.remote else {
return Ok(false);
};
Ok(remote.expire(&counter.key(), self.redis_ttl).await?)
}
pub fn get_local(&self, counter: &SpendCounter) -> Result<Option<Number>, MemoryError> {
self.memory.get_cache(&counter.key())
}
pub fn set_local(
&self,
counter: &SpendCounter,
spend: Number,
ttl: Option<Duration>,
) -> Result<(), MemoryError> {
self.memory.set_cache(counter.key(), spend, ttl).map(drop)
}
pub fn delete_local(&self, counter: &SpendCounter) -> Result<(), MemoryError> {
self.memory.delete_cache(&counter.key())
}
/// `DualCache.async_increment_cache(local_only=True)`.
pub fn increment_local(
&self,
counter: &SpendCounter,
amount: Number,
) -> Result<Number, MemoryError> {
let key = counter.key();
let _guard = lock(&self.updates);
let total = self.memory.get_cache(&key)?.unwrap_or(Number::Int(0)) + amount;
self.memory
.set_cache(key, total, Some(LOCAL_INCREMENT_TTL))?;
Ok(total)
}
pub fn snapshot(&self) -> Result<Vec<(String, Number)>, MemoryError> {
self.memory.entries()
}
pub fn flush_local(&self) -> Result<(), MemoryError> {
self.memory.flush_cache()
}
/// `invalidate_spend_counter`: memory always, Redis best effort.
pub async fn invalidate(&self, counters: &[SpendCounter]) {
if let Err(error) = self.delete(counters).await {
tracing::debug!(%error, "dropping spend counters from Redis failed");
}
}
fn declare_increments(
&self,
batch: &RedisBatch,
increments: &[(SpendCounter, Number)],
) -> Vec<(String, Declared<f64>)> {
increments
.iter()
.map(|(counter, amount)| {
let key = counter.key();
let total =
batch.incr_by_float(&key, amount.as_f64(), Expiry::Refresh(self.redis_ttl));
(key, total)
})
.collect()
}
}
/// Reads declared on a batch, answered once it goes out.
pub struct DeclaredReads(Vec<Declared<Option<Vec<u8>>>>);
impl DeclaredReads {
/// Flushes the batch if nobody has yet. `None` when the read or a decode failed.
pub async fn resolve(self) -> Option<Vec<Option<Number>>> {
let replies = join_all(self.0.into_iter().map(Declared::resolve)).await;
replies
.into_iter()
.map(|reply| {
reply
.map_err(Error::from)
.and_then(|raw| decode_reply(raw).map_err(Error::from))
})
.collect::<Result<Vec<_>, _>>()
.inspect_err(|error| {
tracing::debug!(%error, "spend counter batch read failed, falling back to per-key reads");
})
.ok()
}
}
fn lock(mutex: &Mutex<()>) -> MutexGuard<'_, ()> {
mutex.lock().unwrap_or_else(PoisonError::into_inner)
}

View file

@ -0,0 +1,484 @@
mod support;
use std::{sync::Arc, time::Duration};
use litellm_redis_client::{Operation, RedisClient, Ttl};
use litellm_redis_store::{
Number, RequestRedis,
router_cooldowns::{self, RouterCooldowns},
spend_counters::{LostSeed, Reading, SpendCounter, SpendCounters},
};
use litellm_redis_testing::{RedisDeployment, Topology};
use rstest::rstest;
use support::{connect, connect_observed, within};
use tokio::runtime::Handle;
const TTL: Duration = Duration::from_secs(60);
const LONG_DEADLINE: Duration = Duration::from_secs(60);
const READS: [(&str, Number); 4] = [
("7.25", Number::Float(7.25)),
("3", Number::Int(3)),
("0.0", Number::Float(0.0)),
("1e-05", Number::Float(1e-05)),
];
const SETS: [(Number, &str); 4] = [
(Number::Float(0.0), "0.0"),
(Number::Int(12), "12"),
(Number::Float(1e-05), "1e-05"),
(Number::Float(0.30000000000000004), "0.30000000000000004"),
];
const SET_MAXES: [(Number, &str); 3] = [
(Number::Float(2.5), "2.5"),
(Number::Int(100), "100"),
(Number::Float(1e-05), "1e-05"),
];
fn team(id: &str) -> SpendCounter {
SpendCounter::Team {
team_id: id.to_owned(),
window: None,
}
}
fn store(client: Option<RedisClient>) -> Arc<SpendCounters> {
Arc::new(SpendCounters::new(client, TTL))
}
/// Mirrored by `tests/unit/redis/stores/test_spend_counters.py`, which checks the same keys and
/// bytes against Python's own key builders and `RedisCache` codec.
#[rstest]
#[case(
SpendCounter::Key {
token: "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b".to_owned(),
window: None,
},
"spend:key:88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"
)]
#[case(
SpendCounter::Key {
token: "88dc28d0".to_owned(),
window: Some("30d".to_owned()),
},
"spend:key:88dc28d0:window:30d"
)]
#[case(
SpendCounter::Team {
team_id: "team-a".to_owned(),
window: None,
},
"spend:team:team-a"
)]
#[case(
SpendCounter::Team {
team_id: "team-a".to_owned(),
window: Some("1h".to_owned()),
},
"spend:team:team-a:window:1h"
)]
#[case(
SpendCounter::TeamMember {
user_id: "user:with:colons".to_owned(),
team_id: "team-a".to_owned(),
},
"spend:team_member:user:with:colons:team-a"
)]
#[case(SpendCounter::User("user@example.com".to_owned()), "spend:user:user@example.com")]
#[case(SpendCounter::EndUser("customer é 1".to_owned()), "spend:end_user:customer é 1")]
#[case(SpendCounter::Org("org-1".to_owned()), "spend:org:org-1")]
#[case(SpendCounter::Tag("prod:eu".to_owned()), "spend:tag:prod:eu")]
#[case(SpendCounter::Project("proj-1".to_owned()), "spend:project:proj-1")]
#[case(
SpendCounter::ModelAccessGroup("gpt-group".to_owned()),
"spend:model_access_group:gpt-group"
)]
#[case(SpendCounter::ShadowEval("job-7".to_owned()), "spend:shadow_eval:job-7")]
fn every_counter_formats_and_parses_back_to_mains_key(
#[case] counter: SpendCounter,
#[case] key: &str,
) {
assert_eq!(counter.key(), key);
assert_eq!(SpendCounter::parse(key).unwrap(), counter, "{key}");
}
#[rstest]
#[case("spend_db_floor:spend:key:abc")]
#[case("virtual_key_spend:abc:gpt:1d")]
#[case("spend:unknown:abc")]
#[case("spend:team_member:no-team")]
#[case("spend:")]
fn keys_outside_the_spend_counter_keyspace_are_refused(#[case] key: &str) {
assert!(SpendCounter::parse(key).is_err(), "{key}");
}
#[rstest]
#[tokio::test]
async fn reads_values_python_wrote_and_a_clean_miss_is_a_redis_answer(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = store(Some(connect(&deployment)));
let counters = (0..READS.len())
.map(|index| team(&format!("t{index}")))
.collect::<Vec<_>>();
for ((bytes, _), counter) in READS.iter().zip(&counters) {
raw.set(&counter.key(), *bytes, TTL).await.unwrap();
}
let expected = READS
.iter()
.map(|(_, spend)| Some(*spend))
.collect::<Vec<_>>();
for (counter, value) in counters.iter().zip(&expected) {
assert_eq!(store.read(counter).await.unwrap(), Reading::Redis(*value));
}
assert_eq!(store.read_many(&counters).await, Some(expected));
assert_eq!(
store.read(&team("absent")).await.unwrap(),
Reading::Redis(None)
);
}
#[rstest]
#[tokio::test]
async fn resets_and_repairs_write_mains_bytes_with_the_redis_ttl(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = store(Some(connect(&deployment)));
let counter = team("t");
for (spend, bytes) in SETS {
store.set(&counter, spend, None).await.unwrap();
assert_eq!(
raw.get::<String>(&counter.key()).await.unwrap(),
Some(bytes.to_owned())
);
}
for (spend, bytes) in SET_MAXES {
raw.del(&[&counter.key()]).await.unwrap();
let repaired = store.set_max(&counter, spend).await.unwrap();
assert_eq!(
raw.get::<String>(&counter.key()).await.unwrap(),
Some(bytes.to_owned())
);
assert_eq!(repaired, Some(Number::Float(spend.as_f64())));
}
store
.set(&counter, Number::Int(5), Some(Duration::from_secs(30)))
.await
.unwrap();
let Ttl::Expires(remaining) = raw.ttl(&counter.key()).await.unwrap() else {
panic!("a reset counter has no TTL");
};
assert!(within(remaining, 30), "{remaining:?}");
}
#[rstest]
#[tokio::test]
async fn set_max_only_ever_raises_a_counter(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = store(Some(connect(&deployment)));
let counter = SpendCounter::Key {
token: "hashed".to_owned(),
window: None,
};
store.increment(&counter, Number::Float(9.5)).await.unwrap();
let lower = store.set_max(&counter, Number::Float(4.0)).await.unwrap();
let higher = store.set_max(&counter, Number::Float(12.0)).await.unwrap();
assert_eq!(lower, Some(Number::Float(9.5)));
assert_eq!(higher, Some(Number::Float(12.0)));
assert_eq!(
raw.get::<String>(&counter.key()).await.unwrap(),
Some("12.0".to_owned())
);
let Ttl::Expires(remaining) = raw.ttl(&counter.key()).await.unwrap() else {
panic!("a repaired counter has no TTL");
};
assert!(within(remaining, 60), "{remaining:?}");
assert_eq!(
store.get_local(&counter).unwrap(),
Some(Number::Float(12.0))
);
}
#[rstest]
#[tokio::test]
async fn increments_refresh_the_ttl_and_mirror_the_shared_total_into_memory(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let (client, trips) = connect_observed(&deployment);
let other_pod = store(Some(connect(&deployment)));
let store = store(Some(client));
let [first, second] = [team("a"), team("b")];
other_pod
.increment(&first, Number::Float(1.25))
.await
.unwrap();
raw.expire(&first.key(), Duration::from_secs(5))
.await
.unwrap();
let totals = store
.increment_many(&[
(first.clone(), Number::Float(0.5)),
(second.clone(), Number::Int(2)),
])
.await
.unwrap();
assert_eq!(totals, [Number::Float(1.75), Number::Float(2.0)]);
assert_eq!(trips.seen(), [Operation::Pipeline { ops: 2 }]);
let Ttl::Expires(remaining) = raw.ttl(&first.key()).await.unwrap() else {
panic!("an incremented counter has no TTL");
};
assert!(within(remaining, 60), "{remaining:?}");
assert_eq!(
raw.get::<String>(&first.key()).await.unwrap(),
Some("1.75".to_owned())
);
assert_eq!(store.get_local(&first).unwrap(), Some(Number::Float(1.75)));
assert_eq!(store.get_local(&second).unwrap(), Some(Number::Float(2.0)));
}
#[rstest]
#[tokio::test]
async fn a_post_call_increment_from_a_thread_without_a_runtime_lands_in_redis_and_memory(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = store(Some(connect(&deployment)));
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
let counter = team("t");
let declared = std::thread::spawn({
let (store, request, counter) = (Arc::clone(&store), request.clone(), counter.clone());
move || store.increment_many_post_call(&request, &[(counter, Number::Float(3.5))])
})
.join()
.expect("a Python thread has no Tokio context, and declaring there must not panic");
request.flush_post().await;
assert!(declared);
let mirrored =
litellm_redis_testing::eventually(async || store.get_local(&counter).unwrap()).await;
assert_eq!(mirrored, Number::Float(3.5));
assert_eq!(
raw.get::<String>(&counter.key()).await.unwrap(),
Some("3.5".to_owned())
);
}
#[rstest]
#[tokio::test]
async fn post_call_increments_wait_for_the_post_batch_then_land_in_memory(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let (client, trips) = connect_observed(&deployment);
let store = store(Some(client));
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
let counter = team("t");
assert!(store.increment_many_post_call(&request, &[(counter.clone(), Number::Float(3.5))]));
request.flush_pre().await;
assert!(trips.seen().is_empty());
assert_eq!(raw.get::<String>(&counter.key()).await.unwrap(), None);
request.flush_post().await;
tokio::time::sleep(Duration::from_millis(100)).await;
assert_eq!(trips.seen(), [Operation::Pipeline { ops: 1 }]);
assert_eq!(
raw.get::<String>(&counter.key()).await.unwrap(),
Some("3.5".to_owned())
);
assert_eq!(store.get_local(&counter).unwrap(), Some(Number::Float(3.5)));
}
#[rstest]
#[tokio::test]
async fn the_budget_read_rides_the_same_pre_call_pipeline_as_the_cooldown_read(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let (client, trips) = connect_observed(&deployment);
let cooldowns = RouterCooldowns::new(Some(client.clone()), router_cooldowns::POLICY);
let spend = store(Some(client));
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
let counters = [
SpendCounter::Key {
token: "hashed".to_owned(),
window: None,
},
team("t"),
];
raw.set(&counters[0].key(), "4.5", TTL).await.unwrap();
let active = cooldowns
.declare_active(&request, &["dep-a".to_owned()])
.unwrap();
let budget = spend
.declare_read_many_pre_call(&request, &counters)
.unwrap();
request.flush_pre().await;
assert_eq!(trips.seen(), [Operation::Pipeline { ops: 3 }]);
assert_eq!(cooldowns.resolve_active(active).await.unwrap(), vec![]);
assert_eq!(
budget.resolve().await,
Some(vec![Some(Number::Float(4.5)), None])
);
assert_eq!(trips.seen().len(), 1);
}
#[rstest]
#[tokio::test]
async fn a_seed_lets_the_first_pod_win_and_others_answer_with_its_value(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let first = store(Some(connect(&deployment)));
let second = store(Some(connect(&deployment)));
let counter = SpendCounter::Org("org".to_owned());
let won = first
.seed(&counter, Number::Float(10.0), LostSeed::Answer)
.await
.unwrap();
let lost = second
.seed(&counter, Number::Float(99.0), LostSeed::Increment)
.await
.unwrap();
assert_eq!([won, lost], [Number::Float(10.0), Number::Float(10.0)]);
assert_eq!(
raw.get::<String>(&counter.key()).await.unwrap(),
Some("10.0".to_owned())
);
let Ttl::Expires(remaining) = raw.ttl(&counter.key()).await.unwrap() else {
panic!("a seeded counter has no TTL");
};
assert!(within(remaining, 60), "{remaining:?}");
assert_eq!(
second.get_local(&counter).unwrap(),
Some(Number::Float(10.0))
);
}
#[rstest]
#[tokio::test]
async fn delete_drops_memory_and_redis_and_refresh_ttl_reports_whether_redis_held_it(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = store(Some(connect(&deployment)));
let counters = [team("a"), team("b")];
for counter in &counters {
store.increment(counter, Number::Int(1)).await.unwrap();
}
raw.expire(&counters[0].key(), Duration::from_secs(5))
.await
.unwrap();
assert!(store.refresh_ttl(&counters[0]).await.unwrap());
let Ttl::Expires(remaining) = raw.ttl(&counters[0].key()).await.unwrap() else {
panic!("a refreshed counter has no TTL");
};
assert!(within(remaining, 60), "{remaining:?}");
store.delete(&counters).await.unwrap();
for counter in &counters {
assert_eq!(raw.get::<String>(&counter.key()).await.unwrap(), None);
assert_eq!(store.get_local(counter).unwrap(), None);
}
assert!(!store.refresh_ttl(&counters[0]).await.unwrap());
}
#[rstest]
#[tokio::test]
async fn a_redis_outage_reads_memory_and_a_failed_increment_drops_the_counter() {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
let store = store(Some(connect(&deployment)));
let [kept, failed] = [team("kept"), team("failed")];
store.increment(&kept, Number::Float(2.0)).await.unwrap();
store.increment(&failed, Number::Float(3.0)).await.unwrap();
deployment.stop().await.unwrap();
let read = store.read(&kept).await.unwrap();
let increment = store.increment(&failed, Number::Float(1.0)).await;
assert_eq!(read, Reading::Memory(Some(Number::Float(2.0))));
assert!(increment.is_err());
assert_eq!(store.get_local(&failed).unwrap(), None);
assert_eq!(store.read_many(&[kept]).await, None);
}
#[rstest]
#[tokio::test]
async fn a_failed_post_call_increment_drops_the_counter_from_memory() {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
let store = store(Some(connect(&deployment)));
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
let counter = team("t");
store.increment(&counter, Number::Float(1.0)).await.unwrap();
assert!(store.increment_many_post_call(&request, &[(counter.clone(), Number::Float(1.0))]));
deployment.stop().await.unwrap();
request.flush_post().await;
tokio::time::sleep(Duration::from_millis(200)).await;
assert_eq!(store.get_local(&counter).unwrap(), None);
}
#[tokio::test]
async fn without_redis_counters_live_in_memory() {
let store = store(None);
let counter = team("t");
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
let totals = store
.increment_many(&[
(counter.clone(), Number::Int(2)),
(counter.clone(), Number::Float(0.5)),
])
.await
.unwrap();
let seeded = store
.seed(&counter, Number::Float(1.0), LostSeed::Answer)
.await
.unwrap();
let raised = store.set_max(&counter, Number::Float(7.0)).await.unwrap();
assert_eq!(totals, [Number::Int(2), Number::Float(2.5)]);
assert_eq!(seeded, Number::Float(2.5));
assert_eq!(raised, None);
assert_eq!(
store.read(&counter).await.unwrap(),
Reading::Memory(Some(Number::Float(7.0)))
);
assert!(!store.increment_many_post_call(&request, &[(counter.clone(), Number::Int(1))]));
assert!(
store
.declare_read_many_pre_call(&request, &[counter])
.is_none()
);
}

View file

@ -40,6 +40,9 @@ from litellm.litellm_core_utils.llm_judge import (
)
from litellm.litellm_core_utils.redact_messages import should_redact_message_logging
from litellm.llms.base_llm.base_utils import type_to_response_format_param
from litellm.redis.stores.spend_counters import (
shadow_eval_counter,
)
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.management_endpoints.auto_router_endpoints import ShadowEvalDirection
from litellm.types.utils import SHADOW_EVAL_JUDGE_CALL_ORIGIN, SHADOW_EVAL_ROUTER_CALL_ORIGIN
@ -587,7 +590,7 @@ def _leg_eval_spend(sums: Mapping[str, object]) -> float:
def _job_spend_counter_key(job_id: str) -> str:
return f"spend:shadow_eval:{job_id}"
return shadow_eval_counter(job_id)
async def _job_spend_from_counter(counter_key: str, fallback_spend: float, max_budget: float) -> float:

View file

@ -130,10 +130,8 @@ from litellm.proxy.common_utils.user_api_key_cache import (
get_management_object_ttl,
model_access_group_cache_key,
model_access_group_registry_cache_key,
model_access_group_spend_counter_key,
object_permission_cache_key,
project_cache_key,
project_spend_counter_key,
tag_cache_key,
tag_registry_cache_key,
team_membership_auth_cache_key,
@ -148,8 +146,21 @@ from litellm.proxy.guardrails.tool_name_extraction import (
from litellm.proxy.route_llm_request import route_request
from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start
from litellm.proxy.spend_tracking.carried_budget_state import carry_organization_budget_state
from litellm.proxy.spend_tracking.spend_counter_batch import SPEND_COUNTERS_TARGET
from litellm.proxy.utils import PrismaClient, ProxyLogging, log_db_metrics
from litellm.redis.stores.spend_counters import (
SPEND_COUNTERS_TARGET,
end_user_counter,
key_counter,
key_window_counter,
model_access_group_counter,
org_counter,
project_counter,
tag_counter,
team_counter,
team_member_counter,
team_window_counter,
user_counter,
)
from litellm.repositories.budget_repository import BudgetRepository
from litellm.repositories.object_permission_repository import ObjectPermissionRepository
from litellm.repositories.organization_repository import OrganizationRepository
@ -1417,7 +1428,7 @@ async def common_checks(
user_budget: Final = user_object.max_budget
user_spend: Final = await get_current_spend(
counter_key=f"spend:user:{user_object.user_id}",
counter_key=user_counter(user_object.user_id),
fallback_spend=user_object.spend or 0.0,
max_budget=user_budget,
)
@ -1965,7 +1976,7 @@ async def check_end_user_budget(
from litellm.proxy.proxy_server import get_current_spend
end_user_spend: Final = await get_current_spend(
counter_key=f"spend:end_user:{end_user_obj.user_id}",
counter_key=end_user_counter(end_user_obj.user_id),
fallback_spend=end_user_obj.spend or 0.0,
max_budget=end_user_budget,
fallback_authoritative=True,
@ -3319,37 +3330,40 @@ async def invalidate_team_member_spend_state(
await await_inflight_team_membership_load(user_id=user_id, team_id=team_id)
if new_spend is not None:
from litellm.proxy.proxy_server import SPEND_DB_FLOOR_CACHE_TTL_SECONDS, spend_counter_cache
from litellm.proxy.proxy_server import (
SPEND_DB_FLOOR_CACHE_TTL_SECONDS,
spend_counter_cache,
spend_counter_store,
)
spend_counter_key: Final = f"spend:team_member:{user_id}:{team_id}"
spend_counter_key: Final = team_member_counter(user_id, team_id)
spend_db_floor_key: Final = f"spend_db_floor:{spend_counter_key}"
spend_counter_cache.in_memory_cache.set_cache(key=spend_counter_key, value=new_spend, ttl=60)
if spend_counter_cache.redis_cache is not None:
try:
await spend_counter_store.set(spend_counter_key, new_spend, ttl=60)
except Exception as e: # noqa: BLE001 # fall back to deleting the stale entry before giving up
verbose_proxy_logger.warning(
"Failed to set spend counter %s in Redis after reset: %s; deleting it instead so the next "
"read reseeds from the DB rather than keeping the stale pre-reset value authoritative",
spend_counter_key,
e,
)
try:
await spend_counter_cache.redis_cache.async_set_cache(key=spend_counter_key, value=new_spend, ttl=60)
except Exception as e: # noqa: BLE001 # fall back to deleting the stale entry before giving up
await spend_counter_store.delete((spend_counter_key,))
except Exception: # noqa: BLE001 # stale value now authoritative in Redis; surface instead of reporting success
verbose_proxy_logger.warning(
"Failed to set spend counter %s in Redis after reset: %s; deleting it instead so the next "
"read reseeds from the DB rather than keeping the stale pre-reset value authoritative",
"Failed to delete stale spend counter %s in Redis after a failed reset write",
spend_counter_key,
e,
exc_info=True,
)
try:
await spend_counter_cache.redis_cache.async_delete_cache(key=spend_counter_key)
except Exception: # noqa: BLE001 # stale value now authoritative in Redis; surface instead of reporting success
verbose_proxy_logger.warning(
"Failed to delete stale spend counter %s in Redis after a failed reset write",
spend_counter_key,
exc_info=True,
)
raise HTTPException(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
detail={
"error": "Spend was reset in the database, but Redis is unreachable and still "
"holds the pre-reset counter. Retry once Redis is reachable."
},
) from e
raise HTTPException(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
detail={
"error": "Spend was reset in the database, but Redis is unreachable and still "
"holds the pre-reset counter. Retry once Redis is reachable."
},
) from e
spend_counter_store.set_local(spend_counter_key, new_spend, ttl=60)
spend_counter_cache.in_memory_cache.set_cache(
key=spend_db_floor_key,
@ -6702,7 +6716,7 @@ async def virtual_key_max_budget_check(
from litellm.proxy.proxy_server import get_current_spend
fallback_spend: Final = valid_token.spend or 0.0
counter_key: Final = f"spend:key:{valid_token.token}"
counter_key: Final = key_counter(str(valid_token.token))
# Read spend from cross-pod counter (Redis-first) or cached object (fallback)
spend: Final = await get_current_spend(
@ -6791,7 +6805,7 @@ async def _virtual_key_multi_budget_check(
for window in valid_token.budget_limits:
w: dict = window if isinstance(window, dict) else window.model_dump()
counter_key = f"spend:key:{valid_token.token}:window:{w['budget_duration']}"
counter_key = key_window_counter(str(valid_token.token), str(w["budget_duration"]))
window_spend = await get_current_spend(
counter_key=counter_key,
fallback_spend=0.0,
@ -7098,7 +7112,7 @@ async def _check_team_member_budget(
from litellm.proxy.proxy_server import get_current_spend
team_member_spend = await get_current_spend(
counter_key=f"spend:team_member:{valid_token.user_id}:{team_object.team_id}",
counter_key=team_member_counter(valid_token.user_id, team_object.team_id),
fallback_spend=team_member_spend,
max_budget=team_member_budget,
)
@ -7210,7 +7224,7 @@ async def team_max_budget_check(
# Read spend from cross-pod counter (Redis-first) or cached object (fallback)
spend: Final = await get_current_spend(
counter_key=f"spend:team:{team_object.team_id}",
counter_key=team_counter(team_object.team_id),
fallback_spend=team_object.spend or 0.0,
max_budget=team_object.max_budget,
)
@ -7263,7 +7277,7 @@ async def _team_multi_budget_check(
for window in team_object.budget_limits:
w: dict = window if isinstance(window, dict) else window.model_dump()
counter_key = f"spend:team:{team_object.team_id}:window:{w['budget_duration']}"
counter_key = team_window_counter(team_object.team_id, str(w["budget_duration"]))
window_spend = await get_current_spend(
counter_key=counter_key,
fallback_spend=0.0,
@ -7382,7 +7396,7 @@ async def _project_max_budget_check(
from litellm.proxy.proxy_server import get_current_spend
project_spend: Final = await get_current_spend(
counter_key=project_spend_counter_key(project_object.project_id),
counter_key=project_counter(project_object.project_id),
fallback_spend=project_object.spend or 0.0,
max_budget=max_budget,
)
@ -7589,7 +7603,7 @@ async def organization_max_budget_check(
from litellm.proxy.proxy_server import get_current_spend
org_spend: Final = await get_current_spend(
counter_key=f"spend:org:{org_id}",
counter_key=org_counter(org_id),
fallback_spend=org_table.spend or 0.0,
max_budget=org_max_budget,
)
@ -7679,7 +7693,7 @@ async def tag_max_budget_check_for_tags(
from litellm.proxy.proxy_server import get_current_spend
tag_spend = await get_current_spend(
counter_key=f"spend:tag:{tag_name}",
counter_key=tag_counter(tag_name),
fallback_spend=tag_object.spend or 0.0,
max_budget=tag_object.litellm_budget_table.max_budget,
fallback_authoritative=True,
@ -7734,7 +7748,7 @@ async def _model_access_group_max_budget_check(
continue
group_spend = await get_current_spend(
counter_key=model_access_group_spend_counter_key(group),
counter_key=model_access_group_counter(group),
fallback_spend=budget.spend,
max_budget=budget.max_budget,
fallback_authoritative=True,

View file

@ -44,6 +44,10 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo
_is_model_cost_zero, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
is_model_cost_zero, # pyright: ignore[reportPrivateUsage] # the zero-cost predicate the auth-time budget checks use; no public equivalent
)
from litellm.redis.stores.spend_counters import (
key_counter,
user_counter,
)
from litellm.router_api import RegistryReader, RoutingContext
from litellm.types.llms.base import LiteLLMBaseModel
@ -133,7 +137,7 @@ async def is_token_within_budget_for_model(
key_budget: Final = valid_token.max_budget
if key_budget is not None and valid_token.token is not None:
key_spend: Final = await _counter_spend(
counter_key=f"spend:key:{valid_token.token}",
counter_key=key_counter(valid_token.token),
fallback_spend=valid_token.spend or 0.0,
max_budget=key_budget,
)
@ -149,7 +153,7 @@ async def is_token_within_budget_for_model(
and (valid_token.team_id is None or _applies_user_budget_to_team_keys())
):
user_spend: Final = await _counter_spend(
counter_key=f"spend:user:{valid_token.user_id}",
counter_key=user_counter(valid_token.user_id),
fallback_spend=valid_token.user_spend or 0.0,
max_budget=user_budget,
)

View file

@ -26,7 +26,6 @@ import litellm
from litellm._internal_context import service_target
from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm._service_logger import ServiceLogging
from litellm.caching.redis_cache import RedisCache
from litellm.constants import (
CLIENT_REQUESTED_MODEL_SCOPE_KEY,
GLOBAL_PROXY_SPEND_CACHE_KEY,
@ -180,6 +179,10 @@ from litellm.proxy.utils import (
ProxyLogging,
normalize_route_for_root_path,
)
from litellm.redis.stores.spend_counters import (
SpendCounterStore,
team_member_counter,
)
from litellm.repositories.table_repositories import JWTKeyMappingRepository, TeamMembershipRepository
from litellm.repositories.user_repository import UserRepository
from litellm.repositories.verification_token_repository import VerificationTokenRepository
@ -3057,7 +3060,7 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e
team_member_spend = valid_token.team_member_spend
if valid_token.user_id is not None and valid_token.team_id is not None:
team_member_spend = await get_current_spend(
counter_key=f"spend:team_member:{valid_token.user_id}:{valid_token.team_id}",
counter_key=team_member_counter(valid_token.user_id, valid_token.team_id),
fallback_spend=team_member_spend,
max_budget=team_member_budget,
)
@ -4226,10 +4229,10 @@ async def _authorize_authenticated_request(
return None
def _spend_counter_redis_cache() -> RedisCache | None:
from litellm.proxy.proxy_server import spend_counter_cache
def _spend_counter_store() -> SpendCounterStore:
from litellm.proxy.proxy_server import spend_counter_store
return spend_counter_cache.redis_cache
return spend_counter_store
def _identity_cache_keys(api_key: str, *, end_user_id: str | None, key_is_resolved: bool) -> tuple[str, ...]:
@ -4332,7 +4335,7 @@ async def user_api_key_auth(
with (
phase_span(f"auth {route}"),
service_target(AUTH_OBJECTS_TARGET),
spend_counter_batch_scope(_spend_counter_redis_cache()),
spend_counter_batch_scope(_spend_counter_store()),
):
try:
user_api_key_auth_obj: Final = await user_api_key_auth_builder(

View file

@ -3,7 +3,7 @@ import json
from collections.abc import Mapping, Sequence
from dataclasses import asdict, dataclass
from types import MappingProxyType
from typing import TYPE_CHECKING, Final
from typing import TYPE_CHECKING, Final, Protocol
from pydantic import TypeAdapter
@ -19,6 +19,13 @@ from litellm.proxy.common_utils.config_sync_pubsub import ( # noqa: F401 # leg
)
from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET
class SpendCounterCopies(Protocol):
"""This worker's memory copies of the spend counters, which a broadcast SET or delete reaches too."""
def apply_broadcast(self, cache_key: str, new_value: float | None, ttl: float | None) -> None: ...
if TYPE_CHECKING:
from litellm.caching.dual_cache import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
@ -232,17 +239,19 @@ async def evict_and_broadcast(cache_keys: Sequence[str], user_api_key_cache: "Us
class AuthCacheInvalidationSubscriber:
__slots__ = ("_additional_in_memory_caches", "_redis_cache", "_task", "_user_api_key_cache")
__slots__ = ("_additional_in_memory_caches", "_redis_cache", "_spend_counters", "_task", "_user_api_key_cache")
def __init__(
self,
redis_cache: "RedisCache",
user_api_key_cache: "UserApiKeyCache",
additional_in_memory_caches: Sequence["InMemoryCache"] = (),
spend_counters: SpendCounterCopies | None = None,
) -> None:
self._redis_cache = redis_cache
self._user_api_key_cache = user_api_key_cache
self._additional_in_memory_caches = tuple(additional_in_memory_caches)
self._spend_counters: Final = spend_counters
self._task: asyncio.Task[None] | None = None
def start(self) -> None:
@ -297,6 +306,8 @@ class AuthCacheInvalidationSubscriber:
parsed: Final = _message_from_data(data)
if parsed is None:
return
if self._spend_counters is not None:
self._spend_counters.apply_broadcast(parsed.cache_key, parsed.new_value, parsed.ttl)
if parsed.new_value is not None:
for additional_cache in self._additional_in_memory_caches:
additional_cache.set_cache(parsed.cache_key, parsed.new_value, ttl=parsed.ttl)

View file

@ -15,7 +15,6 @@ from typing_extensions import assert_never
import litellm
from litellm._internal_context import with_service_target
from litellm._logging import verbose_proxy_logger
from litellm.caching.dual_cache import DualCache
from litellm.constants import (
GLOBAL_PROXY_SPEND_CACHE_KEY,
LITELLM_PROXY_BUDGET_NAME,
@ -42,17 +41,29 @@ from litellm.proxy.common_utils.user_api_key_cache import (
AUTH_OBJECTS_TARGET,
end_user_cache_key,
model_access_group_cache_key,
model_access_group_spend_counter_key,
project_cache_key,
project_spend_counter_key,
tag_cache_key,
)
from litellm.proxy.db.budget_window_spend_writer import roll_window_spend_row
from litellm.proxy.db.db_span import db_span, db_spanned
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import POD_LOCK_TARGET, PodLockManager
from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry
from litellm.proxy.spend_tracking.spend_counter_batch import SPEND_COUNTERS_TARGET
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.redis.stores.spend_counters import (
SPEND_COUNTERS_TARGET,
SpendCounterStore,
end_user_counter,
key_counter,
key_window_counter,
model_access_group_counter,
org_counter,
project_counter,
tag_counter,
team_counter,
team_member_counter,
team_window_counter,
user_counter,
)
from litellm.repositories.organization_repository import OrganizationRepository
from litellm.repositories.prisma_protocols import PrismaBatch, SpendLinkedTable
from litellm.repositories.project_repository import ProjectRepository
@ -157,7 +168,7 @@ def _row_carried_spend(row: _BudgetLinkedRow, caps: Mapping[str, float]) -> floa
def _team_membership_counter_key(row: _TeamMembershipRow) -> str:
return f"spend:team_member:{row.user_id}:{row.team_id}"
return team_member_counter(row.user_id, row.team_id)
def _team_membership_cache_keys(row: _TeamMembershipRow) -> tuple[str, ...]:
@ -165,7 +176,7 @@ def _team_membership_cache_keys(row: _TeamMembershipRow) -> tuple[str, ...]:
def _key_counter_key(row: _KeyRow) -> str:
return f"spend:key:{row.token}"
return key_counter(row.token)
def _key_cache_keys(row: _KeyRow) -> tuple[str, ...]:
@ -173,7 +184,7 @@ def _key_cache_keys(row: _KeyRow) -> tuple[str, ...]:
def _org_counter_key(row: _OrgRow) -> str:
return f"spend:org:{row.organization_id}"
return org_counter(row.organization_id)
def _org_cache_keys(row: _OrgRow) -> tuple[str, ...]:
@ -184,7 +195,7 @@ def _org_cache_keys(row: _OrgRow) -> tuple[str, ...]:
def _tag_counter_key(row: _TagRow) -> str:
return f"spend:tag:{row.tag_name}"
return tag_counter(row.tag_name)
def _tag_cache_keys(row: _TagRow) -> tuple[str, ...]:
@ -192,7 +203,7 @@ def _tag_cache_keys(row: _TagRow) -> tuple[str, ...]:
def _model_access_group_counter_key(row: _ModelAccessGroupRow) -> str:
return model_access_group_spend_counter_key(row.access_group_name)
return model_access_group_counter(row.access_group_name)
def _model_access_group_cache_keys(row: _ModelAccessGroupRow) -> tuple[str, ...]:
@ -200,7 +211,7 @@ def _model_access_group_cache_keys(row: _ModelAccessGroupRow) -> tuple[str, ...]
def _project_counter_key(row: _ProjectRow) -> str:
return project_spend_counter_key(row.project_id)
return project_counter(row.project_id)
def _project_cache_keys(row: _ProjectRow) -> tuple[str, ...]:
@ -208,7 +219,7 @@ def _project_cache_keys(row: _ProjectRow) -> tuple[str, ...]:
def _enduser_counter_key(row: _EndUserRow) -> str:
return f"spend:end_user:{row.user_id}"
return end_user_counter(row.user_id)
def _enduser_cache_keys(row: _EndUserRow) -> tuple[str, ...]:
@ -427,7 +438,7 @@ class _WindowSource:
table: str
id_column: str
entity_type: Litellm_EntityType
counter_prefix: str
window_counter: Callable[[str, str], str]
log_subject: str
retry_subject: str
write: Callable[[PrismaClient, str, str], Awaitable[None]]
@ -452,7 +463,7 @@ _WINDOW_SOURCES: Final[tuple[_WindowSource, ...]] = (
table="LiteLLM_VerificationToken",
id_column="token",
entity_type=Litellm_EntityType.KEY,
counter_prefix="spend:key",
window_counter=key_window_counter,
log_subject="keys",
retry_subject="key",
write=_write_key_windows,
@ -461,7 +472,7 @@ _WINDOW_SOURCES: Final[tuple[_WindowSource, ...]] = (
table="LiteLLM_TeamTable",
id_column="team_id",
entity_type=Litellm_EntityType.TEAM,
counter_prefix="spend:team",
window_counter=team_window_counter,
log_subject="teams",
retry_subject="team",
write=_write_team_windows,
@ -607,22 +618,16 @@ class ResetBudgetJob:
commit opens a window where get_current_spend reads 0 from Redis
while the DB still holds the pre-reset value, allowing bypass.
"""
try:
from litellm.proxy.proxy_server import spend_counter_cache
from litellm.proxy.proxy_server import spend_counter_store
spend_counter_cache.in_memory_cache.delete_cache(key=counter_key)
if spend_counter_cache.redis_cache is not None:
try:
await spend_counter_cache.redis_cache.async_delete_cache(key=counter_key)
except Exception as redis_err:
verbose_proxy_logger.warning(
"Failed to reset spend counter %s in Redis: %s. "
"Budget may be over-enforced until counter expires.",
counter_key,
redis_err,
)
try:
await spend_counter_store.delete((counter_key,))
except Exception as e:
verbose_proxy_logger.warning("Failed to reset spend counter %s: %s", counter_key, e)
verbose_proxy_logger.warning(
"Failed to reset spend counter %s in Redis: %s. Budget may be over-enforced until counter expires.",
counter_key,
e,
)
@staticmethod
async def _invalidate_global_proxy_spend_cache() -> None:
@ -671,11 +676,11 @@ class ResetBudgetJob:
if not keys:
return
try:
from litellm.proxy.proxy_server import spend_counter_cache, user_api_key_cache
from litellm.proxy.proxy_server import spend_counter_store, user_api_key_cache
match cache:
case "spend counter":
await spend_counter_cache.async_delete_cache_keys(keys)
await spend_counter_store.delete(keys)
case "user_api_key_cache":
await user_api_key_cache.async_delete_cache_keys(keys)
case _:
@ -1175,7 +1180,7 @@ class ResetBudgetJob:
for k in updated_keys:
token = getattr(k.row, "token", None)
if token:
await self._invalidate_spend_counter(f"spend:key:{token}")
await self._invalidate_spend_counter(key_counter(str(token)))
end_time = time.time()
outcome: Final = _ChunkOutcome(
@ -1286,7 +1291,7 @@ class ResetBudgetJob:
for u in updated_users:
user_id = getattr(u.row, "user_id", None)
if user_id:
await self._invalidate_spend_counter(f"spend:user:{user_id}")
await self._invalidate_spend_counter(user_counter(str(user_id)))
if user_id == LITELLM_PROXY_BUDGET_NAME:
await self._invalidate_global_proxy_spend_cache()
@ -1401,7 +1406,7 @@ class ResetBudgetJob:
for t in updated_teams:
team_id = getattr(t.row, "team_id", None)
if team_id:
await self._invalidate_spend_counter(f"spend:team:{team_id}")
await self._invalidate_spend_counter(team_counter(str(team_id)))
end_time = time.time()
outcome: Final = _ChunkOutcome(
@ -1464,7 +1469,7 @@ class ResetBudgetJob:
async def _reset_expired_window(
window: dict,
counter_key: str,
spend_counter_cache: DualCache,
spend_counter_store: SpendCounterStore,
now: datetime,
reset_settings: BudgetResetSettings,
prisma_client: PrismaClient,
@ -1478,13 +1483,11 @@ class ResetBudgetJob:
reset_at: Final = datetime.fromisoformat(reset_at_str.replace("Z", "+00:00")).replace(tzinfo=None)
if reset_at > now:
return False
new_value: Final = await ResetBudgetJob._window_carried_spend(window, counter_key, spend_counter_cache)
spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=new_value)
if spend_counter_cache.redis_cache is not None:
try:
await spend_counter_cache.redis_cache.async_set_cache(key=counter_key, value=new_value)
except Exception as redis_err:
verbose_proxy_logger.warning("Failed to reset Redis counter %s: %s", counter_key, redis_err)
new_value: Final = await ResetBudgetJob._window_carried_spend(window, counter_key, spend_counter_store)
try:
await spend_counter_store.set(counter_key, new_value)
except Exception as redis_err:
verbose_proxy_logger.warning("Failed to reset Redis counter %s: %s", counter_key, redis_err)
budget_duration: Final = window["budget_duration"]
next_reset_at: Final = compute_budget_reset_at(budget_duration=budget_duration, settings=reset_settings)
window["reset_at"] = next_reset_at.isoformat()
@ -1538,7 +1541,7 @@ class ResetBudgetJob:
@staticmethod
@with_service_target(SPEND_COUNTERS_TARGET)
async def _window_carried_spend(
window: Mapping[str, object], counter_key: str, spend_counter_cache: DualCache
window: Mapping[str, object], counter_key: str, spend_counter_store: SpendCounterStore
) -> float:
"""Per-window spend lives only in the counter, so the carried overage is
read from it before the reset overwrites it."""
@ -1549,11 +1552,11 @@ class ResetBudgetJob:
if cap is None:
return 0.0
try:
current: Final = await spend_counter_cache.async_get_cache(key=counter_key)
current, _ = await spend_counter_store.read(counter_key)
except Exception as e: # noqa: BLE001 # an unreadable counter falls back to a plain zero reset
verbose_proxy_logger.warning("Failed to read spend counter %s for rollover: %s", counter_key, e)
return 0.0
if not isinstance(current, (int, float)):
if current is None:
return 0.0
return _carried_spend(float(current), cap)
@ -1563,12 +1566,12 @@ class ResetBudgetJob:
reset_at <= now. Only the expired windows are reset; other windows are untouched.
"""
from litellm.proxy.proxy_server import spend_counter_cache
from litellm.proxy.proxy_server import spend_counter_store
now: Final = datetime.utcnow()
for source in _WINDOW_SOURCES:
try:
await self._reset_windows_for(source=source, now=now, spend_counter_cache=spend_counter_cache)
await self._reset_windows_for(source=source, now=now, spend_counter_store=spend_counter_store)
except Exception as e:
verbose_proxy_logger.exception("Failed to reset budget windows for %s: %s", source.log_subject, e)
@ -1576,7 +1579,7 @@ class ResetBudgetJob:
self,
source: _WindowSource,
now: datetime,
spend_counter_cache: DualCache,
spend_counter_store: SpendCounterStore,
) -> None:
"""Walk one table's windowed rows a page at a time, to the end.
@ -1598,7 +1601,7 @@ class ResetBudgetJob:
source=source,
cursor=cursor,
now=now,
spend_counter_cache=spend_counter_cache,
spend_counter_store=spend_counter_store,
)
if next_cursor is None:
return
@ -1609,7 +1612,7 @@ class ResetBudgetJob:
source: _WindowSource,
cursor: str,
now: datetime,
spend_counter_cache: DualCache,
spend_counter_store: SpendCounterStore,
) -> str | None:
"""Reset one page of windows; return the next cursor, or None when drained."""
rows: Final = await self._with_db_retry(
@ -1628,11 +1631,11 @@ class ResetBudgetJob:
windows: list[dict[str, object]] = raw if isinstance(raw, list) else json.loads(raw)
changed = False
for window in windows:
counter_key = f"{source.counter_prefix}:{row_id}:window:{window['budget_duration']}"
counter_key = source.window_counter(row_id, str(window["budget_duration"]))
if await ResetBudgetJob._reset_expired_window(
window,
counter_key,
spend_counter_cache,
spend_counter_store,
now,
self.reset_settings,
prisma_client=self.prisma_client,

View file

@ -343,25 +343,10 @@ def model_access_group_registry_cache_key() -> str:
return "model_access_group_registry"
def model_access_group_spend_counter_key(access_group_name: str) -> str:
"""Spend counter key for one model access group; shared so its four owners cannot drift.
The reservation path writes it up front, the cost callback writes it after the call, auth
reads it to enforce ``max_budget``, and the reset job clears it on rollover. A copy that
drifts in any one of them silently resets or reads a counter nobody else touches, which shows
up as a budget that never trips or never resets.
"""
return f"spend:model_access_group:{access_group_name}"
def project_cache_key(project_id: str) -> str:
return f"project_id:{project_id}"
def project_spend_counter_key(project_id: str) -> str:
return f"spend:project:{project_id}"
#: Cached under ``end_user_restricted_registry_cache_key`` when the restricted set exceeds
#: ``END_USER_RESTRICTED_REGISTRY_MAX_SIZE``: registry unusable, fall back to the per-id fetch.
END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL: Final = "__end_user_restricted_registry_overflow__"

View file

@ -26,10 +26,10 @@ from litellm.litellm_core_utils.duration_parser import duration_in_seconds
from litellm.proxy._types import Litellm_EntityType
from litellm.proxy.db.db_lookup_gate import bounded_db_lookup, db_lookup_gate
from litellm.proxy.spend_tracking.spend_counter_batch import (
SPEND_COUNTERS_TARGET,
read_batched_spend_counter,
record_spend_counter_value,
)
from litellm.redis.stores.spend_counters import END_USER_COUNTER_PREFIX, SPEND_COUNTERS_TARGET, SpendCounterStore
from litellm.repositories.organization_repository import OrganizationRepository
from litellm.repositories.project_repository import ProjectRepository
from litellm.repositories.table_repositories import (
@ -47,7 +47,6 @@ from litellm.repositories.verification_token_repository import (
if TYPE_CHECKING:
from prisma.types import LiteLLM_EndUserTableWhereUniqueInput
from litellm.caching.dual_cache import DualCache
from litellm.proxy.utils import PrismaClient
@ -58,8 +57,6 @@ _WINDOW_SPEND_ENTITY_TYPES: Final[Mapping[str, str]] = MappingProxyType(
}
)
END_USER_COUNTER_PREFIX: Final = "spend:end_user:"
_WINDOW_SPEND_LOG_FIELDS: Final[Mapping[str, str]] = MappingProxyType(
{
"Key": "api_key",
@ -114,13 +111,11 @@ class SpendCounterReseed:
@staticmethod
@with_service_target(SPEND_COUNTERS_TARGET)
async def increment_in_memory(spend_counter_cache: "DualCache", counter_key: str, increment: float) -> float | None:
async def increment_in_memory(spend_counter_store: SpendCounterStore, counter_key: str, increment: float) -> float:
"""Apply local deltas after an in-flight reseed establishes the spend balance."""
lock: Final = await SpendCounterReseed._get_lock(counter_key)
async with lock:
return await spend_counter_cache.async_increment_cache(
key=counter_key, value=increment, local_only=True, refresh_ttl=True
)
return spend_counter_store.increment_local(counter_key, increment)
@staticmethod
async def from_db(prisma_client: Optional["PrismaClient"], counter_key: str) -> float | None:
@ -222,7 +217,7 @@ class SpendCounterReseed:
@with_service_target(SPEND_COUNTERS_TARGET)
async def coalesced(
prisma_client: Optional["PrismaClient"],
spend_counter_cache: "DualCache",
spend_counter_store: SpendCounterStore,
counter_key: str,
require_cache_warm: bool = False,
) -> float | None:
@ -236,24 +231,9 @@ class SpendCounterReseed:
"""
lock: Final = await SpendCounterReseed._get_lock(counter_key)
async with lock:
batched: Final = await SpendCounterReseed._read_active_batch(counter_key)
if batched is not None and batched[0] is not None:
return batched[0]
# Re-check after acquiring the lock. Skip in-memory on a clean
# Redis miss - in-memory is per-pod-stale.
redis_clean_miss = batched is not None
if spend_counter_cache.redis_cache is not None and not redis_clean_miss:
try:
val = await spend_counter_cache.redis_cache.async_get_cache(key=counter_key)
if val is not None:
return float(val)
redis_clean_miss = True
except Exception:
pass
if not redis_clean_miss:
val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
if val is not None:
return float(val)
cached: Final = await SpendCounterReseed._read_cached(spend_counter_store, counter_key)
if cached is not None:
return cached
db_spend: Final = await SpendCounterReseed.from_db(prisma_client, counter_key)
if db_spend is None:
@ -264,29 +244,8 @@ class SpendCounterReseed:
# Redis key with db_spend; concurrent seeders read the winner's
# value. INCRBYFLOAT-of-db_spend from N pods would multiply the
# counter (N x db_spend) and trigger spurious budget alerts.
current_value: float = float(db_spend)
try:
if spend_counter_cache.redis_cache is not None:
seeded: Final = await spend_counter_cache.redis_cache.async_set_cache(
key=counter_key,
value=db_spend,
nx=True,
)
if seeded:
current_value = float(db_spend)
else:
cached: Final = await spend_counter_cache.redis_cache.async_get_cache(key=counter_key)
current_value = float(cached) if cached is not None else float(db_spend)
spend_counter_cache.in_memory_cache.set_cache(
key=counter_key,
value=current_value,
)
record_spend_counter_value(counter_key, current_value)
else:
cached_spend: Final = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
seeded_spend: Final = max(db_spend, float(cached_spend)) if cached_spend is not None else db_spend
spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=seeded_spend)
return seeded_spend
seeded: Final = await spend_counter_store.seed(counter_key, db_spend, increment_if_lost=False)
except Exception:
verbose_proxy_logger.exception(
"SpendCounterReseed.coalesced: failed to warm counter %s",
@ -294,7 +253,20 @@ class SpendCounterReseed:
)
if require_cache_warm:
raise
return current_value
return float(db_spend)
if spend_counter_store.has_redis:
record_spend_counter_value(counter_key, float(seeded))
return float(seeded)
@staticmethod
async def _read_cached(spend_counter_store: SpendCounterStore, counter_key: str) -> float | None:
"""Re-check after acquiring the lock. A clean Redis miss, in the request's MGET or a GET, skips memory,
which only holds this pod's writes."""
batched: Final = await SpendCounterReseed._read_active_batch(counter_key)
if batched is not None:
return batched[0]
value, _ = await spend_counter_store.read(counter_key)
return None if value is None else float(value)
@staticmethod
async def window_from_table(
@ -425,7 +397,7 @@ class SpendCounterReseed:
@with_service_target(SPEND_COUNTERS_TARGET)
async def coalesced_window(
prisma_client: Optional["PrismaClient"],
spend_counter_cache: "DualCache",
spend_counter_store: SpendCounterStore,
counter_key: str,
entity_type: str,
entity_id: str,
@ -434,22 +406,9 @@ class SpendCounterReseed:
) -> float | None:
lock: Final = await SpendCounterReseed._get_lock(counter_key)
async with lock:
batched: Final = await SpendCounterReseed._read_active_batch(counter_key)
if batched is not None and batched[0] is not None:
return batched[0]
redis_clean_miss = batched is not None
if spend_counter_cache.redis_cache is not None and not redis_clean_miss:
try:
val = await spend_counter_cache.redis_cache.async_get_cache(key=counter_key)
if val is not None:
return float(val)
redis_clean_miss = True
except Exception:
pass
if not redis_clean_miss:
val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
if val is not None:
return float(val)
cached: Final = await SpendCounterReseed._read_cached(spend_counter_store, counter_key)
if cached is not None:
return cached
window_spend: Final = await SpendCounterReseed.window_from_db(
prisma_client=prisma_client,
@ -461,39 +420,13 @@ class SpendCounterReseed:
if window_spend is None:
return None
try:
if spend_counter_cache.redis_cache is not None:
seeded: Final = await spend_counter_cache.redis_cache.async_set_cache(
key=counter_key,
value=window_spend,
nx=True,
)
if seeded:
current_value = window_spend
else:
current_cached_value = await spend_counter_cache.redis_cache.async_get_cache(key=counter_key)
if current_cached_value is None:
current_value = await spend_counter_cache.redis_cache.async_increment(
key=counter_key,
value=window_spend,
)
else:
current_value = float(current_cached_value)
spend_counter_cache.in_memory_cache.set_cache(
key=counter_key,
value=current_value,
)
record_spend_counter_value(counter_key, float(current_value))
else:
cached_spend: Final = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
seeded_spend: Final = (
max(window_spend, float(cached_spend)) if cached_spend is not None else window_spend
)
spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=seeded_spend)
return seeded_spend
seeded: Final = await spend_counter_store.seed(counter_key, window_spend, increment_if_lost=True)
except Exception:
verbose_proxy_logger.exception(
"SpendCounterReseed.coalesced_window: failed to warm counter %s",
counter_key,
)
raise
return current_value
if spend_counter_store.has_redis:
record_spend_counter_value(counter_key, float(seeded))
return float(seeded)

View file

@ -1357,7 +1357,7 @@ async def _update_database_and_spend_counters(
"""The reservation is reconciled before the spend is persisted, from its own read. One spend counter batch then
spans the database write and the counter update, so the post-call counters are read with a single MGET after the
write and their increments leave in a single pipeline."""
from litellm.proxy.proxy_server import spend_counter_cache
from litellm.proxy.proxy_server import spend_counter_store
from litellm.proxy.spend_tracking.budget_reservation import get_reserved_counter_keys
if budget_reservation is not None:
@ -1376,7 +1376,7 @@ async def _update_database_and_spend_counters(
model_access_groups=model_access_groups,
project_id=project_id,
)
with spend_counter_batch_scope(spend_counter_cache.redis_cache, counter_keys=counter_keys):
with spend_counter_batch_scope(spend_counter_store, counter_keys=counter_keys):
return await _update_database_and_spend_counters_in_batch(
proxy_logging_obj=proxy_logging_obj,
spend_writer=spend_writer,

View file

@ -85,6 +85,9 @@ from litellm.proxy.management_helpers.object_permission_utils import ( # noqa:
)
from litellm.proxy.management_helpers.utils import management_endpoint_wrapper
from litellm.proxy.utils import handle_exception_on_proxy, hash_password
from litellm.redis.stores.spend_counters import (
user_counter,
)
from litellm.repositories.organization_repository import OrganizationRepository
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.table_repositories import (
@ -1444,7 +1447,7 @@ async def _invalidate_user_spend_counter_if_changed(
if non_default_values.get("spend") is not None:
from litellm.proxy.proxy_server import invalidate_spend_counter
await invalidate_spend_counter(counter_key=f"spend:user:{non_default_values['user_id']}")
await invalidate_spend_counter(counter_key=user_counter(str(non_default_values["user_id"])))
def _clears_object_permission(user_request: UpdateUserRequest) -> bool:

View file

@ -30,7 +30,7 @@ from pydantic import TypeAdapter
from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm._internal_context import service_target, with_service_target
from litellm._internal_context import with_service_target
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
from litellm.caching.dual_cache import DualCache
@ -157,7 +157,6 @@ from litellm.proxy.mcp_server.outbound_credentials.sso_assertion_store import (
)
from litellm.proxy.search_endpoints.search_tool_registry import rotate_search_tools_master_key
from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start
from litellm.proxy.spend_tracking.spend_counter_batch import SPEND_COUNTERS_TARGET
from litellm.proxy.spend_tracking.spend_tracking_utils import ( # noqa: F401 # legacy module exports
_is_master_key, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
is_master_key,
@ -171,6 +170,11 @@ from litellm.proxy.utils import (
is_valid_api_key,
require_enterprise_license,
)
from litellm.redis.stores.spend_counters import (
SPEND_COUNTERS_TARGET,
key_counter,
key_window_counter,
)
from litellm.repositories.base_repository import BaseRepository, record_to_dict
from litellm.repositories.budget_repository import BudgetRepository
from litellm.repositories.config_repository import ConfigParam, ConfigRepository
@ -4263,21 +4267,18 @@ async def update_key_fn(
)
if data.spend is not None:
from litellm.proxy.proxy_server import spend_counter_cache
from litellm.proxy.proxy_server import spend_counter_store
counter_key: Final = f"spend:key:{hash_token_if_needed(key)}"
spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=data.spend, ttl=60)
if spend_counter_cache.redis_cache is not None:
try:
with service_target(SPEND_COUNTERS_TARGET):
await spend_counter_cache.redis_cache.async_set_cache(key=counter_key, value=data.spend, ttl=60)
except Exception as redis_err:
verbose_proxy_logger.warning(
"Failed to update spend counter %s in Redis after key spend update: %s. "
"Budget checks may use stale value until counter expires.",
counter_key,
redis_err,
)
counter_key: Final = key_counter(hash_token_if_needed(key))
try:
await spend_counter_store.set(counter_key, data.spend, ttl=60)
except Exception as redis_err:
verbose_proxy_logger.warning(
"Failed to update spend counter %s in Redis after key spend update: %s. "
"Budget checks may use stale value until counter expires.",
counter_key,
redis_err,
)
asyncio.create_task(
KeyManagementEventHooks.async_key_updated_hook(
@ -4901,7 +4902,7 @@ async def _budget_window_usage(
if not isinstance(duration, str) or not duration:
return None
spend: Final = await get_current_spend(
counter_key=f"spend:key:{api_key_hash}:window:{duration}",
counter_key=key_window_counter(api_key_hash, duration),
fallback_spend=0.0,
max_budget=_window_max_budget(window),
window_entity_type="Key",
@ -6820,19 +6821,16 @@ async def _set_spend_counter_with_floor_and_broadcast(counter_key: str, value: f
Without the broadcast, a worker that already cached the pre-reset key object
or floor marker keeps enforcing against it until its own TTL expires.
"""
from litellm.proxy.proxy_server import SPEND_DB_FLOOR_CACHE_TTL_SECONDS, spend_counter_cache
from litellm.proxy.proxy_server import SPEND_DB_FLOOR_CACHE_TTL_SECONDS, spend_counter_cache, spend_counter_store
spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=value, ttl=60)
if spend_counter_cache.redis_cache is not None:
try:
await spend_counter_cache.redis_cache.async_set_cache(key=counter_key, value=value, ttl=60)
except Exception as redis_err:
verbose_proxy_logger.warning(
"Failed to update spend counter %s in Redis: %s. "
"Budget checks may use stale value until counter expires.",
counter_key,
redis_err,
)
try:
await spend_counter_store.set(counter_key, value, ttl=60)
except Exception as redis_err:
verbose_proxy_logger.warning(
"Failed to update spend counter %s in Redis: %s. Budget checks may use stale value until counter expires.",
counter_key,
redis_err,
)
floor_key: Final = f"spend_db_floor:{counter_key}"
spend_counter_cache.in_memory_cache.set_cache(key=floor_key, value=value, ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS)
@ -6942,7 +6940,7 @@ async def _reset_key_budget_windows(
for window in reset_windows:
duration = window.get("budget_duration")
if isinstance(duration, str) and duration:
counter_key = f"spend:key:{hashed_api_key}:window:{duration}"
counter_key = key_window_counter(hashed_api_key, duration)
await _set_spend_counter_with_floor_and_broadcast(counter_key=counter_key, value=0.0)
@ -7016,7 +7014,7 @@ async def reset_key_spend_fn(
# budget_limits windows, so get_current_spend() returns the correct
# amount for every enforcement check immediately instead of the stale
# pre-reset value.
_counter_key: Final = f"spend:key:{hashed_api_key}"
_counter_key: Final = key_counter(hashed_api_key)
await _set_spend_counter_with_floor_and_broadcast(counter_key=_counter_key, value=reset_to)
await _reset_key_budget_windows(
prisma_client=prisma_client,

View file

@ -31,6 +31,9 @@ from litellm.proxy.management_endpoints.common_daily_activity import (
get_daily_activity,
)
from litellm.proxy.management_helpers.utils import handle_budget_for_entity
from litellm.redis.stores.spend_counters import (
tag_counter,
)
from litellm.repositories.daily_activity_buckets import api_key_from_bucket
from litellm.repositories.model_repository import ModelRepository
from litellm.repositories.table_repositories import (
@ -319,7 +322,7 @@ async def new_tag(
)
await _evict_tag_cache_keys((tag_cache_key(tag.name), tag_registry_cache_key()))
await invalidate_spend_counter(counter_key=f"spend:tag:{tag.name}")
await invalidate_spend_counter(counter_key=tag_counter(tag.name))
# Update models with new tag
if tag.models:

View file

@ -202,7 +202,6 @@ from litellm.router_strategy.auto_router.auto_router_tuning_baseline import (
tuning_limit_violation,
)
from litellm.rust_bridge.catalog import Target
from litellm.types.caching import RedisPipelineIncrementOperation
from litellm.types.utils import (
PRICING_OVERRIDES_KEY,
CountContextManagement,
@ -326,10 +325,7 @@ from litellm._logging import (
verbose_router_logger,
)
from litellm.caching.caching import DualCache, RedisCache
from litellm.caching.redis_batch import (
active_post_call_redis_batch,
drain_post_call_redis_batches,
)
from litellm.caching.redis_batch import drain_post_call_redis_batches
from litellm.caching.redis_cache import RedisCircuitBreakerOpenError, is_redis_timeout_failure
from litellm.caching.redis_cluster_cache import RedisClusterCache
from litellm.constants import (
@ -590,9 +586,7 @@ from litellm.proxy.common_utils.user_api_key_cache import (
end_user_cache_key,
get_management_object_ttl,
model_access_group_cache_key,
model_access_group_spend_counter_key,
project_cache_key,
project_spend_counter_key,
tag_cache_key,
)
from litellm.proxy.common_utils.validation_error_body import public_validation_errors
@ -859,6 +853,23 @@ from litellm.proxy.telemetry.middleware import TelemetryMiddleware
from litellm.proxy.telemetry.runtime import TelemetryRuntime
from litellm.proxy.telemetry.settings import describe_errors as describe_telemetry_errors
from litellm.proxy.telemetry.settings import load_settings as load_telemetry_settings
from litellm.redis.stores._connection import native_client
from litellm.redis.stores.spend_counters import (
SPEND_COUNTERS_TARGET,
SpendCounterStore,
end_user_counter,
is_unavailable,
key_counter,
key_window_counter,
model_access_group_counter,
org_counter,
project_counter,
tag_counter,
team_counter,
team_member_counter,
team_window_counter,
user_counter,
)
from litellm.router import router_arguments
from litellm.router_config import build_router_config
@ -945,7 +956,6 @@ from litellm.proxy.spend_tracking.spend_capture_rate import (
run_scheduled_spend_capture_rate_check,
)
from litellm.proxy.spend_tracking.spend_counter_batch import (
SPEND_COUNTERS_TARGET,
PendingSpendIncrement,
active_spend_counter_batch,
forget_spend_counter,
@ -3080,6 +3090,7 @@ user_api_key_cache: UserApiKeyCache = UserApiKeyCache(
default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value
)
spend_counter_cache: Final = DualCache(default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value)
spend_counter_store: Final = SpendCounterStore(None)
cli_sso_session_cache: Final = DualCache(default_in_memory_ttl=CLI_SSO_SESSION_TTL_SECONDS)
saml_sso_cache: Final = DualCache(default_in_memory_ttl=SAML_AUTHN_REQUEST_TTL_SECONDS)
model_max_budget_limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=spend_counter_cache)
@ -3420,28 +3431,18 @@ async def _repair_stale_spend_counter(counter_key: str, db_spend: float) -> None
in-memory copy is guarded by a read-compare-write with no await in between,
so it is atomic within the worker.
"""
cached: Final = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
needs_update = True
if cached is not None:
try:
needs_update = float(cached) < db_spend
except (TypeError, ValueError):
needs_update = True
if needs_update:
spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=db_spend)
if spend_counter_cache.redis_cache is not None:
forget_spend_counter(counter_key)
try:
repaired: Final = await spend_counter_cache.redis_cache.async_set_max(key=counter_key, value=db_spend)
except Exception:
verbose_proxy_logger.debug(
"Unable to repair stale spend counter %s in Redis",
counter_key,
exc_info=True,
)
return
if repaired is not None:
record_spend_counter_value(counter_key, repaired)
forget_spend_counter(counter_key)
try:
repaired: Final = await spend_counter_store.set_max(counter_key, db_spend)
except Exception:
verbose_proxy_logger.debug(
"Unable to repair stale spend counter %s in Redis",
counter_key,
exc_info=True,
)
return
if repaired is not None:
record_spend_counter_value(counter_key, repaired)
async def reseed_spend_counter_from_db(counter_key: str) -> bool:
@ -3538,19 +3539,8 @@ async def read_spend_counter_cache_value(counter_key: str) -> tuple[float | None
if batched is not None:
return batched
if spend_counter_cache.redis_cache is not None:
try:
redis_val: Final = await spend_counter_cache.redis_cache.async_get_cache(key=counter_key)
return (float(redis_val) if redis_val is not None else None), True
except Exception as e:
verbose_proxy_logger.debug(
"get_current_spend: Redis read failed for %s, falling back to in-memory: %s",
counter_key,
e,
)
in_memory_val: Final = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
return (float(in_memory_val) if in_memory_val is not None else None), False
value, authoritative = await spend_counter_store.read(counter_key)
return (float(value) if value is not None else None), authoritative
async def _read_spend_counter_estimate(counter_key: str, fallback_spend: float) -> tuple[float, bool]:
@ -3565,7 +3555,7 @@ async def _read_spend_counter_estimate(counter_key: str, fallback_spend: float)
# Reseed from DB - fallback_spend lags cross-pod, would allow bypass.
db_spend: Final = await SpendCounterReseed.coalesced(
prisma_client=prisma_client,
spend_counter_cache=spend_counter_cache,
spend_counter_store=spend_counter_store,
counter_key=counter_key,
)
if db_spend is not None:
@ -3591,7 +3581,7 @@ async def increment_spend_counters(
"""
Atomically increment spend counters for budget enforcement.
Uses spend_counter_cache (DualCache with Redis backend when available)
Uses spend_counter_store (Redis when available, this worker's memory behind it)
so counters are shared across all pods. Budget check functions read
from these counters via get_current_spend() (Redis-first).
@ -3599,7 +3589,7 @@ async def increment_spend_counters(
updated before the next request's auth check runs.
"""
with spend_counter_batch_scope(
spend_counter_cache.redis_cache,
spend_counter_store,
counter_keys=post_call_counter_keys(
token=token,
team_id=team_id,
@ -3675,7 +3665,7 @@ async def _increment_spend_counters_batched(
hashed_token: Final = (
hash_token(token=key_token) if isinstance(key_token, str) and key_token.startswith("sk-") else key_token
)
key_counter_key: Final = f"spend:key:{hashed_token}"
key_counter_key: Final = key_counter(hashed_token)
key_pending: Final[tuple[PendingSpendIncrement, ...]] = (
()
if key_counter_key in reserved_counter_keys
@ -3695,18 +3685,18 @@ async def _increment_spend_counters_batched(
key_window_reset_at = (
window.get("reset_at") if isinstance(window, dict) else getattr(window, "reset_at", None)
)
key_window_counter: Final = f"spend:key:{hashed_token}:window:{duration}"
key_window_counter_key: Final = key_window_counter(hashed_token, str(duration))
key_window_start = get_budget_window_start(window)
pending_window: Final = (
await _prepare_window_spend_counter_increment(
counter_key=key_window_counter,
counter_key=key_window_counter_key,
entity_type="Key",
entity_id=hashed_token,
window_duration=duration,
window_start=key_window_start,
increment=cost,
)
if key_window_counter not in reserved_counter_keys
if key_window_counter_key not in reserved_counter_keys
else None
)
await _enqueue_window_spend_row_update(
@ -3734,7 +3724,7 @@ async def _increment_spend_counters_batched(
return key_pending + tuple(item for item in window_pending if item is not None)
async def _team_scope(scope_team_id: str) -> tuple[PendingSpendIncrement | BaseException, ...]:
team_counter_key: Final = f"spend:team:{scope_team_id}"
team_counter_key: Final = team_counter(scope_team_id)
team_pending: Final[tuple[PendingSpendIncrement, ...]] = (
()
if team_counter_key in reserved_counter_keys
@ -3754,18 +3744,18 @@ async def _increment_spend_counters_batched(
team_window_reset_at = (
window.get("reset_at") if isinstance(window, dict) else getattr(window, "reset_at", None)
)
team_window_counter: Final = f"spend:team:{scope_team_id}:window:{duration}"
team_window_counter_key: Final = team_window_counter(scope_team_id, str(duration))
team_window_start = get_budget_window_start(window)
pending_window: Final = (
await _prepare_window_spend_counter_increment(
counter_key=team_window_counter,
counter_key=team_window_counter_key,
entity_type="Team",
entity_id=scope_team_id,
window_duration=duration,
window_start=team_window_start,
increment=cost,
)
if team_window_counter not in reserved_counter_keys
if team_window_counter_key not in reserved_counter_keys
else None
)
await _enqueue_window_spend_row_update(
@ -3795,7 +3785,7 @@ async def _increment_spend_counters_batched(
async def _team_member_scope(
scope_user_id: str, scope_team_id: str
) -> tuple[PendingSpendIncrement | BaseException, ...]:
team_member_counter_key: Final = f"spend:team_member:{scope_user_id}:{scope_team_id}"
team_member_counter_key: Final = team_member_counter(scope_user_id, scope_team_id)
if team_member_counter_key in reserved_counter_keys:
return ()
return (
@ -3807,7 +3797,7 @@ async def _increment_spend_counters_batched(
)
async def _user_scope(scope_user_id: str) -> tuple[PendingSpendIncrement | BaseException, ...]:
user_counter_key: Final = f"spend:user:{scope_user_id}"
user_counter_key: Final = user_counter(scope_user_id)
if user_counter_key in reserved_counter_keys:
return ()
return (
@ -3943,7 +3933,7 @@ async def _prepare_end_user_and_tag_spend_increments(
coro
for coro in (
_prepare_unreserved_spend_counter_increment(
counter_key=f"spend:end_user:{end_user_id}",
counter_key=end_user_counter(end_user_id),
source_cache_key=end_user_cache_key(end_user_id),
increment=response_cost,
reserved_counter_keys=reserved_counter_keys,
@ -3952,7 +3942,7 @@ async def _prepare_end_user_and_tag_spend_increments(
else None,
*(
_prepare_unreserved_spend_counter_increment(
counter_key=f"spend:tag:{tag_name}",
counter_key=tag_counter(tag_name),
source_cache_key=tag_cache_key(tag_name),
increment=response_cost,
reserved_counter_keys=reserved_counter_keys,
@ -3988,7 +3978,7 @@ async def _prepare_model_access_group_spend_increments(
results: Final = await asyncio.gather(
*(
_prepare_unreserved_spend_counter_increment(
counter_key=model_access_group_spend_counter_key(group),
counter_key=model_access_group_counter(group),
source_cache_key=model_access_group_cache_key(group),
increment=response_cost,
reserved_counter_keys=reserved_counter_keys,
@ -4009,7 +3999,7 @@ async def _prepare_org_spend_increment(
return ()
pending: Final = await _prepare_unreserved_spend_counter_increment(
counter_key=f"spend:org:{org_id}",
counter_key=org_counter(org_id),
source_cache_key=[f"org_id:{org_id}:with_budget", f"org_id:{org_id}"],
increment=response_cost,
reserved_counter_keys=reserved_counter_keys,
@ -4026,7 +4016,7 @@ async def _prepare_project_spend_increment(
return ()
pending: Final = await _prepare_unreserved_spend_counter_increment(
counter_key=project_spend_counter_key(project_id),
counter_key=project_counter(project_id),
source_cache_key=project_cache_key(project_id),
increment=response_cost,
reserved_counter_keys=reserved_counter_keys,
@ -4061,7 +4051,7 @@ async def _prepare_spend_counter_increment(
pipelined Redis call.
On first access per pod:
1. Check spend_counter_cache (in-memory -> Redis via DualCache)
1. Check spend_counter_store (Redis, then this worker's memory)
2. If not found, reseed from the DB via `SpendCounterReseed.coalesced`.
Falls back to the cached object's `.spend` via user_api_key_cache
only if prisma is unavailable, since that value can lag the flusher.
@ -4166,7 +4156,7 @@ async def ensure_spend_counter_initialized(
# Shares the per-counter lock with get_current_spend.
db_spend: Final = await SpendCounterReseed.coalesced(
prisma_client=prisma_client,
spend_counter_cache=spend_counter_cache,
spend_counter_store=spend_counter_store,
counter_key=counter_key,
require_cache_warm=True,
)
@ -4207,7 +4197,7 @@ async def ensure_window_spend_counter_initialized(
window_spend: Final = await SpendCounterReseed.coalesced_window(
prisma_client=prisma_client,
spend_counter_cache=spend_counter_cache,
spend_counter_store=spend_counter_store,
counter_key=counter_key,
entity_type=entity_type,
entity_id=entity_id,
@ -4233,29 +4223,13 @@ async def _is_spend_counter_cache_warm(counter_key: str) -> bool:
batched_value, _ = batched
if batched_value is None:
return False
spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=batched_value)
spend_counter_store.set_local(counter_key, batched_value)
return True
if spend_counter_cache.redis_cache is not None:
try:
current_value: Final[object] = await spend_counter_cache.redis_cache.async_get_cache(
key=counter_key,
)
if current_value is None:
return False
spend_counter_cache.in_memory_cache.set_cache(
key=counter_key,
value=current_value,
)
return True
except Exception as e:
verbose_proxy_logger.debug(
"Unable to read Redis spend counter %s before initialization, falling back to in-memory: %s",
counter_key,
e,
)
return spend_counter_cache.in_memory_cache.get_cache(key=counter_key) is not None
current_value, from_redis = await spend_counter_store.read(counter_key)
if from_redis and current_value is not None:
spend_counter_store.set_local(counter_key, current_value)
return current_value is not None
async def increment_spend_counter(counter_key: str, increment: float):
@ -4267,10 +4241,10 @@ async def increment_spend_counter(counter_key: str, increment: float):
@with_service_target(SPEND_COUNTERS_TARGET)
async def refresh_spend_counter_ttl(counter_key: str) -> bool:
if spend_counter_cache.redis_cache is None:
if not spend_counter_store.has_redis:
return False
try:
return await spend_counter_cache.redis_cache.async_refresh_ttl(key=counter_key)
return await spend_counter_store.refresh_ttl(counter_key)
except Exception as e:
verbose_proxy_logger.debug("spend counter TTL refresh skipped for %s: %s", counter_key, e)
return False
@ -4278,25 +4252,17 @@ async def refresh_spend_counter_ttl(counter_key: str) -> bool:
@with_service_target(SPEND_COUNTERS_TARGET)
async def increment_spend_counter_cache(counter_key: str, increment: float) -> float | None:
if spend_counter_cache.redis_cache is not None:
if spend_counter_store.has_redis:
try:
current_value: Final = await spend_counter_cache.redis_cache.async_increment(
key=counter_key,
value=increment,
refresh_ttl=True,
)
current_value: Final = await spend_counter_store.increment(counter_key, increment)
except Exception:
await invalidate_spend_counter(counter_key=counter_key)
forget_spend_counter(counter_key)
raise
spend_counter_cache.in_memory_cache.set_cache(
key=counter_key,
value=current_value,
)
record_spend_counter_value(counter_key, float(current_value))
return current_value
return await SpendCounterReseed.increment_in_memory(
spend_counter_cache=spend_counter_cache, counter_key=counter_key, increment=increment
spend_counter_store=spend_counter_store, counter_key=counter_key, increment=increment
)
@ -4306,10 +4272,8 @@ _increment_spend_counter_cache: Final = increment_spend_counter_cache
@with_service_target(SPEND_COUNTERS_TARGET)
async def invalidate_spend_counter(counter_key: str) -> None:
forget_spend_counter(counter_key)
spend_counter_cache.in_memory_cache.delete_cache(key=counter_key)
if spend_counter_cache.redis_cache is not None:
with contextlib.suppress(Exception):
await spend_counter_cache.redis_cache.async_delete_cache(key=counter_key)
with contextlib.suppress(Exception):
await spend_counter_store.delete((counter_key,))
_invalidate_spend_counter: Final = invalidate_spend_counter
@ -4321,51 +4285,22 @@ async def _apply_spend_counter_increments(pending: Sequence[PendingSpendIncremen
try:
await increment_spend_counters_pipeline(pending=pending)
except Exception as e:
if isinstance(e, RedisCircuitBreakerOpenError) or is_redis_timeout_failure(e):
if isinstance(e, RedisCircuitBreakerOpenError) or is_redis_timeout_failure(e) or is_unavailable(e):
return
raise
def _defer_spend_counter_increments(pending: Sequence[PendingSpendIncrement]) -> bool:
"""Post-call increments ride the request's post-call pipeline with the other counters. Each counter's
new value lands in memory when the pipeline settles; a failed one is invalidated so no reader trusts a
"""Post-call increments ride the request's native post-call pipeline with the other counters. Each counter's
new value lands in memory when the pipeline settles; a failed one is dropped everywhere so no reader trusts a
counter whose increment may not have applied, as ``increment_spend_counters_pipeline`` does."""
redis_cache: Final = spend_counter_cache.redis_cache
if redis_cache is None or not pending:
return False
batch: Final = active_post_call_redis_batch(redis_cache)
if batch is None:
return False
ttl: Final = redis_cache.get_ttl()
with service_target(SPEND_COUNTERS_TARGET):
for item in pending:
batch.increment(item.counter_key, item.increment, ttl).on_settled(_settle_spend_counter_increment(item))
return True
def _settle_spend_counter_increment(item: PendingSpendIncrement) -> Callable[[asyncio.Future[float]], Awaitable[None]]:
async def settle(future: asyncio.Future[float]) -> None:
if not future.cancelled() and future.exception() is None:
current_value: Final = float(future.result())
spend_counter_cache.in_memory_cache.set_cache(key=item.counter_key, value=current_value)
record_spend_counter_value(item.counter_key, current_value)
return
if future.cancelled():
if spend_counter_cache.in_memory_cache.get_cache(key=item.counter_key) is not None:
spend_counter_cache.in_memory_cache.increment_cache(key=item.counter_key, value=item.increment)
return
verbose_proxy_logger.warning(
"Spend counter %s increment did not land in the post-call pipeline; invalidating it", item.counter_key
)
await invalidate_spend_counter(counter_key=item.counter_key)
return settle
return spend_counter_store.increment_many_post_call(tuple((item.counter_key, item.increment) for item in pending))
async def increment_spend_counters_pipeline(pending: Sequence[PendingSpendIncrement]) -> tuple[float | None, ...]:
"""One INCRBYFLOAT+EXPIRE pipeline for every pending counter, returning each counter's new value in order; on
failure every counter is invalidated before the error propagates, so no caller can read a half-applied batch."""
if spend_counter_cache.redis_cache is None:
if not spend_counter_store.has_redis:
return await run_spend_counter_pipeline(pending=pending)
try:
return await run_spend_counter_pipeline(pending=pending)
@ -4380,26 +4315,21 @@ async def run_spend_counter_pipeline(pending: Sequence[PendingSpendIncrement]) -
happens to counters whose increment may or may not have landed when the pipeline fails."""
if not pending:
return ()
redis_cache: Final = spend_counter_cache.redis_cache
if redis_cache is None:
if not spend_counter_store.has_redis:
return tuple(
[
await SpendCounterReseed.increment_in_memory(
spend_counter_cache=spend_counter_cache, counter_key=item.counter_key, increment=item.increment
spend_counter_store=spend_counter_store, counter_key=item.counter_key, increment=item.increment
)
for item in pending
]
)
ttl: Final = redis_cache.get_ttl()
increment_list: Final = [
RedisPipelineIncrementOperation(key=item.counter_key, increment_value=item.increment, ttl=ttl)
for item in pending
]
results: Final = await redis_cache.async_increment_pipeline(increment_list=increment_list)
for item, current_value in zip(pending, results or ()):
spend_counter_cache.in_memory_cache.set_cache(key=item.counter_key, value=current_value)
record_spend_counter_value(item.counter_key, float(current_value))
return tuple(float(current_value) for current_value in results or ())
totals: Final = await spend_counter_store.increment_many(
tuple((item.counter_key, item.increment) for item in pending)
)
for item, current_value in zip(pending, totals, strict=True):
record_spend_counter_value(item.counter_key, current_value)
return totals
_CachedSpendObjectT = TypeVar(
@ -5284,6 +5214,7 @@ def _attach_redis_usage_cache(redis_cache: RedisCache, enable_redis_auth_cache:
redis_cache,
default_redis_ttl=litellm.default_redis_ttl,
)
spend_counter_store.attach_redis(native_client(redis_cache), redis_ttl=redis_cache.default_ttl)
cli_sso_session_cache.attach_redis_cache(
redis_cache,
default_redis_ttl=CLI_SSO_SESSION_TTL_SECONDS,
@ -8859,6 +8790,7 @@ class ProxyConfig:
redis_cache=redis_cache,
user_api_key_cache=user_api_key_cache,
additional_in_memory_caches=(spend_counter_cache.in_memory_cache, byok_credential_cache),
spend_counters=spend_counter_store,
)
self.auth_cache_invalidation_subscriber = subscriber
subscriber.start()

View file

@ -35,9 +35,7 @@ from litellm.proxy.common_utils.user_api_key_cache import (
UserApiKeyCache,
end_user_cache_key,
model_access_group_cache_key,
model_access_group_spend_counter_key,
project_cache_key,
project_spend_counter_key,
tag_cache_key,
team_membership_reservation_cache_key,
)
@ -50,6 +48,17 @@ from litellm.proxy.spend_tracking.input_tokens import (
)
from litellm.proxy.spend_tracking.spend_counter_batch import PendingSpendIncrement, spend_counter_batch_scope
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.redis.stores.spend_counters import (
end_user_counter,
key_counter,
model_access_group_counter,
org_counter,
project_counter,
tag_counter,
team_counter,
team_member_counter,
user_counter,
)
from litellm.router import Router
from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget
from litellm.types.router import DeploymentTypedDict
@ -128,15 +137,16 @@ def _start_reservation_lease_renewal(budget_reservation: Mapping[str, object], c
"""A reservation lives inside spend counter keys that expire on their Redis TTL. Renew the TTL
while the request is in flight so a request longer than the TTL does not drop its
reservation and admit concurrent requests against the DB floor on any worker."""
from litellm.proxy.proxy_server import spend_counter_cache
from litellm.proxy.proxy_server import spend_counter_store
if spend_counter_cache.redis_cache is None or not counter_keys:
redis_ttl: Final = spend_counter_store.redis_ttl
if redis_ttl is None or not counter_keys:
return
task: Final = asyncio.create_task(
_renew_reservation_lease(
budget_reservation=budget_reservation,
counter_keys=counter_keys,
interval=spend_counter_cache.redis_cache.default_ttl / 2,
interval=redis_ttl / 2,
request_task=asyncio.current_task(),
)
)
@ -174,7 +184,7 @@ def _key_reservation_should_release_for_throttle(counter_key: str, valid_token:
"""
if valid_token is None:
return False
return counter_key == f"spend:key:{valid_token.token}" and should_throttle_budget_exceeded(valid_token)
return counter_key == key_counter(str(valid_token.token)) and should_throttle_budget_exceeded(valid_token)
async def _apply_over_budget_reservation_policy(
@ -489,7 +499,7 @@ async def _get_budget_counters(
if valid_token.max_budget is not None and valid_token.max_budget > 0:
counters.append(
_BudgetCounter(
counter_key=f"spend:key:{valid_token.token}",
counter_key=key_counter(valid_token.token),
source_cache_key=valid_token.token,
max_budget=float(valid_token.max_budget),
fallback_spend=float(valid_token.spend or 0.0),
@ -499,7 +509,7 @@ async def _get_budget_counters(
)
counters.extend(
_get_budget_limit_counters(
entity_prefix=f"spend:key:{valid_token.token}",
entity_prefix=key_counter(valid_token.token),
entity_type="Key",
entity_id=valid_token.token,
budget_limits=valid_token.budget_limits,
@ -512,7 +522,7 @@ async def _get_budget_counters(
if team_object.max_budget is not None and team_object.max_budget > 0:
counters.append(
_BudgetCounter(
counter_key=f"spend:team:{team_id}",
counter_key=team_counter(team_id),
source_cache_key=f"team_id:{team_id}",
max_budget=float(team_object.max_budget),
fallback_spend=float(team_object.spend or 0.0),
@ -522,7 +532,7 @@ async def _get_budget_counters(
)
counters.extend(
_get_budget_limit_counters(
entity_prefix=f"spend:team:{team_id}",
entity_prefix=team_counter(team_id),
entity_type="Team",
entity_id=team_id,
budget_limits=team_object.budget_limits,
@ -540,7 +550,7 @@ async def _get_budget_counters(
):
counters.append(
_BudgetCounter(
counter_key=f"spend:user:{user_object.user_id}",
counter_key=user_counter(user_object.user_id),
source_cache_key=user_object.user_id,
max_budget=float(user_object.max_budget),
fallback_spend=float(user_object.spend or 0.0),
@ -623,7 +633,7 @@ async def _get_end_user_budget_counter(
return None
return _BudgetCounter(
counter_key=f"spend:end_user:{end_user_id}",
counter_key=end_user_counter(end_user_id),
source_cache_key=source_cache_key,
max_budget=max_budget,
fallback_spend=fallback_spend,
@ -663,7 +673,7 @@ async def _get_tag_budget_counters(
continue
counters.append(
_BudgetCounter(
counter_key=f"spend:tag:{tag_name}",
counter_key=tag_counter(tag_name),
source_cache_key=tag_cache_key(tag_name),
max_budget=max_budget,
fallback_spend=_to_float(_get_value(tag_object, "spend")) or 0.0,
@ -705,7 +715,7 @@ def _model_access_group_counter(group: str, budget: ModelAccessGroupBudget | Non
if budget is None or budget.max_budget is None or budget.max_budget <= 0:
return None
return _BudgetCounter(
counter_key=model_access_group_spend_counter_key(group),
counter_key=model_access_group_counter(group),
source_cache_key=model_access_group_cache_key(group),
max_budget=budget.max_budget,
fallback_spend=budget.spend,
@ -769,7 +779,7 @@ async def _get_team_member_budget_counter(
team_member_spend = cast(LiteLLM_TeamMembership, team_membership).spend if team_membership is not None else 0.0
return _BudgetCounter(
counter_key=f"spend:team_member:{valid_token.user_id}:{team_object.team_id}",
counter_key=team_member_counter(valid_token.user_id, team_object.team_id),
source_cache_key=membership_cache_key,
max_budget=float(team_member_budget),
fallback_spend=float(team_member_spend or 0.0),
@ -808,7 +818,7 @@ async def _get_org_budget_counter(
org_spend: Final = _to_float(_get_value(org_table, "spend")) or 0.0
return _BudgetCounter(
counter_key=f"spend:org:{org_id}",
counter_key=org_counter(org_id),
source_cache_key=f"org_id:{org_id}:with_budget",
max_budget=org_max_budget,
fallback_spend=org_spend,
@ -841,7 +851,7 @@ async def _get_project_budget_counter(
return None
return _BudgetCounter(
counter_key=project_spend_counter_key(valid_token.project_id),
counter_key=project_counter(valid_token.project_id),
source_cache_key=source_cache_key,
max_budget=project_max_budget,
fallback_spend=_to_float(_get_value(project_object, "spend")) or 0.0,
@ -1096,9 +1106,9 @@ async def _get_current_counter_value(counter: _BudgetCounter) -> float:
def _counters_batch_scope(counter_keys: frozenset[str]) -> spend_counter_batch_scope:
"""Each counter is read once, then written, so one MGET up front serves every read in the loop."""
from litellm.proxy.proxy_server import spend_counter_cache
from litellm.proxy.proxy_server import spend_counter_store
return spend_counter_batch_scope(spend_counter_cache.redis_cache, counter_keys=counter_keys)
return spend_counter_batch_scope(spend_counter_store, counter_keys=counter_keys)
@dataclass(frozen=True, slots=True)

View file

@ -7,21 +7,29 @@ from dataclasses import dataclass
from types import MappingProxyType, TracebackType
from typing import Final
from pydantic import TypeAdapter
from litellm._internal_context import service_target
from litellm._logging import verbose_proxy_logger
from litellm.caching.redis_batch import BatchResult, RedisBatch, active_request_redis_batch
from litellm.caching.redis_cache import RedisCache
from litellm.caching.redis_batch import NativeBatchResult, RequestRedisBatches, active_request_redis_batches
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_utils.user_api_key_cache import (
model_access_group_spend_counter_key,
project_spend_counter_key,
from litellm.redis.stores.spend_counters import (
SpendCounterStore,
end_user_counter,
key_counter,
model_access_group_counter,
org_counter,
project_counter,
tag_counter,
team_counter,
team_member_counter,
user_counter,
)
_CounterValues: Final = TypeAdapter(dict[str, float | None])
_NO_VALUES: Final[Mapping[str, float | None]] = MappingProxyType({})
SPEND_COUNTERS_TARGET: Final = "spend_counters"
@dataclass(frozen=True, slots=True)
class _InflightRead:
counters: tuple[str, ...]
result: NativeBatchResult[Sequence[float | None] | None]
@dataclass(frozen=True, slots=True)
@ -35,23 +43,23 @@ class SpendCounterBatch:
``async_batch_get_cache`` maps a clean miss to ``None`` and drops keys only when Redis failed, so an absent
key means "read it yourself" and a present ``None`` is an authoritative miss.
Inside a ``request_redis_batch_scope`` the MGET rides the request's pipeline instead: the batch's flush
hook declares whatever is bound but unread, so whoever flushes first (the auth object prefetch, usually)
carries the spend counters in the same round trip."""
Inside a ``request_redis_batch_scope`` the read rides the request's native pre-call pipeline instead: the
native flush hook declares whatever is bound but unread, so whichever native read the request awaits first
(the armed cooldown read, or this batch's own) carries the spend counters in the same round trip."""
__slots__ = ("_fetched", "_inflight", "_keys", "_loaded", "_lock", "_open", "_redis_cache", "_request_batch")
__slots__ = ("_fetched", "_inflight", "_keys", "_loaded", "_lock", "_open", "_request", "_store")
def __init__(self, redis_cache: RedisCache) -> None:
self._redis_cache: Final = redis_cache
def __init__(self, store: SpendCounterStore) -> None:
self._store: Final = store
self._lock: Final = asyncio.Lock()
self._open = True
self._keys: frozenset[str] = frozenset()
self._fetched: frozenset[str] = frozenset()
self._loaded: Mapping[str, float | None] = _NO_VALUES
self._inflight: Final[list[BatchResult[Mapping[str, object]]]] = [] # mutable-ok: drained by _load
self._request_batch: Final[RedisBatch | None] = active_request_redis_batch(redis_cache)
if self._request_batch is not None:
self._request_batch.add_flush_hook(self._declare_pending)
self._inflight: Final[list[_InflightRead]] = [] # mutable-ok: drained by _load
self._request: Final[RequestRedisBatches | None] = active_request_redis_batches()
if self._request is not None:
self._request.add_native_flush_hook(self._declare_pending)
@property
def counter_keys(self) -> frozenset[str]:
@ -96,7 +104,7 @@ class SpendCounterBatch:
async def _load(self) -> Mapping[str, float | None]:
async with self._lock:
if self._request_batch is not None:
if self._request is not None:
self._declare_pending()
await self._collect_inflight()
return self._loaded
@ -108,34 +116,33 @@ class SpendCounterBatch:
return self._loaded
def _declare_pending(self) -> None:
"""Flush hook: put every bound-but-unread counter on the request pipeline that is about to go out."""
if self._request_batch is None or not self._open:
"""Native flush hook: put every bound-but-unread counter on the native pre-call pipeline about to go out."""
if self._request is None or not self._open:
return
pending: Final = self._keys - self._fetched
if pending:
self._fetched = self._fetched | pending
with service_target(SPEND_COUNTERS_TARGET):
self._inflight.append(self._request_batch.mget(sorted(pending)))
pending: Final = tuple(sorted(self._keys - self._fetched))
if not pending:
return
result: Final = self._store.declare_read_many(self._request, pending)
if result is None:
return
self._fetched = self._fetched | frozenset(pending)
self._inflight.append(_InflightRead(counters=pending, result=result))
async def _collect_inflight(self) -> None:
results: Final = tuple(self._inflight)
reads: Final = tuple(self._inflight)
self._inflight.clear()
for result in results:
for read in reads:
try:
fetched: Mapping[str, float | None] = _CounterValues.validate_python(await result)
values: Sequence[float | None] | None = await read.result
except Exception as e: # noqa: BLE001 # per-key reads take over and apply their own Redis fallback
verbose_proxy_logger.debug("spend counter batch read failed, falling back to per-key reads: %s", e)
continue
self._loaded = MappingProxyType({**fetched, **self._loaded})
if values is not None:
self._loaded = MappingProxyType({**dict(zip(read.counters, values, strict=True)), **self._loaded})
async def _fetch(self, keys: frozenset[str]) -> Mapping[str, float | None]:
try:
with service_target(SPEND_COUNTERS_TARGET):
values: Final = await self._redis_cache.async_batch_get_cache(key_list=sorted(keys)) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # untyped cache API
return _CounterValues.validate_python(values) # pyright: ignore[reportUnknownArgumentType] # untyped cache API
except Exception as e: # noqa: BLE001 # per-key reads take over and apply their own Redis fallback
verbose_proxy_logger.debug("spend counter batch read failed, falling back to per-key reads: %s", e)
return _NO_VALUES
fetched: Final = await self._store.read_many(sorted(keys))
return _NO_VALUES if fetched is None else fetched
_active_batch: Final[ContextVar[SpendCounterBatch | None]] = ContextVar("spend_counter_batch", default=None)
@ -149,21 +156,21 @@ class spend_counter_batch_scope:
"""Reads inside the scope share one MGET for the keys bound here or by ``bind_*`` calls inside it.
Opened inside a scope whose batch is still open, it binds into that batch so both phases share the MGET."""
__slots__ = ("_counter_keys", "_redis_cache", "_token")
__slots__ = ("_counter_keys", "_store", "_token")
def __init__(self, redis_cache: RedisCache | None, counter_keys: frozenset[str] = frozenset()) -> None:
self._redis_cache: Final = redis_cache
def __init__(self, store: SpendCounterStore, counter_keys: frozenset[str] = frozenset()) -> None:
self._store: Final = store
self._counter_keys: Final = counter_keys
self._token: Token[SpendCounterBatch | None] | None = None
def __enter__(self) -> None:
if self._redis_cache is None:
if not self._store.has_redis:
return
outer: Final = _active_batch.get()
if outer is not None and outer.is_open:
outer.bind(self._counter_keys)
return
batch: Final = SpendCounterBatch(self._redis_cache)
batch: Final = SpendCounterBatch(self._store)
batch.bind(self._counter_keys)
self._token = _active_batch.set(batch)
@ -191,19 +198,19 @@ def _iter_entity_counter_keys(
"""Only string ids name a counter; anything else (None, or an unresolved placeholder in synthetic
logging payloads) simply has no counter to bind."""
if isinstance(token, str):
yield f"spend:key:{token}"
yield key_counter(token)
if isinstance(team_id, str):
yield f"spend:team:{team_id}"
yield team_counter(team_id)
if isinstance(user_id, str):
yield f"spend:team_member:{user_id}:{team_id}"
yield team_member_counter(user_id, team_id)
if isinstance(user_id, str):
yield f"spend:user:{user_id}"
yield user_counter(user_id)
if isinstance(end_user_id, str):
yield f"spend:end_user:{end_user_id}"
yield end_user_counter(end_user_id)
if isinstance(org_id, str):
yield f"spend:org:{org_id}"
yield org_counter(org_id)
if isinstance(project_id, str):
yield project_spend_counter_key(project_id)
yield project_counter(project_id)
def admission_counter_keys(token: UserAPIKeyAuth, end_user_id: str | None) -> frozenset[str]:
@ -240,11 +247,9 @@ def post_call_counter_keys(
end_user_id=end_user_id,
)
)
tag_keys: Final = frozenset(f"spend:tag:{tag}" for tag in tags or () if tag and isinstance(tag, str))
tag_keys: Final = frozenset(tag_counter(tag) for tag in tags or () if tag and isinstance(tag, str))
group_keys: Final = frozenset(
model_access_group_spend_counter_key(group)
for group in model_access_groups or ()
if group and isinstance(group, str)
model_access_group_counter(group) for group in model_access_groups or () if group and isinstance(group, str)
)
return entity_keys | tag_keys | group_keys

View file

@ -0,0 +1,225 @@
from __future__ import annotations
from collections.abc import Awaitable, Callable, Mapping, Sequence
from typing import Final, Literal, Protocol, TypeAlias
from litellm._internal_context import service_target, with_service_target
from litellm.caching.redis_batch import (
NativeBatchResult,
NativePending,
RequestRedisBatches,
active_request_redis_batches,
)
from litellm.rust_bridge._native import (
NativeRedisClient,
NativeRequestRedis,
NativeSpendCounterError,
NativeSpendCounters,
)
SPEND_COUNTERS_TARGET: Final = "spend_counters"
SPEND_COUNTER_PREFIX: Final = "spend:"
END_USER_COUNTER_PREFIX: Final = "spend:end_user:"
SpendIncrement: TypeAlias = tuple[str, float]
SpendReading: TypeAlias = tuple[float | None, bool]
"""(value, authoritative): authoritative when Redis answered, a clean miss included."""
def key_counter(token: str) -> str:
return f"spend:key:{token}"
def key_window_counter(token: str, budget_duration: str) -> str:
return f"spend:key:{token}:window:{budget_duration}"
def team_counter(team_id: str) -> str:
return f"spend:team:{team_id}"
def team_window_counter(team_id: str, budget_duration: str) -> str:
return f"spend:team:{team_id}:window:{budget_duration}"
def team_member_counter(user_id: str, team_id: str) -> str:
return f"spend:team_member:{user_id}:{team_id}"
def user_counter(user_id: str) -> str:
return f"spend:user:{user_id}"
def end_user_counter(end_user_id: str) -> str:
return f"{END_USER_COUNTER_PREFIX}{end_user_id}"
def org_counter(org_id: str) -> str:
return f"spend:org:{org_id}"
def tag_counter(tag: str) -> str:
return f"spend:tag:{tag}"
def project_counter(project_id: str) -> str:
return f"spend:project:{project_id}"
def model_access_group_counter(access_group: str) -> str:
return f"spend:model_access_group:{access_group}"
def shadow_eval_counter(job_id: str) -> str:
return f"spend:shadow_eval:{job_id}"
def is_unavailable(error: BaseException) -> bool:
"""Redis timed out or its breaker is open, the failures a post-call increment drops."""
return isinstance(error, NativeSpendCounterError) and error.args[:1] in (("timeout",), ("breaker_open",))
class SpendCountersHandle(Protocol):
def read(self, counter: str) -> Awaitable[tuple[Literal["redis", "memory"], float | None]]: ...
def read_many(self, counters: Sequence[str]) -> Awaitable[Sequence[float | None] | None]: ...
def declare_read_many(
self, request: NativeRequestRedis, counters: Sequence[str]
) -> NativePending[Sequence[float | None] | None] | None: ...
def increment(self, counter: str, amount: float) -> Awaitable[float]: ...
def increment_many(self, increments: Sequence[SpendIncrement]) -> Awaitable[Sequence[float]]: ...
def increment_many_post_call(self, request: NativeRequestRedis, increments: Sequence[SpendIncrement]) -> bool: ...
def set_max(self, counter: str, spend: float) -> Awaitable[float | None]: ...
def seed(self, counter: str, spend: float, *, increment_if_lost: bool) -> Awaitable[float]: ...
def set(self, counter: str, spend: float, ttl_seconds: float | None = None) -> Awaitable[None]: ...
def delete(self, counters: Sequence[str]) -> Awaitable[None]: ...
def refresh_ttl(self, counter: str) -> Awaitable[bool]: ...
def get_local(self, counter: str) -> float | None: ...
def set_local(self, counter: str, spend: float, ttl_seconds: float | None = None) -> None: ...
def delete_local(self, counter: str) -> None: ...
def increment_local(self, counter: str, amount: float) -> float: ...
def flush_local(self) -> None: ...
def snapshot(self) -> Mapping[str, float]: ...
HandleFactory: TypeAlias = Callable[[NativeRedisClient | None, int], SpendCountersHandle]
class SpendCounterStore:
"""The ``spend:*`` budget counters: Redis shared by every worker, and this worker's memory copy that answers
only when Redis cannot. Counters are named by the builders above."""
def __init__(
self, client: NativeRedisClient | None, *, redis_ttl: int = 60, handle: HandleFactory = NativeSpendCounters
) -> None:
self._build: Final = handle
self._handle = handle(client, redis_ttl)
self._redis_ttl = None if client is None else redis_ttl
def attach_redis(self, client: NativeRedisClient | None, *, redis_ttl: int) -> None:
if client is None or self._redis_ttl is not None:
return
self._handle = self._build(client, redis_ttl)
self._redis_ttl = redis_ttl
@property
def redis_ttl(self) -> int | None:
"""The TTL every Redis write gives a counter, None when there is no Redis."""
return self._redis_ttl
@property
def has_redis(self) -> bool:
return self._redis_ttl is not None
@with_service_target(SPEND_COUNTERS_TARGET)
async def read(self, counter: str) -> SpendReading:
source, value = await self._handle.read(counter)
return value, source == "redis"
@with_service_target(SPEND_COUNTERS_TARGET)
async def read_many(self, counters: Sequence[str]) -> Mapping[str, float | None] | None:
"""None when there is no Redis or the read failed, so each caller reads on its own."""
values: Final = await self._handle.read_many(counters)
return None if values is None else dict(zip(counters, values, strict=True))
def declare_read_many(
self, request: RequestRedisBatches, counters: Sequence[str]
) -> NativeBatchResult[Sequence[float | None] | None] | None:
"""``read_many`` declared on the request's native pre-call batch; None when there is no Redis."""
if not self.has_redis:
return None
with service_target(SPEND_COUNTERS_TARGET):
pending: Final = self._handle.declare_read_many(request.native(), counters)
return None if pending is None else request.native_result(pending)
@with_service_target(SPEND_COUNTERS_TARGET)
async def increment(self, counter: str, amount: float) -> float:
"""Raises when the increment may not have landed, after dropping the counter everywhere."""
return await self._handle.increment(counter, amount)
@with_service_target(SPEND_COUNTERS_TARGET)
async def increment_many(self, increments: Sequence[SpendIncrement]) -> tuple[float, ...]:
return tuple(await self._handle.increment_many(increments))
def increment_many_post_call(self, increments: Sequence[SpendIncrement]) -> bool:
"""Declares the increments on the request's post-call batch, False when there is no Redis or request.
Each total lands in memory once the batch goes out; a counter whose increment failed is dropped everywhere."""
request: Final = active_request_redis_batches()
if request is None or not self.has_redis or not increments:
return False
with service_target(SPEND_COUNTERS_TARGET):
return self._handle.increment_many_post_call(request.native_post_call(), increments)
@with_service_target(SPEND_COUNTERS_TARGET)
async def set_max(self, counter: str, spend: float) -> float | None:
"""Raise memory and Redis to ``spend``, never lowering either; what Redis holds after, None without Redis."""
return await self._handle.set_max(counter, spend)
@with_service_target(SPEND_COUNTERS_TARGET)
async def seed(self, counter: str, spend: float, *, increment_if_lost: bool) -> float:
"""Seed a cold counter with SET NX so one worker wins, answering with what the counter then holds."""
return await self._handle.seed(counter, spend, increment_if_lost=increment_if_lost)
@with_service_target(SPEND_COUNTERS_TARGET)
async def set(self, counter: str, spend: float, ttl: float | None = None) -> None:
"""Memory, then Redis; raises when Redis failed, with memory already written."""
await self._handle.set(counter, spend, ttl)
@with_service_target(SPEND_COUNTERS_TARGET)
async def delete(self, counters: Sequence[str]) -> None:
"""Memory, then Redis in one round trip; raises when Redis failed, with memory already dropped."""
await self._handle.delete(counters)
@with_service_target(SPEND_COUNTERS_TARGET)
async def refresh_ttl(self, counter: str) -> bool:
return await self._handle.refresh_ttl(counter)
def get_local(self, counter: str) -> float | None:
return self._handle.get_local(counter)
def set_local(self, counter: str, spend: float, ttl: float | None = None) -> None:
self._handle.set_local(counter, spend, ttl)
def delete_local(self, counter: str) -> None:
self._handle.delete_local(counter)
def increment_local(self, counter: str, amount: float) -> float:
return self._handle.increment_local(counter, amount)
def flush_local(self) -> None:
self._handle.flush_local()
def snapshot(self) -> Mapping[str, float]:
return self._handle.snapshot()
def apply_broadcast(self, cache_key: str, new_value: float | None, ttl: float | None) -> None:
"""A cross-worker auth cache invalidation: a SET writes this worker's copy, a delete drops it. Keys outside
the spend counter keyspace are not this store's."""
if not cache_key.startswith(SPEND_COUNTER_PREFIX):
return
try:
if new_value is None:
self.delete_local(cache_key)
else:
self.set_local(cache_key, new_value, ttl)
except ValueError:
return

View file

@ -12,6 +12,7 @@ class RustBridgeDeclined(Exception): ...
class RustUpstreamError(Exception): ...
class ForkedAfterNativeRuntimeStarted(RuntimeError): ...
class ProcessReservedForForking(RuntimeError): ...
class NativeSpendCounterError(Exception): ...
def trace_encode_error(message: str) -> bytes: ...
def bundled_model_cost_map() -> bytes: ...
@ -84,6 +85,33 @@ class NativeRouterUsage:
def flush_local_counters(self) -> None: ...
def snapshot(self) -> dict[str, float]: ...
@final
class NativeDeclaredSpendCounters:
def resolve(self) -> Future[list[float | None] | None]: ...
@final
class NativeSpendCounters:
def __new__(cls, client: NativeRedisClient | None, redis_ttl_seconds: int = 60, /) -> NativeSpendCounters: ...
def read(self, counter: str) -> Future[tuple[Literal["redis", "memory"], float | None]]: ...
def read_many(self, counters: Sequence[str]) -> Future[list[float | None] | None]: ...
def declare_read_many(
self, request: NativeRequestRedis, counters: Sequence[str]
) -> NativeDeclaredSpendCounters | None: ...
def increment(self, counter: str, amount: float) -> Future[float]: ...
def increment_many(self, increments: Sequence[tuple[str, float]]) -> Future[list[float]]: ...
def increment_many_post_call(self, request: NativeRequestRedis, increments: Sequence[tuple[str, float]]) -> bool: ...
def set_max(self, counter: str, spend: float) -> Future[float | None]: ...
def seed(self, counter: str, spend: float, *, increment_if_lost: bool) -> Future[float]: ...
def set(self, counter: str, spend: float, ttl_seconds: float | None = None) -> Future[None]: ...
def delete(self, counters: Sequence[str]) -> Future[None]: ...
def refresh_ttl(self, counter: str) -> Future[bool]: ...
def get_local(self, counter: str) -> float | None: ...
def set_local(self, counter: str, spend: float, ttl_seconds: float | None = None) -> None: ...
def delete_local(self, counter: str) -> None: ...
def increment_local(self, counter: str, amount: float) -> float: ...
def flush_local(self) -> None: ...
def snapshot(self) -> dict[str, float]: ...
@final
class NativeDatabase:
@staticmethod
@ -274,11 +302,14 @@ __all__ = [
"NativeClickHouseSpendStorage",
"NativeDatabase",
"NativeDeclaredCooldowns",
"NativeDeclaredSpendCounters",
"NativeDiagnosticProcessor",
"NativeRedisClient",
"NativeRequestRedis",
"NativeRouterCooldowns",
"NativeRouterUsage",
"NativeSpendCounterError",
"NativeSpendCounters",
"ProcessReservedForForking",
"ResponsesWebSocketConnection",
"RustBridgeDeclined",

View file

@ -13,6 +13,8 @@ from collections.abc import Awaitable, Callable, Mapping, Sequence
import pytest
from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters
import litellm
from litellm.caching.caching import Cache
from litellm.caching.dual_cache import DualCache
@ -444,76 +446,23 @@ async def test_the_success_handler_closes_the_post_call_batch_after_the_last_cal
@pytest.mark.asyncio
async def test_spend_counter_increments_ride_the_pipeline_and_settle_into_memory(monkeypatch):
async def test_spend_counter_increments_ride_the_native_post_call_batch(monkeypatch):
from litellm.proxy import proxy_server
client = FakeClient(_ok_replies)
spend_cache = DualCache()
spend_cache.attach_redis_cache(PostCallFakeRedisCache(client))
monkeypatch.setattr(proxy_server, "spend_counter_cache", spend_cache)
natives = FakeNatives()
spend = FakeSpendCounters(redis={"spend:key:k1": 7.0})
monkeypatch.setattr(proxy_server, "spend_counter_store", spend.store())
pending = [PendingSpendIncrement("spend:key:k1", 0.5), PendingSpendIncrement("spend:team:t1", 0.5)]
with request_redis_batch_scope():
with request_redis_batch_scope(native_factory=natives): # pyright: ignore[reportArgumentType] # test double
await proxy_server._apply_spend_counter_increments(pending)
assert client.pipelines == []
declared = list(spend.post_call)
assert natives.built[0].events == []
await flush_post_call_redis_batches()
assert [c for c in client.pipelines[0].commands if c[0] == "INCRBYFLOAT"] == [
("INCRBYFLOAT", "spend:key:k1", 0.5),
("INCRBYFLOAT", "spend:team:t1", 0.5),
]
assert spend_cache.in_memory_cache.get_cache("spend:key:k1") == 7.5
@pytest.mark.asyncio
async def test_a_spend_counter_whose_increment_failed_is_invalidated_not_trusted(monkeypatch):
from litellm.proxy import proxy_server
def replies(command: tuple[object, ...]) -> object:
if command[0] == "INCRBYFLOAT" and command[1] == "spend:key:k1":
return Exception("OOM")
return _ok_replies(command)
redis_cache = PostCallFakeRedisCache(FakeClient(replies))
spend_cache = DualCache()
spend_cache.attach_redis_cache(redis_cache)
spend_cache.in_memory_cache.set_cache("spend:key:k1", 3.0)
spend_cache.in_memory_cache.set_cache("spend:team:t1", 3.0)
monkeypatch.setattr(proxy_server, "spend_counter_cache", spend_cache)
with request_redis_batch_scope():
await proxy_server._apply_spend_counter_increments(
[PendingSpendIncrement("spend:key:k1", 0.5), PendingSpendIncrement("spend:team:t1", 0.5)]
)
await flush_post_call_redis_batches()
assert spend_cache.in_memory_cache.get_cache("spend:key:k1") is None
assert redis_cache.alone == [("DEL", "spend:key:k1")]
assert spend_cache.in_memory_cache.get_cache("spend:team:t1") == 7.5
@pytest.mark.asyncio
async def test_a_cancelled_post_call_flush_keeps_the_shared_spend_counter_and_counts_the_spend_locally(monkeypatch):
from litellm.proxy import proxy_server
redis_cache = PostCallFakeRedisCache(
FakeClient(_ok_replies, fail=asyncio.CancelledError()) # pyright: ignore[reportArgumentType] # a cancel raised mid-pipeline
)
spend_cache = DualCache()
spend_cache.attach_redis_cache(redis_cache)
spend_cache.in_memory_cache.set_cache("spend:key:k1", 3.0)
monkeypatch.setattr(proxy_server, "spend_counter_cache", spend_cache)
with request_redis_batch_scope():
await proxy_server._apply_spend_counter_increments(
[PendingSpendIncrement("spend:key:k1", 0.5), PendingSpendIncrement("spend:team:t1", 0.5)]
)
with pytest.raises(asyncio.CancelledError):
await flush_post_call_redis_batches()
assert redis_cache.alone == [], "a cancel says nothing about the shared counter, so Redis keeps it"
assert spend_cache.in_memory_cache.get_cache("spend:key:k1") == 3.5, "the local copy counts the cancelled spend"
assert spend_cache.in_memory_cache.get_cache("spend:team:t1") is None, "an absent local copy is not seeded"
assert declared == [(natives.built[0], (("spend:key:k1", 0.5), ("spend:team:t1", 0.5)))]
assert natives.built[0].events == ["flush_post"]
assert spend.commands == []
@pytest.mark.asyncio

View file

@ -3,6 +3,8 @@ from typing import Final
from unittest.mock import MagicMock
import pytest
from litellm.redis.stores.spend_counters import SpendCounterStore
from pydantic import BaseModel
import litellm.proxy.proxy_server as proxy_server
@ -70,6 +72,7 @@ def proxy_auth_state(monkeypatch: pytest.MonkeyPatch) -> ProxyAuthState:
"prisma_client": MagicMock(),
"user_api_key_cache": cache,
"spend_counter_cache": DualCache(),
"spend_counter_store": SpendCounterStore(None),
"proxy_logging_obj": ProxyLogging(user_api_key_cache=cache),
"general_settings": {},
"user_custom_auth": None,

View file

@ -398,7 +398,7 @@ async def test_a_sub_call_reserves_nothing_against_the_callers_budget(
proxy_auth_state.serve(_classifier_router(cost_per_token=0.001))
caller: Final = UserAPIKeyAuth(api_key="sk-live", spend=0.0, max_budget=10.0)
counter_key: Final = f"spend:key:{caller.token}"
await proxy_server.spend_counter_cache.async_set_cache(counter_key, 0.0)
proxy_server.spend_counter_store.set_local(counter_key, 0.0)
request: Final = Request(
scope={"type": "http", "method": "POST", "path": "/v1/messages", "headers": [], "query_string": b""}
)
@ -410,7 +410,7 @@ async def test_a_sub_call_reserves_nothing_against_the_callers_budget(
await run_centralized_common_checks(caller, request, priced_body, "/v1/messages", sub_call=sub_call)
counter: Final = await proxy_server.spend_counter_cache.async_get_cache(counter_key)
counter: Final = proxy_server.spend_counter_store.get_local(counter_key)
assert (counter > 0, caller.budget_reservation is not None) == (reserved, reserved)

View file

@ -105,6 +105,8 @@ from litellm.utils import get_utc_datetime
from litellm.llms.base_llm.managed_resources.utils import encode_unified_id, generate_unified_id_string
from litellm.vector_stores.vector_store_registry import VectorStoreRegistry
from tests._master_key import MASTER_KEY
from litellm.redis.stores.spend_counters import SpendCounterStore
from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters
def _rendered_log_message(call):
@ -9749,18 +9751,17 @@ def _project_with_budget(spend: float, max_budget: float | None):
async def test_project_max_budget_check_blocks_when_live_spend_reaches_the_budget(
counter_spend, db_spend, max_budget, blocks
):
from litellm.caching.dual_cache import DualCache
from litellm.proxy.auth.auth_checks import _project_max_budget_check
real_spend_counter_cache = DualCache()
real_spend_counter_store = SpendCounterStore(None)
if counter_spend is not None:
real_spend_counter_cache.in_memory_cache.set_cache(key="spend:project:p-budget", value=counter_spend)
real_spend_counter_store.set_local("spend:project:p-budget", counter_spend)
valid_token = UserAPIKeyAuth(api_key="hashed-key", project_id="p-budget", team_id="t-1", user_id="u-1")
proxy_logging_obj = MagicMock()
proxy_logging_obj.budget_alerts = AsyncMock()
with patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock
"litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache
with patch( # test-quality-ok: injects a real store for the module global, not a behavior mock
"litellm.proxy.proxy_server.spend_counter_store", real_spend_counter_store
):
if not blocks:
await _project_max_budget_check(
@ -10596,13 +10597,19 @@ async def test_invalidate_team_member_spend_state_sets_the_spend_counter_and_cle
await real_cache.async_set_cache(key="team_membership:user-1:team-1", value="stale-membership")
real_spend_counter_cache = DualCache()
real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:user-1:team-1", value=999.0)
real_spend_counter_cache.in_memory_cache.set_cache(
key="spend_db_floor:spend:team_member:user-1:team-1", value=999.0
)
real_spend_counter_store = SpendCounterStore(None)
real_spend_counter_store.set_local("spend:team_member:user-1:team-1", 999.0)
with patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock
"litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache
with (
patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock
"litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache
),
patch( # test-quality-ok: injects a real store for the module global, not a behavior mock
"litellm.proxy.proxy_server.spend_counter_store", real_spend_counter_store
),
):
await invalidate_team_member_spend_state(
user_id="user-1",
@ -10613,7 +10620,7 @@ async def test_invalidate_team_member_spend_state_sets_the_spend_counter_and_cle
assert await real_cache.async_get_cache(key="team-1_user-1") is None
assert await real_cache.async_get_cache(key="team_membership:user-1:team-1") is None
assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-1:team-1") == 0.0
assert real_spend_counter_store.get_local("spend:team_member:user-1:team-1") == 0.0
assert (
real_spend_counter_cache.in_memory_cache.get_cache(key="spend_db_floor:spend:team_member:user-1:team-1") == 0.0
), "the DB-floor marker kept the pre-reset value; a stale-floor read can raise the counter right back up"
@ -10626,17 +10633,16 @@ async def test_invalidate_team_member_spend_state_leaves_the_live_spend_counter_
case would force the next read to reseed from the DB's own spend column, which lags the live
counter via periodic batch writes, briefly UNDER-enforcing the raised cap against a spend
value lower than what was actually tracked (regression: PR #37971 Bugbot finding)."""
from litellm.caching.dual_cache import DualCache
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
real_cache = UserApiKeyCache()
await real_cache.async_set_cache(key="team-1_user-1", value="stale-membership")
real_spend_counter_cache = DualCache()
real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:user-1:team-1", value=999.0)
real_spend_counter_store = SpendCounterStore(None)
real_spend_counter_store.set_local("spend:team_member:user-1:team-1", 999.0)
with patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock
"litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache
with patch( # test-quality-ok: injects a real store for the module global, not a behavior mock
"litellm.proxy.proxy_server.spend_counter_store", real_spend_counter_store
):
await invalidate_team_member_spend_state(
user_id="user-1",
@ -10645,7 +10651,7 @@ async def test_invalidate_team_member_spend_state_leaves_the_live_spend_counter_
)
assert await real_cache.async_get_cache(key="team-1_user-1") is None
assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-1:team-1") == 999.0
assert real_spend_counter_store.get_local("spend:team_member:user-1:team-1") == 999.0
@pytest.mark.asyncio
@ -10653,17 +10659,13 @@ async def test_invalidate_team_member_spend_state_sets_new_spend_instead_of_dele
"""/key/{key}/reset_spend SETs its counter to the reset value rather than deleting it, so a
worker's next read reflects it directly instead of falling back through a DB reseed. A reset
caller passing new_spend must match that precedent, not merely delete the counter."""
from litellm.caching.dual_cache import DualCache
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
real_cache = UserApiKeyCache()
real_spend_counter_cache = DualCache()
fake_redis_cache = MagicMock()
fake_redis_cache.async_set_cache = AsyncMock()
real_spend_counter_cache.redis_cache = fake_redis_cache
spend = FakeSpendCounters()
with patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock
"litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache
with patch( # test-quality-ok: injects a store over a fake handle for the module global
"litellm.proxy.proxy_server.spend_counter_store", spend.store()
):
await invalidate_team_member_spend_state(
user_id="user-1",
@ -10672,8 +10674,8 @@ async def test_invalidate_team_member_spend_state_sets_new_spend_instead_of_dele
new_spend=2.5,
)
assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-1:team-1") == 2.5
fake_redis_cache.async_set_cache.assert_awaited_once_with(key="spend:team_member:user-1:team-1", value=2.5, ttl=60)
assert spend.memory == {"spend:team_member:user-1:team-1": 2.5}
assert spend.commands == ["SET spend:team_member:user-1:team-1 2.5 60"]
@pytest.mark.asyncio
@ -10682,18 +10684,13 @@ async def test_invalidate_team_member_spend_state_deletes_redis_counter_when_set
first), so a failed Redis SET would otherwise leave the OLD pre-reset value authoritative
for every worker even though the reset reported success. On a failed SET, the stale Redis
entry must be deleted instead, so the next read clean-misses and reseeds from the DB."""
from litellm.caching.dual_cache import DualCache
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
real_cache = UserApiKeyCache()
real_spend_counter_cache = DualCache()
fake_redis_cache = MagicMock()
fake_redis_cache.async_set_cache = AsyncMock(side_effect=ConnectionError("redis down"))
fake_redis_cache.async_delete_cache = AsyncMock()
real_spend_counter_cache.redis_cache = fake_redis_cache
spend = FakeSpendCounters(redis={"spend:team_member:user-1:team-1": 9.0}, failing=frozenset({"set"}))
with patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock
"litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache
with patch( # test-quality-ok: injects a store over a fake handle for the module global
"litellm.proxy.proxy_server.spend_counter_store", spend.store()
):
await invalidate_team_member_spend_state(
user_id="user-1",
@ -10702,7 +10699,9 @@ async def test_invalidate_team_member_spend_state_deletes_redis_counter_when_set
new_spend=2.5,
)
fake_redis_cache.async_delete_cache.assert_awaited_once_with(key="spend:team_member:user-1:team-1")
assert spend.commands == ["DEL spend:team_member:user-1:team-1"]
assert spend.redis == {}
assert spend.memory == {"spend:team_member:user-1:team-1": 2.5}
@pytest.mark.asyncio
@ -10712,19 +10711,14 @@ async def test_invalidate_team_member_spend_state_raises_503_when_both_redis_wri
member, so the reset must surface a 503 instead (regression: PR #37971 Greptile finding)."""
from fastapi import HTTPException
from litellm.caching.dual_cache import DualCache
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
real_cache = UserApiKeyCache()
real_spend_counter_cache = DualCache()
fake_redis_cache = MagicMock()
fake_redis_cache.async_set_cache = AsyncMock(side_effect=ConnectionError("redis down"))
fake_redis_cache.async_delete_cache = AsyncMock(side_effect=ConnectionError("redis still down"))
real_spend_counter_cache.redis_cache = fake_redis_cache
spend = FakeSpendCounters(failing=frozenset({"set", "delete"}))
with (
patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock
"litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache
patch( # test-quality-ok: injects a store over a fake handle for the module global
"litellm.proxy.proxy_server.spend_counter_store", spend.store()
),
pytest.raises(HTTPException) as exc_info,
):
@ -10776,13 +10770,17 @@ async def test_invalidate_team_member_spend_state_broadcasts_the_spend_counter_t
remote_user_api_key_cache = UserApiKeyCache()
remote_spend_counter_in_memory_cache = InMemoryCache()
remote_spend_counter_in_memory_cache.set_cache("spend:team_member:user-1:team-1", 999.0)
remote_spend_counter_in_memory_cache.set_cache("spend_db_floor:spend:team_member:user-1:team-1", 999.0)
remote_spend_counter_store = SpendCounterStore(None)
remote_spend_counter_store.set_local("spend:team_member:user-1:team-1", 999.0)
with (
patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock
"litellm.proxy.proxy_server.spend_counter_cache", local_spend_counter_cache
),
patch( # test-quality-ok: injects a real store for the module global, not a behavior mock
"litellm.proxy.proxy_server.spend_counter_store", SpendCounterStore(None)
),
patch( # test-quality-ok: injects a fake pub/sub-capable redis cache; no live redis in this unit test
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache",
return_value=_FakeRedisCache(),
@ -10804,13 +10802,14 @@ async def test_invalidate_team_member_spend_state_broadcasts_the_spend_counter_t
redis_cache=_FakeRedisCache(),
user_api_key_cache=remote_user_api_key_cache,
additional_in_memory_caches=(remote_spend_counter_in_memory_cache,),
spend_counters=remote_spend_counter_store,
)
for cache_key in ("spend:team_member:user-1:team-1", "spend_db_floor:spend:team_member:user-1:team-1"):
remote_subscriber._apply_message( # pyright: ignore[reportPrivateUsage] # exercising the real cross-worker message handler, not a public API
{"type": "message", "data": _published_message_for(cache_key)}
)
assert remote_spend_counter_in_memory_cache.get_cache("spend:team_member:user-1:team-1") == 0.0
assert remote_spend_counter_store.get_local("spend:team_member:user-1:team-1") == 0.0
assert remote_spend_counter_in_memory_cache.get_cache("spend_db_floor:spend:team_member:user-1:team-1") == 0.0, (
"the DB-floor marker was not broadcast; a remote worker can re-raise the counter off its stale floor"
)
@ -10847,12 +10846,16 @@ async def test_invalidate_team_member_spend_state_self_delivered_broadcast_does_
return _RecordingRedisClient()
local_spend_counter_cache = DualCache()
local_spend_counter_store = SpendCounterStore(None)
local_user_api_key_cache = UserApiKeyCache()
with (
patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock
"litellm.proxy.proxy_server.spend_counter_cache", local_spend_counter_cache
),
patch( # test-quality-ok: injects a real store for the module global, not a behavior mock
"litellm.proxy.proxy_server.spend_counter_store", local_spend_counter_store
),
patch( # test-quality-ok: injects a fake pub/sub-capable redis cache; no live redis in this unit test
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache",
return_value=_FakeRedisCache(),
@ -10869,13 +10872,14 @@ async def test_invalidate_team_member_spend_state_self_delivered_broadcast_does_
redis_cache=_FakeRedisCache(),
user_api_key_cache=local_user_api_key_cache,
additional_in_memory_caches=(local_spend_counter_cache.in_memory_cache,),
spend_counters=local_spend_counter_store,
)
for _, message in published:
own_subscriber._apply_message( # pyright: ignore[reportPrivateUsage] # exercising the real cross-worker message handler, not a public API
{"type": "message", "data": message}
)
assert local_spend_counter_cache.in_memory_cache.get_cache("spend:team_member:user-1:team-1") == 0.0, (
assert local_spend_counter_store.get_local("spend:team_member:user-1:team-1") == 0.0, (
"the handler's self-delivered broadcast erased the post-reset spend counter"
)
assert (

View file

@ -34,16 +34,16 @@ from litellm.proxy.common_utils.user_api_key_cache import (
NO_TEAM_MEMBERSHIP_SENTINEL,
UserApiKeyCache,
model_access_group_registry_cache_key,
model_access_group_spend_counter_key,
team_membership_reservation_cache_key,
)
from litellm.proxy.utils import ProxyLogging
from litellm.redis.stores.spend_counters import model_access_group_counter
TEAM_ID = "team-1"
USER_ID = "user-1"
ORG_ID = "org-1"
BUDGETED_GROUPS = ("tier-a", "tier-b", "claude-tier")
MODEL_ACCESS_GROUP_COUNTER_KEY = model_access_group_spend_counter_key("tier-a")
MODEL_ACCESS_GROUP_COUNTER_KEY = model_access_group_counter("tier-a")
MODEL_LIST = [
{
@ -462,7 +462,7 @@ async def test_an_over_budget_group_blocks_even_when_another_matched_group_is_fi
_MagBudgetRow("tier-b", max_budget=1.0),
spend_by_counter_key={
MODEL_ACCESS_GROUP_COUNTER_KEY: 1.0,
model_access_group_spend_counter_key("tier-b"): 5.0,
model_access_group_counter("tier-b"): 5.0,
},
)

View file

@ -20,6 +20,9 @@ from unittest.mock import ANY, AsyncMock, MagicMock, patch
import fastapi
import orjson
import pytest
from litellm.redis.stores.spend_counters import SpendCounterStore
from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters
from fastapi import HTTPException, Request, status
from starlette.datastructures import UploadFile
@ -6981,8 +6984,8 @@ async def test_centralized_common_checks_reserves_request_end_user_budget():
)
attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None)
counter_cache = DualCache()
attrs["spend_counter_cache"] = counter_cache
counter_cache = SpendCounterStore(None)
attrs["spend_counter_store"] = counter_cache
originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs}
try:
for k, v in attrs.items():
@ -7026,7 +7029,7 @@ async def test_centralized_common_checks_reserves_request_end_user_budget():
"applied_adjustment": 0.0,
}
]
assert counter_cache.in_memory_cache.get_cache(key="spend:end_user:alice") == pytest.approx(0.6)
assert counter_cache.get_local("spend:end_user:alice") == pytest.approx(0.6)
@pytest.mark.asyncio
@ -12894,12 +12897,12 @@ async def test_admission_and_budget_reservation_read_the_key_spend_counter_with_
async def _reservation_reads_spend(**kwargs):
reads.append(("reservation", await read_batched_spend_counter("spend:key:hashed")))
redis = MagicMock()
redis.async_batch_get_cache = AsyncMock(return_value={"spend:key:hashed": 4.0})
spend = FakeSpendCounters(redis={"spend:key:hashed": 4.0})
store = spend.store()
attrs = {
**_proxy_attrs_for_centralized_checks(user_custom_auth=None),
"prisma_client": MagicMock(),
"spend_counter_cache": MagicMock(redis_cache=redis),
"spend_counter_store": store,
}
originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs}
try:
@ -12914,7 +12917,7 @@ async def test_admission_and_budget_reservation_read_the_key_spend_counter_with_
"litellm.proxy.spend_tracking.budget_reservation.reserve_budget_for_request",
side_effect=_reservation_reads_spend,
),
spend_counter_batch_scope(redis),
spend_counter_batch_scope(store),
):
await run_centralized_common_checks(
user_api_key_auth_obj=token,
@ -12932,8 +12935,8 @@ async def test_admission_and_budget_reservation_read_the_key_spend_counter_with_
("reservation", (4.0, True)),
("after admission", None),
], "admission and reservation share one snapshot, and read-then-write callers go to Redis once it closes"
assert redis.async_batch_get_cache.await_count == 1
assert "spend:key:hashed" in redis.async_batch_get_cache.await_args.kwargs["key_list"]
mgets = [command.split()[1:] for command in spend.commands if command.startswith("MGET")]
assert len(mgets) == 1 and "spend:key:hashed" in mgets[0]
def test_identity_prefetch_keys_match_what_auth_reads_for_the_request():

View file

@ -20,6 +20,7 @@ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
publish_auth_cache_invalidation,
)
from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET, UserApiKeyCache
from litellm.redis.stores.spend_counters import SpendCounterStore
class _RecordingRedisClient(Redis):
@ -271,6 +272,36 @@ async def test_subscriber_deletes_additional_in_memory_cache_entry_on_message()
assert spend_counter_in_memory_cache.get_cache("spend:team_member:u-1:t-1") is None
@pytest.mark.asyncio
async def test_subscriber_writes_and_drops_the_native_spend_counter_copies_a_broadcast_names() -> None:
"""The spend counters' memory tier lives in the native store, not in a Python InMemoryCache, so a broadcast
has to reach it there: a SET carries the reset value into it, a delete drops it, and keys outside the spend
counter keyspace leave it alone."""
store = SpendCounterStore(None)
store.set_local("spend:team_member:u-1:t-1", 999.0)
store.set_local("spend:key:k-1", 7.0)
pubsub = _QueuePubSub(
initial_messages=[
_invalidation_message("spend:team_member:u-1:t-1"),
{"type": "message", "data": json.dumps({"cache_key": "spend:key:k-1", "new_value": 0.0, "ttl": 60})},
{"type": "message", "data": json.dumps({"cache_key": "spend_db_floor:spend:key:k-1", "new_value": 3.0})},
{"type": "message", "data": json.dumps({"cache_key": "spend:unknown:x", "new_value": 3.0})},
]
)
subscriber = AuthCacheInvalidationSubscriber(
redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[pubsub])),
user_api_key_cache=UserApiKeyCache(),
spend_counters=store,
)
subscriber.start()
try:
assert await asyncio.wait_for(pubsub.drained.wait(), timeout=5)
finally:
await subscriber.stop()
assert store.snapshot() == {"spend:key:k-1": 0.0}
@pytest.mark.asyncio
async def test_subscriber_ignores_malformed_messages() -> None:
cache = UserApiKeyCache()

View file

@ -30,6 +30,7 @@ from litellm.proxy.common_utils.reset_budget_job import (
from litellm.proxy.common_utils.timezone_utils import BudgetResetSettings
from litellm.proxy.utils import PrismaClient
from tests.unit.proxy.db.fake_prisma_engine import engine_call
from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters
# Mock classes for testing
@ -1045,7 +1046,7 @@ def _make_reset_budget_windows_job(
team_rows: List[Dict[str, Any]],
):
"""Build a ResetBudgetJob with a fully-mocked prisma client and a fake
`litellm.proxy.proxy_server` module exposing a stub `spend_counter_cache`.
`litellm.proxy.proxy_server` module exposing a stub `spend_counter_store`.
Returns (job, prisma_client_mock, spend_counter_cache_mock).
"""
@ -1065,14 +1066,12 @@ def _make_reset_budget_windows_job(
prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=None)
# Stub out litellm.proxy.proxy_server so the in-function
# `from litellm.proxy.proxy_server import spend_counter_cache` resolves
# `from litellm.proxy.proxy_server import spend_counter_store` resolves
# without importing the real (heavy) module.
spend_counter_cache = MagicMock()
spend_counter_cache.in_memory_cache.set_cache = MagicMock()
spend_counter_cache.redis_cache = None # skip the async redis branch
spend_counter_cache = FakeSpendCounters(redis=None)
fake_module = types.ModuleType("litellm.proxy.proxy_server")
fake_module.spend_counter_cache = spend_counter_cache
fake_module.spend_counter_store = spend_counter_cache.store()
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_module)
job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client)
@ -1129,7 +1128,7 @@ def test_reset_budget_windows_resets_expired_key_window(monkeypatch):
assert new_reset_at > now
# The spend counter for this key+window was cleared.
spend_counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:key:sk-expired:window:1d", value=0.0)
assert spend_counter_cache.memory["spend:key:sk-expired:window:1d"] == 0.0
def _window_spend_rolls(prisma_client):
@ -1268,7 +1267,7 @@ def test_reset_budget_windows_survives_a_failed_window_spend_roll(monkeypatch):
asyncio.run(job.reset_budget_windows())
prisma_client.db.litellm_verificationtoken.update.assert_awaited_once()
spend_counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:key:sk-expired:window:1d", value=0.0)
assert spend_counter_cache.memory["spend:key:sk-expired:window:1d"] == 0.0
def test_reset_budget_windows_skips_unexpired_key_window(monkeypatch):
@ -1311,7 +1310,7 @@ def test_reset_budget_windows_resets_expired_team_window(monkeypatch):
assert call_kwargs["where"] == {"team_id": "team-expired"}
assert "budget_limits" in call_kwargs["data"]
spend_counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:team:team-expired:window:30d", value=0.0)
assert spend_counter_cache.memory["spend:team:team-expired:window:30d"] == 0.0
def test_reset_budget_windows_handles_string_budget_limits(monkeypatch):
@ -1372,11 +1371,8 @@ def test_reset_budget_windows_query_error_does_not_break_team_path(monkeypatch):
prisma_client.db.query_raw = AsyncMock(side_effect=fake_query_raw)
prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=None)
spend_counter_cache = MagicMock()
spend_counter_cache.in_memory_cache.set_cache = MagicMock()
spend_counter_cache.redis_cache = None
fake_module = types.ModuleType("litellm.proxy.proxy_server")
fake_module.spend_counter_cache = spend_counter_cache
fake_module.spend_counter_store = FakeSpendCounters(redis=None).store()
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_module)
job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client)
@ -1392,17 +1388,13 @@ def test_reset_budget_windows_query_error_does_not_break_team_path(monkeypatch):
def _make_counter_invalidation_job(monkeypatch):
"""Stub spend_counter_cache (and user_api_key_cache) so we can observe
"""Stub spend_counter_store (and user_api_key_cache) so we can observe
invalidation calls.
Both caches are looked up via ``from litellm.proxy.proxy_server import
<name>`` inside the reset job, so we publish them on a fake module.
"""
spend_counter_cache = MagicMock()
spend_counter_cache.in_memory_cache.set_cache = MagicMock()
spend_counter_cache.redis_cache = MagicMock()
spend_counter_cache.redis_cache.async_set_cache = AsyncMock()
spend_counter_cache.redis_cache.async_delete_cache = AsyncMock()
spend_counter_cache = FakeSpendCounters()
user_api_key_cache = MagicMock()
user_api_key_cache.async_delete_cache = AsyncMock()
@ -1410,24 +1402,18 @@ def _make_counter_invalidation_job(monkeypatch):
# Batch deletes fan out to the same per-key calls the real DualCache makes,
# so an assertion reads "this key was invalidated" whether the caller went
# one key at a time or a page at a time.
async def _delete_counter_keys(keys):
for key in keys:
spend_counter_cache.in_memory_cache.delete_cache(key=key)
await spend_counter_cache.redis_cache.async_delete_cache(key=key)
async def _delete_management_keys(keys):
for key in keys:
await user_api_key_cache.async_delete_cache(key=key)
spend_counter_cache.async_delete_cache_keys = AsyncMock(side_effect=_delete_counter_keys)
user_api_key_cache.async_delete_cache_keys = AsyncMock(side_effect=_delete_management_keys)
fake_module = types.ModuleType("litellm.proxy.proxy_server")
fake_module.spend_counter_cache = spend_counter_cache
fake_module.spend_counter_store = spend_counter_cache.store()
fake_module.user_api_key_cache = user_api_key_cache
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_module)
spend_counter_cache.user_api_key_cache = user_api_key_cache
spend_counter_cache.user_api_key_cache = user_api_key_cache # pyright: ignore[reportAttributeAccessIssue] # the test reads both caches off one handle
return spend_counter_cache
@ -1452,8 +1438,7 @@ def test_reset_budget_for_keys_invalidates_redis_counter(reset_budget_job, mock_
asyncio.run(reset_budget_job.reset_budget_for_litellm_keys())
counter_cache.in_memory_cache.delete_cache.assert_any_call(key="spend:key:sk-abc")
counter_cache.redis_cache.async_delete_cache.assert_any_await(key="spend:key:sk-abc")
assert "spend:key:sk-abc" in counter_cache.deleted()
def test_reset_budget_for_users_invalidates_redis_counter(reset_budget_job, mock_prisma_client, monkeypatch):
@ -1477,8 +1462,7 @@ def test_reset_budget_for_users_invalidates_redis_counter(reset_budget_job, mock
asyncio.run(reset_budget_job.reset_budget_for_litellm_users())
counter_cache.in_memory_cache.delete_cache.assert_any_call(key="spend:user:alice")
counter_cache.redis_cache.async_delete_cache.assert_any_await(key="spend:user:alice")
assert "spend:user:alice" in counter_cache.deleted()
def test_reset_budget_for_proxy_budget_row_invalidates_global_spend_cache(
@ -1562,8 +1546,7 @@ def test_reset_budget_for_teams_invalidates_redis_counter(reset_budget_job, mock
asyncio.run(reset_budget_job.reset_budget_for_litellm_teams())
counter_cache.in_memory_cache.delete_cache.assert_any_call(key="spend:team:team-x")
counter_cache.redis_cache.async_delete_cache.assert_any_await(key="spend:team:team-x")
assert "spend:team:team-x" in counter_cache.deleted()
def test_reset_does_not_zero_counter_when_db_write_fails(monkeypatch):
@ -1623,7 +1606,7 @@ def test_reset_does_not_zero_counter_when_db_write_fails(monkeypatch):
# assert_not_called() instead of iterating call_args_list, because the
# latter is vacuously true when the list is empty (would pass even if
# the bypass were re-introduced via a different code path).
counter_cache.in_memory_cache.delete_cache.assert_not_called()
assert counter_cache.deleted() == frozenset()
def test_reset_budget_for_keys_writes_only_spend_and_reset_at(reset_budget_job, mock_prisma_client):
@ -1727,8 +1710,7 @@ def test_budget_table_reset_invalidates_counters_and_management_cache(
asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table())
counter_cache.in_memory_cache.delete_cache.assert_any_call(key=counter_key)
counter_cache.redis_cache.async_delete_cache.assert_any_await(key=counter_key)
assert counter_key in counter_cache.deleted()
deleted = {call.kwargs.get("key") for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list}
assert cache_keys <= deleted
@ -1766,8 +1748,7 @@ def test_budget_table_reset_invalidates_enduser_counter_and_cache(reset_budget_j
asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table())
counter_cache.in_memory_cache.delete_cache.assert_any_call(key="spend:end_user:customer-42")
counter_cache.redis_cache.async_delete_cache.assert_any_await(key="spend:end_user:customer-42")
assert "spend:end_user:customer-42" in counter_cache.deleted()
deleted: Final = {call.kwargs.get("key") for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list}
assert "end_user_id:customer-42" in deleted
@ -1805,13 +1786,10 @@ def test_enduser_invalidation_is_paged_and_batched(reset_budget_job, mock_prisma
f"cust-{RESET_BUDGET_JOB_BATCH_SIZE * 2 - 1:06d}",
]
assert counter_cache.async_delete_cache_keys.await_count == 6
assert len(counter_cache.deletes) == 6
assert counter_cache.user_api_key_cache.async_delete_cache_keys.await_count == 6
counter_cache.async_delete_cache.assert_not_called()
invalidated: Final = {
key for call in counter_cache.async_delete_cache_keys.await_args_list for key in call.args[0]
}
invalidated: Final = counter_cache.deleted()
assert invalidated == {f"spend:end_user:cust-{i:06d}" for i in range(population)}
evicted: Final = {
key for call in counter_cache.user_api_key_cache.async_delete_cache_keys.await_args_list for key in call.args[0]
@ -1861,7 +1839,7 @@ def test_a_failed_counter_batch_still_evicts_the_management_cache(
until they expire.
"""
counter_cache: Final = _make_counter_invalidation_job(monkeypatch)
counter_cache.async_delete_cache_keys = AsyncMock(side_effect=RuntimeError("redis unavailable"))
counter_cache.fail = True
mock_prisma_client.data["budget"] = [_budget_row(budget_id="budget-1")]
mock_prisma_client.db.litellm_endusertable.set_find_many_results(
[type("EndUser", (), {"user_id": "customer-42", "spend": 5.0, "budget_id": "budget-1"})]
@ -1932,7 +1910,7 @@ def test_access_groups_are_untouched_when_no_budget_is_due(reset_budget_job, moc
assert mock_prisma_client.db.litellm_modelaccessgroupbudgettable.find_many_calls == []
assert _batch_writes(mock_prisma_client, "model_access_group") == []
counter_cache.in_memory_cache.delete_cache.assert_not_called()
assert counter_cache.deleted() == frozenset()
counter_cache.user_api_key_cache.async_delete_cache.assert_not_awaited()
@ -1951,7 +1929,7 @@ def test_budget_table_reset_invalidates_every_access_group_not_just_the_first(
deleted = {call.kwargs.get("key") for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list}
assert deleted == {"model_access_group:group-a", "model_access_group:group-b", "model_access_group:group-c"}
for name in ("group-a", "group-b", "group-c"):
counter_cache.in_memory_cache.delete_cache.assert_any_call(key=f"spend:model_access_group:{name}")
assert f"spend:model_access_group:{name}" in counter_cache.deleted()
def test_project_reset_zeroes_spend_on_due_tiers(reset_budget_job, mock_prisma_client, monkeypatch):
@ -2001,7 +1979,7 @@ def test_budget_cascade_carries_access_group_overage_when_rollover_enabled(
} in writes
assert _replay_spend_writes(writes, 15.0) == 5.0
assert _replay_spend_writes(writes, 8.0) == 0
counter_cache.in_memory_cache.delete_cache.assert_any_call(key="spend:model_access_group:gpt-4-group")
assert "spend:model_access_group:gpt-4-group" in counter_cache.deleted()
# ---------------------------------------------------------------------------
@ -2092,7 +2070,7 @@ def test_budget_reset_at_is_not_advanced_when_the_cascade_fails(db_factory, monk
assert prisma_client.db.batch_calls == [], "a failed cascade must not persist any write"
assert prisma_client.db.batchers[0].committed is False
assert prisma_client.updated_data["budget"] == [], "budget_reset_at must not be advanced outside the transaction"
counter_cache.in_memory_cache.delete_cache.assert_not_called()
assert counter_cache.deleted() == frozenset()
counter_cache.user_api_key_cache.async_delete_cache.assert_not_awaited()
@ -2133,7 +2111,7 @@ def test_caches_are_invalidated_only_after_the_transaction_commits(monkeypatch):
cap while the DB still holds the over-budget spend."""
events = []
counter_cache = _make_counter_invalidation_job(monkeypatch)
counter_cache.in_memory_cache.delete_cache.side_effect = lambda **kwargs: events.append("counter")
counter_cache.on_delete = lambda _: events.append("counter")
job, _ = _job_with_expired_budget(OrderRecordingDB(events))
@ -2623,10 +2601,8 @@ def _make_leader_election_job(monkeypatch, pod_lock_manager):
prisma_client = MockPrismaClient()
prisma_client.db.query_raw = AsyncMock(return_value=[])
spend_counter_cache = MagicMock()
spend_counter_cache.redis_cache = None
fake_module = types.ModuleType("litellm.proxy.proxy_server")
fake_module.spend_counter_cache = spend_counter_cache
fake_module.spend_counter_store = FakeSpendCounters(redis=None).store()
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_module)
job = ResetBudgetJob(
@ -2756,10 +2732,8 @@ def _paginating_window_job(monkeypatch, pages_by_table: Dict[str, List[List[Dict
prisma_client.db.litellm_verificationtoken.update = AsyncMock(return_value=None)
prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=None)
spend_counter_cache = MagicMock()
spend_counter_cache.redis_cache = None
fake_module = types.ModuleType("litellm.proxy.proxy_server")
fake_module.spend_counter_cache = spend_counter_cache
fake_module.spend_counter_store = FakeSpendCounters(redis=None).store()
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_module)
job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client)
@ -2826,10 +2800,8 @@ def test_reset_budget_windows_survives_one_table_failing(monkeypatch):
return []
prisma_client.db.query_raw = AsyncMock(side_effect=fake_query_raw)
spend_counter_cache = MagicMock()
spend_counter_cache.redis_cache = None
fake_module = types.ModuleType("litellm.proxy.proxy_server")
fake_module.spend_counter_cache = spend_counter_cache
fake_module.spend_counter_store = FakeSpendCounters(redis=None).store()
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_module)
job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client)
@ -2897,10 +2869,8 @@ def _cursor_paginating_window_job(monkeypatch, key_rows: List[Dict[str, Any]]):
prisma_client.db.litellm_verificationtoken.update = AsyncMock(return_value=None)
prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=None)
spend_counter_cache = MagicMock()
spend_counter_cache.redis_cache = None
fake_module = types.ModuleType("litellm.proxy.proxy_server")
fake_module.spend_counter_cache = spend_counter_cache
fake_module.spend_counter_store = FakeSpendCounters(redis=None).store()
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_module)
job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client)
@ -3322,7 +3292,7 @@ def test_direct_reset_carries_overage_when_rollover_enabled(
assert writes[0]["data"]["spend"] == {"decrement": 100.0}
assert writes[0]["data"]["budget_reset_at"] > now
counter_prefix = {"key": "spend:key", "user": "spend:user", "team": "spend:team"}[table]
counter_cache.in_memory_cache.delete_cache.assert_any_call(key=f"{counter_prefix}:{id_value}")
assert f"{counter_prefix}:{id_value}" in counter_cache.deleted()
def test_direct_reset_zeroes_under_budget_row_even_with_rollover(
@ -3341,7 +3311,7 @@ def test_direct_reset_zeroes_under_budget_row_even_with_rollover(
asyncio.run(reset_budget_job.reset_budget_for_litellm_keys())
assert _batch_writes(mock_prisma_client, "key")[0]["data"]["spend"] == {"decrement": 40.0}
counter_cache.in_memory_cache.delete_cache.assert_any_call(key="spend:key:tok-under")
assert "spend:key:tok-under" in counter_cache.deleted()
def test_direct_reset_zeroes_row_without_max_budget_even_with_rollover(
@ -3394,7 +3364,7 @@ def test_budget_cascade_carries_overage_per_tier_when_rollover_enabled(
"where": {"budget_id": "budget-roll", "spend": {"gt": 0, "lte": 10.0}},
"data": {"spend": 0},
} in membership_writes
counter_cache.in_memory_cache.delete_cache.assert_any_call(key="spend:team_member:member-1:team-1")
assert "spend:team_member:member-1:team-1" in counter_cache.deleted()
def test_budget_cascade_floors_membership_total_spend_before_zeroing_spend(
@ -3479,8 +3449,7 @@ def test_budget_cascade_carries_default_tier_enduser_counter_when_rollover_enabl
asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table())
counter_cache.in_memory_cache.delete_cache.assert_any_call(key="spend:end_user:enduser-implicit")
counter_cache.redis_cache.async_delete_cache.assert_any_await(key="spend:end_user:enduser-implicit")
assert "spend:end_user:enduser-implicit" in counter_cache.deleted()
deleted: Final = {call.kwargs.get("key") for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list}
assert "end_user_id:enduser-implicit" in deleted
@ -3561,12 +3530,12 @@ def test_window_reset_carries_counter_overage_when_rollover_enabled(rollover_ena
job, prisma_client, spend_counter_cache = _make_reset_budget_windows_job(
monkeypatch, key_rows=key_rows, team_rows=[]
)
spend_counter_cache.async_get_cache = AsyncMock(return_value=130.0)
spend_counter_cache.memory["spend:key:sk-roll:window:1d"] = 130.0
asyncio.run(job.reset_budget_windows())
prisma_client.db.litellm_verificationtoken.update.assert_awaited_once()
spend_counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:key:sk-roll:window:1d", value=30.0)
assert spend_counter_cache.memory["spend:key:sk-roll:window:1d"] == 30.0
def test_window_reset_zeroes_counter_when_rollover_disabled(monkeypatch):
@ -3581,12 +3550,11 @@ def test_window_reset_zeroes_counter_when_rollover_disabled(monkeypatch):
job, prisma_client, spend_counter_cache = _make_reset_budget_windows_job(
monkeypatch, key_rows=key_rows, team_rows=[]
)
spend_counter_cache.async_get_cache = AsyncMock(return_value=130.0)
spend_counter_cache.memory["spend:key:sk-off:window:1d"] = 130.0
asyncio.run(job.reset_budget_windows())
spend_counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:key:sk-off:window:1d", value=0.0)
spend_counter_cache.async_get_cache.assert_not_awaited()
assert spend_counter_cache.memory["spend:key:sk-off:window:1d"] == 0.0
def _apply_spend_payload(db_spend: float, spend_field: dict[str, float]) -> float:
@ -3718,10 +3686,9 @@ def test_reset_deletes_spend_counter_instead_of_seeding(reset_budget_job, mock_p
asyncio.run(reset_budget_job.reset_budget_for_litellm_users())
counter_cache.in_memory_cache.delete_cache.assert_any_call(key="spend:user:carol")
counter_cache.redis_cache.async_delete_cache.assert_any_await(key="spend:user:carol")
counter_cache.in_memory_cache.set_cache.assert_not_called()
counter_cache.redis_cache.async_set_cache.assert_not_awaited()
assert "spend:user:carol" in counter_cache.deleted()
assert counter_cache.memory == {}
assert [command for command in counter_cache.commands if command.startswith("SET ")] == []
@pytest.mark.asyncio

View file

@ -21,6 +21,7 @@ from litellm.caching.in_memory_cache import InMemoryCache
from litellm.constants import PROXY_DB_LOOKUP_MAX_CONCURRENCY
from litellm.proxy.db.db_lookup_gate import LoopBoundSemaphore, db_lookup_stall_tracker
from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed
from litellm.redis.stores.spend_counters import SpendCounterStore
WINDOW_START = datetime(2026, 8, 1, tzinfo=timezone.utc)
@ -97,13 +98,13 @@ class _PausedSpendTable:
async def _reseed_with_paused_table(
table: _PausedSpendTable, cache: DualCache, counter_key: str, window: bool
table: _PausedSpendTable, cache: SpendCounterStore, counter_key: str, window: bool
) -> float | None:
prisma: Final = SimpleNamespace(db=SimpleNamespace(litellm_usertable=table, litellm_budgetwindowspend=table))
if window:
return await SpendCounterReseed.coalesced_window(
prisma_client=prisma,
spend_counter_cache=cache,
spend_counter_store=cache,
counter_key=counter_key,
entity_type="Team",
entity_id="team-1",
@ -112,7 +113,7 @@ async def _reseed_with_paused_table(
)
return await SpendCounterReseed.coalesced(
prisma_client=prisma,
spend_counter_cache=cache,
spend_counter_store=cache,
counter_key=counter_key,
)
@ -291,12 +292,12 @@ async def test_window_from_db_without_a_duration_skips_the_row_lookup():
@pytest.mark.asyncio
async def test_coalesced_window_seeds_a_cold_counter_from_the_row():
prisma = _FakePrismaClient(row=_row(WINDOW_START, 4.5), spend_logs_total=100.0)
cache = DualCache()
cache = SpendCounterStore(None)
counter_key = "spend:key:tok-1:window:30d"
result = await SpendCounterReseed.coalesced_window(
prisma_client=prisma,
spend_counter_cache=cache,
spend_counter_store=cache,
counter_key=counter_key,
entity_type="Key",
entity_id="tok-1",
@ -305,7 +306,7 @@ async def test_coalesced_window_seeds_a_cold_counter_from_the_row():
)
assert result == 4.5
assert cache.in_memory_cache.get_cache(key=counter_key) == 4.5
assert cache.get_local(counter_key) == 4.5
assert prisma.db.litellm_spendlogs.call_count == 0
@ -316,19 +317,19 @@ async def test_cold_reseed_does_not_add_database_spend_to_concurrent_cache(
window: bool,
concurrent_spend: float,
) -> None:
cache: Final = DualCache(in_memory_cache=InMemoryCache())
cache: Final = SpendCounterStore(None)
counter_key: Final = "spend:team:team-1:window:1d" if window else "spend:user:user-1"
db_spend: Final = 989.01459411
table: Final = _PausedSpendTable(db_spend)
reseed_task: Final = asyncio.create_task(_reseed_with_paused_table(table, cache, counter_key, window))
await asyncio.wait_for(table.read_started.wait(), timeout=5)
cache.in_memory_cache.set_cache(key=counter_key, value=concurrent_spend)
cache.set_local(counter_key, concurrent_spend)
table.resume_read.set()
result: Final = await asyncio.wait_for(reseed_task, timeout=5)
expected: Final = max(db_spend, concurrent_spend)
assert cache.in_memory_cache.get_cache(key=counter_key) == expected
assert cache.get_local(counter_key) == expected
assert result == expected
@ -341,14 +342,14 @@ async def test_cold_reseed_preserves_concurrent_local_increment(
) -> None:
from litellm.proxy import proxy_server
cache: Final = DualCache(in_memory_cache=InMemoryCache())
cache: Final = SpendCounterStore(None)
counter_key: Final = (
f"spend:team:concurrent-{batch}-{increment}:window:1d"
if window
else f"spend:user:concurrent-{batch}-{increment}"
)
table: Final = _PausedSpendTable(100.0)
monkeypatch.setattr(proxy_server, "spend_counter_cache", cache)
monkeypatch.setattr(proxy_server, "spend_counter_store", cache)
reseed_task: Final = asyncio.create_task(_reseed_with_paused_table(table, cache, counter_key, window))
await asyncio.wait_for(table.read_started.wait(), timeout=5)
@ -363,7 +364,7 @@ async def test_cold_reseed_preserves_concurrent_local_increment(
table.resume_read.set()
await asyncio.wait_for(asyncio.gather(reseed_task, increment_task), timeout=5)
assert cache.in_memory_cache.get_cache(key=counter_key) == 100.0 + increment
assert cache.get_local(counter_key) == 100.0 + increment
@pytest.mark.asyncio

View file

@ -10,6 +10,8 @@ from unittest.mock import AsyncMock, MagicMock, create_autospec, patch
import httpx
import pytest
from litellm.redis.stores.spend_counters import SpendCounterStore
from fastapi import HTTPException
import litellm
@ -2679,9 +2681,9 @@ async def test_async_post_call_failure_hook_records_recovered_partial_spend():
@pytest.fixture()
def live_spend_counters(monkeypatch: pytest.MonkeyPatch) -> DualCache:
counters: Final = DualCache()
monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", counters)
def live_spend_counters(monkeypatch: pytest.MonkeyPatch) -> SpendCounterStore:
counters: Final = SpendCounterStore(None)
monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_store", counters)
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", DualCache())
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
@ -2734,8 +2736,8 @@ async def _fail_request(logger: ProxyDBLogger, failure: Mapping[str, object]) ->
)
async def _counter_values(counters: DualCache) -> dict[str, float]:
return {key: await counters.async_get_cache(key) for key in _LIVE_COUNTERS}
async def _counter_values(counters: SpendCounterStore) -> dict[str, float | None]:
return {key: counters.get_local(key) for key in _LIVE_COUNTERS}
@pytest.mark.asyncio
@ -2754,12 +2756,12 @@ async def _counter_values(counters: DualCache) -> dict[str, float]:
],
)
async def test_failure_spend_rows_raise_the_live_spend_counters_by_their_cost(
live_spend_counters: DualCache, failure: Mapping[str, object], charged: float
live_spend_counters: SpendCounterStore, failure: Mapping[str, object], charged: float
):
"""#35568: budget checks read the live counters, so a failure row's cost has to land there as well as in the
spend log, or an over-budget key keeps passing until the counter is next repaired from the database."""
for key, spend in _LIVE_COUNTERS.items():
await live_spend_counters.async_set_cache(key, spend)
live_spend_counters.set_local(key, spend)
await _fail_request(ProxyDBLogger(), failure)
@ -2769,11 +2771,11 @@ async def test_failure_spend_rows_raise_the_live_spend_counters_by_their_cost(
@pytest.mark.asyncio
async def test_a_completed_cost_failure_raises_the_live_spend_counters_once(live_spend_counters: DualCache):
async def test_a_completed_cost_failure_raises_the_live_spend_counters_once(live_spend_counters: SpendCounterStore):
"""A failure whose cost was already finalized settles through the success path's write, which raises the
counters itself."""
for key, spend in _LIVE_COUNTERS.items():
await live_spend_counters.async_set_cache(key, spend)
live_spend_counters.set_local(key, spend)
await _fail_request(
ProxyDBLogger(), {"standard_logging_object": {"cost_finalized": True, "response_cost": 0.125}}
@ -2786,11 +2788,11 @@ async def test_a_completed_cost_failure_raises_the_live_spend_counters_once(live
@pytest.mark.asyncio
async def test_failure_spend_rows_leave_the_live_spend_counters_alone_while_spend_updates_are_disabled(
live_spend_counters: DualCache, monkeypatch: pytest.MonkeyPatch
live_spend_counters: SpendCounterStore, monkeypatch: pytest.MonkeyPatch
):
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"disable_spend_updates": True})
for key, spend in _LIVE_COUNTERS.items():
await live_spend_counters.async_set_cache(key, spend)
live_spend_counters.set_local(key, spend)
await _fail_request(ProxyDBLogger(), _failure_with_billed_attempts())
@ -2833,7 +2835,7 @@ async def test_a_slow_spend_counter_store_lets_the_failure_out_and_raises_the_co
@pytest.mark.asyncio
async def test_a_failure_spend_row_lands_in_the_key_budget_window_as_of_the_request_start(
live_spend_counters: DualCache, monkeypatch: pytest.MonkeyPatch
live_spend_counters: SpendCounterStore, monkeypatch: pytest.MonkeyPatch
):
"""The window row's one-time seed from the spend log stops at the request start, which keeps the failure row
it already logged from being counted twice."""
@ -2870,7 +2872,7 @@ async def test_a_failure_spend_row_lands_in_the_key_budget_window_as_of_the_requ
@pytest.mark.asyncio
async def test_a_billed_attempt_on_a_deployment_the_router_dropped_debits_every_attributed_access_group(
live_spend_counters: DualCache, monkeypatch: pytest.MonkeyPatch
live_spend_counters: SpendCounterStore, monkeypatch: pytest.MonkeyPatch
):
"""Its spend row names the dropped deployment, which narrows nothing, so the counters must not narrow it to the
final attempt's groups either."""
@ -2896,7 +2898,7 @@ async def test_a_billed_attempt_on_a_deployment_the_router_dropped_debits_every_
)
assert {
group: await live_spend_counters.async_get_cache(f"spend:model_access_group:{group}")
group: live_spend_counters.get_local(f"spend:model_access_group:{group}")
for group in ("premium", "pool")
} == pytest.approx({"premium": 0.25, "pool": 1.25})

View file

@ -108,6 +108,7 @@ from litellm.types.proxy.management_endpoints.key_management_endpoints import (
BulkUpdateKeyResponse,
CustomKeyPolicyRequest,
)
from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters
client = TestClient(app)
@ -384,12 +385,9 @@ async def test_budget_reset_and_expires_at_first_of_month(monkeypatch):
from datetime import datetime, timedelta, timezone
import pytest
from litellm.proxy.management_endpoints.key_management_endpoints import (
generate_key_helper_fn,
)
from litellm.proxy.proxy_server import prisma_client
# Use monkeypatch to set the prisma_client
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
@ -466,9 +464,12 @@ async def test_key_expiration_exact_duration_hours(monkeypatch):
from datetime import datetime, timedelta, timezone
import pytest
from litellm.proxy.management_endpoints.key_management_endpoints import (
generate_key_helper_fn,
)
from litellm.proxy.proxy_server import prisma_client
# Use monkeypatch to set the prisma_client
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
@ -8643,6 +8644,8 @@ async def test_reset_key_spend_success(monkeypatch):
"litellm.proxy.proxy_server.spend_counter_cache",
mock_spend_counter_cache,
)
spend = FakeSpendCounters(redis=None)
monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_store", spend.store())
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
@ -8663,9 +8666,7 @@ async def test_reset_key_spend_success(monkeypatch):
assert response["max_budget"] == 200.0
mock_prisma_client.db.litellm_verificationtoken.update.assert_called_once()
mock_delete_cache.assert_awaited_once()
mock_spend_counter_cache.in_memory_cache.set_cache.assert_any_call(
key=f"spend:key:{hashed_key}", value=50.0, ttl=60
)
assert spend.memory[f"spend:key:{hashed_key}"] == 50.0
# spend_db_floor marker is also set to the reset value (LIT-3803 pattern),
# so a request landing on a pod with a warm pre-reset floor marker cannot
# re-derive and re-apply the stale spend.
@ -8742,6 +8743,10 @@ async def test_reset_key_spend_resets_budget_windows(monkeypatch):
mock_spend_counter_cache,
)
spend = FakeSpendCounters()
monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_store", spend.store())
with (
patch("litellm.proxy.proxy_server.hash_token") as mock_hash_token, # test-quality-ok: no HTTP boundary; same pattern as test_reset_key_spend_success
patch( # test-quality-ok: no HTTP boundary; same pattern as test_reset_key_spend_success
@ -8773,12 +8778,8 @@ async def test_reset_key_spend_resets_budget_windows(monkeypatch):
assert response["spend"] == 0.0
window_counter_key = f"spend:key:{hashed_key}:window:1d"
mock_spend_counter_cache.in_memory_cache.set_cache.assert_any_call(
key=window_counter_key, value=0.0, ttl=60
)
mock_spend_counter_cache.redis_cache.async_set_cache.assert_any_call(
key=window_counter_key, value=0.0, ttl=60
)
assert spend.memory[window_counter_key] == 0.0
assert f"SET {window_counter_key} 0.0 60" in spend.commands
mock_spend_counter_cache.in_memory_cache.set_cache.assert_any_call(
key=f"spend_db_floor:{window_counter_key}", value=0.0, ttl=5
)
@ -8851,6 +8852,10 @@ async def test_reset_key_spend_no_budget_limits_skips_window_reset(monkeypatch):
mock_spend_counter_cache,
)
spend = FakeSpendCounters(redis=None)
monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_store", spend.store())
with (
patch("litellm.proxy.proxy_server.hash_token") as mock_hash_token, # test-quality-ok: no HTTP boundary; same pattern as test_reset_key_spend_success
patch( # test-quality-ok: no HTTP boundary; same pattern as test_reset_key_spend_success
@ -8959,6 +8964,8 @@ async def test_update_key_spend_updates_counter(monkeypatch):
"litellm.proxy.proxy_server.spend_counter_cache",
mock_spend_counter_cache,
)
spend = FakeSpendCounters()
monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_store", spend.store())
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints.delete_cache_key_object"
@ -8982,12 +8989,8 @@ async def test_update_key_spend_updates_counter(monkeypatch):
)
mock_delete_cache.assert_awaited_once()
mock_spend_counter_cache.in_memory_cache.set_cache.assert_called_once_with(
key=f"spend:key:{hashed_key}", value=0.0, ttl=60
)
mock_spend_counter_cache.redis_cache.async_set_cache.assert_awaited_once_with(
key=f"spend:key:{hashed_key}", value=0.0, ttl=60
)
assert spend.memory == {f"spend:key:{hashed_key}": 0.0}
assert spend.commands == [f"SET spend:key:{hashed_key} 0.0 60"]
@pytest.mark.asyncio

View file

@ -104,6 +104,7 @@ from tests.unit.proxy.management_endpoints.jwt_key_mapping_doubles import (
CascadingJWTMappingTable,
JWTMappingRow,
)
from litellm.redis.stores.spend_counters import SpendCounterStore
# Setup TestClient
client = TestClient(app)
@ -16630,7 +16631,8 @@ async def test_reset_team_member_spend_fn_success(monkeypatch):
await real_cache.async_set_cache(key="team-1_member-1", value="stale-membership")
await real_cache.async_set_cache(key="team_membership:member-1:team-1", value="stale-membership")
real_spend_counter_cache = DualCache()
real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:member-1:team-1", value=999.0)
real_spend_counter_store = SpendCounterStore(None)
real_spend_counter_store.set_local("spend:team_member:member-1:team-1", 999.0)
membership_row = LiteLLM_TeamMembership(
user_id="member-1",
@ -16645,6 +16647,7 @@ async def test_reset_team_member_spend_fn_success(monkeypatch):
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", real_cache)
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj)
monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache)
monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_store", real_spend_counter_store)
with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
"litellm.proxy.management_endpoints.team_endpoints.get_team_object",
@ -16668,7 +16671,7 @@ async def test_reset_team_member_spend_fn_success(monkeypatch):
assert "spend >= $3" in statement
assert await real_cache.async_get_cache(key="team-1_member-1") is None
assert await real_cache.async_get_cache(key="team_membership:member-1:team-1") is None
assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:member-1:team-1") == 0.0
assert real_spend_counter_store.get_local("spend:team_member:member-1:team-1") == 0.0
@pytest.mark.asyncio
@ -16683,7 +16686,8 @@ async def test_reset_team_member_spend_fn_stale_spend_returns_409_and_keeps_cach
real_cache = UserApiKeyCache()
await real_cache.async_set_cache(key="team_membership:member-1:team-1", value="cached-membership")
real_spend_counter_cache = DualCache()
real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:member-1:team-1", value=999.0)
real_spend_counter_store = SpendCounterStore(None)
real_spend_counter_store.set_local("spend:team_member:member-1:team-1", 999.0)
membership_row = LiteLLM_TeamMembership(user_id="member-1", team_id="team-1", spend=10.0)
mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row)
@ -16693,6 +16697,7 @@ async def test_reset_team_member_spend_fn_stale_spend_returns_409_and_keeps_cach
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", real_cache)
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache)
monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_store", real_spend_counter_store)
with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
"litellm.proxy.management_endpoints.team_endpoints.get_team_object",
@ -16709,7 +16714,7 @@ async def test_reset_team_member_spend_fn_stale_spend_returns_409_and_keeps_cach
assert exc_info.value.status_code == 409
assert await real_cache.async_get_cache(key="team_membership:member-1:team-1") == "cached-membership"
assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:member-1:team-1") == 999.0
assert real_spend_counter_store.get_local("spend:team_member:member-1:team-1") == 999.0
@pytest.mark.asyncio
@ -17097,7 +17102,8 @@ async def test_team_member_update_invalidates_team_member_spend_state_when_budge
await real_cache.async_set_cache(key="team-1_member-1", value="stale-membership")
await real_cache.async_set_cache(key="team_membership:member-1:team-1", value="stale-membership")
real_spend_counter_cache = DualCache()
real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:member-1:team-1", value=999.0)
real_spend_counter_store = SpendCounterStore(None)
real_spend_counter_store.set_local("spend:team_member:member-1:team-1", 999.0)
team_row = LiteLLM_TeamTable(team_id="team-1", metadata={}, members_with_roles=[])
team_info_response = {
@ -17111,6 +17117,7 @@ async def test_team_member_update_invalidates_team_member_spend_state_when_budge
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", real_cache)
monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache)
monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_store", real_spend_counter_store)
mock_tx = AsyncMock()
mock_prisma_client.tx.return_value.__aenter__ = AsyncMock(return_value=mock_tx)
@ -17136,7 +17143,7 @@ async def test_team_member_update_invalidates_team_member_spend_state_when_budge
assert await real_cache.async_get_cache(key="team-1_member-1") is None
assert await real_cache.async_get_cache(key="team_membership:member-1:team-1") is None
assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:member-1:team-1") == 999.0
assert real_spend_counter_store.get_local("spend:team_member:member-1:team-1") == 999.0
@pytest.mark.asyncio
@ -17150,7 +17157,8 @@ async def test_team_member_update_skips_invalidation_when_no_budget_fields_sent(
real_cache = UserApiKeyCache()
await real_cache.async_set_cache(key="team-1_member-1", value="still-fresh-membership")
real_spend_counter_cache = DualCache()
real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:member-1:team-1", value=1.5)
real_spend_counter_store = SpendCounterStore(None)
real_spend_counter_store.set_local("spend:team_member:member-1:team-1", 1.5)
team_row = LiteLLM_TeamTable(team_id="team-1", metadata={}, members_with_roles=[])
team_info_response = {
@ -17164,6 +17172,7 @@ async def test_team_member_update_skips_invalidation_when_no_budget_fields_sent(
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", real_cache)
monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache)
monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_store", real_spend_counter_store)
mock_tx = AsyncMock()
mock_prisma_client.tx.return_value.__aenter__ = AsyncMock(return_value=mock_tx)
@ -17188,7 +17197,7 @@ async def test_team_member_update_skips_invalidation_when_no_budget_fields_sent(
)
assert await real_cache.async_get_cache(key="team-1_member-1") == "still-fresh-membership"
assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:member-1:team-1") == 1.5
assert real_spend_counter_store.get_local("spend:team_member:member-1:team-1") == 1.5
@pytest.mark.asyncio

View file

@ -36,6 +36,7 @@ from litellm.proxy.spend_tracking.budget_reservation import (
)
from litellm.proxy.utils import ProxyLogging
from litellm.utils import Rules, function_setup
from litellm.redis.stores.spend_counters import SpendCounterStore
KEY_TOKEN: Final = "hashed-release-middleware-key"
COUNTER_KEY: Final = f"spend:key:{KEY_TOKEN}"
@ -43,11 +44,11 @@ CHAT_BODY: Final = {"model": "gpt-4o", "messages": [{"role": "user", "content":
@pytest.fixture
def spend_counter_cache(monkeypatch: pytest.MonkeyPatch) -> DualCache:
cache: Final = DualCache()
monkeypatch.setattr(proxy_server, "spend_counter_cache", cache)
def spend_counter_store(monkeypatch: pytest.MonkeyPatch) -> SpendCounterStore:
store: Final = SpendCounterStore(None)
monkeypatch.setattr(proxy_server, "spend_counter_store", store)
monkeypatch.setattr(proxy_server, "prisma_client", None)
return cache
return store
@pytest.fixture
@ -135,14 +136,14 @@ async def _post(app: ASGIApp) -> None:
await app(scope, receive, send)
def _counter(spend_counter_cache: DualCache) -> float | None:
return spend_counter_cache.in_memory_cache.get_cache(key=COUNTER_KEY)
def _counter(spend_counter_store: SpendCounterStore) -> float | None:
return spend_counter_store.get_local(COUNTER_KEY)
@pytest.mark.asyncio
async def test_unbound_reservation_is_released_after_the_response(spend_counter_cache: DualCache):
async def test_unbound_reservation_is_released_after_the_response(spend_counter_store: SpendCounterStore):
reservation: Final = await _reserve()
assert _counter(spend_counter_cache) == pytest.approx(reservation["reserved_cost"])
assert _counter(spend_counter_store) == pytest.approx(reservation["reserved_cost"])
async def handler(request: Request) -> Response:
request.state.budget_reservation = reservation
@ -150,12 +151,12 @@ async def test_unbound_reservation_is_released_after_the_response(spend_counter_
await _post(_app(handler))
assert _counter(spend_counter_cache) == pytest.approx(0.0)
assert _counter(spend_counter_store) == pytest.approx(0.0)
assert reservation["finalized"] is True
@pytest.mark.asyncio
async def test_unbound_reservation_is_released_when_the_handler_raises(spend_counter_cache: DualCache):
async def test_unbound_reservation_is_released_when_the_handler_raises(spend_counter_store: SpendCounterStore):
reservation: Final = await _reserve()
async def handler(request: Request) -> Response:
@ -165,13 +166,13 @@ async def test_unbound_reservation_is_released_when_the_handler_raises(spend_cou
with pytest.raises(RuntimeError):
await _post(_app(handler))
assert _counter(spend_counter_cache) == pytest.approx(0.0)
assert _counter(spend_counter_store) == pytest.approx(0.0)
assert reservation["finalized"] is True
@pytest.mark.asyncio
async def test_reservation_seen_only_by_the_proxy_pre_call_logging_object_is_released(
spend_counter_cache: DualCache, no_callbacks: None
spend_counter_store: SpendCounterStore, no_callbacks: None
):
reservation: Final = await _reserve()
@ -185,13 +186,13 @@ async def test_reservation_seen_only_by_the_proxy_pre_call_logging_object_is_rel
await _post(_app(handler))
assert _counter(spend_counter_cache) == pytest.approx(0.0)
assert _counter(spend_counter_store) == pytest.approx(0.0)
assert reservation["finalized"] is True
@pytest.mark.asyncio
async def test_reservation_of_a_failed_call_is_released_after_the_error_response(
spend_counter_cache: DualCache, no_callbacks: None
spend_counter_store: SpendCounterStore, no_callbacks: None
):
reservation: Final = await _reserve()
refused: Final = litellm.AuthenticationError(message="bad key", llm_provider="openai", model="gpt-4o")
@ -207,13 +208,13 @@ async def test_reservation_of_a_failed_call_is_released_after_the_error_response
await _post(_app(handler))
assert _counter(spend_counter_cache) == pytest.approx(0.0)
assert _counter(spend_counter_store) == pytest.approx(0.0)
assert reservation["finalized"] is True
@pytest.mark.asyncio
async def test_reservation_claimed_by_a_completed_call_is_left_for_the_callback(
spend_counter_cache: DualCache, no_callbacks: None
spend_counter_store: SpendCounterStore, no_callbacks: None
):
reservation: Final = await _reserve()
reserved_cost: Final = reservation["reserved_cost"]
@ -226,13 +227,13 @@ async def test_reservation_claimed_by_a_completed_call_is_left_for_the_callback(
await _post(_app(handler))
assert _counter(spend_counter_cache) == pytest.approx(reserved_cost)
assert _counter(spend_counter_store) == pytest.approx(reserved_cost)
assert reservation["finalized"] is False
@pytest.mark.asyncio
async def test_reservation_claimed_by_a_streaming_call_is_left_for_the_callback_that_finishes_after_the_response(
spend_counter_cache: DualCache, no_callbacks: None
spend_counter_store: SpendCounterStore, no_callbacks: None
):
reservation: Final = await _reserve()
reserved_cost: Final = reservation["reserved_cost"]
@ -251,19 +252,19 @@ async def test_reservation_claimed_by_a_streaming_call_is_left_for_the_callback_
await _post(_app(handler))
assert _counter(spend_counter_cache) == pytest.approx(reserved_cost)
assert _counter(spend_counter_store) == pytest.approx(reserved_cost)
assert reservation["finalized"] is False
actual_cost: Final = reserved_cost / 4
await reconcile_budget_reservation(budget_reservation=reservation, actual_cost=actual_cost)
assert _counter(spend_counter_cache) == pytest.approx(actual_cost)
assert _counter(spend_counter_store) == pytest.approx(actual_cost)
assert reservation["finalized"] is True
@pytest.mark.asyncio
async def test_unbound_reservation_of_a_websocket_session_is_released_when_the_socket_closes(
spend_counter_cache: DualCache,
spend_counter_store: SpendCounterStore,
):
reservation: Final = await _reserve()
@ -283,7 +284,7 @@ async def test_unbound_reservation_of_a_websocket_session_is_released_when_the_s
)
await middleware({"type": "websocket", "path": "/deepgram/v1/listen", "headers": []}, receive, send)
assert _counter(spend_counter_cache) == pytest.approx(0.0)
assert _counter(spend_counter_store) == pytest.approx(0.0)
assert reservation["finalized"] is True

View file

@ -6,6 +6,8 @@ from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
from litellm.redis.stores.spend_counters import SpendCounterStore
from pydantic import TypeAdapter
import litellm
@ -754,10 +756,10 @@ async def test_completed_passthrough_cancellation_preserves_accounting(
)
monkeypatch.setattr(litellm, "apply_search_free_tier", True)
monkeypatch.setattr(proxy_server, "general_settings", {})
counters: Final = DualCache()
counters: Final = SpendCounterStore(None)
for key in ("spend:key:pass-through-key", "spend:team:pass-through-team"):
counters.set_cache(key=key, value=1.9)
monkeypatch.setattr(proxy_server, "spend_counter_cache", counters)
counters.set_local(key, 1.9)
monkeypatch.setattr(proxy_server, "spend_counter_store", counters)
reservation: Final = {
"reserved_cost": 0.4,
"callback_bound": False,
@ -856,4 +858,4 @@ async def test_completed_passthrough_cancellation_preserves_accounting(
assert (payload["prompt_tokens"], payload["completion_tokens"]) == (10, 5)
assert reservation["finalized"] is True
for key in ("spend:key:pass-through-key", "spend:team:pass-through-team"):
assert await counters.async_get_cache(key=key) == pytest.approx(1.98)
assert counters.get_local(key) == pytest.approx(1.98)

File diff suppressed because it is too large Load diff

View file

@ -34,6 +34,7 @@ from litellm.proxy.spend_tracking.budget_reservation import (
from litellm.proxy.utils import ProxyLogging
from litellm.router import Router
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
from litellm.redis.stores.spend_counters import SpendCounterStore
TOKEN_COUNTING_ROUTES: Final = (
"/responses/input_tokens",
@ -106,11 +107,11 @@ TINY_BUDGET_KEY_TOKEN: Final = "hashed-count-tokens-key"
@pytest.fixture
def spend_counter_cache(monkeypatch: pytest.MonkeyPatch) -> DualCache:
cache: Final = DualCache()
monkeypatch.setattr(proxy_server, "spend_counter_cache", cache)
def spend_counter_store(monkeypatch: pytest.MonkeyPatch) -> SpendCounterStore:
store: Final = SpendCounterStore(None)
monkeypatch.setattr(proxy_server, "spend_counter_store", store)
monkeypatch.setattr(proxy_server, "prisma_client", None)
return cache
return store
async def _reserve_for_tiny_budget_key(route: str, request_body: dict[str, object]) -> dict[str, object] | None:
@ -130,13 +131,13 @@ async def _reserve_for_tiny_budget_key(route: str, request_body: dict[str, objec
@pytest.mark.asyncio
@pytest.mark.parametrize(("route", "request_body"), COUNT_TOKENS_REQUESTS)
async def test_repeated_token_counting_never_touches_a_tiny_budget(
spend_counter_cache: DualCache, route: str, request_body: dict[str, object]
spend_counter_store: SpendCounterStore, route: str, request_body: dict[str, object]
):
counter_key: Final = f"spend:key:{TINY_BUDGET_KEY_TOKEN}"
assert await _reserve_for_tiny_budget_key(route, request_body) is None
assert await _reserve_for_tiny_budget_key(route, request_body) is None
assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) is None
assert spend_counter_store.get_local(counter_key) is None
completion: Final = await _reserve_for_tiny_budget_key(
"/v1/messages", {"model": "claude-sonnet-5", "max_tokens": 16, "messages": ANTHROPIC_MESSAGES}
@ -145,7 +146,7 @@ async def test_repeated_token_counting_never_touches_a_tiny_budget(
reserved_cost: Final = completion["reserved_cost"]
assert isinstance(reserved_cost, float)
assert reserved_cost > 0
assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(reserved_cost)
assert spend_counter_store.get_local(counter_key) == pytest.approx(reserved_cost)
BEDROCK_SONNET: Final = "us.anthropic.claude-sonnet-4-6"
@ -320,22 +321,22 @@ async def test_reservation_starts_unbound_to_any_callback():
@pytest.mark.asyncio
async def test_release_unbound_budget_reservation_frees_the_counter(spend_counter_cache: DualCache):
async def test_release_unbound_budget_reservation_frees_the_counter(spend_counter_store: SpendCounterStore):
counter_key: Final = f"spend:key:{TINY_BUDGET_KEY_TOKEN}"
reservation: Final = await _reserve_for_tiny_budget_key(
"/v1/chat/completions", {"model": "gpt-4o", "messages": [{"role": "user", "content": "hello"}]}
)
assert reservation is not None
assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(reservation["reserved_cost"])
assert spend_counter_store.get_local(counter_key) == pytest.approx(reservation["reserved_cost"])
await release_unbound_budget_reservation(reservation)
assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(0.0)
assert spend_counter_store.get_local(counter_key) == pytest.approx(0.0)
assert reservation["finalized"] is True
@pytest.mark.asyncio
async def test_release_unbound_budget_reservation_leaves_a_bound_one_to_its_callback(spend_counter_cache: DualCache):
async def test_release_unbound_budget_reservation_leaves_a_bound_one_to_its_callback(spend_counter_store: SpendCounterStore):
counter_key: Final = f"spend:key:{TINY_BUDGET_KEY_TOKEN}"
reservation: Final = await _reserve_for_tiny_budget_key(
"/v1/chat/completions", {"model": "gpt-4o", "messages": [{"role": "user", "content": "hello"}]}
@ -345,12 +346,12 @@ async def test_release_unbound_budget_reservation_leaves_a_bound_one_to_its_call
await release_unbound_budget_reservation(reservation)
assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(reservation["reserved_cost"])
assert spend_counter_store.get_local(counter_key) == pytest.approx(reservation["reserved_cost"])
assert reservation["finalized"] is False
@pytest.mark.asyncio
async def test_reservation_holds_cost_against_cached_org_and_project_budgets(spend_counter_cache: DualCache):
async def test_reservation_holds_cost_against_cached_org_and_project_budgets(spend_counter_store: SpendCounterStore):
cache: Final = UserApiKeyCache()
await cache.async_set_cache(
key="org_id:org-budgeted:with_budget",
@ -404,8 +405,8 @@ async def test_reservation_holds_cost_against_cached_org_and_project_budgets(spe
"applied_adjustment": 0.0,
},
]
assert spend_counter_cache.in_memory_cache.get_cache(key="spend:org:org-budgeted") == pytest.approx(reserved_cost)
assert spend_counter_cache.in_memory_cache.get_cache(key="spend:project:project-budgeted") == pytest.approx(
assert spend_counter_store.get_local("spend:org:org-budgeted") == pytest.approx(reserved_cost)
assert spend_counter_store.get_local("spend:project:project-budgeted") == pytest.approx(
reserved_cost
)

View file

@ -19,38 +19,8 @@ import pytest
from litellm.caching import DualCache
from litellm.proxy import proxy_server
class _FlakyRedisCache:
def __init__(self) -> None:
self._store: dict = {}
self._increment_calls = 0
async def async_increment(self, key, value, **kwargs):
self._increment_calls += 1
if self._increment_calls == 1:
raise Exception("Redis timeout")
self._store[key] = float(self._store.get(key, 0.0)) + float(value)
return self._store[key]
async def async_get_cache(self, key, *args, **kwargs):
return self._store.get(key)
async def async_delete_cache(self, key, *args, **kwargs):
self._store.pop(key, None)
async def async_set_cache(self, key, value, *args, **kwargs):
self._store[key] = float(value)
return True
async def async_increment_pipeline(self, increment_list, **kwargs):
results = []
for op in increment_list:
results.append(await self.async_increment(op["key"], op["increment_value"]))
return results
def get_ttl(self, **kwargs):
return None
from litellm.rust_bridge._native import NativeSpendCounterError
from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters
@pytest.mark.asyncio
@ -62,15 +32,13 @@ async def test_direct_increment_runs_when_reservation_reconcile_hits_redis_failu
reserved_cost = 0.5
response_cost = 1.0
flaky_redis = _FlakyRedisCache()
flaky_redis._store[counter_key] = reserved_cost
spend = FakeSpendCounters(
redis={counter_key: reserved_cost}, memory={counter_key: reserved_cost}, fail_pipelines=True
)
monkeypatch.setattr(proxy_server, "prisma_client", None)
monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache())
monkeypatch.setattr(proxy_server.spend_counter_cache, "redis_cache", flaky_redis)
proxy_server.spend_counter_cache.in_memory_cache.set_cache(
key=counter_key, value=reserved_cost
)
monkeypatch.setattr(proxy_server, "spend_counter_store", spend.store())
budget_reservation = {
"reserved_cost": reserved_cost,
@ -86,7 +54,7 @@ async def test_direct_increment_runs_when_reservation_reconcile_hits_redis_failu
],
}
with pytest.raises(Exception, match="Redis timeout"):
with pytest.raises(NativeSpendCounterError, match="transport"):
await proxy_server.increment_spend_counters(
token=hashed_token,
team_id=None,
@ -95,5 +63,5 @@ async def test_direct_increment_runs_when_reservation_reconcile_hits_redis_failu
budget_reservation=budget_reservation,
)
assert await flaky_redis.async_get_cache(key=counter_key) is None
assert proxy_server.spend_counter_cache.in_memory_cache.get_cache(key=counter_key) is None
assert spend.redis == {}
assert spend.memory == {}

View file

@ -10,6 +10,7 @@ from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.proxy.spend_tracking.cost_finalization import finalize_response_cost
from litellm.types.utils import ModelResponse, PromptTokensDetailsWrapper, Usage
from litellm.redis.stores.spend_counters import SpendCounterStore
NOW: Final = datetime(2026, 10, 20, 12, 30, tzinfo=timezone.utc)
MODEL: Final = "gemini-finalization-test"
@ -321,8 +322,8 @@ async def test_guardrail_json_response_uses_finalized_header_when_cost_is_known(
logging_obj: Logging, monkeypatch: pytest.MonkeyPatch, known_cost: bool
) -> None:
import httpx
from fastapi import Response
from fastapi import Response
from litellm.caching.dual_cache import DualCache
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.proxy.utils import ProxyLogging
@ -366,8 +367,8 @@ async def test_guardrail_json_response_uses_finalized_header_when_cost_is_known(
async def test_malformed_passthrough_logging_preserves_response_dispatch_contract(
logging_obj: Logging, inline: bool
) -> None:
import httpx
import httpx
from litellm.llms.anthropic.common_utils import AnthropicError
from litellm.proxy.pass_through_endpoints.success_handler import PassThroughEndpointLogging
from litellm.types.passthrough_endpoints.pass_through_endpoints import PassthroughStandardLoggingPayload
@ -450,18 +451,17 @@ async def test_completed_call_keeps_its_charge_when_post_call_processing_fails(
disable_spend_updates: bool,
) -> None:
from unittest.mock import create_autospec
import anyio
from fastapi import Request, Response
from litellm.proxy._types import ProxyException
from litellm.caching.caching import DualCache
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy import proxy_server
from litellm.proxy._types import ProxyException
from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger
from litellm.proxy.proxy_server import ProxyConfig
from litellm.proxy import proxy_server
from litellm.proxy.utils import ProxyLogging
scope: Final = anyio.CancelScope()
@ -505,11 +505,11 @@ async def test_completed_call_keeps_its_charge_when_post_call_processing_fails(
}
processor: Final = ProxyBaseLLMRequestProcessing(data=data)
proxy_logging: Final = ProxyLogging(user_api_key_cache=DualCache())
counters: Final = DualCache()
counters: Final = SpendCounterStore(None)
for key in ("spend:key:completed-key", "spend:team:completed-team"):
counters.set_cache(key=key, value=1.5 + reserved_cost)
counters.set_cache(key="spend:tag:User-Agent: budget-client", value=1.5)
monkeypatch.setattr(proxy_server, "spend_counter_cache", counters)
counters.set_local(key, 1.5 + reserved_cost)
counters.set_local("spend:tag:User-Agent: budget-client", 1.5)
monkeypatch.setattr(proxy_server, "spend_counter_store", counters)
reservation: Final = (
{
"reserved_cost": reserved_cost,
@ -547,7 +547,7 @@ async def test_completed_call_keeps_its_charge_when_post_call_processing_fails(
if disable_spend_updates:
assert writer.update_database.await_count == 0
for key in ("spend:key:completed-key", "spend:team:completed-team", "spend:tag:User-Agent: budget-client"):
assert await counters.async_get_cache(key=key) == pytest.approx(1.5)
assert counters.get_local(key) == pytest.approx(1.5)
if auth.budget_reservation is not None:
assert auth.budget_reservation["finalized"] is True
return
@ -562,12 +562,10 @@ async def test_completed_call_keeps_its_charge_when_post_call_processing_fails(
assert not logging_obj.async_success_log_enqueued
if auth.budget_reservation is not None:
assert auth.budget_reservation["finalized"] is True
assert await counters.async_get_cache(key="spend:key:completed-key") == pytest.approx(2.0 + guardrail_cost)
assert await counters.async_get_cache(key="spend:team:completed-team") == pytest.approx(2.0 + guardrail_cost)
assert counters.get_local("spend:key:completed-key") == pytest.approx(2.0 + guardrail_cost)
assert counters.get_local("spend:team:completed-team") == pytest.approx(2.0 + guardrail_cost)
assert await counters.async_get_cache(key="spend:tag:User-Agent: budget-client") == pytest.approx(
2.0 + guardrail_cost
)
assert counters.get_local("spend:tag:User-Agent: budget-client") == pytest.approx(2.0 + guardrail_cost)
@pytest.mark.asyncio

View file

@ -22,6 +22,10 @@ from litellm.proxy.spend_tracking.spend_counter_batch import (
release_spend_counter_batch,
spend_counter_batch_scope,
)
from litellm.rust_bridge._native import NativeSpendCounterError
from litellm.caching.redis_batch import request_redis_batch_scope
from tests.unit.caching.test_redis_batch import FakeNativePending, FakeNatives
from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters
TOKEN = UserAPIKeyAuth(token="hashed", team_id="team", user_id="user", org_id="org")
TOKEN_KEYS = frozenset(
@ -36,50 +40,6 @@ TOKEN_KEYS = frozenset(
)
class CountingRedis(RedisCache):
def __init__(self, store: dict[str, object] | None = None, fail: bool = False) -> None:
self.store: dict[str, object] = dict(store or {})
self.fail = fail
self.commands: list[str] = []
async def async_get_cache(self, key: str, **kwargs: object) -> object:
if self.fail:
raise ConnectionError("redis down")
self.commands.append(f"GET {key}")
return self.store.get(key)
async def async_batch_get_cache(self, key_list: Sequence[str], **kwargs: object) -> dict[str, object]:
if self.fail:
raise ConnectionError("redis down")
self.commands.append(f"MGET {' '.join(key_list)}")
return {key: self.store.get(key) for key in key_list}
def get_ttl(self, **kwargs: object) -> int | None:
return None
async def async_increment(self, key: str, value: float, **kwargs: object) -> float:
self.commands.append(f"INCRBYFLOAT {key} {value}")
return self._incr(key, value)
async def async_increment_pipeline(
self, increment_list: Sequence[Mapping[str, object]], **kwargs: object
) -> list[float]:
self.commands.append(f"PIPELINE {' '.join(str(op['key']) for op in increment_list)}")
return [self._incr(str(op["key"]), float(str(op["increment_value"]))) for op in increment_list]
def _incr(self, key: str, value: float) -> float:
total = float(str(self.store.get(key, 0.0))) + value
self.store[key] = total
return total
def _spend_counter_cache(redis: RedisCache | None, in_memory: dict[str, float] | None = None) -> MagicMock:
cache = MagicMock()
cache.redis_cache = redis
cache.in_memory_cache.get_cache = MagicMock(side_effect=lambda key: (in_memory or {}).get(key))
return cache
def test_admission_counter_keys_cover_every_entity_the_checks_read():
assert admission_counter_keys(TOKEN, end_user_id="eu") == TOKEN_KEYS
assert admission_counter_keys(UserAPIKeyAuth(token="hashed"), end_user_id=None) == {"spend:key:hashed"}
@ -105,14 +65,14 @@ def test_post_call_counter_keys_skip_ids_that_are_not_strings():
@pytest.mark.asyncio
async def test_bound_counters_share_one_mget_and_a_clean_miss_is_authoritative():
redis = CountingRedis({"spend:key:hashed": 1.5, "spend:team:team": 2.5})
batch = SpendCounterBatch(redis)
spend = FakeSpendCounters(redis={"spend:key:hashed": 1.5, "spend:team:team": 2.5})
batch = SpendCounterBatch(spend.store())
batch.bind(TOKEN_KEYS)
reads = await asyncio.gather(*(batch.read(key) for key in sorted(TOKEN_KEYS)))
assert len(redis.commands) == 1
assert set(redis.commands[0].split()[1:]) == TOKEN_KEYS
assert len(spend.commands) == 1
assert set(spend.commands[0].split()[1:]) == TOKEN_KEYS
assert dict(zip(sorted(TOKEN_KEYS), reads)) == {
"spend:end_user:eu": (None, True),
"spend:key:hashed": (1.5, True),
@ -123,27 +83,74 @@ async def test_bound_counters_share_one_mget_and_a_clean_miss_is_authoritative()
}
@pytest.mark.asyncio
async def test_counters_bound_at_auth_join_the_native_pipeline_the_armed_cooldown_read_flushes():
"""One round trip: whichever native read the request awaits first flushes the native pre-call batch, and the
batch's flush hook puts every bound counter on it before it goes out, so the budget read and the armed cooldown
read leave together and the budget read costs no trip of its own."""
natives = FakeNatives()
spend = FakeSpendCounters(redis={"spend:key:hashed": 3.0})
with (
request_redis_batch_scope(native_factory=natives) as request, # pyright: ignore[reportArgumentType] # test double
spend_counter_batch_scope(spend.store()),
):
bind_admission_counter_keys(TOKEN, end_user_id=None)
native = request.native()
native.events.append("declare cooldowns")
armed_cooldowns = request.native_result(FakeNativePending(native, lambda: ()))
batch = active_spend_counter_batch()
assert batch is not None
assert await armed_cooldowns == ()
budget = await batch.read("spend:key:hashed")
assert budget == (3.0, True)
assert native.events[:3] == ["declare cooldowns", "declare spend_counters", "flush_pre"]
assert native.events.count("declare spend_counters") == 1
assert [declared.counters for declared in spend.declared] == [tuple(sorted(TOKEN_KEYS - {"spend:end_user:eu"}))]
assert spend.commands == [f"MGET {' '.join(sorted(TOKEN_KEYS - {'spend:end_user:eu'}))}"]
@pytest.mark.asyncio
async def test_counters_bound_after_the_native_batch_went_out_ride_the_next_one():
natives = FakeNatives()
spend = FakeSpendCounters(redis={"spend:key:hashed": 3.0, "spend:org:org": 4.0})
with (
request_redis_batch_scope(native_factory=natives), # pyright: ignore[reportArgumentType] # test double
spend_counter_batch_scope(spend.store(), counter_keys=frozenset({"spend:key:hashed"})),
):
batch = active_spend_counter_batch()
assert batch is not None
assert await batch.read("spend:key:hashed") == (3.0, True)
batch.bind(frozenset({"spend:org:org"}))
assert await batch.read("spend:org:org") == (4.0, True)
assert [declared.counters for declared in spend.declared] == [("spend:key:hashed",), ("spend:org:org",)]
@pytest.mark.asyncio
async def test_unbound_counter_and_closed_batch_leave_the_read_to_the_caller():
redis = CountingRedis({"spend:key:hashed": 1.0})
batch = SpendCounterBatch(redis)
spend = FakeSpendCounters(redis={"spend:key:hashed": 1.0})
batch = SpendCounterBatch(spend.store())
batch.bind(frozenset({"spend:key:hashed"}))
assert await batch.read("spend:tag:prod") is None
assert redis.commands == []
assert spend.commands == []
assert await batch.read("spend:key:hashed") == (1.0, True)
batch.close()
batch.bind(frozenset({"spend:org:org"}))
assert await batch.read("spend:key:hashed") is None
assert await batch.read("spend:org:org") is None
assert len(redis.commands) == 1
assert len(spend.commands) == 1
@pytest.mark.asyncio
async def test_keys_bound_after_the_first_read_join_one_more_mget_for_only_the_new_keys():
redis = CountingRedis({"spend:key:hashed": 1.0, "spend:org:org": 9.0})
batch = SpendCounterBatch(redis)
spend = FakeSpendCounters(redis={"spend:key:hashed": 1.0, "spend:org:org": 9.0})
batch = SpendCounterBatch(spend.store())
batch.bind(frozenset({"spend:key:hashed"}))
assert await batch.read("spend:key:hashed") == (1.0, True)
@ -151,13 +158,13 @@ async def test_keys_bound_after_the_first_read_join_one_more_mget_for_only_the_n
assert await batch.read("spend:org:org") == (9.0, True)
assert await batch.read("spend:key:hashed") == (1.0, True)
assert redis.commands == ["MGET spend:key:hashed", "MGET spend:org:org"]
assert spend.commands == ["MGET spend:key:hashed", "MGET spend:org:org"]
@pytest.mark.asyncio
async def test_a_recorded_write_result_answers_later_reads_without_another_redis_read():
redis = CountingRedis({"spend:key:hashed": 1.0})
batch = SpendCounterBatch(redis)
spend = FakeSpendCounters(redis={"spend:key:hashed": 1.0})
batch = SpendCounterBatch(spend.store())
batch.bind(frozenset({"spend:key:hashed"}))
assert await batch.read("spend:key:hashed") == (1.0, True)
@ -166,28 +173,28 @@ async def test_a_recorded_write_result_answers_later_reads_without_another_redis
assert await batch.read("spend:key:hashed") == (3.5, True)
assert await batch.read("spend:org:org") == (7.0, True)
assert redis.commands == ["MGET spend:key:hashed"]
assert spend.commands == ["MGET spend:key:hashed"]
@pytest.mark.asyncio
async def test_a_forgotten_counter_is_read_fresh_from_redis_when_it_is_bound_again():
redis = CountingRedis({"spend:key:hashed": 1.0})
batch = SpendCounterBatch(redis)
spend = FakeSpendCounters(redis={"spend:key:hashed": 1.0})
batch = SpendCounterBatch(spend.store())
batch.bind(frozenset({"spend:key:hashed"}))
batch.record("spend:key:hashed", 3.5)
batch.forget("spend:key:hashed")
assert await batch.read("spend:key:hashed") is None
redis.store["spend:key:hashed"] = 9.0
spend.redis["spend:key:hashed"] = 9.0
batch.bind(frozenset({"spend:key:hashed"}))
assert await batch.read("spend:key:hashed") == (9.0, True)
assert redis.commands == ["MGET spend:key:hashed"]
assert spend.commands == ["MGET spend:key:hashed"]
@pytest.mark.asyncio
async def test_failed_mget_hands_every_counter_back_to_the_caller():
batch = SpendCounterBatch(CountingRedis(fail=True))
batch = SpendCounterBatch(FakeSpendCounters(fail=True).store())
batch.bind(TOKEN_KEYS)
assert await batch.read("spend:key:hashed") is None
@ -195,9 +202,9 @@ async def test_failed_mget_hands_every_counter_back_to_the_caller():
@pytest.mark.asyncio
async def test_non_numeric_counter_payload_hands_the_batch_back_to_the_caller():
redis = CountingRedis()
redis.store["spend:key:hashed"] = "garbage"
batch = SpendCounterBatch(redis)
spend = FakeSpendCounters()
spend.redis["spend:key:hashed"] = "garbage"
batch = SpendCounterBatch(spend.store())
batch.bind(frozenset({"spend:key:hashed"}))
assert await batch.read("spend:key:hashed") is None
@ -205,11 +212,11 @@ async def test_non_numeric_counter_payload_hands_the_batch_back_to_the_caller():
def test_scope_installs_a_batch_only_when_redis_exists_and_release_closes_without_clearing():
assert active_spend_counter_batch() is None
with spend_counter_batch_scope(None):
with spend_counter_batch_scope(FakeSpendCounters(redis=None).store()):
assert active_spend_counter_batch() is None
bind_admission_counter_keys(TOKEN, end_user_id=None)
with spend_counter_batch_scope(CountingRedis()):
with spend_counter_batch_scope(FakeSpendCounters().store()):
batch = active_spend_counter_batch()
assert batch is not None
bind_admission_counter_keys(TOKEN, end_user_id="eu")
@ -223,11 +230,11 @@ def test_scope_installs_a_batch_only_when_redis_exists_and_release_closes_withou
@pytest.mark.asyncio
async def test_get_current_spend_inside_the_scope_costs_one_mget_for_all_admission_counters(monkeypatch):
redis = CountingRedis({"spend:key:hashed": 3.0, "spend:team:team": 4.0, "spend:org:org": 5.0})
monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis))
spend = FakeSpendCounters(redis={"spend:key:hashed": 3.0, "spend:team:team": 4.0, "spend:org:org": 5.0})
monkeypatch.setattr(ps, "spend_counter_store", spend.store())
monkeypatch.setattr(ps, "prisma_client", None)
with spend_counter_batch_scope(redis):
with spend_counter_batch_scope(spend.store()):
bind_admission_counter_keys(TOKEN, end_user_id="eu")
key_spend = await ps.get_current_spend(counter_key="spend:key:hashed", fallback_spend=0.0)
team_spend = await ps.get_current_spend(counter_key="spend:team:team", fallback_spend=0.0)
@ -235,61 +242,61 @@ async def test_get_current_spend_inside_the_scope_costs_one_mget_for_all_admissi
user_spend = await ps.get_current_spend(counter_key="spend:user:user", fallback_spend=7.0)
assert (key_spend, team_spend, org_spend, user_spend) == (3.0, 4.0, 5.0, 7.0)
assert [c for c in redis.commands if c.startswith("GET ")] == [], "the cold reseed reuses the MGET miss"
assert [c for c in redis.commands if c.startswith("MGET ")] == [f"MGET {' '.join(sorted(TOKEN_KEYS))}"]
assert [c for c in spend.commands if c.startswith("GET ")] == [], "the cold reseed reuses the MGET miss"
assert [c for c in spend.commands if c.startswith("MGET ")] == [f"MGET {' '.join(sorted(TOKEN_KEYS))}"]
@pytest.mark.asyncio
async def test_get_current_spend_outside_the_scope_still_reads_redis_per_counter(monkeypatch):
redis = CountingRedis({"spend:key:hashed": 3.0})
monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis))
spend = FakeSpendCounters(redis={"spend:key:hashed": 3.0})
monkeypatch.setattr(ps, "spend_counter_store", spend.store())
assert await ps.get_current_spend(counter_key="spend:key:hashed", fallback_spend=0.0) == 3.0
assert redis.commands == ["GET spend:key:hashed"]
assert spend.commands == ["GET spend:key:hashed"]
@pytest.mark.asyncio
async def test_after_release_a_read_goes_to_redis_directly_so_read_then_write_sees_fresh_values(monkeypatch):
redis = CountingRedis({"spend:key:hashed": 3.0})
monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis))
spend = FakeSpendCounters(redis={"spend:key:hashed": 3.0})
monkeypatch.setattr(ps, "spend_counter_store", spend.store())
with spend_counter_batch_scope(redis):
with spend_counter_batch_scope(spend.store()):
bind_admission_counter_keys(TOKEN, end_user_id=None)
assert await ps.read_spend_counter_cache_value("spend:key:hashed") == (3.0, True)
redis.store["spend:key:hashed"] = 8.0
spend.redis["spend:key:hashed"] = 8.0
assert await ps.read_spend_counter_cache_value("spend:key:hashed") == (3.0, True)
release_spend_counter_batch()
assert await ps.read_spend_counter_cache_value("spend:key:hashed") == (8.0, True)
assert [c.split()[0] for c in redis.commands] == ["MGET", "GET"]
assert [c.split()[0] for c in spend.commands] == ["MGET", "GET"]
@pytest.mark.asyncio
async def test_batched_clean_miss_does_not_fall_back_to_the_per_pod_in_memory_copy(monkeypatch):
redis = CountingRedis()
monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis, in_memory={"spend:key:hashed": 99.0}))
spend = FakeSpendCounters(memory={"spend:key:hashed": 99.0})
monkeypatch.setattr(ps, "spend_counter_store", spend.store())
with spend_counter_batch_scope(redis):
with spend_counter_batch_scope(spend.store()):
bind_admission_counter_keys(TOKEN, end_user_id=None)
assert await ps.read_spend_counter_cache_value("spend:key:hashed") == (None, True)
@pytest.mark.asyncio
async def test_batched_redis_failure_falls_back_to_the_per_pod_in_memory_copy(monkeypatch):
redis = CountingRedis(fail=True)
monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis, in_memory={"spend:key:hashed": 99.0}))
spend = FakeSpendCounters(fail=True, memory={"spend:key:hashed": 99.0})
monkeypatch.setattr(ps, "spend_counter_store", spend.store())
with spend_counter_batch_scope(redis):
with spend_counter_batch_scope(spend.store()):
bind_admission_counter_keys(TOKEN, end_user_id=None)
assert await ps.read_spend_counter_cache_value("spend:key:hashed") == (99.0, False)
@pytest.mark.asyncio
async def test_scope_is_per_task_so_concurrent_requests_do_not_share_a_batch():
redis = CountingRedis({"spend:key:a": 1.0, "spend:key:b": 2.0})
spend = FakeSpendCounters(redis={"spend:key:a": 1.0, "spend:key:b": 2.0})
async def request(token: str) -> tuple[float | None, bool] | None:
with spend_counter_batch_scope(redis):
with spend_counter_batch_scope(spend.store()):
bind_admission_counter_keys(UserAPIKeyAuth(token=token), end_user_id=None)
batch = active_spend_counter_batch()
assert batch is not None
@ -297,18 +304,18 @@ async def test_scope_is_per_task_so_concurrent_requests_do_not_share_a_batch():
return await batch.read(f"spend:key:{token}")
assert await asyncio.gather(request("a"), request("b")) == [(1.0, True), (2.0, True)]
assert sorted(redis.commands) == ["MGET spend:key:a", "MGET spend:key:b"]
assert sorted(spend.commands) == ["MGET spend:key:a", "MGET spend:key:b"]
@pytest.mark.asyncio
async def test_batch_reads_never_touch_a_prisma_client_when_redis_answers(monkeypatch):
redis = CountingRedis({"spend:key:hashed": 3.0})
spend = FakeSpendCounters(redis={"spend:key:hashed": 3.0})
prisma = MagicMock()
prisma.db.litellm_verificationtoken.find_unique = AsyncMock()
monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis))
monkeypatch.setattr(ps, "spend_counter_store", spend.store())
monkeypatch.setattr(ps, "prisma_client", prisma)
with spend_counter_batch_scope(redis):
with spend_counter_batch_scope(spend.store()):
bind_admission_counter_keys(TOKEN, end_user_id=None)
assert await ps.get_current_spend(counter_key="spend:key:hashed", fallback_spend=0.0) == 3.0
@ -323,16 +330,16 @@ def _reseed_prisma(spend: float) -> MagicMock:
@pytest.mark.asyncio
async def test_reseed_reuses_the_admission_mget_instead_of_its_own_get():
redis = CountingRedis({"spend:key:hashed": 7.5})
spend = FakeSpendCounters(redis={"spend:key:hashed": 7.5})
prisma = _reseed_prisma(spend=1.0)
cache = _spend_counter_cache(redis)
cache = spend.store()
with spend_counter_batch_scope(redis):
with spend_counter_batch_scope(spend.store()):
bind_admission_counter_keys(TOKEN, end_user_id=None)
value = await SpendCounterReseed.coalesced(prisma, cache, counter_key="spend:key:hashed")
assert value == 7.5
assert redis.commands == [
assert spend.commands == [
"MGET spend:key:hashed spend:org:org spend:team:team spend:team_member:user:team spend:user:user"
]
assert prisma.db.mock_calls == []
@ -340,30 +347,27 @@ async def test_reseed_reuses_the_admission_mget_instead_of_its_own_get():
@pytest.mark.asyncio
async def test_reseed_treats_a_batched_clean_miss_as_authoritative_and_seeds_from_the_db():
redis = CountingRedis()
redis.async_set_cache = AsyncMock(return_value=True)
spend = FakeSpendCounters(memory={"spend:key:hashed": 99.0})
prisma = _reseed_prisma(spend=2.25)
cache = _spend_counter_cache(redis, in_memory={"spend:key:hashed": 99.0})
cache = spend.store()
with spend_counter_batch_scope(redis):
with spend_counter_batch_scope(spend.store()):
bind_admission_counter_keys(TOKEN, end_user_id=None)
value = await SpendCounterReseed.coalesced(prisma, cache, counter_key="spend:key:hashed")
assert value == 2.25
assert [c for c in redis.commands if c.startswith("GET")] == []
redis.async_set_cache.assert_awaited_once_with(key="spend:key:hashed", value=2.25, nx=True)
assert [c for c in spend.commands if c.startswith("GET")] == []
assert [c for c in spend.commands if c.startswith("SETNX")] == ["SETNX spend:key:hashed 2.25"]
@pytest.mark.asyncio
async def test_reseed_outside_the_scope_still_re_checks_redis_itself():
redis = CountingRedis({"spend:key:hashed": 4.0})
spend = FakeSpendCounters(redis={"spend:key:hashed": 4.0})
value = await SpendCounterReseed.coalesced(
_reseed_prisma(spend=1.0), _spend_counter_cache(redis), "spend:key:hashed"
)
value = await SpendCounterReseed.coalesced(_reseed_prisma(spend=1.0), spend.store(), "spend:key:hashed")
assert value == 4.0
assert redis.commands == ["GET spend:key:hashed"]
assert spend.commands == ["GET spend:key:hashed"]
POST_CALL_KEYS = TOKEN_KEYS | {"spend:tag:prod", "spend:model_access_group:premium"}
@ -371,8 +375,8 @@ POST_CALL_KEYS = TOKEN_KEYS | {"spend:tag:prod", "spend:model_access_group:premi
@pytest.mark.asyncio
async def test_post_call_increment_for_every_entity_costs_one_mget_and_one_pipeline(monkeypatch):
redis = CountingRedis({key: 1.0 for key in POST_CALL_KEYS})
monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis))
spend = FakeSpendCounters(redis={key: 1.0 for key in POST_CALL_KEYS})
monkeypatch.setattr(ps, "spend_counter_store", spend.store())
monkeypatch.setattr(ps, "prisma_client", None)
await ps.increment_spend_counters(
@ -386,26 +390,25 @@ async def test_post_call_increment_for_every_entity_costs_one_mget_and_one_pipel
response_cost=0.5,
)
assert [c.split()[0] for c in redis.commands] == ["MGET", "PIPELINE"], redis.commands
assert set(redis.commands[0].split()[1:]) == POST_CALL_KEYS
assert set(redis.commands[1].split()[1:]) == POST_CALL_KEYS
assert {key: redis.store[key] for key in POST_CALL_KEYS} == {key: 1.5 for key in POST_CALL_KEYS}
assert [c.split()[0] for c in spend.commands] == ["MGET", "PIPELINE"], spend.commands
assert set(spend.commands[0].split()[1:]) == POST_CALL_KEYS
assert set(spend.commands[1].split()[1:]) == POST_CALL_KEYS
assert {key: spend.redis[key] for key in POST_CALL_KEYS} == {key: 1.5 for key in POST_CALL_KEYS}
@pytest.mark.asyncio
async def test_post_call_cold_counters_seed_from_the_mget_miss_without_a_second_read(monkeypatch):
redis = CountingRedis({"spend:key:hashed": 1.0})
redis.async_set_cache = AsyncMock(return_value=True)
spend = FakeSpendCounters(redis={"spend:key:hashed": 1.0})
prisma = MagicMock()
prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=MagicMock(spend=4.0))
monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis))
monkeypatch.setattr(ps, "spend_counter_store", spend.store())
monkeypatch.setattr(ps, "prisma_client", prisma)
await ps.increment_spend_counters(token="hashed", team_id="team", user_id=None, response_cost=0.5)
assert [c.split()[0] for c in redis.commands] == ["MGET", "PIPELINE"], redis.commands
redis.async_set_cache.assert_awaited_once_with(key="spend:team:team", value=4.0, nx=True)
assert redis.store["spend:key:hashed"] == 1.5
assert [c.split()[0] for c in spend.commands] == ["MGET", "SETNX", "PIPELINE"], spend.commands
assert spend.commands[1] == "SETNX spend:team:team 4.0"
assert spend.redis["spend:key:hashed"] == 1.5
RESERVED_KEYS = frozenset(
@ -425,8 +428,8 @@ def _reservation(reserved_cost: float, counter_keys: frozenset[str] = RESERVED_K
@pytest.mark.asyncio
async def test_post_call_with_a_reservation_costs_one_mget_and_one_pipeline_for_reconcile_and_increments(monkeypatch):
redis = CountingRedis({key: 1.0 for key in POST_CALL_KEYS})
monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis))
spend = FakeSpendCounters(redis={key: 1.0 for key in POST_CALL_KEYS})
monkeypatch.setattr(ps, "spend_counter_store", spend.store())
monkeypatch.setattr(ps, "prisma_client", None)
reservation = _reservation(reserved_cost=0.4)
@ -442,10 +445,10 @@ async def test_post_call_with_a_reservation_costs_one_mget_and_one_pipeline_for_
budget_reservation=reservation,
)
assert [c.split()[0] for c in redis.commands] == ["MGET", "PIPELINE"], redis.commands
assert set(redis.commands[0].split()[1:]) == POST_CALL_KEYS, "reconcile and warm checks share the MGET"
assert set(redis.commands[1].split()[1:]) == POST_CALL_KEYS, "reconcile adjustments ride the increment pipeline"
assert {key: round(redis.store[key], 6) for key in POST_CALL_KEYS} == {
assert [c.split()[0] for c in spend.commands] == ["MGET", "PIPELINE"], spend.commands
assert set(spend.commands[0].split()[1:]) == POST_CALL_KEYS, "reconcile and warm checks share the MGET"
assert set(spend.commands[1].split()[1:]) == POST_CALL_KEYS, "reconcile adjustments ride the increment pipeline"
assert {key: round(spend.redis[key], 6) for key in POST_CALL_KEYS} == {
key: (1.1 if key in RESERVED_KEYS else 1.5) for key in POST_CALL_KEYS
}
assert [round(entry["applied_adjustment"], 6) for entry in reservation["entries"]] == [0.1] * len(RESERVED_KEYS)
@ -454,44 +457,35 @@ async def test_post_call_with_a_reservation_costs_one_mget_and_one_pipeline_for_
@pytest.mark.asyncio
async def test_a_stale_counter_repair_updates_the_open_batch_instead_of_forcing_a_second_mget(monkeypatch):
redis = CountingRedis({"spend:key:hashed": 1.0, "spend:team:team": 1.0})
spend = FakeSpendCounters(redis={"spend:key:hashed": 1.0, "spend:team:team": 1.0})
monkeypatch.setattr(ps, "spend_counter_store", spend.store())
async def set_max(key: str, value: float, **kwargs: object) -> float:
redis.commands.append(f"SETMAX {key} {value}")
redis.store[key] = max(float(str(redis.store.get(key, 0.0))), value)
return float(str(redis.store[key]))
redis.async_set_max = set_max
monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis))
with spend_counter_batch_scope(redis, counter_keys=frozenset({"spend:key:hashed", "spend:team:team"})):
with spend_counter_batch_scope(spend.store(), counter_keys=frozenset({"spend:key:hashed", "spend:team:team"})):
assert await ps.read_spend_counter_cache_value(counter_key="spend:team:team") == (1.0, True)
await ps._repair_stale_spend_counter(counter_key="spend:team:team", db_spend=4.0)
assert await ps.read_spend_counter_cache_value(counter_key="spend:team:team") == (4.0, True)
assert await ps.read_spend_counter_cache_value(counter_key="spend:key:hashed") == (1.0, True)
assert [c.split()[0] for c in redis.commands] == ["MGET", "SETMAX"], redis.commands
assert [c.split()[0] for c in spend.commands] == ["MGET", "SETMAX"], spend.commands
@pytest.mark.asyncio
async def test_reconcile_settles_a_flushed_counter_on_its_own_after_the_shared_pipeline(monkeypatch):
from litellm.proxy.spend_tracking.budget_reservation import reconcile_budget_reservation
redis = CountingRedis({key: 1.0 for key in RESERVED_KEYS - {"spend:team:team"}})
redis.async_set_max = AsyncMock(return_value=4.0)
spend = FakeSpendCounters(redis={key: 1.0 for key in RESERVED_KEYS - {"spend:team:team"}})
prisma = MagicMock()
prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=MagicMock(spend=4.0))
monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis))
monkeypatch.setattr(ps, "spend_counter_store", spend.store())
monkeypatch.setattr(ps, "prisma_client", prisma)
reservation = _reservation(reserved_cost=0.4)
await reconcile_budget_reservation(budget_reservation=reservation, actual_cost=0.5)
assert [c.split()[0] for c in redis.commands] == ["MGET", "PIPELINE", "INCRBYFLOAT"], redis.commands
assert set(redis.commands[1].split()[1:]) == RESERVED_KEYS - {"spend:team:team"}
assert redis.commands[2] == "INCRBYFLOAT spend:team:team 0.5"
redis.async_set_max.assert_awaited_once()
assert redis.async_set_max.await_args.kwargs["key"] == "spend:team:team"
assert [c.split()[0] for c in spend.commands] == ["MGET", "PIPELINE", "SETMAX", "INCRBYFLOAT"], spend.commands
assert set(spend.commands[1].split()[1:]) == RESERVED_KEYS - {"spend:team:team"}
assert spend.commands[2] == "SETMAX spend:team:team 4.0"
assert spend.commands[3] == "INCRBYFLOAT spend:team:team 0.5"
assert all(round(entry["applied_adjustment"], 6) == 0.1 for entry in reservation["entries"])
@ -499,8 +493,8 @@ async def test_reconcile_settles_a_flushed_counter_on_its_own_after_the_shared_p
async def test_pre_call_resize_against_an_inconsistent_counter_writes_nothing_and_denies(monkeypatch):
from litellm.proxy.spend_tracking.budget_reservation import _resize_applied_reservation
redis = CountingRedis({"spend:key:hashed": 1.0})
monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis))
spend = FakeSpendCounters(redis={"spend:key:hashed": 1.0})
monkeypatch.setattr(ps, "spend_counter_store", spend.store())
monkeypatch.setattr(ps, "prisma_client", None)
entries = _reservation(reserved_cost=0.4, counter_keys=frozenset({"spend:key:hashed", "spend:team:team"}))[
"entries"
@ -509,25 +503,19 @@ async def test_pre_call_resize_against_an_inconsistent_counter_writes_nothing_an
with pytest.raises(RuntimeError, match="spend:team:team"):
await _resize_applied_reservation(entries=entries, current_reserved_cost=0.4, new_reserved_cost=0.9)
assert [c.split()[0] for c in redis.commands] == ["MGET"], redis.commands
assert redis.store["spend:key:hashed"] == 1.0
assert [c.split()[0] for c in spend.commands] == ["MGET"], spend.commands
assert spend.redis["spend:key:hashed"] == 1.0
assert all("applied_adjustment" not in entry for entry in entries)
@pytest.mark.asyncio
async def test_a_failed_post_call_pipeline_invalidates_every_counter_it_carried_and_stamps_nothing(monkeypatch):
redis = CountingRedis({key: 1.0 for key in POST_CALL_KEYS})
redis.async_delete_cache = AsyncMock()
async def _pipeline(increment_list: Sequence[Mapping[str, object]], **kwargs: object) -> list[float]:
raise ConnectionError("redis down")
redis.async_increment_pipeline = _pipeline # pyright: ignore[reportAttributeAccessIssue] # instance override
monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis))
spend = FakeSpendCounters(redis={key: 1.0 for key in POST_CALL_KEYS}, fail_pipelines=True)
monkeypatch.setattr(ps, "spend_counter_store", spend.store())
monkeypatch.setattr(ps, "prisma_client", None)
reservation = _reservation(reserved_cost=0.4)
with pytest.raises(ConnectionError):
with pytest.raises(NativeSpendCounterError):
await ps.increment_spend_counters(
token="hashed",
team_id="team",
@ -538,24 +526,22 @@ async def test_a_failed_post_call_pipeline_invalidates_every_counter_it_carried_
budget_reservation=reservation,
)
assert [c.split()[0] for c in redis.commands] == ["MGET"], redis.commands
assert {call.kwargs["key"] for call in redis.async_delete_cache.await_args_list} == RESERVED_KEYS | {
"spend:user:user"
}
assert [c.split()[0] for c in spend.commands if not c.startswith("DEL ")] == ["MGET"], spend.commands
assert {c.split()[1] for c in spend.commands if c.startswith("DEL ")} == RESERVED_KEYS | {"spend:user:user"}
assert all("applied_adjustment" not in entry for entry in reservation["entries"])
assert {key: redis.store[key] for key in POST_CALL_KEYS} == {key: 1.0 for key in POST_CALL_KEYS}
assert spend.redis == {"spend:tag:prod": 1.0, "spend:model_access_group:premium": 1.0}
def test_a_scope_opened_inside_an_open_scope_joins_its_batch_and_a_closed_one_gets_its_own():
redis = CountingRedis()
with spend_counter_batch_scope(redis, counter_keys=frozenset({"spend:key:a"})):
spend = FakeSpendCounters()
with spend_counter_batch_scope(spend.store(), counter_keys=frozenset({"spend:key:a"})):
outer = active_spend_counter_batch()
assert outer is not None
with spend_counter_batch_scope(redis, counter_keys=frozenset({"spend:key:b"})):
with spend_counter_batch_scope(spend.store(), counter_keys=frozenset({"spend:key:b"})):
assert active_spend_counter_batch() is outer
assert outer.counter_keys == {"spend:key:a", "spend:key:b"}
release_spend_counter_batch()
with spend_counter_batch_scope(redis, counter_keys=frozenset({"spend:key:c"})):
with spend_counter_batch_scope(spend.store(), counter_keys=frozenset({"spend:key:c"})):
inner = active_spend_counter_batch()
assert inner is not outer
assert inner is not None and inner.counter_keys == {"spend:key:c"}
@ -568,14 +554,13 @@ async def test_reservation_inside_the_admission_scope_reuses_its_mget_and_reserv
from litellm.proxy._types import LiteLLM_TeamTableCachedObj
from litellm.proxy.spend_tracking.budget_reservation import reserve_budget_for_request
redis = CountingRedis({"spend:key:hashed": 1.0, "spend:team:team": 2.0})
redis.default_ttl = 3600
monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis))
spend = FakeSpendCounters(redis={"spend:key:hashed": 1.0, "spend:team:team": 2.0})
monkeypatch.setattr(ps, "spend_counter_store", spend.store())
monkeypatch.setattr(ps, "prisma_client", None)
monkeypatch.setattr("litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", lambda **_: 0.5)
token = UserAPIKeyAuth(token="hashed", team_id="team", max_budget=10.0)
with spend_counter_batch_scope(redis, counter_keys=admission_counter_keys(token, end_user_id=None)):
with spend_counter_batch_scope(spend.store(), counter_keys=admission_counter_keys(token, end_user_id=None)):
reservation = await reserve_budget_for_request(
request_body={"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hi"}]},
route="/chat/completions",
@ -589,10 +574,10 @@ async def test_reservation_inside_the_admission_scope_reuses_its_mget_and_reserv
)
assert reservation is not None
assert [c.split()[0] for c in redis.commands] == ["MGET", "PIPELINE"], redis.commands
assert set(redis.commands[0].split()[1:]) == {"spend:key:hashed", "spend:team:team"}
assert redis.commands[1] == "PIPELINE spend:key:hashed spend:team:team"
assert redis.store == {"spend:key:hashed": 1.5, "spend:team:team": 2.5}
assert [c.split()[0] for c in spend.commands] == ["MGET", "PIPELINE"], spend.commands
assert set(spend.commands[0].split()[1:]) == {"spend:key:hashed", "spend:team:team"}
assert spend.commands[1] == "PIPELINE spend:key:hashed spend:team:team"
assert spend.redis == {"spend:key:hashed": 1.5, "spend:team:team": 2.5}
assert [entry["counter_key"] for entry in reservation["entries"]] == ["spend:key:hashed", "spend:team:team"]
@ -602,15 +587,8 @@ async def test_a_failed_reservation_pipeline_drops_every_counter_and_reserves_no
from litellm.proxy._types import LiteLLM_TeamTableCachedObj
from litellm.proxy.spend_tracking.budget_reservation import reserve_budget_for_request
redis = CountingRedis({"spend:key:hashed": 1.0, "spend:team:team": 2.0})
redis.default_ttl = 3600
redis.async_delete_cache = AsyncMock()
async def _pipeline(increment_list: Sequence[Mapping[str, object]], **kwargs: object) -> list[float]:
raise ConnectionError("redis down")
redis.async_increment_pipeline = _pipeline # pyright: ignore[reportAttributeAccessIssue] # instance override
monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis))
spend = FakeSpendCounters(redis={"spend:key:hashed": 1.0, "spend:team:team": 2.0}, fail_pipelines=True)
monkeypatch.setattr(ps, "spend_counter_store", spend.store())
monkeypatch.setattr(ps, "prisma_client", None)
monkeypatch.setattr("litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", lambda **_: 0.5)
token = UserAPIKeyAuth(token="hashed", team_id="team", max_budget=10.0)
@ -628,24 +606,21 @@ async def test_a_failed_reservation_pipeline_drops_every_counter_and_reserves_no
)
assert reservation is None
assert {call.kwargs["key"] for call in redis.async_delete_cache.await_args_list} == {
"spend:key:hashed",
"spend:team:team",
}
assert redis.store == {"spend:key:hashed": 1.0, "spend:team:team": 2.0}
assert {c.split()[1] for c in spend.commands if c.startswith("DEL ")} == {"spend:key:hashed", "spend:team:team"}
assert spend.redis == {}
@pytest.mark.asyncio
async def test_post_call_lifecycle_reads_the_counters_after_the_db_update_and_writes_one_pipeline(monkeypatch):
from litellm.proxy.hooks.proxy_track_cost_callback import _update_database_and_spend_counters
redis = CountingRedis({key: 1.0 for key in POST_CALL_KEYS})
monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis))
spend = FakeSpendCounters(redis={key: 1.0 for key in POST_CALL_KEYS})
monkeypatch.setattr(ps, "spend_counter_store", spend.store())
monkeypatch.setattr(ps, "prisma_client", None)
proxy_logging_obj = MagicMock()
async def _update_database(**kwargs: object) -> bool:
redis.commands.append("DB")
spend.commands.append("DB")
return True
proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(side_effect=_update_database)
@ -671,11 +646,11 @@ async def test_post_call_lifecycle_reads_the_counters_after_the_db_update_and_wr
assert charged is True
proxy_logging_obj.db_spend_update_writer.update_database.assert_awaited_once()
assert [c.split()[0] for c in redis.commands] == ["MGET", "DB", "MGET", "PIPELINE"], redis.commands
assert set(redis.commands[0].split()[1:]) == RESERVED_KEYS
assert set(redis.commands[2].split()[1:]) == POST_CALL_KEYS
assert set(redis.commands[3].split()[1:]) == POST_CALL_KEYS
assert {key: round(redis.store[key], 6) for key in POST_CALL_KEYS} == {
assert [c.split()[0] for c in spend.commands] == ["MGET", "DB", "MGET", "PIPELINE"], spend.commands
assert set(spend.commands[0].split()[1:]) == RESERVED_KEYS
assert set(spend.commands[2].split()[1:]) == POST_CALL_KEYS
assert set(spend.commands[3].split()[1:]) == POST_CALL_KEYS
assert {key: round(spend.redis[key], 6) for key in POST_CALL_KEYS} == {
key: (1.1 if key in RESERVED_KEYS else 1.5) for key in POST_CALL_KEYS
}
assert [round(entry["applied_adjustment"], 6) for entry in reservation["entries"]] == [0.1] * len(RESERVED_KEYS)
@ -683,19 +658,18 @@ async def test_post_call_lifecycle_reads_the_counters_after_the_db_update_and_wr
assert active_spend_counter_batch() is None
def _reservation_fixture(monkeypatch, redis: CountingRedis) -> None:
redis.default_ttl = 3600
monkeypatch.setattr(ps, "spend_counter_cache", _spend_counter_cache(redis))
def _reservation_fixture(monkeypatch, spend: FakeSpendCounters) -> None:
monkeypatch.setattr(ps, "spend_counter_store", spend.store())
monkeypatch.setattr(ps, "prisma_client", None)
monkeypatch.setattr("litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", lambda **_: 0.5)
async def _reserve(redis: CountingRedis, token: UserAPIKeyAuth, team_max_budget: float) -> dict | None:
async def _reserve(spend: FakeSpendCounters, token: UserAPIKeyAuth, team_max_budget: float) -> dict | None:
from litellm.caching.dual_cache import DualCache
from litellm.proxy._types import LiteLLM_TeamTableCachedObj
from litellm.proxy.spend_tracking.budget_reservation import reserve_budget_for_request
with spend_counter_batch_scope(redis, counter_keys=admission_counter_keys(token, end_user_id=None)):
with spend_counter_batch_scope(spend.store(), counter_keys=admission_counter_keys(token, end_user_id=None)):
return await reserve_budget_for_request(
request_body={"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hi"}]},
route="/chat/completions",
@ -713,27 +687,27 @@ async def _reserve(redis: CountingRedis, token: UserAPIKeyAuth, team_max_budget:
async def test_a_rejected_counter_is_charged_alone_so_the_counters_after_it_are_never_touched(monkeypatch):
"""Only counters the admission MGET says still fit the estimate share the reservation pipeline; a counter that
does not is charged on its own first, so its rejection never inflates a sibling counter, not even briefly."""
redis = CountingRedis({"spend:key:hashed": 10.0, "spend:team:team": 2.0})
_reservation_fixture(monkeypatch, redis)
spend = FakeSpendCounters(redis={"spend:key:hashed": 10.0, "spend:team:team": 2.0})
_reservation_fixture(monkeypatch, spend)
token = UserAPIKeyAuth(token="hashed", team_id="team", max_budget=10.0)
with pytest.raises(litellm.BudgetExceededError):
await _reserve(redis, token, team_max_budget=20.0)
await _reserve(spend, token, team_max_budget=20.0)
writes = [c for c in redis.commands if not c.startswith("MGET")]
assert writes and all("spend:team:team" not in c for c in writes), redis.commands
assert redis.store == {"spend:key:hashed": 10.0, "spend:team:team": 2.0}
writes = [c for c in spend.commands if not c.startswith("MGET")]
assert writes and all("spend:team:team" not in c for c in writes), spend.commands
assert spend.redis == {"spend:key:hashed": 10.0, "spend:team:team": 2.0}
@pytest.mark.asyncio
async def test_a_resized_reservation_is_carried_at_its_resized_cost_to_the_counters_charged_after_it(monkeypatch):
redis = CountingRedis({"spend:key:hashed": 9.8, "spend:team:team": 2.0})
_reservation_fixture(monkeypatch, redis)
spend = FakeSpendCounters(redis={"spend:key:hashed": 9.8, "spend:team:team": 2.0})
_reservation_fixture(monkeypatch, spend)
token = UserAPIKeyAuth(token="hashed", team_id="team", max_budget=10.0)
reservation = await _reserve(redis, token, team_max_budget=2.1)
reservation = await _reserve(spend, token, team_max_budget=2.1)
assert reservation is not None
assert reservation["reserved_cost"] == pytest.approx(0.1)
assert redis.store["spend:key:hashed"] == pytest.approx(9.9)
assert redis.store["spend:team:team"] == pytest.approx(2.1)
assert spend.redis["spend:key:hashed"] == pytest.approx(9.9)
assert spend.redis["spend:team:team"] == pytest.approx(2.1)

File diff suppressed because it is too large Load diff

File diff suppressed because it is too large Load diff

View file

@ -15,7 +15,8 @@ import litellm
import litellm.proxy.proxy_server as ps
from litellm.caching.caching import RedisCache
from litellm.caching.dual_cache import DualCache
from litellm.redis.stores.spend_counters import SpendCounterStore
from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters
# ---------------------------------------------------------------------------
# Helpers
@ -31,6 +32,7 @@ class _FakeRedisCache(RedisCache):
def __init__(self): # super().__init__ skipped intentionally
self._store = {}
self.default_ttl = 60
def set_cache(self, key, value, **kwargs): # type: ignore[override]
# Enforce Redis JSON-serializable payload contract.
@ -44,6 +46,10 @@ class _FakeRedisCache(RedisCache):
return json.loads(raw)
def _spend_counter_store() -> SpendCounterStore:
return SpendCounterStore(None, handle=lambda client, _: FakeSpendCounters(redis=None if client is None else {}))
@contextmanager
def _patched_init_cache(litellm_settings: dict, cache_params: dict):
"""
@ -54,7 +60,7 @@ def _patched_init_cache(litellm_settings: dict, cache_params: dict):
_FakeRedisCache (passes the isinstance guard in _init_cache).
3. Extracts enable_redis_auth_cache from litellm_settings and passes it
as the second argument to _init_cache (matching production behaviour).
4. Yields (user_api_key_cache, spend_counter_cache, cli_sso_session_cache)
4. Yields (user_api_key_cache, spend_counter_store, cli_sso_session_cache)
after calling _init_cache, then restores everything.
"""
fake_redis = _FakeRedisCache()
@ -63,7 +69,7 @@ def _patched_init_cache(litellm_settings: dict, cache_params: dict):
mock_litellm_cache.cache = fake_redis
fresh_user_cache = DualCache()
fresh_spend_cache = DualCache()
fresh_spend_store = _spend_counter_store()
fresh_cli_sso_cache = DualCache()
fresh_config_cache = DualCache()
@ -71,7 +77,8 @@ def _patched_init_cache(litellm_settings: dict, cache_params: dict):
with (
patch.object(ps, "user_api_key_cache", fresh_user_cache),
patch.object(ps, "spend_counter_cache", fresh_spend_cache),
patch.object(ps, "spend_counter_cache", DualCache()),
patch.object(ps, "spend_counter_store", fresh_spend_store),
patch.object(ps, "cli_sso_session_cache", fresh_cli_sso_cache),
patch.object(ps, "saml_sso_cache", DualCache()),
patch.object(ps, "litellm_config_cache", fresh_config_cache),
@ -81,7 +88,7 @@ def _patched_init_cache(litellm_settings: dict, cache_params: dict):
):
litellm.cache = None
ps.ProxyConfig()._init_cache(cache_params, enable_redis_auth_cache)
yield fresh_user_cache, fresh_spend_cache, fresh_cli_sso_cache
yield fresh_user_cache, fresh_spend_store, fresh_cli_sso_cache
# ---------------------------------------------------------------------------
@ -97,8 +104,7 @@ class TestRedisAuthCacheFlag:
cache_params={"type": "redis", "host": "localhost", "port": 6379},
) as (user_cache, _, _cli_sso_cache):
assert user_cache.redis_cache is not None, (
"Redis should be attached to user_api_key_cache when "
"enable_redis_auth_cache=True"
"Redis should be attached to user_api_key_cache when enable_redis_auth_cache=True"
)
def test_flag_false_leaves_user_api_key_cache_in_memory_only(self):
@ -108,8 +114,7 @@ class TestRedisAuthCacheFlag:
cache_params={"type": "redis", "host": "localhost", "port": 6379},
) as (user_cache, _, _cli_sso_cache):
assert user_cache.redis_cache is None, (
"user_api_key_cache must remain in-memory-only when "
"enable_redis_auth_cache=False"
"user_api_key_cache must remain in-memory-only when enable_redis_auth_cache=False"
)
def test_flag_absent_leaves_user_api_key_cache_in_memory_only(self):
@ -123,21 +128,16 @@ class TestRedisAuthCacheFlag:
"enable_redis_auth_cache is absent from litellm_settings"
)
def test_spend_counter_cache_always_gets_redis_regardless_of_flag(self):
"""spend_counter_cache must receive Redis regardless of the auth-cache flag."""
def test_spend_counter_store_always_gets_redis_regardless_of_flag(self):
"""The spend counter store must receive Redis regardless of the auth-cache flag."""
for flag_value in (True, False, None):
ls = (
{"enable_redis_auth_cache": flag_value}
if flag_value is not None
else {}
)
ls = {"enable_redis_auth_cache": flag_value} if flag_value is not None else {}
with _patched_init_cache(
litellm_settings=ls,
cache_params={"type": "redis", "host": "localhost", "port": 6379},
) as (_, spend_cache, _cli_sso_cache):
assert spend_cache.redis_cache is not None, (
f"spend_counter_cache must always get Redis "
f"(enable_redis_auth_cache={flag_value!r})"
) as (_, spend_store, _cli_sso_cache):
assert spend_store.has_redis, (
f"spend_counter_store must always get Redis (enable_redis_auth_cache={flag_value!r})"
)
def test_flag_false_spend_gets_redis_but_user_cache_does_not(self):
@ -145,8 +145,8 @@ class TestRedisAuthCacheFlag:
with _patched_init_cache(
litellm_settings={"enable_redis_auth_cache": False},
cache_params={"type": "redis", "host": "localhost", "port": 6379},
) as (user_cache, spend_cache, _cli_sso_cache):
assert spend_cache.redis_cache is not None
) as (user_cache, spend_store, _cli_sso_cache):
assert spend_store.has_redis
assert user_cache.redis_cache is None
def test_cli_sso_session_cache_always_gets_redis_regardless_of_flag(self):
@ -157,18 +157,13 @@ class TestRedisAuthCacheFlag:
session" bug)
"""
for flag_value in (True, False, None):
ls = (
{"enable_redis_auth_cache": flag_value}
if flag_value is not None
else {}
)
ls = {"enable_redis_auth_cache": flag_value} if flag_value is not None else {}
with _patched_init_cache(
litellm_settings=ls,
cache_params={"type": "redis", "host": "localhost", "port": 6379},
) as (_, _, cli_sso_cache):
assert cli_sso_cache.redis_cache is not None, (
f"cli_sso_session_cache must always get Redis "
f"(enable_redis_auth_cache={flag_value!r})"
f"cli_sso_session_cache must always get Redis (enable_redis_auth_cache={flag_value!r})"
)
def test_flag_absent_still_shares_the_model_budget_counters_over_redis(self):
@ -188,6 +183,7 @@ class TestRedisAuthCacheFlag:
ps.litellm_config_cache,
)
with ExitStack() as detached:
detached.enter_context(patch.object(ps, "spend_counter_store", _spend_counter_store()))
for cache in touched_caches:
detached.enter_context(patch.object(cache, "redis_cache", None))
ps._attach_redis_usage_cache(fake_redis, enable_redis_auth_cache=False)

View file

@ -0,0 +1,245 @@
"""A stand-in for ``NativeSpendCounters`` with the native store's semantics: Redis first, this worker's memory
behind it, every Redis round trip logged in ``commands``."""
from __future__ import annotations
import asyncio
from collections.abc import Awaitable, Callable, Sequence
from dataclasses import dataclass, field
from typing import Final, Literal, TypeVar
from litellm.redis.stores.spend_counters import SpendCounterStore, SpendIncrement
from litellm.rust_bridge._native import NativeRedisClient, NativeSpendCounterError
T = TypeVar("T")
def _done(value: T) -> asyncio.Future[T]:
future: Final[asyncio.Future[T]] = asyncio.get_running_loop().create_future()
future.set_result(value)
return future
def _failed(error: BaseException) -> asyncio.Future[T]:
future: Final[asyncio.Future[T]] = asyncio.get_running_loop().create_future()
future.set_exception(error)
return future
def unavailable(kind: str = "transport") -> NativeSpendCounterError:
return NativeSpendCounterError(kind, "Redis did not answer")
def _is_number(value: object) -> bool:
"""What the native codec decodes; anything else reads as a failed Redis read."""
return value is None or (isinstance(value, (int, float)) and not isinstance(value, bool))
def _check(counter: str) -> None:
if not counter.startswith("spend:"):
raise ValueError(f"{counter!r} is not a spend counter")
@dataclass
class DeclaredRead:
counters: tuple[str, ...]
owner: FakeSpendCounters
request: object
async def resolve(self) -> Sequence[float | None] | None:
events: Final = getattr(self.request, "events", None)
if events is not None and "flush_pre" not in events:
raise ConnectionError("resolved before its batch was flushed")
return self.owner.mget(self.counters)
@dataclass
class FakeSpendCounters:
"""``redis`` is None for a worker without Redis, and ``redis_default`` answers reads of a key it lacks. ``fail``
makes every Redis call fail with ``failure``, ``fail_pipelines`` only the increment pipelines and ``failing`` only
the named writes ("set", "delete")."""
redis: dict[str, object] | None = field(default_factory=dict)
memory: dict[str, float] = field(default_factory=dict)
commands: list[str] = field(default_factory=list)
fail: bool = False
fail_pipelines: bool = False
failing: frozenset[str] = frozenset()
failure: str = "transport"
redis_default: float | None = None
redis_ttl: float = 60
batch_reads: bool = True
read_hook: Callable[[str], Awaitable[None]] | None = None
on_delete: Callable[[tuple[str, ...]], None] | None = None
deletes: list[tuple[str, ...]] = field(default_factory=list)
post_call: list[tuple[object, tuple[SpendIncrement, ...]]] = field(default_factory=list)
declared: list[DeclaredRead] = field(default_factory=list)
def store(self) -> SpendCounterStore:
client: Final = None if self.redis is None else NativeRedisClient.connect({"host": "127.0.0.1", "port": 1})
return SpendCounterStore(client, redis_ttl=self.redis_ttl, handle=lambda _, __: self)
def mget(self, counters: Sequence[str]) -> list[float | None] | None:
if self.redis is None or self.fail or not self.batch_reads:
return None
self.commands.append(f"MGET {' '.join(counters)}")
values: Final = [self.redis.get(counter, self.redis_default) for counter in counters]
return values if all(_is_number(value) for value in values) else None
async def read(self, counter: str) -> tuple[Literal["redis", "memory"], float | None]:
_check(counter)
if self.read_hook is not None:
await self.read_hook(counter)
if self.redis is not None and not self.fail:
self.commands.append(f"GET {counter}")
value: Final = self.redis.get(counter, self.redis_default)
if _is_number(value):
return "redis", value
return "memory", self.memory.get(counter)
def read_many(self, counters: Sequence[str]) -> asyncio.Future[Sequence[float | None] | None]:
return _done(self.mget(counters))
def declare_read_many(self, request: object, counters: Sequence[str]) -> DeclaredRead | None:
if self.redis is None or not self.batch_reads:
return None
declared: Final = DeclaredRead(tuple(counters), self, request)
self.declared.append(declared)
events: Final = getattr(request, "events", None)
if isinstance(events, list):
events.append("declare spend_counters")
return declared
def increment(self, counter: str, amount: float) -> asyncio.Future[float]:
_check(counter)
if self.redis is None:
return _done(self.increment_local(counter, amount))
if self.fail:
self.memory.pop(counter, None)
return _failed(unavailable(self.failure))
self.commands.append(f"INCRBYFLOAT {counter} {amount}")
return _done(self._incr(counter, amount))
def increment_many(self, increments: Sequence[SpendIncrement]) -> asyncio.Future[Sequence[float]]:
if self.redis is None:
return _done([self.increment_local(counter, amount) for counter, amount in increments])
if self.fail or self.fail_pipelines:
return _failed(unavailable(self.failure))
self.commands.append(f"PIPELINE {' '.join(counter for counter, _ in increments)}")
return _done([self._incr(counter, amount) for counter, amount in increments])
def increment_many_post_call(self, request: object, increments: Sequence[SpendIncrement]) -> bool:
if self.redis is None:
return False
self.post_call.append((request, tuple(increments)))
return True
def settle_post_call(self) -> None:
"""What the native task does once the post-call batch goes out."""
pending: Final = tuple(self.post_call)
self.post_call.clear()
for _, increments in pending:
if self.fail:
for counter, _ in increments:
self.memory.pop(counter, None)
continue
self.commands.append(f"POST PIPELINE {' '.join(counter for counter, _ in increments)}")
for counter, amount in increments:
self._incr(counter, amount)
def set_max(self, counter: str, spend: float) -> asyncio.Future[float | None]:
_check(counter)
if self.memory.get(counter) is None or self.memory[counter] < spend:
self.memory[counter] = spend
if self.redis is None:
return _done(None)
if self.fail:
return _failed(unavailable(self.failure))
self.commands.append(f"SETMAX {counter} {spend}")
current: Final = self.redis.get(counter)
if current is None or current < spend:
self.redis[counter] = spend
return _done(float(self.redis[counter]))
def seed(self, counter: str, spend: float, *, increment_if_lost: bool) -> asyncio.Future[float]:
_check(counter)
if self.redis is None:
cached: Final = self.memory.get(counter)
seeded: Final = spend if cached is None or cached <= spend else float(cached)
self.memory[counter] = seeded
return _done(seeded)
if self.fail:
return _failed(unavailable(self.failure))
self.commands.append(f"SETNX {counter} {spend}")
current: Final = self.redis.setdefault(counter, spend)
self.memory[counter] = float(current)
return _done(float(current))
def set(self, counter: str, spend: float, ttl_seconds: float | None = None) -> asyncio.Future[None]:
_check(counter)
self.memory[counter] = spend
if self.redis is None:
return _done(None)
if self.fail or "set" in self.failing:
return _failed(unavailable(self.failure))
self.commands.append(f"SET {counter} {spend} {ttl_seconds}")
self.redis[counter] = spend
return _done(None)
def deleted(self) -> frozenset[str]:
"""Every counter a delete reached, memory and Redis alike."""
return frozenset(counter for batch in self.deletes for counter in batch)
def delete(self, counters: Sequence[str]) -> asyncio.Future[None]:
for counter in counters:
_check(counter)
self.memory.pop(counter, None)
self.deletes.append(tuple(counters))
if self.on_delete is not None:
self.on_delete(tuple(counters))
if self.redis is None:
return _done(None)
if self.fail or "delete" in self.failing:
return _failed(unavailable(self.failure))
self.commands.append(f"DEL {' '.join(counters)}")
for counter in counters:
self.redis.pop(counter, None)
return _done(None)
def refresh_ttl(self, counter: str) -> asyncio.Future[bool]:
_check(counter)
if self.redis is None:
return _done(False)
if self.fail:
return _failed(unavailable(self.failure))
self.commands.append(f"EXPIRE {counter}")
return _done(counter in self.redis)
def get_local(self, counter: str) -> float | None:
_check(counter)
return self.memory.get(counter)
def set_local(self, counter: str, spend: float, ttl_seconds: float | None = None) -> None:
_check(counter)
self.memory[counter] = spend
def delete_local(self, counter: str) -> None:
_check(counter)
self.memory.pop(counter, None)
def increment_local(self, counter: str, amount: float) -> float:
_check(counter)
self.memory[counter] = self.memory.get(counter, 0) + amount
return self.memory[counter]
def flush_local(self) -> None:
self.memory.clear()
def snapshot(self) -> dict[str, float]:
return dict(self.memory)
def _incr(self, counter: str, amount: float) -> float:
assert self.redis is not None
self.redis[counter] = float(self.redis.get(counter, 0.0)) + amount
self.memory[counter] = self.redis[counter]
return self.redis[counter]

View file

@ -0,0 +1,181 @@
import json
from typing import Final
import pytest
from litellm._internal_context import current_service_target
from litellm.caching.redis_batch import request_redis_batch_scope
from litellm.caching.redis_cache import RedisCache
from litellm.redis.stores.spend_counters import (
SPEND_COUNTERS_TARGET,
SpendCounterStore,
end_user_counter,
is_unavailable,
key_counter,
key_window_counter,
model_access_group_counter,
org_counter,
project_counter,
shadow_eval_counter,
tag_counter,
team_counter,
team_member_counter,
team_window_counter,
user_counter,
)
from litellm.rust_bridge._native import NativeRedisClient, NativeSpendCounterError, NativeSpendCounters
from tests.unit.caching.test_redis_batch import FakeNatives
from .spend_counters_double import FakeSpendCounters
_KEYS: Final = (
(
key_counter("88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"),
"spend:key:88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b",
),
(key_window_counter("88dc28d0", "30d"), "spend:key:88dc28d0:window:30d"),
(team_counter("team-a"), "spend:team:team-a"),
(team_window_counter("team-a", "1h"), "spend:team:team-a:window:1h"),
(team_member_counter("user:with:colons", "team-a"), "spend:team_member:user:with:colons:team-a"),
(user_counter("user@example.com"), "spend:user:user@example.com"),
(end_user_counter("customer é 1"), "spend:end_user:customer é 1"),
(org_counter("org-1"), "spend:org:org-1"),
(tag_counter("prod:eu"), "spend:tag:prod:eu"),
(project_counter("proj-1"), "spend:project:proj-1"),
(model_access_group_counter("gpt-group"), "spend:model_access_group:gpt-group"),
(shadow_eval_counter("job-7"), "spend:shadow_eval:job-7"),
)
@pytest.mark.parametrize(("built", "key"), _KEYS, ids=[key for _, key in _KEYS])
def test_python_builds_every_key_the_native_store_formats(built: str, key: str) -> None:
store: Final = NativeSpendCounters(None)
store.set_local(built, 1.5)
assert built == key
assert store.snapshot() == {key: 1.5}
@pytest.mark.parametrize(
("spend", "wire"), [(0.0, "0.0"), (12, "12"), (1e-05, "1e-05"), (0.30000000000000004, "0.30000000000000004")]
)
def test_a_reset_writes_what_redis_cache_set_cache_writes(spend: float, wire: str) -> None:
assert json.dumps(spend) == wire
@pytest.mark.parametrize(("spend", "wire"), [(2.5, "2.5"), (100, "100"), (1e-05, "1e-05")])
def test_a_repair_writes_what_redis_cache_set_max_writes(spend: float, wire: str) -> None:
assert str(spend) == wire
@pytest.mark.parametrize(("wire", "spend"), [("7.25", 7.25), ("3", 3), ("0.0", 0.0), ("1e-05", 1e-05)])
def test_a_read_decodes_what_redis_cache_decodes(wire: str, spend: float) -> None:
decoded: Final = RedisCache._get_cache_logic( # pyright: ignore[reportPrivateUsage] # main's codec is the contract
RedisCache.__new__(RedisCache), wire.encode()
)
assert decoded == spend
assert type(decoded) is type(spend)
@pytest.mark.parametrize("key", ["spend_db_floor:spend:key:k", "virtual_key_spend:k:gpt:1d", "spend:unknown:k"])
def test_the_native_store_refuses_keys_outside_the_spend_counter_keyspace(key: str) -> None:
with pytest.raises(ValueError, match="is not a spend counter"):
NativeSpendCounters(None).get_local(key)
@pytest.mark.asyncio
async def test_without_redis_the_native_store_counts_and_answers_from_memory() -> None:
store: Final = SpendCounterStore(None)
first: Final = await store.increment("spend:key:k", 2)
second: Final = await store.increment_many((("spend:key:k", 0.5), ("spend:team:t", 1.0)))
seeded: Final = await store.seed("spend:org:o", 4.0, increment_if_lost=False)
repaired: Final = await store.set_max("spend:key:k", 9.0)
assert (first, second, seeded, repaired) == (2, (2.5, 1.0), 4.0, None)
assert await store.read("spend:key:k") == (9.0, False)
assert await store.read_many(("spend:key:k",)) is None
assert not store.has_redis
assert not store.increment_many_post_call((("spend:key:k", 1.0),))
def test_a_broadcast_writes_and_drops_only_spend_counter_copies() -> None:
store: Final = SpendCounterStore(None)
store.set_local("spend:key:k", 5.0)
store.set_local("spend:team:t", 5.0)
store.apply_broadcast("spend:key:k", 0.0, 60)
store.apply_broadcast("spend:team:t", None, None)
store.apply_broadcast("spend_db_floor:spend:key:k", 1.0, 5)
store.apply_broadcast("spend:unknown:x", 1.0, None)
assert store.snapshot() == {"spend:key:k": 0.0}
@pytest.mark.parametrize(
"kind, dropped",
[("timeout", True), ("breaker_open", True), ("transport", False), ("rejected", False), ("memory", False)],
)
def test_only_a_timeout_or_an_open_breaker_lets_a_post_call_increment_drop(kind: str, dropped: bool) -> None:
assert is_unavailable(NativeSpendCounterError(kind, "message")) is dropped
assert not is_unavailable(TimeoutError(kind))
def test_redis_attaches_once_and_its_ttl_is_the_lease_the_store_reports() -> None:
built: Final[list[tuple[NativeRedisClient | None, int]]] = []
def handle(client: NativeRedisClient | None, redis_ttl: int) -> FakeSpendCounters:
built.append((client, redis_ttl))
return FakeSpendCounters(redis=None if client is None else {})
store: Final = SpendCounterStore(None, handle=handle)
first: Final = NativeRedisClient.connect({"host": "127.0.0.1", "port": 1})
second: Final = NativeRedisClient.connect({"host": "127.0.0.2", "port": 1})
store.attach_redis(None, redis_ttl=10)
store.attach_redis(first, redis_ttl=30)
store.attach_redis(second, redis_ttl=90)
assert built == [(None, 60), (first, 30)]
assert store.has_redis
assert store.redis_ttl == 30
@pytest.mark.asyncio
async def test_post_call_increments_go_on_the_requests_native_post_call_batch_under_the_spend_counters_target() -> None:
natives: Final = FakeNatives()
targets: Final[list[str | None]] = []
class Recording(FakeSpendCounters):
def increment_many_post_call(self, request: object, increments: object) -> bool:
targets.append(current_service_target())
return super().increment_many_post_call(request, increments) # pyright: ignore[reportArgumentType] # recorded as given
spend: Final = Recording()
store: Final = spend.store()
outside: Final = store.increment_many_post_call((("spend:key:k", 1.0),))
with request_redis_batch_scope(native_factory=natives): # pyright: ignore[reportArgumentType] # test double
inside: Final = store.increment_many_post_call((("spend:key:k", 1.0),))
empty: Final = store.increment_many_post_call(())
assert (outside, inside, empty) == (False, True, False)
assert spend.post_call == [(natives.built[0], (("spend:key:k", 1.0),))]
assert targets == [SPEND_COUNTERS_TARGET]
@pytest.mark.asyncio
async def test_an_unreachable_redis_fails_writes_as_transport_and_serves_reads_from_memory() -> None:
store: Final = NativeSpendCounters(
NativeRedisClient.connect({"host": "127.0.0.1", "port": 1, "socket_timeout": 0.2})
)
store.set_local("spend:key:k", 4.0)
with pytest.raises(NativeSpendCounterError) as failure:
await store.increment("spend:key:k", 1.0)
assert failure.value.args[0] == "transport"
assert not is_unavailable(failure.value)
assert await store.read("spend:key:k") == ("memory", None)