From f135feeccd5e648d1fc75967d7c4d810f3772ec8 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 3 Nov 2026 15:55:33 +0900 Subject: [PATCH] feat(redis): move the spend counter keyspace onto the Rust store 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 The spend store sends the commands main does. One MGET for a batched read, INCRBYFLOAT with EXPIRE on its own pipeline for increment_many, async_set_max as a plain EVAL, and the reset job's deletes as multi-key DEL in chunks of 1000. --- litellm-rust/crates/python-bridge/src/lib.rs | 7 + .../python-bridge/src/redis/store/mod.rs | 1 + .../src/redis/store/spend_counters.rs | 329 +++++++ .../crates/redis-client/src/client.rs | 21 + litellm-rust/crates/redis-store/src/lib.rs | 1 + .../crates/redis-store/src/spend_counters.rs | 554 ++++++++++++ .../redis-store/tests/spend_counters.rs | 534 +++++++++++ .../crates/redis-store/tests/support/mod.rs | 26 + litellm/integrations/shadow_eval_logger.py | 5 +- litellm/proxy/auth/auth_checks.py | 90 +- litellm/proxy/auth/fallback_budget.py | 8 +- litellm/proxy/auth/user_api_key_auth.py | 15 +- .../auth_cache_invalidation_pubsub.py | 15 +- .../proxy/common_utils/reset_budget_job.py | 105 +-- .../proxy/common_utils/user_api_key_cache.py | 15 - litellm/proxy/db/spend_counter_reseed.py | 127 +-- .../proxy/hooks/proxy_track_cost_callback.py | 4 +- .../internal_user_endpoints.py | 5 +- .../key_management_endpoints.py | 60 +- .../tag_management_endpoints.py | 5 +- litellm/proxy/proxy_server.py | 228 ++--- .../spend_tracking/budget_reservation.py | 48 +- .../spend_tracking/spend_counter_batch.py | 115 +-- litellm/redis/stores/spend_counters.py | 225 +++++ litellm/rust_bridge/_native.pyi | 31 + .../test_request_redis_batch_post_call.py | 75 +- .../unit/llms/anthropic/messages/conftest.py | 3 + .../messages/test_caller_model_checks.py | 4 +- ...st_auth_checks_object_access_and_lookup.py | 92 +- .../auth/test_model_access_group_budgets.py | 6 +- .../test_user_api_key_auth_request_flow.py | 21 +- .../test_auth_cache_invalidation_pubsub.py | 31 + .../common_utils/test_reset_budget_job.py | 117 +-- .../proxy/db/test_spend_counter_reseed.py | 25 +- .../hooks/test_proxy_track_cost_callback.py | 30 +- .../test_key_management_endpoints.py | 39 +- .../test_team_endpoints.py | 25 +- ...t_budget_reservation_release_middleware.py | 45 +- .../test_success_handler.py | 10 +- .../proxy/proxy_server/test_spend_counters.py | 642 ++++--------- .../spend_tracking/test_budget_reservation.py | 31 +- .../test_budget_reservation_redis_failure.py | 50 +- .../spend_tracking/test_cost_finalization.py | 23 +- .../test_spend_counter_batch.py | 400 ++++----- tests/unit/proxy/test_budget_reservation.py | 611 ++++--------- ...test_proxy_server_endpoints_and_startup.py | 844 ++++++------------ .../unit/proxy/test_redis_auth_cache_flag.py | 54 +- .../redis/stores/spend_counters_double.py | 245 +++++ .../unit/redis/stores/test_spend_counters.py | 181 ++++ 49 files changed, 3679 insertions(+), 2499 deletions(-) create mode 100644 litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs create mode 100644 litellm-rust/crates/redis-store/src/spend_counters.rs create mode 100644 litellm-rust/crates/redis-store/tests/spend_counters.rs create mode 100644 litellm/redis/stores/spend_counters.py create mode 100644 tests/unit/redis/stores/spend_counters_double.py create mode 100644 tests/unit/redis/stores/test_spend_counters.py diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 24522dd4e04..4bbf55c6cbf 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -47,6 +47,10 @@ mod _native { #[pymodule_export] use crate::redis::store::{NativePending, router_cooldowns::NativeRouterCooldowns}; #[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}; @@ -124,6 +128,9 @@ mod tests { "NativePending", "NativeRouterCooldowns", "NativeRouterUsage", + "NativeDeclaredSpendCounters", + "NativeSpendCounterError", + "NativeSpendCounters", "native_redis_breaker_states", "NativeClickHouseSpendConfig", "NativeClickHouseSpendStorage", diff --git a/litellm-rust/crates/python-bridge/src/redis/store/mod.rs b/litellm-rust/crates/python-bridge/src/redis/store/mod.rs index f4cd89b342d..8a5e463b7f0 100644 --- a/litellm-rust/crates/python-bridge/src/redis/store/mod.rs +++ b/litellm-rust/crates/python-bridge/src/redis/store/mod.rs @@ -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; mod value; use std::{ diff --git a/litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs b/litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs new file mode 100644 index 00000000000..cf573858d6a --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs @@ -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 { + let key: String = value.extract()?; + SpendCounter::parse(&key) + .map(Self) + .map_err(|error| PyValueError::new_err(error.to_string())) + } +} + +fn counters(counters: Vec) -> Vec { + 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) -> Option { + seconds.map(|seconds| Duration::from_secs_f64(seconds.max(0.0))) +} + +fn numbers(values: Vec>) -> Vec> { + 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); + +#[pymethods] +impl NativeSpendCounters { + #[new] + #[pyo3(signature = (client, redis_ttl_seconds = 60))] + fn new(client: Option>, 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> { + 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, + ) -> PyResult> { + 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, + ) -> Option { + 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> { + 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> { + 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::>()) + }) + } + + /// 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> { + 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> { + 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, + ) -> PyResult> { + 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, + ) -> PyResult> { + 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> { + 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> { + 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, + ) -> 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 { + 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> { + 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>); + +#[pymethods] +impl NativeDeclaredSpendCounters { + fn resolve<'py>(&self, py: Python<'py>) -> PyResult> { + 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)) + }) + } +} diff --git a/litellm-rust/crates/redis-client/src/client.rs b/litellm-rust/crates/redis-client/src/client.rs index d615f616d19..69c0121d5e3 100644 --- a/litellm-rust/crates/redis-client/src/client.rs +++ b/litellm-rust/crates/redis-client/src/client.rs @@ -233,6 +233,27 @@ impl RedisClient { self.0.subscriber.clone().messages(channel.to_owned()) } + /// A script sent with its source as `EVAL` every time, like redis-py's `Redis.eval`, for the + /// scripts `RedisCache` runs that way rather than through `register_script`. + pub async fn eval_source( + &self, + source: &str, + keys: &[&str], + args: impl ToRedisArgs, + ) -> Result { + let keys = self.keys(keys); + let command = redis::cmd("EVAL") + .arg(source) + .arg(keys.len()) + .arg(keys) + .arg(args) + .to_owned(); + self.guarded(Operation::Script, async { + command.query_async(&mut self.0.connection.clone()).await + }) + .await + } + pub fn pipeline(&self) -> Pipeline { Pipeline::new(self.0.namespace.clone()) } diff --git a/litellm-rust/crates/redis-store/src/lib.rs b/litellm-rust/crates/redis-store/src/lib.rs index c87f54163b3..87749369287 100644 --- a/litellm-rust/crates/redis-store/src/lib.rs +++ b/litellm-rust/crates/redis-store/src/lib.rs @@ -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; diff --git a/litellm-rust/crates/redis-store/src/spend_counters.rs b/litellm-rust/crates/redis-store/src/spend_counters.rs new file mode 100644 index 00000000000..d2b2b9a24a5 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/spend_counters.rs @@ -0,0 +1,554 @@ +use std::{ + sync::{Arc, Mutex, MutexGuard, PoisonError}, + time::Duration, +}; + +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}; + +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, sent as `EVAL` like it: a counter only ever rises to +/// the authoritative spend. +const SET_MAX: &str = 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" +); +/// `DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE`, the chunk `DualCache.async_delete_cache_keys` deletes at once. +const DELETE_CHUNK: usize = 1000; + +/// 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, + }, + Team { + team_id: String, + window: Option, + }, + 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| 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 { + 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), + Memory(Option), +} + +/// 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::from_value(&python_value(raw)?).ok_or(DecodeError::Shape("a spend counter is a number")) +} + +fn decode_reply(raw: Option>) -> Result, 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, + remote: Option, + redis_ttl: Duration, + updates: Mutex<()>, +} + +impl SpendCounters { + pub fn new(client: Option, redis_ttl: Duration) -> Self { + Self::with_clock(client, redis_ttl, system_clock()) + } + + pub fn with_clock(client: Option, 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 { + if let Some(remote) = &self.remote { + match remote + .get::>(&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>> { + 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 { + self.declare_read_many(|client| request.pre(client), counters) + } + + fn declare_read_many( + &self, + batch_for: impl FnOnce(&RedisClient) -> RedisBatch, + counters: &[SpendCounter], + ) -> Option { + let batch = batch_for(self.remote.as_ref()?); + let keys: Vec = counters.iter().map(SpendCounter::key).collect(); + let keys: Vec<&str> = keys.iter().map(String::as_str).collect(); + Some(DeclaredReads(batch.mget::>(&keys))) + } + + /// `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 { + let Some(remote) = &self.remote else { + return Ok(self.increment_local(counter, amount)?); + }; + let key = counter.key(); + match remote + .incr_by_float(&key, amount.into(), 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, Error> { + let Some(remote) = &self.remote else { + return increments + .iter() + .map(|(counter, amount)| Ok(self.increment_local(counter, *amount)?)) + .collect(); + }; + let pipeline = increments + .iter() + .fold(remote.pipeline(), |pipeline, (counter, amount)| { + pipeline.incr_by_float(&counter.key(), (*amount).into(), Some(self.redis_ttl)) + }); + let replies = remote.run(&pipeline).await; + let mut totals = Vec::with_capacity(replies.len()); + for ((counter, _), reply) in increments.iter().zip(replies) { + let total = Number::Float(float_reply(reply?)?); + self.memory.set_cache(counter.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, + 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::>(); + 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, 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> = remote + .eval_source( + 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::().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 { + 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::>(&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.into(), 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, + ) -> 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 keys: Vec = counters.iter().map(SpendCounter::key).collect(); + for chunk in keys.chunks(DELETE_CHUNK) { + let chunk: Vec<&str> = chunk.iter().map(String::as_str).collect(); + remote.del(&chunk).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 { + 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, MemoryError> { + self.memory.get_cache(&counter.key()) + } + + pub fn set_local( + &self, + counter: &SpendCounter, + spend: Number, + ttl: Option, + ) -> 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 { + 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, 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)> { + increments + .iter() + .map(|(counter, amount)| { + let key = counter.key(); + let total = batch.incr_by_float(&key, (*amount).into(), Some(self.redis_ttl)); + (key, total) + }) + .collect() + } +} + +/// One `MGET` declared on a batch, answered once it goes out. +pub struct DeclaredReads(Declared>>>); + +impl DeclaredReads { + /// Flushes the batch if nobody has yet. `None` when the read or a decode failed. + pub async fn resolve(self) -> Option>> { + self.0 + .resolve() + .await + .map_err(Error::from) + .and_then(|replies| { + replies + .into_iter() + .map(|raw| decode_reply(raw).map_err(Error::from)) + .collect::, _>>() + }) + .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) +} + +fn float_reply(value: litellm_redis_client::Value) -> Result { + let total = match &value { + litellm_redis_client::Value::BulkString(bytes) => std::str::from_utf8(bytes) + .ok() + .and_then(|text| text.parse().ok()), + litellm_redis_client::Value::Double(total) => Some(*total), + _ => None, + }; + Ok(total.ok_or(DecodeError::Shape("INCRBYFLOAT did not answer a number"))?) +} diff --git a/litellm-rust/crates/redis-store/tests/spend_counters.rs b/litellm-rust/crates/redis-store/tests/spend_counters.rs new file mode 100644 index 00000000000..af24e6c1b63 --- /dev/null +++ b/litellm-rust/crates/redis-store/tests/spend_counters.rs @@ -0,0 +1,534 @@ +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) -> Arc { + 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::>(); + 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::>(); + + 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::(&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::(&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::(&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::(&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::(&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::(&counter.key()).await.unwrap(), None); + + request.flush_post().await; + tokio::time::sleep(Duration::from_millis(100)).await; + + let increment = match topology { + Topology::Standalone => Operation::Pipeline { ops: 1 }, + Topology::Cluster => Operation::Command("INCRBYFLOAT"), + }; + assert_eq!(trips.seen(), [increment]); + assert_eq!( + raw.get::(&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; + + let expected = match topology { + Topology::Standalone => vec![Operation::Pipeline { ops: 2 }], + Topology::Cluster => vec![Operation::Command("MGET"); 2], + }; + assert_eq!(trips.seen(), expected); + 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(), expected.len()); +} + +#[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::(&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::(&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() + ); +} + +// main's RedisCache: async_increment(refresh_ttl=True), MGET for a batched read, async_set_max as +// a plain EVAL, SET EX NX for a seed and one multi-key DEL per chunk +#[rstest] +#[tokio::test] +async fn spend_counters_send_the_commands_main_sends( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let store = store(Some(connect(&deployment))); + let counters = [team("a"), team("b")]; + deployment + .cli_on_every_node(&["CONFIG", "RESETSTAT"]) + .await + .unwrap(); + + store + .increment(&counters[0], Number::Float(1.5)) + .await + .unwrap(); + store.read_many(&counters).await.unwrap(); + store + .set_max(&counters[0], Number::Float(9.0)) + .await + .unwrap(); + store + .seed(&counters[1], Number::Float(2.0), LostSeed::Answer) + .await + .unwrap(); + store.delete(&counters).await.unwrap(); + + // commandstats also counts the GET, SET and EXPIRE the set_max script runs inside Redis + let calls = async |command| support::calls(&deployment, command).await; + assert_eq!(calls("incrbyfloat").await, 1); + assert_eq!(calls("expire").await, 2); + assert_eq!(calls("eval").await, 1); + assert_eq!(calls("evalsha").await, 0); + assert_eq!(calls("set").await, 2); + assert!(calls("mget").await > 0); + assert_eq!(calls("get").await, 1); + assert!(calls("del").await > 0); +} diff --git a/litellm-rust/crates/redis-store/tests/support/mod.rs b/litellm-rust/crates/redis-store/tests/support/mod.rs index 70bd1743da6..81dbc68ad21 100644 --- a/litellm-rust/crates/redis-store/tests/support/mod.rs +++ b/litellm-rust/crates/redis-store/tests/support/mod.rs @@ -189,3 +189,29 @@ pub fn within(ttl: Duration, whole_seconds: u64) -> bool { ttl <= Duration::from_secs(whole_seconds) && ttl > Duration::from_secs(whole_seconds.saturating_sub(5)) } + +/// How many times each command reached the servers since the last `CONFIG RESETSTAT`, so a test +/// can pin the commands a store sends against the ones main's Python sends. +pub async fn calls(deployment: &RedisDeployment, command: &str) -> u64 { + let stats = match deployment { + RedisDeployment::Standalone(server) => server.cli(&["INFO", "commandstats"]).await.unwrap(), + RedisDeployment::Cluster(cluster) => { + let mut stats = String::new(); + for (_, port) in cluster.nodes() { + stats.push_str( + &cluster + .cli_on_node(port, &["INFO", "commandstats"]) + .await + .unwrap(), + ); + } + stats + } + }; + let prefix = format!("cmdstat_{command}:calls="); + stats + .lines() + .filter_map(|line| line.trim().strip_prefix(&prefix)) + .filter_map(|rest| rest.split(',').next()?.parse::().ok()) + .sum() +} diff --git a/litellm/integrations/shadow_eval_logger.py b/litellm/integrations/shadow_eval_logger.py index b208ab76ead..f683108a51a 100644 --- a/litellm/integrations/shadow_eval_logger.py +++ b/litellm/integrations/shadow_eval_logger.py @@ -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: diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index af60a315c60..4db340a0d48 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -133,10 +133,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, @@ -151,8 +149,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 @@ -1426,7 +1437,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, ) @@ -1974,7 +1985,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, @@ -3328,37 +3339,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, @@ -6722,7 +6736,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( @@ -6811,7 +6825,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, @@ -7118,7 +7132,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, ) @@ -7230,7 +7244,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, ) @@ -7283,7 +7297,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, @@ -7402,7 +7416,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, ) @@ -7609,7 +7623,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, ) @@ -7699,7 +7713,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, @@ -7754,7 +7768,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, diff --git a/litellm/proxy/auth/fallback_budget.py b/litellm/proxy/auth/fallback_budget.py index 5065c69e512..9aef59dc8cb 100644 --- a/litellm/proxy/auth/fallback_budget.py +++ b/litellm/proxy/auth/fallback_budget.py @@ -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, ) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 82135cb1b74..0a15b9aa62c 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -27,7 +27,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, @@ -182,6 +181,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 @@ -3097,7 +3100,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, ) @@ -4266,10 +4269,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, ...]: @@ -4372,7 +4375,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( diff --git a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py index edc8048ee73..39d300bb898 100644 --- a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py +++ b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py @@ -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) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 59fd00cec43..540944b08a6 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -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, diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index d61ab8c8518..011b4e744ea 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -360,25 +360,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__" diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 721999b1868..4ed3ac30be7 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -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) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 90887422567..ce4d5ff8128 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -1281,7 +1281,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: @@ -1300,7 +1300,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, diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index d2aac9cae38..a78d653df54 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -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 ( @@ -1471,7 +1474,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: diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index c8895c01f7b..5f0e8b010f8 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -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, diff --git a/litellm/proxy/management_endpoints/tag_management_endpoints.py b/litellm/proxy/management_endpoints/tag_management_endpoints.py index 6ace23d016b..3ec92165a84 100644 --- a/litellm/proxy/management_endpoints/tag_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tag_management_endpoints.py @@ -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: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 87bc78730b5..08902bce1d6 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -205,7 +205,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, CallTypes, @@ -332,10 +331,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 ( @@ -615,9 +611,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 @@ -882,6 +876,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 @@ -968,7 +979,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, @@ -3139,6 +3149,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) @@ -3449,28 +3460,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: @@ -3567,19 +3568,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]: @@ -3594,7 +3584,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: @@ -3620,7 +3610,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). @@ -3628,7 +3618,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, @@ -3749,7 +3739,11 @@ async def _increment_spend_counters_batched( window_reset_at: Final = ( window.get("reset_at") if isinstance(window, dict) else getattr(window, "reset_at", None) ) - window_counter: Final = f"spend:{entity_type.value}:{entity_id}:window:{duration}" + window_counter: Final = ( + key_window_counter(entity_id, str(duration)) + if entity_type is Litellm_EntityType.KEY + else team_window_counter(entity_id, str(duration)) + ) window_start: Final = get_budget_window_start(window) pending_window: Final = ( await _prepare_window_spend_counter_increment( @@ -3811,7 +3805,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 @@ -3827,7 +3821,7 @@ async def _increment_spend_counters_batched( return key_pending + await _cached_window_increments(Litellm_EntityType.KEY, hashed_token, hashed_token) async def _team_scope(scope_team_id: str) -> tuple[_SpendScopeItem, ...]: - 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 @@ -3847,7 +3841,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 ( @@ -3859,7 +3853,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 ( @@ -4022,7 +4016,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, @@ -4031,7 +4025,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, @@ -4067,7 +4061,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, @@ -4088,7 +4082,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, @@ -4105,7 +4099,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, @@ -4140,7 +4134,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. @@ -4245,7 +4239,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, ) @@ -4286,7 +4280,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, @@ -4312,29 +4306,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): @@ -4346,10 +4324,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 @@ -4357,25 +4335,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 ) @@ -4385,10 +4355,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 @@ -4400,51 +4368,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) @@ -4459,26 +4398,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( @@ -5363,6 +5297,9 @@ def _attach_redis_usage_cache(redis_cache: RedisCache, enable_redis_auth_cache: redis_cache, default_redis_ttl=litellm.default_redis_ttl, ) + native_settings: Final = redis_cache.native_settings + if native_settings is not None: + spend_counter_store.attach_redis(native_client(native_settings), redis_ttl=native_settings.default_ttl) cli_sso_session_cache.attach_redis_cache( redis_cache, default_redis_ttl=CLI_SSO_SESSION_TTL_SECONDS, @@ -8997,6 +8934,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() diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index dcfd6148f54..0b14f1baa1c 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -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) diff --git a/litellm/proxy/spend_tracking/spend_counter_batch.py b/litellm/proxy/spend_tracking/spend_counter_batch.py index 7e0821e091e..e15f1566c11 100644 --- a/litellm/proxy/spend_tracking/spend_counter_batch.py +++ b/litellm/proxy/spend_tracking/spend_counter_batch.py @@ -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 diff --git a/litellm/redis/stores/spend_counters.py b/litellm/redis/stores/spend_counters.py new file mode 100644 index 00000000000..18e4a0777aa --- /dev/null +++ b/litellm/redis/stores/spend_counters.py @@ -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 diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index a24320de6e6..565e51daf4e 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -14,6 +14,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: ... @@ -91,6 +92,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 @@ -280,12 +308,15 @@ __all__ = [ "NativeClickHouseSpendConfig", "NativeClickHouseSpendStorage", "NativeDatabase", + "NativeDeclaredSpendCounters", "NativeDiagnosticProcessor", "NativePending", "NativeRedisClient", "NativeRequestRedis", "NativeRouterCooldowns", "NativeRouterUsage", + "NativeSpendCounterError", + "NativeSpendCounters", "ProcessReservedForForking", "ResponsesWebSocketConnection", "RustBridgeDeclined", diff --git a/tests/unit/caching/test_request_redis_batch_post_call.py b/tests/unit/caching/test_request_redis_batch_post_call.py index b3932285220..0fae8a06b4c 100644 --- a/tests/unit/caching/test_request_redis_batch_post_call.py +++ b/tests/unit/caching/test_request_redis_batch_post_call.py @@ -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 diff --git a/tests/unit/llms/anthropic/messages/conftest.py b/tests/unit/llms/anthropic/messages/conftest.py index 108ce8fed98..57e554e33d2 100644 --- a/tests/unit/llms/anthropic/messages/conftest.py +++ b/tests/unit/llms/anthropic/messages/conftest.py @@ -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, diff --git a/tests/unit/llms/anthropic/messages/test_caller_model_checks.py b/tests/unit/llms/anthropic/messages/test_caller_model_checks.py index 751f7cc3934..59ce07f2e54 100644 --- a/tests/unit/llms/anthropic/messages/test_caller_model_checks.py +++ b/tests/unit/llms/anthropic/messages/test_caller_model_checks.py @@ -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) diff --git a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py index 024d00d5d95..ddcfe36deb6 100644 --- a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py +++ b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py @@ -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 ( diff --git a/tests/unit/proxy/auth/test_model_access_group_budgets.py b/tests/unit/proxy/auth/test_model_access_group_budgets.py index 7506bd031d9..1650ccd7fb5 100644 --- a/tests/unit/proxy/auth/test_model_access_group_budgets.py +++ b/tests/unit/proxy/auth/test_model_access_group_budgets.py @@ -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, }, ) diff --git a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py index 17e6c004566..14044e36e1f 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py @@ -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 @@ -13076,12 +13079,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: @@ -13096,7 +13099,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, @@ -13114,8 +13117,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(): diff --git a/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py b/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py index 2dc98b5ca6b..fa2f189703a 100644 --- a/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py +++ b/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py @@ -25,6 +25,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, budget_windows_absent_cache_key, ) +from litellm.redis.stores.spend_counters import SpendCounterStore class _RecordingRedisClient(Redis): @@ -298,6 +299,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() diff --git a/tests/unit/proxy/common_utils/test_reset_budget_job.py b/tests/unit/proxy/common_utils/test_reset_budget_job.py index 183dd7919c1..4a4490afc33 100644 --- a/tests/unit/proxy/common_utils/test_reset_budget_job.py +++ b/tests/unit/proxy/common_utils/test_reset_budget_job.py @@ -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 `` 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 diff --git a/tests/unit/proxy/db/test_spend_counter_reseed.py b/tests/unit/proxy/db/test_spend_counter_reseed.py index ab931277313..b4e68f8297a 100644 --- a/tests/unit/proxy/db/test_spend_counter_reseed.py +++ b/tests/unit/proxy/db/test_spend_counter_reseed.py @@ -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 diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index c6366a40504..e6c22c9edbf 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -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 @@ -2694,9 +2696,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) @@ -2749,8 +2751,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 @@ -2769,12 +2771,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) @@ -2784,11 +2786,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}} @@ -2801,11 +2803,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()) @@ -2848,7 +2850,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.""" @@ -2885,7 +2887,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.""" @@ -2911,7 +2913,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}) diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index b453f985512..77dc6eddf51 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -112,6 +112,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) @@ -388,12 +389,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) @@ -470,9 +468,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) @@ -8647,6 +8648,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, @@ -8667,9 +8670,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. @@ -8746,6 +8747,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 @@ -8777,12 +8782,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 ) @@ -8855,6 +8856,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 @@ -8966,6 +8971,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" @@ -8989,12 +8996,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 diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index 697d9810284..d0ea7e2391b 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -105,6 +105,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) @@ -16635,7 +16636,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", @@ -16650,6 +16652,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", @@ -16673,7 +16676,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 @@ -16688,7 +16691,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) @@ -16698,6 +16702,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", @@ -16714,7 +16719,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 @@ -17102,7 +17107,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 = { @@ -17116,6 +17122,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) @@ -17141,7 +17148,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 @@ -17155,7 +17162,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 = { @@ -17169,6 +17177,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) @@ -17193,7 +17202,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 diff --git a/tests/unit/proxy/middleware/test_budget_reservation_release_middleware.py b/tests/unit/proxy/middleware/test_budget_reservation_release_middleware.py index f37a20dff8b..aed4f74032b 100644 --- a/tests/unit/proxy/middleware/test_budget_reservation_release_middleware.py +++ b/tests/unit/proxy/middleware/test_budget_reservation_release_middleware.py @@ -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 diff --git a/tests/unit/proxy/pass_through_endpoints/test_success_handler.py b/tests/unit/proxy/pass_through_endpoints/test_success_handler.py index 4934573a777..be8584b7188 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_success_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/test_success_handler.py @@ -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) diff --git a/tests/unit/proxy/proxy_server/test_spend_counters.py b/tests/unit/proxy/proxy_server/test_spend_counters.py index d273ef8f42f..53638fa0289 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counters.py +++ b/tests/unit/proxy/proxy_server/test_spend_counters.py @@ -52,42 +52,14 @@ from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, end_u from litellm.proxy.management_endpoints.team_endpoints import refresh_cached_team from .conftest import normalize +from litellm.rust_bridge._native import NativeSpendCounterError +from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters -def _make_spend_counter_cache( - *, - redis_get_value=None, - redis_get_side_effect=None, - redis_increment_value=None, - redis_increment_side_effect=None, - in_memory_value=None, - with_redis: bool = True, -): - cache = MagicMock() - cache.in_memory_cache = MagicMock() - cache.in_memory_cache.get_cache = MagicMock(return_value=in_memory_value) - cache.in_memory_cache.set_cache = MagicMock() - cache.in_memory_cache.delete_cache = MagicMock() - if with_redis: - cache.redis_cache = MagicMock() - cache.redis_cache.async_get_cache = AsyncMock(return_value=redis_get_value, side_effect=redis_get_side_effect) - cache.redis_cache.async_increment = AsyncMock( - return_value=redis_increment_value, - side_effect=redis_increment_side_effect, - ) - cache.redis_cache.async_delete_cache = AsyncMock() - cache.redis_cache.async_set_cache = AsyncMock() - cache.redis_cache.async_set_max = AsyncMock() - cache.redis_cache.async_increment_pipeline = AsyncMock(return_value=None) - cache.redis_cache.get_ttl = MagicMock(return_value=None) - else: - cache.redis_cache = None - cache.async_increment_cache = AsyncMock(return_value=redis_increment_value) - cache.async_get_cache = AsyncMock(return_value=None) - cache.async_set_cache = AsyncMock() - cache.async_delete_cache = AsyncMock() - cache.async_set_cache_pipeline = AsyncMock() - return cache +def _spend_counters(monkeypatch: pytest.MonkeyPatch, **fake: object) -> FakeSpendCounters: + spend: Final = FakeSpendCounters(**fake) # pyright: ignore[reportArgumentType] # test double fields + monkeypatch.setattr(ps, "spend_counter_store", spend.store()) + return spend def _make_user_api_key_cache(get_value=None, get_side_effect=None): @@ -105,30 +77,17 @@ def _make_user_api_key_cache(get_value=None, get_side_effect=None): @pytest.mark.asyncio async def test_get_current_spend_reads_redis_first(monkeypatch): - fake_cache = _make_spend_counter_cache(redis_get_value=42.5) - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + spend = _spend_counters(monkeypatch, redis={"spend:key:abc": 42.5}, memory={"spend:key:abc": 1.0}) result = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=0.0) - observed = { - "value": result, - "redis_called": fake_cache.redis_cache.async_get_cache.called, - "in_memory_called": fake_cache.in_memory_cache.get_cache.called, - } - assert normalize(observed) == { - "value": 42.5, - "redis_called": True, - "in_memory_called": False, - } + assert result == 42.5 + assert spend.commands == ["GET spend:key:abc"] @pytest.mark.asyncio async def test_get_current_spend_redis_error_falls_back_to_in_memory(monkeypatch): - fake_cache = _make_spend_counter_cache( - redis_get_side_effect=RuntimeError("redis down"), - in_memory_value=17.0, - ) - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + _spend_counters(monkeypatch, fail=True, memory={"spend:key:abc": 17.0}) result = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=99.0) assert result == 17.0 @@ -139,8 +98,7 @@ async def test_get_current_spend_floors_stale_low_counter_against_db(monkeypatch """A Redis counter left stale-low by a Redis restart must not admit a key whose authoritative DB spend is already over budget. With max_budget set, get_current_spend re-checks the DB and returns the higher recorded spend.""" - fake_cache = _make_spend_counter_cache(redis_get_value=2.0) - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + spend = _spend_counters(monkeypatch, redis={"spend:key:abc": 2.0}) from_db = AsyncMock(return_value=12.0) monkeypatch.setattr(ps.SpendCounterReseed, "from_db", from_db) @@ -155,15 +113,15 @@ async def test_get_current_spend_floors_stale_low_counter_against_db(monkeypatch # the stale counter is repaired up to the authoritative DB value via a # monotonic set-max so other workers read the corrected total, and a # concurrent increment cannot be clobbered - fake_cache.redis_cache.async_set_max.assert_awaited_once_with(key="spend:key:abc", value=12.0) + assert spend.commands == ["GET spend:key:abc", "SETMAX spend:key:abc 12.0"] + assert spend.redis == {"spend:key:abc": 12.0} @pytest.mark.asyncio async def test_get_current_spend_no_db_recheck_when_counter_healthy(monkeypatch): """A healthy counter (at or above the caller's recorded spend) is trusted without a DB read, so under-budget traffic stays off the DB path.""" - fake_cache = _make_spend_counter_cache(redis_get_value=5.0) - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + _spend_counters(monkeypatch, redis={"spend:key:abc": 5.0}) from_db = AsyncMock(return_value=99.0) monkeypatch.setattr(ps.SpendCounterReseed, "from_db", from_db) @@ -181,8 +139,7 @@ async def test_get_current_spend_no_db_recheck_when_counter_healthy(monkeypatch) async def test_get_current_spend_no_floor_without_max_budget(monkeypatch): """Without max_budget the read-time DB floor is skipped: callers that only read spend (alerts, soft budgets) keep the cheap counter-only behavior.""" - fake_cache = _make_spend_counter_cache(redis_get_value=2.0) - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + _spend_counters(monkeypatch, redis={"spend:key:abc": 2.0}) from_db = AsyncMock(return_value=12.0) monkeypatch.setattr(ps.SpendCounterReseed, "from_db", from_db) @@ -197,8 +154,7 @@ async def test_get_current_spend_floor_admits_after_reset(monkeypatch): """Right after a weekly reset the counter is 0 while the per-worker cached spend can still be last week's value. The DB floor reads the reset spend (0) and admits, so reset keys are not over-blocked.""" - fake_cache = _make_spend_counter_cache(redis_get_value=0.0) - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + spend = _spend_counters(monkeypatch, redis={"spend:key:abc": 0.0}) from_db = AsyncMock(return_value=0.0) monkeypatch.setattr(ps.SpendCounterReseed, "from_db", from_db) @@ -211,17 +167,14 @@ async def test_get_current_spend_floor_admits_after_reset(monkeypatch): assert result == 0.0 assert from_db.await_count == 1 # counter already matches the DB (reset to 0); nothing to repair, so no write - fake_cache.redis_cache.async_set_max.assert_not_called() + assert spend.commands == ["GET spend:key:abc"] @pytest.mark.asyncio async def test_get_current_spend_floor_caches_db_read(monkeypatch): """A persistently stale-low counter must not drive a DB read per request: the authoritative spend is cached in-process and reused within the window.""" - cache = ps.DualCache() - cache.redis_cache = MagicMock() - cache.redis_cache.async_get_cache = AsyncMock(return_value=2.0) - monkeypatch.setattr(ps, "spend_counter_cache", cache) + _spend_counters(monkeypatch, redis={"spend:key:abc": 2.0}) from_db = AsyncMock(return_value=12.0) monkeypatch.setattr(ps.SpendCounterReseed, "from_db", from_db) @@ -240,8 +193,7 @@ async def test_get_current_spend_floors_end_user_tag_against_fallback(monkeypatc none to read without a DB client. When such a counter is stale-low, enforcement falls back to the caller's recorded spend (loaded fresh in auth) instead of trusting the stale counter.""" - fake_cache: Final = _make_spend_counter_cache(redis_get_value=2.0) - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + spend: Final = _spend_counters(monkeypatch, redis_default=2.0) monkeypatch.setattr(ps, "prisma_client", None) monkeypatch.setattr(ps.SpendCounterReseed, "from_db", AsyncMock(return_value=None)) @@ -253,7 +205,7 @@ async def test_get_current_spend_floors_end_user_tag_against_fallback(monkeypatc assert result == 20.0 # no DB row to repair against, so the shared counter is left untouched - fake_cache.redis_cache.async_set_max.assert_not_called() + assert [command for command in spend.commands if command.startswith("SETMAX")] == [] def _make_prisma_with_end_user_row(spend: float | None): @@ -270,8 +222,7 @@ async def test_get_current_spend_end_user_floor_admits_after_a_reset_on_a_stale_ evicts the cached end-user object only on the worker that ran the reset. Every other worker still passes the pre-reset spend as fallback_spend, and that stale copy must not out-vote the reset row.""" - fake_cache: Final = _make_spend_counter_cache(redis_get_value=0.0) - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + spend: Final = _spend_counters(monkeypatch, redis={"spend:end_user:customer-42": 0.0}) prisma: Final = _make_prisma_with_end_user_row(spend=0.0) monkeypatch.setattr(ps, "prisma_client", prisma) @@ -284,7 +235,7 @@ async def test_get_current_spend_end_user_floor_admits_after_a_reset_on_a_stale_ assert result == 0.0 prisma.db.litellm_endusertable.find_unique.assert_awaited_once_with(where={"user_id": "customer-42"}) - fake_cache.redis_cache.async_set_max.assert_not_called() + assert [command for command in spend.commands if command.startswith("SETMAX")] == [] @pytest.mark.asyncio @@ -292,8 +243,7 @@ async def test_get_current_spend_end_user_floor_repairs_a_stale_low_counter(monk """After a Redis restart the end-user counter can sit below the recorded spend; the row wins and the shared counter is raised so other workers stop admitting on the stale value.""" - fake_cache: Final = _make_spend_counter_cache(redis_get_value=2.0) - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + spend: Final = _spend_counters(monkeypatch, redis={"spend:end_user:customer-42": 2.0}) monkeypatch.setattr(ps, "prisma_client", _make_prisma_with_end_user_row(spend=12.0)) result: Final = await ps.get_current_spend( @@ -303,13 +253,13 @@ async def test_get_current_spend_end_user_floor_repairs_a_stale_low_counter(monk ) assert result == 12.0 - fake_cache.redis_cache.async_set_max.assert_awaited_once_with(key="spend:end_user:customer-42", value=12.0) + assert spend.redis == {"spend:end_user:customer-42": 12.0} + assert spend.commands[-1] == "SETMAX spend:end_user:customer-42 12.0" @pytest.mark.asyncio async def test_get_current_spend_end_user_without_a_row_keeps_the_cached_spend(monkeypatch): - fake_cache: Final = _make_spend_counter_cache(redis_get_value=0.0) - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + spend: Final = _spend_counters(monkeypatch, redis_default=0.0) monkeypatch.setattr(ps, "prisma_client", _make_prisma_with_end_user_row(spend=None)) result: Final = await ps.get_current_spend( @@ -319,7 +269,7 @@ async def test_get_current_spend_end_user_without_a_row_keeps_the_cached_spend(m ) assert result == 20.0 - fake_cache.redis_cache.async_set_max.assert_not_called() + assert [command for command in spend.commands if command.startswith("SETMAX")] == [] @pytest.mark.asyncio @@ -329,13 +279,12 @@ async def test_get_current_spend_floors_window_against_spend_logs(monkeypatch): window spend, even though the caller's fallback is 0.""" from datetime import datetime, timezone - fake_cache = _make_spend_counter_cache(redis_get_value=2.0) - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + counter_key = "spend:key:tok:window:7d" + spend = _spend_counters(monkeypatch, redis={counter_key: 2.0}) monkeypatch.setattr(ps.SpendCounterReseed, "from_db", AsyncMock(return_value=None)) wfsl = AsyncMock(return_value=15.0) monkeypatch.setattr(ps.SpendCounterReseed, "window_from_spend_logs", wfsl) - counter_key = "spend:key:tok:window:7d" result = await ps.get_current_spend( counter_key=counter_key, fallback_spend=0.0, @@ -347,7 +296,8 @@ async def test_get_current_spend_floors_window_against_spend_logs(monkeypatch): assert result == 15.0 assert wfsl.await_count == 1 - fake_cache.redis_cache.async_set_max.assert_awaited_once_with(key=counter_key, value=15.0) + assert spend.redis == {counter_key: 15.0} + assert spend.commands[-1] == f"SETMAX {counter_key} 15.0" def _make_window_spend_prisma(row=None, spend_logs_total=0.0): @@ -371,11 +321,10 @@ async def test_get_current_spend_floors_window_against_maintained_row(monkeypatc row=SimpleNamespace(window_start=window_start, spend=15.0), spend_logs_total=100.0, ) - fake_cache = _make_spend_counter_cache(redis_get_value=2.0) - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + counter_key = "spend:key:tok:window:7d" + spend = _spend_counters(monkeypatch, redis={counter_key: 2.0}) monkeypatch.setattr(ps, "prisma_client", fake_prisma) - counter_key = "spend:key:tok:window:7d" result = await ps.get_current_spend( counter_key=counter_key, fallback_spend=0.0, @@ -388,7 +337,8 @@ async def test_get_current_spend_floors_window_against_maintained_row(monkeypatc assert result == 15.0 fake_prisma.db.litellm_spendlogs.group_by.assert_not_awaited() - fake_cache.redis_cache.async_set_max.assert_awaited_once_with(key=counter_key, value=15.0) + assert spend.redis == {counter_key: 15.0} + assert spend.commands[-1] == f"SETMAX {counter_key} 15.0" @pytest.mark.asyncio @@ -403,8 +353,7 @@ async def test_get_current_spend_floors_window_against_logs_when_row_stale(monke row=SimpleNamespace(window_start=window_start - timedelta(days=7), spend=999.0), spend_logs_total=15.0, ) - fake_cache = _make_spend_counter_cache(redis_get_value=2.0) - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + _spend_counters(monkeypatch, redis_default=2.0) monkeypatch.setattr(ps, "prisma_client", fake_prisma) result = await ps.get_current_spend( @@ -428,8 +377,7 @@ async def test_get_current_spend_fail_closed_rejects_when_unverifiable(monkeypat rather than admitted on an unverifiable budget.""" from fastapi import HTTPException - fake_cache = _make_spend_counter_cache(redis_get_side_effect=RuntimeError("redis down")) - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + _spend_counters(monkeypatch, fail=True) monkeypatch.setattr(ps, "general_settings", {"fail_closed_budget_enforcement": True}) monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) @@ -442,8 +390,7 @@ async def test_get_current_spend_fail_closed_rejects_when_unverifiable(monkeypat async def test_get_current_spend_fail_closed_off_admits_when_unverifiable(monkeypatch): """Default (flag off): an unverifiable read keeps the existing behavior and admits using the cached fallback — no new rejection.""" - fake_cache = _make_spend_counter_cache(redis_get_side_effect=RuntimeError("redis down")) - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + _spend_counters(monkeypatch, fail=True) monkeypatch.setattr(ps, "general_settings", {}) monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) @@ -455,8 +402,7 @@ async def test_get_current_spend_fail_closed_off_admits_when_unverifiable(monkey async def test_get_current_spend_fail_closed_admits_when_redis_verified(monkeypatch): """Fail-closed only rejects unverifiable reads: a value served by Redis is authoritative, so an under-budget request is admitted normally.""" - fake_cache = _make_spend_counter_cache(redis_get_value=1.0) - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + _spend_counters(monkeypatch, redis={"spend:key:abc": 1.0}) monkeypatch.setattr(ps, "general_settings", {"fail_closed_budget_enforcement": True}) result = await ps.get_current_spend(counter_key="spend:key:abc", fallback_spend=1.0, max_budget=10.0) @@ -468,8 +414,7 @@ async def test_get_current_spend_fail_closed_allows_authoritative_fallback(monke """End-user/tag callers pass fallback_authoritative=True (their spend is loaded fresh from the DB in auth), so fail-closed does not reject them even when the counter path is unreadable.""" - fake_cache = _make_spend_counter_cache(redis_get_side_effect=RuntimeError("redis down")) - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + _spend_counters(monkeypatch, fail=True) monkeypatch.setattr(ps, "general_settings", {"fail_closed_budget_enforcement": True}) monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) @@ -487,8 +432,7 @@ async def test_get_current_spend_strict_floors_when_fallback_also_stale(monkeypa """Strict mode closes the both-stale gap: when the counter AND the caller's cached spend are both stale-low (cheap guard would skip), strict mode still re-checks the authoritative DB and enforces against it.""" - fake_cache = _make_spend_counter_cache(redis_get_value=0.00001) - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + _spend_counters(monkeypatch, redis_default=0.00001) monkeypatch.setattr(ps, "general_settings", {"fail_closed_budget_enforcement": True}) from_db = AsyncMock(return_value=0.5) monkeypatch.setattr(ps.SpendCounterReseed, "from_db", from_db) @@ -511,16 +455,11 @@ async def test_get_current_spend_strict_floors_when_fallback_also_stale(monkeypa @pytest.mark.asyncio async def test_increment_spend_counters_increments_all_buckets(monkeypatch): - fake_cache = _make_spend_counter_cache(redis_get_value=None, redis_increment_value=5.0) + spend = _spend_counters(monkeypatch) fake_user_cache = _make_user_api_key_cache(get_value=None) - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) - - async def _fake_coalesced(**kwargs): - return None - - monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(side_effect=_fake_coalesced)) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) await ps.increment_spend_counters( token="hashed-tok", @@ -529,30 +468,15 @@ async def test_increment_spend_counters_increments_all_buckets(monkeypatch): response_cost=5.0, ) - pipeline = fake_cache.redis_cache.async_increment_pipeline - pipeline.assert_awaited_once() - increment_list = pipeline.await_args.kwargs["increment_list"] - assert {op["key"] for op in increment_list} == { - "spend:key:hashed-tok", - "spend:team:t1", - "spend:team_member:u1:t1", - "spend:user:u1", - } - assert all(op["increment_value"] == 5.0 for op in increment_list) - observed = { - "redis_increment_called": fake_cache.redis_cache.async_increment.called, - "pipeline_calls": pipeline.await_count, - "user_cache_used": fake_user_cache.async_get_cache.called, - } - assert normalize(observed) == { - "redis_increment_called": False, - "pipeline_calls": 1, - "user_cache_used": True, - } + keys = {"spend:key:hashed-tok", "spend:team:t1", "spend:team_member:u1:t1", "spend:user:u1"} + assert [command.split()[0] for command in spend.commands] == ["MGET", "PIPELINE"] + assert set(spend.commands[1].split()[1:]) == keys + assert spend.redis == dict.fromkeys(keys, 5.0) + assert fake_user_cache.async_get_cache.called is True class _ConcurrencyProbe: - """Stand-in for redis_cache.async_get_cache that pins concurrency. + """Stand-in for the per-counter warm-check read that pins concurrency. Each warm-check read registers itself as in-flight and blocks on ``release`` until the test lets it proceed. ``all_arrived`` fires once ``expected`` @@ -569,16 +493,15 @@ class _ConcurrencyProbe: self.release = asyncio.Event() self.keys: list[str] = [] - async def async_get_cache(self, *, key, **kwargs): + async def read(self, counter: str) -> None: self.in_flight += 1 self.max_in_flight = max(self.max_in_flight, self.in_flight) - self.keys.append(key) + self.keys.append(counter) if self.in_flight >= self.expected: self.all_arrived.set() if not self.release.is_set(): await self.release.wait() self.in_flight -= 1 - return 1.0 @pytest.mark.asyncio @@ -586,24 +509,21 @@ async def test_increment_spend_counters_runs_scopes_concurrently(monkeypatch): """The six independent scopes (key, team, team_member, user, end_user+tags, org) must prepare their increments concurrently. The probe only fires once all eight warm-check reads (one per counter: 6 scopes + 2 tags) are - suspended in async_get_cache at the same time, which is impossible if the - awaits are chained sequentially.""" + suspended at the same time, which is impossible if the awaits are chained + sequentially. Batch reads are off so every counter reads on its own.""" probe = _ConcurrencyProbe(expected_concurrency=8) - fake_cache = _make_spend_counter_cache() - fake_cache.redis_cache.async_get_cache = probe.async_get_cache - recorded: dict[str, float] = {} - - async def _record_pipeline(increment_list, **_): - results = [] - for op in increment_list: - recorded[op["key"]] = recorded.get(op["key"], 0.0) + op["increment_value"] - results.append(recorded[op["key"]]) - return results - - fake_cache.redis_cache.async_increment_pipeline = AsyncMock(side_effect=_record_pipeline) - fake_user_cache = _make_user_api_key_cache(get_value=None) - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) + keys = ( + "spend:key:hashed-tok", + "spend:team:t1", + "spend:team_member:u1:t1", + "spend:user:u1", + "spend:end_user:eu1", + "spend:tag:a", + "spend:tag:b", + "spend:org:org1", + ) + spend = _spend_counters(monkeypatch, redis=dict.fromkeys(keys, 1.0), batch_reads=False, read_hook=probe.read) + monkeypatch.setattr(ps, "user_api_key_cache", _make_user_api_key_cache(get_value=None)) monkeypatch.setattr(ps, "prisma_client", None) monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) @@ -634,16 +554,7 @@ async def test_increment_spend_counters_runs_scopes_concurrently(monkeypatch): probe.release.set() await task - assert recorded == { - "spend:key:hashed-tok": 5.0, - "spend:team:t1": 5.0, - "spend:team_member:u1:t1": 5.0, - "spend:user:u1": 5.0, - "spend:end_user:eu1": 5.0, - "spend:tag:a": 5.0, - "spend:tag:b": 5.0, - "spend:org:org1": 5.0, - } + assert spend.redis == dict.fromkeys(keys, 6.0) @pytest.mark.asyncio @@ -656,21 +567,8 @@ async def test_increment_spend_counters_skips_reserved_counter_keys(monkeypatch) reserved = {"spend:key:hashed-tok", "spend:org:org1"} monkeypatch.setattr(br, "get_reserved_counter_keys", MagicMock(return_value=set(reserved))) monkeypatch.setattr(br, "reconcile_budget_reservation", AsyncMock(return_value=())) - - recorded: dict[str, float] = {} - - async def _record_pipeline(increment_list, **_): - results = [] - for op in increment_list: - recorded[op["key"]] = recorded.get(op["key"], 0.0) + op["increment_value"] - results.append(recorded[op["key"]]) - return results - - fake_cache = _make_spend_counter_cache(redis_get_value=None) - fake_cache.redis_cache.async_increment_pipeline = AsyncMock(side_effect=_record_pipeline) - fake_user_cache = _make_user_api_key_cache(get_value=None) - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) + spend = _spend_counters(monkeypatch) + monkeypatch.setattr(ps, "user_api_key_cache", _make_user_api_key_cache(get_value=None)) monkeypatch.setattr(ps, "prisma_client", None) monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) @@ -687,7 +585,7 @@ async def test_increment_spend_counters_skips_reserved_counter_keys(monkeypatch) ) assert reservation["finalized"] is True - assert recorded == { + assert spend.redis == { "spend:team:t1": 5.0, "spend:team_member:u1:t1": 5.0, "spend:user:u1": 5.0, @@ -708,37 +606,25 @@ async def test_increment_spend_counters_failing_scope_propagates_after_siblings_ enforcement.""" warmed_keys: list[str] = [] - async def _warm_check(*, key, **kwargs): - warmed_keys.append(key) - if key == "spend:team:t1": - raise RuntimeError("redis get failed") - return 1.0 + async def _warm_check(counter: str) -> None: + warmed_keys.append(counter) async def _reseed_fails(*, counter_key, **kwargs): if counter_key == "spend:team:t1": raise RuntimeError("reseed failed") - applied: dict[str, float] = {} - - async def _record_pipeline(increment_list, **_): - results = [] - for op in increment_list: - applied[op["key"]] = op["increment_value"] - results.append(op["increment_value"]) - return results - - fake_cache = _make_spend_counter_cache() - fake_cache.redis_cache.async_get_cache = AsyncMock(side_effect=_warm_check) - fake_cache.redis_cache.async_increment_pipeline = AsyncMock(side_effect=_record_pipeline) - fake_user_cache = _make_user_api_key_cache(get_value=None) - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) - monkeypatch.setattr(ps, "prisma_client", None) - monkeypatch.setattr( - ps.SpendCounterReseed, - "coalesced", - AsyncMock(side_effect=_reseed_fails), + keys = ( + "spend:key:hashed-tok", + "spend:team_member:u1:t1", + "spend:user:u1", + "spend:end_user:eu1", + "spend:tag:a", + "spend:org:org1", ) + spend = _spend_counters(monkeypatch, redis=dict.fromkeys(keys, 1.0), batch_reads=False, read_hook=_warm_check) + monkeypatch.setattr(ps, "user_api_key_cache", _make_user_api_key_cache(get_value=None)) + monkeypatch.setattr(ps, "prisma_client", None) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(side_effect=_reseed_fails)) reservation = {"finalized": False} with pytest.raises(RuntimeError, match="reseed failed"): @@ -755,35 +641,18 @@ async def test_increment_spend_counters_failing_scope_propagates_after_siblings_ assert reservation["finalized"] is False # every sibling scope settled (its warm-check ran) before the error propagated - assert set(warmed_keys) == { - "spend:key:hashed-tok", - "spend:team:t1", - "spend:team_member:u1:t1", - "spend:user:u1", - "spend:end_user:eu1", - "spend:tag:a", - "spend:org:org1", - } + assert set(warmed_keys) == {*keys, "spend:team:t1"} # the surviving scopes' increments were still applied, in one pipeline call - fake_cache.redis_cache.async_increment_pipeline.assert_awaited_once() - assert applied == { - "spend:key:hashed-tok": 5.0, - "spend:team_member:u1:t1": 5.0, - "spend:user:u1": 5.0, - "spend:end_user:eu1": 5.0, - "spend:tag:a": 5.0, - "spend:org:org1": 5.0, - } - fake_cache.redis_cache.async_increment.assert_not_awaited() + assert [command.split()[0] for command in spend.commands if not command.startswith("GET")] == ["PIPELINE"] + assert spend.redis == dict.fromkeys(keys, 6.0) @pytest.mark.asyncio async def test_increment_spend_counters_zero_cost_is_noop_finalizes_reservation( monkeypatch, ): - fake_cache = _make_spend_counter_cache() + spend = _spend_counters(monkeypatch) fake_user_cache = _make_user_api_key_cache() - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) reservation = {"finalized": False} @@ -797,28 +666,26 @@ async def test_increment_spend_counters_zero_cost_is_noop_finalizes_reservation( ) assert reservation == {"finalized": True} - assert fake_cache.redis_cache.async_increment.called is False - fake_cache.redis_cache.async_increment_pipeline.assert_not_awaited() + assert spend.commands == [] @pytest.mark.asyncio async def test_increment_spend_counters_pipelines_all_scopes_in_one_redis_call( monkeypatch, ): - """Every scope's increment must go out in a single async_increment_pipeline - call, not one INCRBYFLOAT round-trip per scope.""" - counter_cache = ps.DualCache() - fake_redis = AsyncMock() - fake_redis.async_get_cache = AsyncMock(return_value=1.0) # counters warm - - async def _pipeline(increment_list, **_): - return [1.5] * len(increment_list) - - fake_redis.async_increment_pipeline = AsyncMock(side_effect=_pipeline) - fake_redis.async_increment = AsyncMock() - fake_redis.get_ttl = MagicMock(return_value=None) - counter_cache.redis_cache = fake_redis - monkeypatch.setattr(ps, "spend_counter_cache", counter_cache) + """Every scope's increment must go out in one pipeline, not one INCRBYFLOAT + round-trip per scope, and each new total lands in this worker's memory.""" + expected_keys = { + "spend:key:hashed", + "spend:team:team-1", + "spend:team_member:user-1:team-1", + "spend:user:user-1", + "spend:end_user:eu-1", + "spend:tag:tag-a", + "spend:tag:tag-b", + "spend:org:org-1", + } + spend = _spend_counters(monkeypatch, redis=dict.fromkeys(expected_keys, 1.0)) monkeypatch.setattr(ps, "user_api_key_cache", ps.DualCache()) monkeypatch.setattr(ps, "prisma_client", None) @@ -832,23 +699,9 @@ async def test_increment_spend_counters_pipelines_all_scopes_in_one_redis_call( tags=["tag-a", "tag-b"], ) - fake_redis.async_increment_pipeline.assert_awaited_once() - assert fake_redis.async_increment.await_count == 0 - increment_list = fake_redis.async_increment_pipeline.await_args.kwargs["increment_list"] - expected_keys = { - "spend:key:hashed", - "spend:team:team-1", - "spend:team_member:user-1:team-1", - "spend:user:user-1", - "spend:end_user:eu-1", - "spend:tag:tag-a", - "spend:tag:tag-b", - "spend:org:org-1", - } - assert {op["key"] for op in increment_list} == expected_keys - assert all(op["increment_value"] == 0.5 for op in increment_list) - for key in expected_keys: - assert counter_cache.in_memory_cache.get_cache(key=key) == 1.5 + assert [command.split()[0] for command in spend.commands] == ["MGET", "PIPELINE"] + assert set(spend.commands[1].split()[1:]) == expected_keys + assert spend.memory == dict.fromkeys(expected_keys, 1.5) _PENDING_SPEND_COUNTER_KEYS: tuple[str, ...] = ( @@ -863,21 +716,17 @@ _PENDING_SPEND_COUNTER_KEYS: tuple[str, ...] = ( ) -def _warm_counters_behind_a_failing_pipeline(monkeypatch, pipeline_error: Exception) -> AsyncMock: - counter_cache = ps.DualCache() - for key in _PENDING_SPEND_COUNTER_KEYS: - counter_cache.in_memory_cache.set_cache(key=key, value=1.0) - fake_redis = AsyncMock() - fake_redis.async_get_cache = AsyncMock(return_value=1.0) # counters warm - fake_redis.async_increment_pipeline = AsyncMock(side_effect=pipeline_error) - fake_redis.async_increment = AsyncMock() - fake_redis.async_delete_cache = AsyncMock() - fake_redis.get_ttl = MagicMock(return_value=None) - counter_cache.redis_cache = fake_redis - monkeypatch.setattr(ps, "spend_counter_cache", counter_cache) +def _warm_counters_behind_a_failing_pipeline(monkeypatch, failure: str) -> FakeSpendCounters: + spend: Final = _spend_counters( + monkeypatch, + redis=dict.fromkeys(_PENDING_SPEND_COUNTER_KEYS, 1.0), + memory=dict.fromkeys(_PENDING_SPEND_COUNTER_KEYS, 1.0), + fail_pipelines=True, + failure=failure, + ) monkeypatch.setattr(ps, "user_api_key_cache", ps.DualCache()) monkeypatch.setattr(ps, "prisma_client", None) - return fake_redis + return spend async def _increment_every_spend_counter() -> None: @@ -892,12 +741,13 @@ async def _increment_every_spend_counter() -> None: ) -def _assert_every_pending_counter_invalidated(fake_redis: AsyncMock) -> None: - assert fake_redis.async_increment.await_count == 0 - deleted_keys = {call.kwargs["key"] for call in fake_redis.async_delete_cache.await_args_list} - assert deleted_keys == set(_PENDING_SPEND_COUNTER_KEYS) - for key in _PENDING_SPEND_COUNTER_KEYS: - assert ps.spend_counter_cache.in_memory_cache.get_cache(key=key) is None +def _assert_every_pending_counter_invalidated(spend: FakeSpendCounters) -> None: + assert [command.split()[0] for command in spend.commands if not command.startswith("DEL ")] == ["MGET"] + assert {command.split()[1] for command in spend.commands if command.startswith("DEL ")} == set( + _PENDING_SPEND_COUNTER_KEYS + ) + assert spend.redis == {} + assert spend.memory == {} @pytest.mark.asyncio @@ -907,33 +757,12 @@ async def test_increment_spend_counters_pipeline_failure_invalidates_all_counter """A failing pipeline must invalidate every pending counter so the next request reseeds from the DB (which already holds this request's cost) instead of trusting a value the write may have partially applied.""" - from redis.exceptions import ConnectionError as RedisConnectionError + spend = _warm_counters_behind_a_failing_pipeline(monkeypatch, failure="transport") - fake_redis = _warm_counters_behind_a_failing_pipeline(monkeypatch, RedisConnectionError("Redis is down")) - - with pytest.raises(RedisConnectionError): + with pytest.raises(NativeSpendCounterError): await _increment_every_spend_counter() - _assert_every_pending_counter_invalidated(fake_redis) - - -@pytest.mark.asyncio -async def test_increment_spend_counters_pool_exhaustion_invalidates_all_counters_without_failing( - monkeypatch, -): - """A cluster node with every connection busy raises MaxConnectionsError at once. That is a busy pool, the - same as the standalone pool's wait timeout, so the counters are invalidated and the request still succeeds.""" - from redis import exceptions as redis_exceptions - - max_connections_error = cast( # cast-ok: types-redis 4.6 stubs predate MaxConnectionsError - "type[Exception]", - redis_exceptions.MaxConnectionsError, # pyright: ignore[reportAttributeAccessIssue] # old stubs - ) - fake_redis = _warm_counters_behind_a_failing_pipeline(monkeypatch, max_connections_error()) - - await _increment_every_spend_counter() - - _assert_every_pending_counter_invalidated(fake_redis) + _assert_every_pending_counter_invalidated(spend) # --------------------------------------------------------------------------- @@ -1013,9 +842,8 @@ async def test_reconcile_budget_reservation_for_counter_update_finalized_reserva async def test_prepare_end_user_and_tag_spend_increments_returns_each_unique_tag( monkeypatch, ): - fake_cache = _make_spend_counter_cache(redis_get_value=1.0) + _spend_counters(monkeypatch, redis_default=1.0) fake_user_cache = _make_user_api_key_cache() - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) @@ -1039,8 +867,7 @@ async def test_prepare_end_user_and_tag_spend_increments_returns_each_unique_tag async def test_prepare_end_user_and_tag_spend_increments_no_end_user_no_tags_invalid_input_noop( monkeypatch, ): - fake_cache = _make_spend_counter_cache() - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + spend = _spend_counters(monkeypatch) pending = await ps._prepare_end_user_and_tag_spend_increments( end_user_id=None, @@ -1050,7 +877,7 @@ async def test_prepare_end_user_and_tag_spend_increments_no_end_user_no_tags_inv ) assert pending == () - assert fake_cache.redis_cache.async_increment.called is False + assert spend.commands == [] # --------------------------------------------------------------------------- @@ -1060,9 +887,8 @@ async def test_prepare_end_user_and_tag_spend_increments_no_end_user_no_tags_inv @pytest.mark.asyncio async def test_prepare_org_spend_increment_returns_pending_when_org_present(monkeypatch): - fake_cache = _make_spend_counter_cache(redis_get_value=1.0) + _spend_counters(monkeypatch, redis_default=1.0) fake_user_cache = _make_user_api_key_cache() - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) @@ -1080,8 +906,7 @@ async def test_prepare_org_spend_increment_returns_pending_when_org_present(monk @pytest.mark.asyncio async def test_prepare_org_spend_increment_no_org_is_noop_invalid_id(monkeypatch): - fake_cache = _make_spend_counter_cache() - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + spend = _spend_counters(monkeypatch) pending = await ps._prepare_org_spend_increment( org_id=None, @@ -1090,7 +915,7 @@ async def test_prepare_org_spend_increment_no_org_is_noop_invalid_id(monkeypatch ) assert pending == () - assert fake_cache.redis_cache.async_increment.called is False + assert spend.commands == [] # --------------------------------------------------------------------------- @@ -1102,8 +927,7 @@ async def test_prepare_org_spend_increment_no_org_is_noop_invalid_id(monkeypatch async def test_prepare_unreserved_spend_counter_increment_skips_reserved_keys( monkeypatch, ): - fake_cache = _make_spend_counter_cache() - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + spend = _spend_counters(monkeypatch) pending = await ps._prepare_unreserved_spend_counter_increment( counter_key="spend:tag:x", @@ -1113,19 +937,17 @@ async def test_prepare_unreserved_spend_counter_increment_skips_reserved_keys( ) assert pending is None - assert fake_cache.redis_cache.async_increment.called is False + assert spend.commands == [] @pytest.mark.asyncio async def test_prepare_unreserved_spend_counter_increment_proceeds_when_not_reserved( monkeypatch, ): - fake_cache = _make_spend_counter_cache(redis_get_value=None) - fake_user_cache = _make_user_api_key_cache() - reseed = AsyncMock(return_value=None) - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) + spend = _spend_counters(monkeypatch) + monkeypatch.setattr(ps, "user_api_key_cache", _make_user_api_key_cache()) monkeypatch.setattr(ps, "prisma_client", None) + reseed = AsyncMock(return_value=None) monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", reseed) pending = await ps._prepare_unreserved_spend_counter_increment( @@ -1138,7 +960,7 @@ async def test_prepare_unreserved_spend_counter_increment_proceeds_when_not_rese assert pending is not None assert pending.counter_key == "spend:tag:y" assert pending.increment == 2.0 - assert fake_cache.redis_cache.async_get_cache.called is True + assert spend.commands == ["GET spend:tag:y"] assert reseed.called is True @@ -1149,10 +971,8 @@ async def test_prepare_unreserved_spend_counter_increment_proceeds_when_not_rese @pytest.mark.asyncio async def test_prepare_spend_counter_increment_warm_cache_skips_reseed(monkeypatch): - fake_cache = _make_spend_counter_cache(redis_get_value=11.0) - fake_user_cache = _make_user_api_key_cache() - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) + spend = _spend_counters(monkeypatch, redis={"spend:key:k": 11.0}) + monkeypatch.setattr(ps, "user_api_key_cache", _make_user_api_key_cache()) monkeypatch.setattr(ps, "prisma_client", None) reseed = AsyncMock(return_value=None) monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", reseed) @@ -1165,16 +985,9 @@ async def test_prepare_spend_counter_increment_warm_cache_skips_reseed(monkeypat assert pending.counter_key == "spend:key:k" assert pending.increment == 3.0 - observed = { - "reseed_called": reseed.called, - "increment_called": fake_cache.redis_cache.async_increment.called, - "in_memory_seeded_from_redis": fake_cache.in_memory_cache.set_cache.called, - } - assert normalize(observed) == { - "reseed_called": False, - "increment_called": False, - "in_memory_seeded_from_redis": True, - } + assert reseed.called is False + assert spend.commands == ["GET spend:key:k"] + assert spend.memory == {"spend:key:k": 11.0} # --------------------------------------------------------------------------- @@ -1186,8 +999,7 @@ async def test_prepare_spend_counter_increment_warm_cache_skips_reseed(monkeypat async def test_prepare_window_spend_counter_increment_returns_pending_when_initialized( monkeypatch, ): - fake_cache = _make_spend_counter_cache(redis_get_value=0.0) - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + _spend_counters(monkeypatch, redis_default=0.0) monkeypatch.setattr(ps, "prisma_client", None) monkeypatch.setattr( ps.SpendCounterReseed, @@ -1213,8 +1025,7 @@ async def test_prepare_window_spend_counter_increment_returns_pending_when_initi async def test_prepare_window_spend_counter_increment_missing_window_start_invalid_skips( monkeypatch, ): - fake_cache = _make_spend_counter_cache() - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + spend = _spend_counters(monkeypatch) pending = await ps._prepare_window_spend_counter_increment( counter_key="spend:key:k:window:1d", @@ -1226,7 +1037,7 @@ async def test_prepare_window_spend_counter_increment_missing_window_start_inval ) assert pending is None - assert fake_cache.redis_cache.async_increment.called is False + assert spend.commands == [] # --------------------------------------------------------------------------- @@ -1248,51 +1059,38 @@ async def test_apply_spend_counter_increments_open_breaker_invalidates_and_retur Re-raising the refusal sent every request through the cost callback's error path, which logged an ERROR and fired the failed-tracking alert once per request for the whole outage. """ - from litellm.caching.redis_cache import RedisCircuitBreakerOpenError - - fake_cache = _make_spend_counter_cache() - fake_cache.redis_cache.async_increment_pipeline = AsyncMock( - side_effect=RedisCircuitBreakerOpenError("Redis circuit breaker is open") + spend = _spend_counters( + monkeypatch, memory={"spend:key:k": 1.0, "spend:team:t": 1.0}, fail_pipelines=True, failure="breaker_open" ) - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) await ps._apply_spend_counter_increments(_two_pending_increments()) - deleted_keys = sorted(call.kwargs["key"] for call in fake_cache.in_memory_cache.delete_cache.call_args_list) - assert deleted_keys == ["spend:key:k", "spend:team:t"] - fake_cache.in_memory_cache.set_cache.assert_not_called() + assert spend.memory == {} + assert spend.commands == ["DEL spend:key:k", "DEL spend:team:t"] @pytest.mark.asyncio async def test_apply_spend_counter_increments_redis_timeout_invalidates_and_returns(monkeypatch): """A Redis timeout invalidates the counters and returns without reaching the cost callback's error path.""" - from redis.exceptions import TimeoutError as RedisTimeoutError - - fake_cache = _make_spend_counter_cache() - fake_cache.redis_cache.async_increment_pipeline = AsyncMock( - side_effect=RedisTimeoutError("Timeout reading from 127.0.0.1:6379") + spend = _spend_counters( + monkeypatch, memory={"spend:key:k": 1.0, "spend:team:t": 1.0}, fail_pipelines=True, failure="timeout" ) - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) await ps._apply_spend_counter_increments(_two_pending_increments()) - deleted_keys = sorted(call.kwargs["key"] for call in fake_cache.in_memory_cache.delete_cache.call_args_list) - assert deleted_keys == ["spend:key:k", "spend:team:t"] - fake_cache.in_memory_cache.set_cache.assert_not_called() + assert spend.memory == {} + assert spend.commands == ["DEL spend:key:k", "DEL spend:team:t"] @pytest.mark.asyncio async def test_apply_spend_counter_increments_other_redis_error_invalidates_and_raises(monkeypatch): - fake_cache = _make_spend_counter_cache() - fake_cache.redis_cache.async_increment_pipeline = AsyncMock(side_effect=ConnectionError("redis down")) - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + spend = _spend_counters(monkeypatch, memory={"spend:key:k": 1.0, "spend:team:t": 1.0}, fail_pipelines=True) - with pytest.raises(ConnectionError, match="redis down"): + with pytest.raises(NativeSpendCounterError, match="transport"): await ps._apply_spend_counter_increments(_two_pending_increments()) - deleted_keys = sorted(call.kwargs["key"] for call in fake_cache.in_memory_cache.delete_cache.call_args_list) - assert deleted_keys == ["spend:key:k", "spend:team:t"] - fake_cache.in_memory_cache.set_cache.assert_not_called() + assert spend.memory == {} + assert spend.commands == ["DEL spend:key:k", "DEL spend:team:t"] # --------------------------------------------------------------------------- @@ -1304,9 +1102,8 @@ async def test_apply_spend_counter_increments_other_redis_error_invalidates_and_ async def test_ensure_spend_counter_initialized_warm_skips_reseed_and_source( monkeypatch, ): - fake_cache = _make_spend_counter_cache(redis_get_value=20.0) + spend = _spend_counters(monkeypatch, redis={"spend:user:u": 20.0}) fake_user_cache = _make_user_api_key_cache() - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) reseed = AsyncMock(return_value=None) @@ -1317,25 +1114,17 @@ async def test_ensure_spend_counter_initialized_warm_skips_reseed_and_source( source_cache_key="u", ) - observed = { - "warm_check_redis": fake_cache.redis_cache.async_get_cache.called, - "reseed_called": reseed.called, - "source_cache_called": fake_user_cache.async_get_cache.called, - } - assert normalize(observed) == { - "warm_check_redis": True, - "reseed_called": False, - "source_cache_called": False, - } + assert spend.commands == ["GET spend:user:u"] + assert reseed.called is False + assert fake_user_cache.async_get_cache.called is False @pytest.mark.asyncio async def test_ensure_spend_counter_initialized_cold_seeds_from_source_cache( monkeypatch, ): - fake_cache = _make_spend_counter_cache(redis_get_value=None, redis_increment_value=7.0) + spend = _spend_counters(monkeypatch) fake_user_cache = _make_user_api_key_cache(get_value={"spend": 7.0}) - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) monkeypatch.setattr(ps, "prisma_client", None) monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) @@ -1345,16 +1134,9 @@ async def test_ensure_spend_counter_initialized_cold_seeds_from_source_cache( source_cache_key="u", ) - observed = { - "source_cache_called": fake_user_cache.async_get_cache.called, - "seed_increment_called": fake_cache.redis_cache.async_increment.called, - "warm_check_done": fake_cache.redis_cache.async_get_cache.called, - } - assert normalize(observed) == { - "source_cache_called": True, - "seed_increment_called": True, - "warm_check_done": True, - } + assert fake_user_cache.async_get_cache.called is True + assert spend.commands == ["GET spend:user:u", "INCRBYFLOAT spend:user:u 7.0"] + assert spend.redis == {"spend:user:u": 7.0} # --------------------------------------------------------------------------- @@ -1409,8 +1191,7 @@ async def test_get_source_cache_base_spend_no_hits_returns_zero_fallback(monkeyp @pytest.mark.asyncio async def test_ensure_window_spend_counter_initialized_warm_returns_true(monkeypatch): - fake_cache = _make_spend_counter_cache(redis_get_value=3.0) - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + spend = _spend_counters(monkeypatch, redis={"spend:key:k:window:1d": 3.0}) monkeypatch.setattr(ps, "prisma_client", None) window_reseed = AsyncMock(return_value=0.0) monkeypatch.setattr(ps.SpendCounterReseed, "coalesced_window", window_reseed) @@ -1423,24 +1204,16 @@ async def test_ensure_window_spend_counter_initialized_warm_returns_true(monkeyp window_start=datetime(2024, 1, 1), ) - observed = { - "initialized": initialized, - "reseed_called": window_reseed.called, - "redis_get_called": fake_cache.redis_cache.async_get_cache.called, - } - assert normalize(observed) == { - "initialized": True, - "reseed_called": False, - "redis_get_called": True, - } + assert initialized is True + assert window_reseed.called is False + assert spend.commands == ["GET spend:key:k:window:1d"] @pytest.mark.asyncio async def test_ensure_window_spend_counter_initialized_db_failure_invalid_returns_false( monkeypatch, ): - fake_cache = _make_spend_counter_cache(redis_get_value=None) - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + _spend_counters(monkeypatch) monkeypatch.setattr(ps, "prisma_client", None) monkeypatch.setattr( ps.SpendCounterReseed, @@ -1466,32 +1239,20 @@ async def test_ensure_window_spend_counter_initialized_db_failure_invalid_return @pytest.mark.asyncio async def test_is_spend_counter_cache_warm_redis_hit_seeds_in_memory(monkeypatch): - fake_cache = _make_spend_counter_cache(redis_get_value=99.0) - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + spend = _spend_counters(monkeypatch, redis={"spend:user:u": 99.0}) result = await ps._is_spend_counter_cache_warm(counter_key="spend:user:u") - observed = { - "result": result, - "redis_get_called": fake_cache.redis_cache.async_get_cache.called, - "in_memory_set_called": fake_cache.in_memory_cache.set_cache.called, - } - assert normalize(observed) == { - "result": True, - "redis_get_called": True, - "in_memory_set_called": True, - } + assert result is True + assert spend.commands == ["GET spend:user:u"] + assert spend.memory == {"spend:user:u": 99.0} @pytest.mark.asyncio async def test_is_spend_counter_cache_warm_redis_error_falls_back_to_in_memory( monkeypatch, ): - fake_cache = _make_spend_counter_cache( - redis_get_side_effect=RuntimeError("redis err"), - in_memory_value=None, - ) - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + _spend_counters(monkeypatch, fail=True) result = await ps._is_spend_counter_cache_warm(counter_key="spend:user:u") assert result is False @@ -1504,35 +1265,25 @@ async def test_is_spend_counter_cache_warm_redis_error_falls_back_to_in_memory( @pytest.mark.asyncio async def test_increment_spend_counter_cache_redis_path_returns_new_value(monkeypatch): - fake_cache = _make_spend_counter_cache(redis_increment_value=44.0) - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + spend = _spend_counters(monkeypatch, redis={"spend:key:k": 40.0}) result = await ps._increment_spend_counter_cache(counter_key="spend:key:k", increment=4.0) - observed = { - "result": result, - "redis_increment_called": fake_cache.redis_cache.async_increment.called, - "in_memory_set_called": fake_cache.in_memory_cache.set_cache.called, - } - assert normalize(observed) == { - "result": 44.0, - "redis_increment_called": True, - "in_memory_set_called": True, - } + assert result == 44.0 + assert spend.commands == ["INCRBYFLOAT spend:key:k 4.0"] + assert spend.memory == {"spend:key:k": 44.0} @pytest.mark.asyncio async def test_increment_spend_counter_cache_redis_error_raises_and_invalidates( monkeypatch, ): - fake_cache = _make_spend_counter_cache(redis_increment_side_effect=RuntimeError("incr fail")) - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + spend = _spend_counters(monkeypatch, memory={"spend:key:k": 3.0}, fail=True) - with pytest.raises(RuntimeError): + with pytest.raises(NativeSpendCounterError): await ps._increment_spend_counter_cache(counter_key="spend:key:k", increment=1.0) - assert fake_cache.in_memory_cache.delete_cache.called is True - assert fake_cache.redis_cache.async_delete_cache.called is True + assert spend.memory == {} # --------------------------------------------------------------------------- @@ -1542,32 +1293,22 @@ async def test_increment_spend_counter_cache_redis_error_raises_and_invalidates( @pytest.mark.asyncio async def test_invalidate_spend_counter_deletes_in_memory_and_redis(monkeypatch): - fake_cache = _make_spend_counter_cache() - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + spend = _spend_counters(monkeypatch, redis={"spend:key:k": 3.0}, memory={"spend:key:k": 3.0}) await ps._invalidate_spend_counter(counter_key="spend:key:k") - observed = { - "in_memory_delete_called": fake_cache.in_memory_cache.delete_cache.called, - "redis_delete_called": fake_cache.redis_cache.async_delete_cache.called, - "delete_args_key": fake_cache.redis_cache.async_delete_cache.call_args.kwargs["key"], - } - assert normalize(observed) == { - "in_memory_delete_called": True, - "redis_delete_called": True, - "delete_args_key": "spend:key:k", - } + assert spend.commands == ["DEL spend:key:k"] + assert spend.redis == {} + assert spend.memory == {} @pytest.mark.asyncio async def test_invalidate_spend_counter_swallows_redis_failure_no_raise(monkeypatch): - fake_cache = _make_spend_counter_cache() - fake_cache.redis_cache.async_delete_cache = AsyncMock(side_effect=RuntimeError("redis down")) - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + spend = _spend_counters(monkeypatch, memory={"spend:key:k": 3.0}, fail=True) await ps._invalidate_spend_counter(counter_key="spend:key:k") - assert fake_cache.in_memory_cache.delete_cache.called is True + assert spend.memory == {} # --------------------------------------------------------------------------- @@ -2013,3 +1754,16 @@ async def test_update_cache_still_adds_spend_to_the_others_when_one_cached_objec 1.25, 10.25, ) + + +@pytest.mark.asyncio +async def test_increment_spend_counters_timeout_invalidates_all_counters_without_failing( + monkeypatch, +): + """A pipeline that timed out (a busy pool or a slow node) invalidates the counters and the request still + succeeds.""" + spend = _warm_counters_behind_a_failing_pipeline(monkeypatch, failure="timeout") + + await _increment_every_spend_counter() + + _assert_every_pending_counter_invalidated(spend) diff --git a/tests/unit/proxy/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py index a87ac7fa7c0..75983abf586 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py @@ -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 ) diff --git a/tests/unit/proxy/spend_tracking/test_budget_reservation_redis_failure.py b/tests/unit/proxy/spend_tracking/test_budget_reservation_redis_failure.py index e0a74d50a6c..918adf25a4d 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation_redis_failure.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation_redis_failure.py @@ -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 == {} diff --git a/tests/unit/proxy/spend_tracking/test_cost_finalization.py b/tests/unit/proxy/spend_tracking/test_cost_finalization.py index ef70a193280..2c342abc028 100644 --- a/tests/unit/proxy/spend_tracking/test_cost_finalization.py +++ b/tests/unit/proxy/spend_tracking/test_cost_finalization.py @@ -11,6 +11,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" @@ -298,8 +299,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 @@ -343,8 +344,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 @@ -481,11 +482,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, @@ -523,7 +524,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 @@ -538,12 +539,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 diff --git a/tests/unit/proxy/spend_tracking/test_spend_counter_batch.py b/tests/unit/proxy/spend_tracking/test_spend_counter_batch.py index 8ba79a45ff3..7fd3b1dd542 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_counter_batch.py +++ b/tests/unit/proxy/spend_tracking/test_spend_counter_batch.py @@ -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) diff --git a/tests/unit/proxy/test_budget_reservation.py b/tests/unit/proxy/test_budget_reservation.py index 6ce8d400daf..c7ecd914197 100644 --- a/tests/unit/proxy/test_budget_reservation.py +++ b/tests/unit/proxy/test_budget_reservation.py @@ -37,7 +37,6 @@ from litellm.proxy.common_utils.reset_budget_job import _model_access_group_coun from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, model_access_group_cache_key, - model_access_group_spend_counter_key, ) from litellm.proxy.spend_tracking.budget_reservation import ( _get_model_access_group_budget_counters, @@ -55,26 +54,28 @@ from litellm.proxy.spend_tracking.input_tokens import ( from litellm.proxy.utils import ProxyLogging from litellm.router import Router from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget +from litellm.redis.stores.spend_counters import model_access_group_counter +from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters, unavailable @pytest.fixture() def spend_counter_state(): import litellm.proxy.proxy_server as ps - original_counter_cache = ps.spend_counter_cache + original_counter_store = ps.spend_counter_store original_key_cache = ps.user_api_key_cache original_prisma_client = ps.prisma_client - counter_cache = DualCache() + counter_cache = FakeSpendCounters(redis=None) key_cache = UserApiKeyCache() - ps.spend_counter_cache = counter_cache + ps.spend_counter_store = counter_cache.store() ps.user_api_key_cache = key_cache ps.prisma_client = None try: yield counter_cache, key_cache finally: - ps.spend_counter_cache = original_counter_cache + ps.spend_counter_store = original_counter_store ps.user_api_key_cache = original_key_cache ps.prisma_client = original_prisma_client @@ -124,16 +125,11 @@ async def test_reservation_still_protects_under_budget_throttled_key( reservation = await _reserve(valid_token, 0.6, key_cache, proxy_logging_obj) assert reservation is not None - assert ( - counter_cache.in_memory_cache.get_cache(key="spend:key:key-throttle-under") - == 0.6 - ) + assert counter_cache.get_local("spend:key:key-throttle-under") == 0.6 @pytest.mark.asyncio -async def test_reservation_does_not_block_over_budget_throttled_key( - spend_counter_state, monkeypatch -): +async def test_reservation_does_not_block_over_budget_throttled_key(spend_counter_state, monkeypatch): """Once an opted-in key is over budget the reservation path must not raise; the rate limiter throttles it instead.""" monkeypatch.setattr(litellm, "budget_exceeded_throttle_percentage", 0.1) @@ -156,16 +152,11 @@ async def test_reservation_does_not_block_over_budget_throttled_key( # increment is released so the counter is not permanently inflated result = await _reserve(valid_token, 0.6, key_cache, proxy_logging_obj) assert result is None - assert ( - counter_cache.in_memory_cache.get_cache(key="spend:key:key-throttle-over") - == 0.6 - ) + assert counter_cache.get_local("spend:key:key-throttle-over") == 0.6 @pytest.mark.asyncio -async def test_reservation_blocks_over_budget_non_throttled_key( - spend_counter_state, monkeypatch -): +async def test_reservation_blocks_over_budget_non_throttled_key(spend_counter_state, monkeypatch): monkeypatch.setattr(litellm, "budget_exceeded_throttle_percentage", 0.1) counter_cache, key_cache = spend_counter_state proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) @@ -257,10 +248,7 @@ async def test_should_shrink_second_key_reservation_to_remaining_budget( proxy_logging_obj=proxy_logging_obj, ) assert reservation is not None - assert ( - counter_cache.in_memory_cache.get_cache(key="spend:key:key-budget-race") - == 0.6 - ) + assert counter_cache.get_local("spend:key:key-budget-race") == 0.6 second_reservation = await reserve_budget_for_request( request_body=_request_body(), @@ -275,9 +263,7 @@ async def test_should_shrink_second_key_reservation_to_remaining_budget( ) assert second_reservation is not None assert second_reservation["reserved_cost"] == pytest.approx(0.4) - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-budget-race" - ) == pytest.approx(1.0) + assert counter_cache.get_local("spend:key:key-budget-race") == pytest.approx(1.0) with pytest.raises(litellm.BudgetExceededError): await reserve_budget_for_request( @@ -292,14 +278,10 @@ async def test_should_shrink_second_key_reservation_to_remaining_budget( proxy_logging_obj=proxy_logging_obj, ) - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-budget-race" - ) == pytest.approx(1.0) + assert counter_cache.get_local("spend:key:key-budget-race") == pytest.approx(1.0) await release_budget_reservation(second_reservation) - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-budget-race" - ) == pytest.approx(0.6) + assert counter_cache.get_local("spend:key:key-budget-race") == pytest.approx(0.6) await release_budget_reservation(reservation) @@ -337,9 +319,7 @@ async def test_should_shrink_second_end_user_reservation_to_remaining_budget( end_user_object=end_user_object, ) assert reservation is not None - assert counter_cache.in_memory_cache.get_cache( - key="spend:end_user:end-user-budget-race" - ) == pytest.approx(0.6) + assert counter_cache.get_local("spend:end_user:end-user-budget-race") == pytest.approx(0.6) second_reservation = await reserve_budget_for_request( request_body=_request_body(), @@ -355,9 +335,7 @@ async def test_should_shrink_second_end_user_reservation_to_remaining_budget( ) assert second_reservation is not None assert second_reservation["reserved_cost"] == pytest.approx(0.4) - assert counter_cache.in_memory_cache.get_cache( - key="spend:end_user:end-user-budget-race" - ) == pytest.approx(1.0) + assert counter_cache.get_local("spend:end_user:end-user-budget-race") == pytest.approx(1.0) with pytest.raises(litellm.BudgetExceededError): await reserve_budget_for_request( @@ -373,14 +351,10 @@ async def test_should_shrink_second_end_user_reservation_to_remaining_budget( end_user_object=end_user_object, ) - assert counter_cache.in_memory_cache.get_cache( - key="spend:end_user:end-user-budget-race" - ) == pytest.approx(1.0) + assert counter_cache.get_local("spend:end_user:end-user-budget-race") == pytest.approx(1.0) await release_budget_reservation(second_reservation) - assert counter_cache.in_memory_cache.get_cache( - key="spend:end_user:end-user-budget-race" - ) == pytest.approx(0.6) + assert counter_cache.get_local("spend:end_user:end-user-budget-race") == pytest.approx(0.6) from litellm.proxy.proxy_server import increment_spend_counters @@ -393,9 +367,7 @@ async def test_should_shrink_second_end_user_reservation_to_remaining_budget( end_user_id="end-user-budget-race", ) - assert counter_cache.in_memory_cache.get_cache( - key="spend:end_user:end-user-budget-race" - ) == pytest.approx(0.2) + assert counter_cache.get_local("spend:end_user:end-user-budget-race") == pytest.approx(0.2) @pytest.mark.asyncio @@ -453,13 +425,8 @@ async def test_should_shrink_second_tag_reservation_to_remaining_budget( "applied_adjustment": 0.0, } ] - assert counter_cache.in_memory_cache.get_cache( - key="spend:tag:tag-budget-race" - ) == pytest.approx(0.6) - assert ( - counter_cache.in_memory_cache.get_cache(key="spend:tag:tag-without-budget") - is None - ) + assert counter_cache.get_local("spend:tag:tag-budget-race") == pytest.approx(0.6) + assert counter_cache.get_local("spend:tag:tag-without-budget") is None second_reservation = await reserve_budget_for_request( request_body=request_body, @@ -474,9 +441,7 @@ async def test_should_shrink_second_tag_reservation_to_remaining_budget( ) assert second_reservation is not None assert second_reservation["reserved_cost"] == pytest.approx(0.4) - assert counter_cache.in_memory_cache.get_cache( - key="spend:tag:tag-budget-race" - ) == pytest.approx(1.0) + assert counter_cache.get_local("spend:tag:tag-budget-race") == pytest.approx(1.0) with pytest.raises(litellm.BudgetExceededError): await reserve_budget_for_request( @@ -491,14 +456,10 @@ async def test_should_shrink_second_tag_reservation_to_remaining_budget( proxy_logging_obj=proxy_logging_obj, ) - assert counter_cache.in_memory_cache.get_cache( - key="spend:tag:tag-budget-race" - ) == pytest.approx(1.0) + assert counter_cache.get_local("spend:tag:tag-budget-race") == pytest.approx(1.0) await release_budget_reservation(second_reservation) - assert counter_cache.in_memory_cache.get_cache( - key="spend:tag:tag-budget-race" - ) == pytest.approx(0.6) + assert counter_cache.get_local("spend:tag:tag-budget-race") == pytest.approx(0.6) from litellm.proxy.proxy_server import increment_spend_counters @@ -511,9 +472,7 @@ async def test_should_shrink_second_tag_reservation_to_remaining_budget( tags=["tag-budget-race"], ) - assert counter_cache.in_memory_cache.get_cache( - key="spend:tag:tag-budget-race" - ) == pytest.approx(0.2) + assert counter_cache.get_local("spend:tag:tag-budget-race") == pytest.approx(0.2) @pytest.mark.asyncio @@ -556,15 +515,9 @@ async def test_should_seed_and_update_end_user_and_tag_counters_without_reservat tags=["paid-tag", "paid-tag", "other-tag", ""], ) - assert counter_cache.in_memory_cache.get_cache( - key="spend:end_user:customer-1" - ) == pytest.approx(4.50) - assert counter_cache.in_memory_cache.get_cache( - key="spend:tag:paid-tag" - ) == pytest.approx(7.50) - assert counter_cache.in_memory_cache.get_cache( - key="spend:tag:other-tag" - ) == pytest.approx(2.50) + assert counter_cache.get_local("spend:end_user:customer-1") == pytest.approx(4.50) + assert counter_cache.get_local("spend:tag:paid-tag") == pytest.approx(7.50) + assert counter_cache.get_local("spend:tag:other-tag") == pytest.approx(2.50) @pytest.mark.asyncio @@ -627,12 +580,8 @@ async def test_should_reserve_team_member_and_org_budget_counters(spend_counter_ proxy_logging_obj=proxy_logging_obj, ) - assert counter_cache.in_memory_cache.get_cache( - key="spend:team_member:user-budget-shared:team-budget-shared" - ) == pytest.approx(0.4) - assert counter_cache.in_memory_cache.get_cache( - key="spend:org:org-budget-shared" - ) == pytest.approx(0.4) + assert counter_cache.get_local("spend:team_member:user-budget-shared:team-budget-shared") == pytest.approx(0.4) + assert counter_cache.get_local("spend:org:org-budget-shared") == pytest.approx(0.4) await release_budget_reservation(reservation) @@ -703,10 +652,8 @@ async def test_should_reserve_project_and_team_member_counters_for_project_scope ) assert reservation is not None - assert counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-proj:team-proj") == pytest.approx( - 0.1 + estimated - ) - assert counter_cache.in_memory_cache.get_cache(key="spend:project:proj-1") == pytest.approx(0.2 + estimated) + assert counter_cache.get_local("spend:team_member:user-proj:team-proj") == pytest.approx(0.1 + estimated) + assert counter_cache.get_local("spend:project:proj-1") == pytest.approx(0.2 + estimated) from litellm.proxy.proxy_server import increment_spend_counters @@ -719,8 +666,8 @@ async def test_should_reserve_project_and_team_member_counters_for_project_scope project_id="proj-1", ) - assert counter_cache.in_memory_cache.get_cache(key="spend:project:proj-1") == pytest.approx(0.25) - assert counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-proj:team-proj") == pytest.approx(0.15) + assert counter_cache.get_local("spend:project:proj-1") == pytest.approx(0.25) + assert counter_cache.get_local("spend:team_member:user-proj:team-proj") == pytest.approx(0.15) @pytest.mark.asyncio @@ -749,7 +696,7 @@ async def test_exhausted_team_member_budget_still_blocks_project_scoped_key(spen ) assert "TeamMember=user-proj:team-proj" in str(exc_info.value) - assert counter_cache.in_memory_cache.get_cache(key="spend:project:proj-1") in (None, pytest.approx(0.0)) + assert counter_cache.get_local("spend:project:proj-1") in (None, pytest.approx(0.0)) @pytest.mark.asyncio @@ -779,7 +726,7 @@ async def test_exhausted_project_budget_blocks_project_scoped_key(spend_counter_ assert "Project=proj-1" in str(exc_info.value) assert exc_info.value.entity_type == Litellm_EntityType.PROJECT.value - assert counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-proj:team-proj") in ( + assert counter_cache.get_local("spend:team_member:user-proj:team-proj") in ( None, pytest.approx(0.0), ) @@ -820,7 +767,7 @@ async def test_should_not_reserve_user_budget_counter_for_team_key(spend_counter proxy_logging_obj=proxy_logging_obj, ) - assert counter_cache.in_memory_cache.get_cache(key="spend:user:user-on-team") is None + assert counter_cache.get_local("spend:user:user-on-team") is None await release_budget_reservation(reservation) @@ -861,7 +808,7 @@ async def test_should_reserve_user_budget_counter_for_team_key_when_flag_enabled apply_user_budget_to_team_keys=True, ) - assert counter_cache.in_memory_cache.get_cache(key="spend:user:user-on-team-flagged") == pytest.approx(0.3) + assert counter_cache.get_local("spend:user:user-on-team-flagged") == pytest.approx(0.3) await release_budget_reservation(reservation) @@ -893,9 +840,7 @@ async def test_should_seed_org_counter_from_with_budget_cache(spend_counter_stat response_cost=0.25, ) - assert counter_cache.in_memory_cache.get_cache( - key="spend:org:org-counter-with-budget" - ) == pytest.approx(2.25) + assert counter_cache.get_local("spend:org:org-counter-with-budget") == pytest.approx(2.25) @pytest.mark.asyncio @@ -924,9 +869,7 @@ async def test_should_seed_org_counter_from_plain_org_cache(spend_counter_state) response_cost=0.25, ) - assert counter_cache.in_memory_cache.get_cache( - key="spend:org:org-counter-plain" - ) == pytest.approx(2.25) + assert counter_cache.get_local("spend:org:org-counter-plain") == pytest.approx(2.25) @pytest.mark.asyncio @@ -940,10 +883,7 @@ async def test_should_cap_known_estimate_to_remaining_budget( spend=0.9, max_budget=1.0, ) - counter_cache.in_memory_cache.set_cache( - key="spend:key:key-budget-known-estimate-cap", - value=0.9, - ) + counter_cache.set_local("spend:key:key-budget-known-estimate-cap", 0.9) with patch( "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", @@ -963,14 +903,10 @@ async def test_should_cap_known_estimate_to_remaining_budget( assert reservation is not None assert reservation["reserved_cost"] == pytest.approx(0.1) - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-budget-known-estimate-cap" - ) == pytest.approx(1.0) + assert counter_cache.get_local("spend:key:key-budget-known-estimate-cap") == pytest.approx(1.0) await release_budget_reservation(reservation) - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-budget-known-estimate-cap" - ) == pytest.approx(0.9) + assert counter_cache.get_local("spend:key:key-budget-known-estimate-cap") == pytest.approx(0.9) @pytest.mark.asyncio @@ -988,10 +924,7 @@ async def test_fail_closed_rejects_known_estimate_exceeding_remaining_budget( spend=0.9, max_budget=1.0, ) - counter_cache.in_memory_cache.set_cache( - key="spend:key:key-budget-known-estimate-fail-closed", - value=0.9, - ) + counter_cache.set_local("spend:key:key-budget-known-estimate-fail-closed", 0.9) with patch( # test-quality-ok: reserve_budget_for_request takes no estimator, so pinning the estimate needs this attribute "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", @@ -1014,9 +947,7 @@ async def test_fail_closed_rejects_known_estimate_exceeding_remaining_budget( assert exc_info.value.current_cost == pytest.approx(0.9) assert exc_info.value.max_budget == pytest.approx(1.0) assert "Current cost: 0.9, Estimated request cost: 0.6, Max budget: 1.0" in str(exc_info.value) - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-budget-known-estimate-fail-closed" - ) == pytest.approx(0.9) + assert counter_cache.get_local("spend:key:key-budget-known-estimate-fail-closed") == pytest.approx(0.9) @pytest.mark.asyncio @@ -1032,10 +963,7 @@ async def test_fail_closed_tolerates_float_noise_when_estimate_exactly_fits( spend=0.1, max_budget=0.3, ) - counter_cache.in_memory_cache.set_cache( - key="spend:key:key-budget-fail-closed-float-noise", - value=0.1, - ) + counter_cache.set_local("spend:key:key-budget-fail-closed-float-noise", 0.1) with patch( # test-quality-ok: reserve_budget_for_request takes no estimator, so pinning the estimate needs this attribute "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", @@ -1056,9 +984,7 @@ async def test_fail_closed_tolerates_float_noise_when_estimate_exactly_fits( assert reservation is not None assert reservation["reserved_cost"] == pytest.approx(0.2) - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-budget-fail-closed-float-noise" - ) == pytest.approx(0.3) + assert counter_cache.get_local("spend:key:key-budget-fail-closed-float-noise") == pytest.approx(0.3) @pytest.mark.asyncio @@ -1699,10 +1625,7 @@ async def test_should_skip_budget_window_with_unparseable_duration( } ], ) - counter_cache.in_memory_cache.set_cache( - key="spend:key:key-budget-malformed-window", - value=0.9, - ) + counter_cache.set_local("spend:key:key-budget-malformed-window", 0.9) with patch( "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", @@ -1721,23 +1644,12 @@ async def test_should_skip_budget_window_with_unparseable_duration( ) assert reservation is not None - assert [entry["counter_key"] for entry in reservation["entries"]] == [ - "spend:key:key-budget-malformed-window" - ] - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-budget-malformed-window" - ) == pytest.approx(1.1) - assert ( - counter_cache.in_memory_cache.get_cache( - key="spend:key:key-budget-malformed-window:window:not-a-duration" - ) - is None - ) + assert [entry["counter_key"] for entry in reservation["entries"]] == ["spend:key:key-budget-malformed-window"] + assert counter_cache.get_local("spend:key:key-budget-malformed-window") == pytest.approx(1.1) + assert counter_cache.get_local("spend:key:key-budget-malformed-window:window:not-a-duration") is None await release_budget_reservation(reservation) - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-budget-malformed-window" - ) == pytest.approx(0.9) + assert counter_cache.get_local("spend:key:key-budget-malformed-window") == pytest.approx(0.9) @pytest.mark.asyncio @@ -1773,12 +1685,7 @@ async def test_should_skip_window_reservation_when_db_baseline_unavailable( ) assert reservation is None - assert ( - counter_cache.in_memory_cache.get_cache( - key="spend:key:key-budget-window-db-unavailable:window:1h" - ) - is None - ) + assert counter_cache.get_local("spend:key:key-budget-window-db-unavailable:window:1h") is None @pytest.mark.asyncio @@ -1794,10 +1701,10 @@ async def test_should_skip_reservation_when_counter_increment_fails( max_budget=1.0, ) - async def fail_increment_cache(*args, **kwargs): + def fail_increment_local(counter: str, amount: float) -> float: raise RuntimeError("counter unavailable") - monkeypatch.setattr(counter_cache, "async_increment_cache", fail_increment_cache) + monkeypatch.setattr(counter_cache, "increment_local", fail_increment_local) with ( patch( @@ -1822,12 +1729,7 @@ async def test_should_skip_reservation_when_counter_increment_fails( assert reservation is None assert mock_warning.call_count >= 1 - assert ( - counter_cache.in_memory_cache.get_cache( - key="spend:key:key-budget-reserve-unavailable" - ) - is None - ) + assert counter_cache.get_local("spend:key:key-budget-reserve-unavailable") is None @pytest.mark.asyncio @@ -1845,10 +1747,10 @@ async def test_should_raise_503_when_counter_increment_fails_and_fail_closed( max_budget=1.0, ) - async def fail_increment_cache(*args, **kwargs): + def fail_increment_local(counter: str, amount: float) -> float: raise RuntimeError("counter unavailable") - monkeypatch.setattr(counter_cache, "async_increment_cache", fail_increment_cache) + monkeypatch.setattr(counter_cache, "increment_local", fail_increment_local) with patch( "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", @@ -1869,12 +1771,7 @@ async def test_should_raise_503_when_counter_increment_fails_and_fail_closed( ) assert exc_info.value.status_code == 503 - assert ( - counter_cache.in_memory_cache.get_cache( - key="spend:key:key-budget-reserve-fail-closed" - ) - is None - ) + assert counter_cache.get_local("spend:key:key-budget-reserve-fail-closed") is None @pytest.mark.asyncio @@ -1916,8 +1813,8 @@ async def test_fail_closed_releases_earlier_counters_before_503( ) assert exc_info.value.status_code == 503 - assert counter_cache.in_memory_cache.get_cache(key="spend:key:key-budget-fail-closed-release") is None - assert counter_cache.in_memory_cache.get_cache(key="spend:key:key-budget-fail-closed-release:window:1h") is None + assert counter_cache.get_local("spend:key:key-budget-fail-closed-release") is None + assert counter_cache.get_local("spend:key:key-budget-fail-closed-release:window:1h") is None @pytest.mark.asyncio @@ -1959,12 +1856,7 @@ async def test_should_skip_reservation_when_counter_initialization_fails( assert reservation is None assert mock_warning.call_count >= 1 - assert ( - counter_cache.in_memory_cache.get_cache( - key="spend:key:key-budget-reserve-init-unavailable" - ) - is None - ) + assert counter_cache.get_local("spend:key:key-budget-reserve-init-unavailable") is None @pytest.mark.asyncio @@ -1981,7 +1873,7 @@ async def test_should_release_tracked_entry_when_reservation_fails_after_increme async def fail_after_increment(pending): for item in pending: - await counter_cache.async_increment_cache(key=item.counter_key, value=item.increment) + counter_cache.increment_local(item.counter_key, item.increment) raise RuntimeError("lost increment response") with ( @@ -2011,9 +1903,7 @@ async def test_should_release_tracked_entry_when_reservation_fails_after_increme ) assert reservation is None - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-budget-reserve-after-increment-failure" - ) == pytest.approx(0.0) + assert counter_cache.get_local("spend:key:key-budget-reserve-after-increment-failure") == pytest.approx(0.0) @pytest.mark.asyncio @@ -2054,12 +1944,8 @@ async def test_should_reconcile_reserved_counter_to_actual_spend( budget_reservation=reservation, ) - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-budget-reconcile" - ) == pytest.approx(0.2) - assert counter_cache.in_memory_cache.get_cache( - key="spend:team:team-without-budget" - ) == pytest.approx(0.2) + assert counter_cache.get_local("spend:key:key-budget-reconcile") == pytest.approx(0.2) + assert counter_cache.get_local("spend:team:team-without-budget") == pytest.approx(0.2) @pytest.mark.asyncio @@ -2091,9 +1977,7 @@ async def test_should_release_reservation_on_failure(spend_counter_state): await release_budget_reservation(reservation) await release_budget_reservation(reservation) - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-budget-release" - ) == pytest.approx(0.0) + assert counter_cache.get_local("spend:key:key-budget-release") == pytest.approx(0.0) @pytest.mark.asyncio @@ -2131,40 +2015,28 @@ async def test_should_retry_partial_release_without_double_decrement( proxy_logging_obj=proxy_logging_obj, ) - original_increment_cache = counter_cache.async_increment_cache + original_increment_local = counter_cache.increment_local fail_next_team_release = True - async def flaky_increment_cache(key, value, *args, **kwargs): + def flaky_increment_local(counter: str, amount: float) -> float: nonlocal fail_next_team_release - if ( - key == "spend:team:team-budget-partial-release" - and value < 0 - and fail_next_team_release - ): + if counter == "spend:team:team-budget-partial-release" and amount < 0 and fail_next_team_release: fail_next_team_release = False raise RuntimeError("simulated counter failure") - return await original_increment_cache(key=key, value=value, *args, **kwargs) + return original_increment_local(counter, amount) - monkeypatch.setattr(counter_cache, "async_increment_cache", flaky_increment_cache) + monkeypatch.setattr(counter_cache, "increment_local", flaky_increment_local) with pytest.raises(RuntimeError): await release_budget_reservation(reservation) - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-budget-partial-release" - ) == pytest.approx(0.0) - assert counter_cache.in_memory_cache.get_cache( - key="spend:team:team-budget-partial-release" - ) == pytest.approx(0.4) + assert counter_cache.get_local("spend:key:key-budget-partial-release") == pytest.approx(0.0) + assert counter_cache.get_local("spend:team:team-budget-partial-release") == pytest.approx(0.4) await release_budget_reservation(reservation) - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-budget-partial-release" - ) == pytest.approx(0.0) - assert counter_cache.in_memory_cache.get_cache( - key="spend:team:team-budget-partial-release" - ) == pytest.approx(0.0) + assert counter_cache.get_local("spend:key:key-budget-partial-release") == pytest.approx(0.0) + assert counter_cache.get_local("spend:team:team-budget-partial-release") == pytest.approx(0.0) @pytest.mark.asyncio @@ -2190,18 +2062,17 @@ async def test_should_preserve_budget_error_and_continue_partial_cleanup( value=team_object, ) - original_increment_cache = counter_cache.async_increment_cache + original_increment_local = counter_cache.increment_local fail_key_cleanup = True - async def flaky_increment_cache(key, value, *args, **kwargs): + def flaky_increment_local(counter: str, amount: float) -> float: nonlocal fail_key_cleanup - if key == "spend:key:key-budget-cleanup-failure" and value < 0: - if fail_key_cleanup: - fail_key_cleanup = False - raise RuntimeError("simulated cleanup failure") - return await original_increment_cache(key=key, value=value, *args, **kwargs) + if counter == "spend:key:key-budget-cleanup-failure" and amount < 0 and fail_key_cleanup: + fail_key_cleanup = False + raise RuntimeError("simulated cleanup failure") + return original_increment_local(counter, amount) - monkeypatch.setattr(counter_cache, "async_increment_cache", flaky_increment_cache) + monkeypatch.setattr(counter_cache, "increment_local", flaky_increment_local) with ( patch( @@ -2225,15 +2096,8 @@ async def test_should_preserve_budget_error_and_continue_partial_cleanup( proxy_logging_obj=proxy_logging_obj, ) - assert ( - counter_cache.in_memory_cache.get_cache( - key="spend:key:key-budget-cleanup-failure" - ) - is None - ) - assert counter_cache.in_memory_cache.get_cache( - key="spend:team:team-budget-cleanup-failure" - ) == pytest.approx(0.3) + assert counter_cache.get_local("spend:key:key-budget-cleanup-failure") is None + assert counter_cache.get_local("spend:team:team-budget-cleanup-failure") == pytest.approx(0.3) mock_log_exception.assert_called() @@ -2262,12 +2126,7 @@ async def test_release_missing_counter_reseeds_from_db_instead_of_failing( await release_budget_reservation(reservation) # counter not driven negative / not corrupted; left absent (no DB to reseed) - assert ( - counter_cache.in_memory_cache.get_cache( - key="spend:key:key-budget-missing-release" - ) - is None - ) + assert counter_cache.get_local("spend:key:key-budget-missing-release") is None assert reservation["finalized"] is True @@ -2279,10 +2138,7 @@ async def test_release_underflow_counter_reseeds_from_db(spend_counter_state): import litellm.proxy.proxy_server as ps counter_cache, _ = spend_counter_state - await counter_cache.async_increment_cache( - key="spend:key:key-budget-underflow-release", - value=0.1, - ) + counter_cache.increment_local("spend:key:key-budget-underflow-release", 0.1) reservation = { "reserved_cost": 0.4, "entries": [ @@ -2299,103 +2155,77 @@ async def test_release_underflow_counter_reseeds_from_db(spend_counter_state): await release_budget_reservation(reservation) # counter reseeded up to the authoritative DB value, not deleted or negated - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-budget-underflow-release" - ) == pytest.approx(0.25) + assert counter_cache.get_local("spend:key:key-budget-underflow-release") == pytest.approx(0.25) assert reservation["finalized"] is True -@pytest.mark.asyncio -async def test_release_non_numeric_counter_reseeds_from_db(spend_counter_state): - """A non-numeric counter value (corrupt/stale) during release is recovered by - reseeding from the DB, not by deleting the counter and raising.""" - import litellm.proxy.proxy_server as ps +class _ExpiringRedis(dict[str, object]): + """A Redis keyspace with real wall-clock key expiry; every write gives the key a fresh TTL.""" - counter_cache, _ = spend_counter_state - counter_cache.in_memory_cache.set_cache( - key="spend:key:key-budget-nonnumeric-release", - value="stale", - ) - reservation = { - "reserved_cost": 0.4, - "entries": [ - { - "counter_key": "spend:key:key-budget-nonnumeric-release", - "reserved_cost": 0.4, - "applied_adjustment": 0.0, - } - ], - "finalized": False, - } - - with patch.object(ps.SpendCounterReseed, "from_db", AsyncMock(return_value=0.5)): - await release_budget_reservation(reservation) - - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-budget-nonnumeric-release" - ) == pytest.approx(0.5) - assert reservation["finalized"] is True - - -class _ExpiringRedisCache: - """In-memory stand-in for RedisCache with real wall-clock key expiry.""" - - def __init__(self, default_ttl: float = 60.0, fail_first_refresh: bool = False) -> None: - self.default_ttl = default_ttl - self.store: dict[str, float] = {} + def __init__(self, ttl: float) -> None: + super().__init__() + self.ttl = ttl self.expires_at: dict[str, float] = {} + + def _evict_expired(self, key: str) -> None: + if self.expires_at.get(key, float("inf")) <= time.monotonic(): + super().pop(key, None) + self.expires_at.pop(key, None) + + def __contains__(self, key: object) -> bool: + if isinstance(key, str): + self._evict_expired(key) + return super().__contains__(key) + + def __getitem__(self, key: str) -> object: + self._evict_expired(key) + return super().__getitem__(key) + + def get(self, key: str, default: object = None) -> object: # pyright: ignore[reportIncompatibleMethodOverride] # dict.get overloads + self._evict_expired(key) + return super().get(key, default) + + def __setitem__(self, key: str, value: object) -> None: + super().__setitem__(key, value) + self.expires_at[key] = time.monotonic() + self.ttl + + def setdefault(self, key: str, default: object = None) -> object: # pyright: ignore[reportIncompatibleMethodOverride] # dict.setdefault overloads + if key not in self: + self[key] = default + return self[key] + + +class _ExpiringSpendCounters(FakeSpendCounters): + """Spend counters whose Redis keys expire, counting the lease renewals that keep them alive.""" + + def __init__(self, ttl: float = 60.0, fail_first_refresh: bool = False) -> None: + super().__init__(redis=_ExpiringRedis(ttl), redis_ttl=ttl) self.refresh_attempts = 0 self.refresh_count = 0 self.fail_first_refresh = fail_first_refresh - def _evict_expired(self, key: str) -> None: - if self.expires_at.get(key, float("inf")) <= time.monotonic(): - self.store.pop(key, None) - self.expires_at.pop(key, None) - - async def async_get_cache(self, key: str, *args: object, **kwargs: object) -> float | None: - self._evict_expired(key) - return self.store.get(key) - - async def async_increment(self, key: str, value: float, **kwargs: object) -> float: - self._evict_expired(key) - self.store[key] = self.store.get(key, 0.0) + float(value) - self.expires_at[key] = time.monotonic() + self.default_ttl - return self.store[key] - - async def async_set_max(self, key: str, value: float, **kwargs: object) -> float: - self._evict_expired(key) - self.store[key] = max(self.store.get(key, float("-inf")), float(value)) - self.expires_at[key] = time.monotonic() + self.default_ttl - return self.store[key] - - async def async_set_cache(self, key: str, value: float, *args: object, **kwargs: object) -> bool: - self.store[key] = float(value) - self.expires_at[key] = time.monotonic() + self.default_ttl - return True - - async def async_delete_cache(self, key: str, *args: object, **kwargs: object) -> None: - self.store.pop(key, None) - self.expires_at.pop(key, None) - - async def async_refresh_ttl(self, key: str, ttl: int | None = None) -> bool: + def refresh_ttl(self, counter: str) -> asyncio.Future[bool]: self.refresh_attempts += 1 + future: asyncio.Future[bool] = asyncio.get_running_loop().create_future() + redis = self.redis + assert isinstance(redis, _ExpiringRedis) if self.fail_first_refresh and self.refresh_attempts == 1: - raise ConnectionError("Redis circuit breaker is open") - self._evict_expired(key) - if key not in self.store: - return False + future.set_exception(unavailable("breaker_open")) + return future + if counter not in redis: + future.set_result(False) + return future self.refresh_count += 1 - self.expires_at[key] = time.monotonic() + (ttl if ttl is not None else self.default_ttl) - return True + redis[counter] = redis[counter] + future.set_result(True) + return future - async def async_increment_pipeline( - self, increment_list: Sequence[RedisPipelineIncrementOperation], **kwargs: object - ) -> list[float]: - return [await self.async_increment(op["key"], op["increment_value"]) for op in increment_list] - def get_ttl(self, **kwargs: object) -> int | None: - return int(self.default_ttl) +def _install(counters: FakeSpendCounters) -> FakeSpendCounters: + import litellm.proxy.proxy_server as ps + + ps.spend_counter_store = counters.store() + return counters @pytest.mark.asyncio @@ -2406,8 +2236,7 @@ async def test_reservation_survives_redis_counter_ttl_while_request_in_flight( (so a concurrent request on any worker still sees it), and renewal must stop once the reservation is reconciled so an idle counter still expires on its own.""" counter_cache, key_cache = spend_counter_state - redis_cache = _ExpiringRedisCache(default_ttl=0.2) - counter_cache.redis_cache = redis_cache + redis_cache = _install(_ExpiringSpendCounters(ttl=0.2)) proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) valid_token = UserAPIKeyAuth(token="key-lease", spend=0.0, max_budget=1.0) counter_key = "spend:key:key-lease" @@ -2416,7 +2245,7 @@ async def test_reservation_survives_redis_counter_ttl_while_request_in_flight( assert reservation is not None await asyncio.sleep(0.5) - assert await redis_cache.async_get_cache(key=counter_key) == pytest.approx(0.6) + assert redis_cache.redis.get(counter_key) == pytest.approx(0.6) concurrent = await _reserve(valid_token, 0.6, key_cache, proxy_logging_obj) assert concurrent is not None assert concurrent["reserved_cost"] == pytest.approx(0.4) @@ -2427,7 +2256,7 @@ async def test_reservation_survives_redis_counter_ttl_while_request_in_flight( refreshes_after_release = redis_cache.refresh_count await asyncio.sleep(0.35) assert redis_cache.refresh_count == refreshes_after_release - assert await redis_cache.async_get_cache(key=counter_key) is None + assert redis_cache.redis.get(counter_key) is None @pytest.mark.asyncio @@ -2437,8 +2266,7 @@ async def test_reservation_lease_keeps_renewing_after_transient_redis_failure( """One failed EXPIRE (Redis blip, open circuit breaker) must not end renewal for the rest of the request.""" counter_cache, key_cache = spend_counter_state - redis_cache = _ExpiringRedisCache(default_ttl=0.2, fail_first_refresh=True) - counter_cache.redis_cache = redis_cache + redis_cache = _install(_ExpiringSpendCounters(ttl=0.2, fail_first_refresh=True)) proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) valid_token = UserAPIKeyAuth(token="key-lease-blip", spend=0.0, max_budget=1.0) @@ -2458,8 +2286,7 @@ async def test_reservation_lease_stops_when_request_task_ends_without_reconcilin cost callbacks) must not keep renewing: the counter falls back to its plain TTL instead of pinning the reservation until the request timeout.""" counter_cache, key_cache = spend_counter_state - redis_cache = _ExpiringRedisCache(default_ttl=0.2) - counter_cache.redis_cache = redis_cache + redis_cache = _install(_ExpiringSpendCounters(ttl=0.2)) proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) valid_token = UserAPIKeyAuth(token="key-lease-orphan", spend=0.0, max_budget=1.0) counter_key = "spend:key:key-lease-orphan" @@ -2470,7 +2297,7 @@ async def test_reservation_lease_stops_when_request_task_ends_without_reconcilin await asyncio.sleep(0.5) assert redis_cache.refresh_count == 0 - assert await redis_cache.async_get_cache(key=counter_key) is None + assert redis_cache.redis.get(counter_key) is None class _TeamMembershipFloorDb: @@ -2493,11 +2320,9 @@ async def test_reconcile_after_redis_counter_expiry_keeps_request_cost_enforced( instead of applying ``actual - reserved`` to the empty key.""" import litellm.proxy.proxy_server as ps - counter_cache, _ = spend_counter_state counter_key = "spend:team_member:user-expiry:team-expiry" - redis_cache = _ExpiringRedisCache() - counter_cache.redis_cache = redis_cache - counter_cache.in_memory_cache.set_cache(key=counter_key, value=0.6) + redis_cache = _install(_ExpiringSpendCounters()) + redis_cache.set_local(counter_key, 0.6) reservation = { "reserved_cost": 0.6, @@ -2524,8 +2349,8 @@ async def test_reconcile_after_redis_counter_expiry_keeps_request_cost_enforced( budget_reservation=reservation, ) - assert redis_cache.store[counter_key] == pytest.approx(0.35) - assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(0.35) + assert redis_cache.redis[counter_key] == pytest.approx(0.35) + assert redis_cache.get_local(counter_key) == pytest.approx(0.35) assert reservation["finalized"] is True @@ -2543,9 +2368,8 @@ async def test_reconcile_before_db_update_does_not_double_count_when_flush_lands counter_cache, _ = spend_counter_state counter_key = "spend:team_member:user-flush:team-flush" - redis_cache = _ExpiringRedisCache() - counter_cache.redis_cache = redis_cache - counter_cache.in_memory_cache.set_cache(key=counter_key, value=0.6) + redis_cache = _install(_ExpiringSpendCounters()) + redis_cache.set_local(counter_key, 0.6) db_floor = _TeamMembershipFloorDb(spend=0.3) ps.prisma_client = SimpleNamespace(db=db_floor) @@ -2565,7 +2389,7 @@ async def test_reconcile_before_db_update_does_not_double_count_when_flush_lands await reconcile_budget_reservation(budget_reservation=reservation, actual_cost=0.05, finalize=False) - assert redis_cache.store[counter_key] == pytest.approx(0.35) + assert redis_cache.redis[counter_key] == pytest.approx(0.35) assert reservation["entries"][0]["applied_adjustment"] == pytest.approx(-0.55) assert reservation["finalized"] is False @@ -2578,15 +2402,10 @@ async def test_reconcile_before_db_update_does_not_double_count_when_flush_lands budget_reservation=reservation, ) - assert redis_cache.store[counter_key] == pytest.approx(0.35) + assert redis_cache.redis[counter_key] == pytest.approx(0.35) assert reservation["finalized"] is True -class _BatchReadingRedisCache(_ExpiringRedisCache): - async def async_batch_get_cache(self, key_list: Sequence[str], **kwargs: object) -> dict[str, float | None]: - return {key: await self.async_get_cache(key) for key in key_list} - - @pytest.mark.asyncio async def test_reserved_counter_deleted_during_spend_write_is_reseeded_instead_of_going_negative( spend_counter_state, @@ -2596,14 +2415,13 @@ async def test_reserved_counter_deleted_during_spend_write_is_reseeded_instead_o counter_cache, _ = spend_counter_state counter_key = "spend:key:key-deleted-mid-write" - redis_cache = _BatchReadingRedisCache() - counter_cache.redis_cache = redis_cache - await redis_cache.async_set_cache(counter_key, 0.6) - counter_cache.in_memory_cache.set_cache(key=counter_key, value=0.6) + redis_cache = _install(_ExpiringSpendCounters()) + redis_cache.redis[counter_key] = 0.6 + redis_cache.set_local(counter_key, 0.6) async def _delete_counter_while_persisting(**kwargs: object) -> bool: - await redis_cache.async_delete_cache(counter_key) - counter_cache.in_memory_cache.delete_cache(key=counter_key) + redis_cache.redis.pop(counter_key, None) + redis_cache.delete_local(counter_key) return True proxy_logging_obj = MagicMock() @@ -2644,7 +2462,7 @@ async def test_reserved_counter_deleted_during_spend_write_is_reseeded_instead_o ) assert charged is True - assert redis_cache.store[counter_key] == pytest.approx(0.35), redis_cache.store + assert redis_cache.redis[counter_key] == pytest.approx(0.35), redis_cache.store assert reservation["finalized"] is True @@ -2653,14 +2471,8 @@ async def test_should_invalidate_reserved_counters_after_persisted_spend_failure spend_counter_state, ): counter_cache, _ = spend_counter_state - await counter_cache.async_increment_cache( - key="spend:key:key-budget-invalidate", - value=0.4, - ) - await counter_cache.async_increment_cache( - key="spend:team:team-budget-invalidate", - value=0.4, - ) + counter_cache.increment_local("spend:key:key-budget-invalidate", 0.4) + counter_cache.increment_local("spend:team:team-budget-invalidate", 0.4) await invalidate_budget_reservation_counters( { @@ -2672,14 +2484,8 @@ async def test_should_invalidate_reserved_counters_after_persisted_spend_failure } ) - assert ( - counter_cache.in_memory_cache.get_cache(key="spend:key:key-budget-invalidate") - is None - ) - assert ( - counter_cache.in_memory_cache.get_cache(key="spend:team:team-budget-invalidate") - is None - ) + assert counter_cache.get_local("spend:key:key-budget-invalidate") is None + assert counter_cache.get_local("spend:team:team-budget-invalidate") is None @pytest.mark.asyncio @@ -2714,12 +2520,8 @@ async def test_should_reserve_all_budgeted_counters(spend_counter_state): proxy_logging_obj=proxy_logging_obj, ) - assert ( - counter_cache.in_memory_cache.get_cache(key="spend:key:key-budget-all") == 0.3 - ) - assert ( - counter_cache.in_memory_cache.get_cache(key="spend:team:team-budget-all") == 0.3 - ) + assert counter_cache.get_local("spend:key:key-budget-all") == 0.3 + assert counter_cache.get_local("spend:team:team-budget-all") == 0.3 @pytest.mark.asyncio @@ -2782,12 +2584,7 @@ async def test_should_not_block_concurrent_team_request_when_first_request_lacks # The team counter must not be pinned at max_budget while the first # request is in flight, otherwise concurrent requests false-positive. - team_counter_after_first = ( - counter_cache.in_memory_cache.get_cache( - key=f"spend:team:{team_object.team_id}" - ) - or 0.0 - ) + team_counter_after_first = counter_cache.get_local(f"spend:team:{team_object.team_id}") or 0.0 assert team_counter_after_first < team_object.max_budget, ( f"Team counter sat at {team_counter_after_first} after one uncapped " f"reservation against a {team_object.max_budget} budget — concurrent " @@ -2846,25 +2643,19 @@ async def test_release_budget_reservation_on_cancel_gives_back_counter( proxy_logging_obj=proxy_logging_obj, ) assert reservation is not None - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-cancel-give-back" - ) == pytest.approx(3.0) + assert counter_cache.get_local("spend:key:key-cancel-give-back") == pytest.approx(3.0) await release_budget_reservation_on_cancel(reservation) # the provider already received the input, so the reservation is reconciled # to the input cost (0.5), not refunded to zero; the worst-case output # reservation (3.0 -> 0.5) is released - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-cancel-give-back" - ) == pytest.approx(0.5) + assert counter_cache.get_local("spend:key:key-cancel-give-back") == pytest.approx(0.5) assert reservation["finalized"] is True # idempotent: a second cancel reconcile must not change the counter again await release_budget_reservation_on_cancel(reservation) - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-cancel-give-back" - ) == pytest.approx(0.5) + assert counter_cache.get_local("spend:key:key-cancel-give-back") == pytest.approx(0.5) @pytest.mark.asyncio @@ -2898,9 +2689,7 @@ async def test_release_budget_reservation_on_cancel_noop_when_finalized( await release_budget_reservation_on_cancel(reservation) # already reconciled by the success/failure path -> must stay untouched - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-cancel-finalized" - ) == pytest.approx(3.0) + assert counter_cache.get_local("spend:key:key-cancel-finalized") == pytest.approx(3.0) async def _reserve_for_stream(counter_cache, key_cache, proxy_logging_obj, token: str): @@ -2927,9 +2716,7 @@ async def _reserve_for_stream(counter_cache, key_cache, proxy_logging_obj, token proxy_logging_obj=proxy_logging_obj, ) assert reservation is not None - assert counter_cache.in_memory_cache.get_cache( - key=f"spend:key:{token}" - ) == pytest.approx(2.0) + assert counter_cache.get_local(f"spend:key:{token}") == pytest.approx(2.0) valid_token.budget_reservation = reservation return valid_token, reservation @@ -2982,9 +2769,7 @@ async def test_streaming_cancel_before_any_chunk_reconciles_to_input_cost( assert received == [] # no chunk delivered, but the provider already received the input, so the # reservation is reconciled to the input cost (0.5), not refunded to zero - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-cancel-no-chunk" - ) == pytest.approx(0.5) + assert counter_cache.get_local("spend:key:key-cancel-no-chunk") == pytest.approx(0.5) assert reservation["finalized"] is True streaming_logging_obj.arelease_max_parallel_requests_on_disconnect.assert_awaited_once() @@ -3016,9 +2801,7 @@ async def test_streaming_cancel_after_chunk_keeps_reservation( assert received == ["data: chunk\n\n"] # a consumed stream must NOT be refunded - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-cancel-after-chunk" - ) == pytest.approx(2.0) + assert counter_cache.get_local("spend:key:key-cancel-after-chunk") == pytest.approx(2.0) assert reservation.get("finalized") is not True streaming_logging_obj.arelease_max_parallel_requests_on_disconnect.assert_awaited_once() @@ -3048,9 +2831,7 @@ async def test_streaming_cancel_after_only_keepalive_pings_reconciles_to_input_c await _drain() assert received == [STREAM_SSE_KEEPALIVE_PING_BYTES] - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-cancel-after-ping" - ) == pytest.approx(0.5) + assert counter_cache.get_local("spend:key:key-cancel-after-ping") == pytest.approx(0.5) assert reservation["finalized"] is True @@ -3119,9 +2900,7 @@ async def test_streaming_cancel_while_holding_back_provider_output_keeps_reserva await asyncio.wait_for(_drain(), timeout=5) assert received and received == [STREAM_SSE_KEEPALIVE_PING_BYTES] * len(received) - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-cancel-held-back" - ) == pytest.approx(2.0) + assert counter_cache.get_local("spend:key:key-cancel-held-back") == pytest.approx(2.0) assert reservation.get("finalized") is not True @@ -3186,9 +2965,7 @@ async def test_streaming_cancel_in_slow_path_before_yield_refunds(spend_counter_ assert received == [] # cancellation happened before any chunk reached the client, but the # provider already received the input -> reconcile to the input cost (0.5) - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-cancel-slowpath" - ) == pytest.approx(0.5) + assert counter_cache.get_local("spend:key:key-cancel-slowpath") == pytest.approx(0.5) assert reservation["finalized"] is True streaming_logging_obj.arelease_max_parallel_requests_on_disconnect.assert_awaited_once() @@ -3216,9 +2993,7 @@ async def test_streaming_disconnect_after_consuming_chunk_keeps_reservation( await generator.aclose() # output was delivered, so the reservation must NOT be refunded - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-disconnect-after-chunk" - ) == pytest.approx(2.0) + assert counter_cache.get_local("spend:key:key-disconnect-after-chunk") == pytest.approx(2.0) assert reservation.get("finalized") is not True streaming_logging_obj.arelease_max_parallel_requests_on_disconnect.assert_awaited_once() @@ -3650,18 +3425,18 @@ async def test_model_access_group_counter_accumulates_across_calls_without_a_res from litellm.proxy.proxy_server import increment_spend_counters - counter_key = model_access_group_spend_counter_key("premium") + counter_key = model_access_group_counter("premium") await increment_spend_counters( token=None, team_id=None, user_id=None, response_cost=0.25, model_access_groups=["premium"] ) - assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(1.25) + assert counter_cache.get_local(counter_key) == pytest.approx(1.25) await increment_spend_counters( token=None, team_id=None, user_id=None, response_cost=0.75, model_access_groups=["premium", "premium", ""] ) - assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(2.0) - assert counter_cache.in_memory_cache.get_cache(key=model_access_group_spend_counter_key("")) is None + assert counter_cache.get_local(counter_key) == pytest.approx(2.0) + assert counter_cache.get_local(model_access_group_counter("")) is None @pytest.mark.asyncio @@ -3671,8 +3446,8 @@ async def test_reserved_model_access_group_is_not_charged_twice(spend_counter_st await _cache_model_access_group_budget(key_cache, "premium", spend=1.0, max_budget=25.0) reservation = await _reserve_for_model_access_groups(key_cache, ["premium"], estimate=0.6) - counter_key = model_access_group_spend_counter_key("premium") - assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(1.6) + counter_key = model_access_group_counter("premium") + assert counter_cache.get_local(counter_key) == pytest.approx(1.6) from litellm.proxy.proxy_server import increment_spend_counters @@ -3687,7 +3462,7 @@ async def test_reserved_model_access_group_is_not_charged_twice(spend_counter_st # 1.0 recorded + the reservation reconciled down to the 0.2 actually spent. A second # increment would land at 1.4. - assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(1.2) + assert counter_cache.get_local(counter_key) == pytest.approx(1.2) @pytest.mark.asyncio @@ -3715,9 +3490,5 @@ async def test_unreserved_model_access_group_is_charged_alongside_a_reserved_one model_access_groups=["premium", "starter", "starter", "premium"], ) - assert counter_cache.in_memory_cache.get_cache( - key=model_access_group_spend_counter_key("premium") - ) == pytest.approx(1.2) - assert counter_cache.in_memory_cache.get_cache( - key=model_access_group_spend_counter_key("starter") - ) == pytest.approx(4.2) + assert counter_cache.get_local(model_access_group_counter("premium")) == pytest.approx(1.2) + assert counter_cache.get_local(model_access_group_counter("starter")) == pytest.approx(4.2) diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 9dd15649d3e..ee1407c895a 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -57,6 +57,9 @@ from litellm.proxy.proxy_server import app, aws_error_type_exception_handler, in from litellm.utils import _invalidate_model_cost_lowercase_map from litellm.proxy.common_utils.banner import LITELLM_BANNER from tests._master_key import MASTER_KEY +from litellm.redis.stores.spend_counters import SpendCounterStore +from litellm.rust_bridge._native import NativeSpendCounterError +from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters PUBLICLY_KNOWN_KEY: Final = "sk-" + "1234" @@ -1950,11 +1953,11 @@ async def test_proxy_startup_after_a_failed_migration_does_not_reuse_the_closed_ async with proxy_startup_event(FastAPI()) as state: assert state["database"] is None - @pytest.mark.asyncio async def test_proxy_startup_names_the_config_file_that_set_the_unsafe_key(monkeypatch, tmp_path): from fastapi import FastAPI + from litellm.proxy.auth.master_key_boot_check import UnsafeMasterKeyError from litellm.proxy.proxy_server import proxy_startup_event @@ -4493,7 +4496,6 @@ async def test_chat_completion_result_no_nested_none_values(): from unittest.mock import AsyncMock, MagicMock, patch from fastapi import Request, Response - from pydantic import BaseModel import litellm from litellm.proxy._types import UserAPIKeyAuth @@ -8855,51 +8857,30 @@ async def test_store_model_in_db_db_failure_graceful(monkeypatch): @pytest.mark.asyncio -async def test_get_current_spend_reads_redis_first(): +async def test_get_current_spend_reads_redis_first(monkeypatch): """get_current_spend should prefer Redis over in-memory.""" - from litellm.caching.dual_cache import DualCache - - counter_cache = DualCache() - - # In-memory has stale value - counter_cache.in_memory_cache.set_cache(key="spend:key:test", value=0.30) - - # Mock Redis with cross-pod authoritative value - mock_redis = AsyncMock() - mock_redis.async_get_cache = AsyncMock(return_value=0.90) - counter_cache.redis_cache = mock_redis - import litellm.proxy.proxy_server as ps - original = ps.spend_counter_cache - ps.spend_counter_cache = counter_cache + spend = FakeSpendCounters(redis={"spend:key:test": 0.90}, memory={"spend:key:test": 0.30}) + monkeypatch.setattr(ps, "spend_counter_store", spend.store()) - try: - from litellm.proxy.proxy_server import get_current_spend + result = await ps.get_current_spend(counter_key="spend:key:test", fallback_spend=0.0) - result = await get_current_spend( - counter_key="spend:key:test", - fallback_spend=0.0, - ) - # Should return Redis value (0.90), not in-memory (0.30) - assert result == 0.90 - mock_redis.async_get_cache.assert_called_once_with(key="spend:key:test") - finally: - ps.spend_counter_cache = original + assert result == 0.90 + assert spend.commands == ["GET spend:key:test"] @pytest.mark.asyncio async def test_get_current_spend_fallback_to_in_memory(): """When Redis is not configured, get_current_spend uses in-memory.""" - from litellm.caching.dual_cache import DualCache - counter_cache = DualCache() # no redis_cache - counter_cache.in_memory_cache.set_cache(key="spend:key:test", value=0.50) + counter_cache = SpendCounterStore(None) + counter_cache.set_local("spend:key:test", 0.50) import litellm.proxy.proxy_server as ps - original = ps.spend_counter_cache - ps.spend_counter_cache = counter_cache + original = ps.spend_counter_store + ps.spend_counter_store = counter_cache try: from litellm.proxy.proxy_server import get_current_spend @@ -8910,7 +8891,7 @@ async def test_get_current_spend_fallback_to_in_memory(): ) assert result == 0.50 finally: - ps.spend_counter_cache = original + ps.spend_counter_store = original @pytest.mark.asyncio @@ -8924,7 +8905,7 @@ async def test_increment_spend_counters_initializes_and_increments(): from litellm.proxy._types import LiteLLM_VerificationTokenView, hash_token key_cache = DualCache() - counter_cache = DualCache() + counter_cache = SpendCounterStore(None) # In production, the auth flow hashes the raw key before it reaches # the cost callback. Simulate that by passing the hashed token. @@ -8941,9 +8922,9 @@ async def test_increment_spend_counters_initializes_and_increments(): import litellm.proxy.proxy_server as ps original_key_cache = ps.user_api_key_cache - original_counter_cache = ps.spend_counter_cache + original_counter_cache = ps.spend_counter_store ps.user_api_key_cache = key_cache - ps.spend_counter_cache = counter_cache + ps.spend_counter_store = counter_cache try: from litellm.proxy.proxy_server import increment_spend_counters @@ -8957,7 +8938,7 @@ async def test_increment_spend_counters_initializes_and_increments(): ) # Counter should be: base(5.0) + increment(0.50) = 5.50 - counter = counter_cache.in_memory_cache.get_cache(key=f"spend:key:{hashed_token}") + counter = counter_cache.get_local(f"spend:key:{hashed_token}") assert counter == 5.50 # Second increment — counter already exists, just increment @@ -8968,11 +8949,11 @@ async def test_increment_spend_counters_initializes_and_increments(): response_cost=0.25, ) - counter = counter_cache.in_memory_cache.get_cache(key=f"spend:key:{hashed_token}") + counter = counter_cache.get_local(f"spend:key:{hashed_token}") assert counter == 5.75 finally: ps.user_api_key_cache = original_key_cache - ps.spend_counter_cache = original_counter_cache + ps.spend_counter_store = original_counter_cache @pytest.mark.asyncio @@ -8982,7 +8963,7 @@ async def test_increment_spend_counters_team_and_member(): from litellm.proxy._types import LiteLLM_TeamTable key_cache = DualCache() - counter_cache = DualCache() + counter_cache = SpendCounterStore(None) # Cached team object team_obj = LiteLLM_TeamTable(team_id="team-1", spend=2.0) @@ -8997,9 +8978,9 @@ async def test_increment_spend_counters_team_and_member(): import litellm.proxy.proxy_server as ps original_key_cache = ps.user_api_key_cache - original_counter_cache = ps.spend_counter_cache + original_counter_cache = ps.spend_counter_store ps.user_api_key_cache = key_cache - ps.spend_counter_cache = counter_cache + ps.spend_counter_store = counter_cache try: from litellm.proxy.proxy_server import increment_spend_counters @@ -9011,141 +8992,73 @@ async def test_increment_spend_counters_team_and_member(): response_cost=0.30, ) - team_counter = counter_cache.in_memory_cache.get_cache(key="spend:team:team-1") + team_counter = counter_cache.get_local("spend:team:team-1") assert team_counter == 2.30 - member_counter = counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-1:team-1") + member_counter = counter_cache.get_local("spend:team_member:user-1:team-1") assert member_counter == 1.30 finally: ps.user_api_key_cache = original_key_cache - ps.spend_counter_cache = original_counter_cache + ps.spend_counter_store = original_counter_cache @pytest.mark.asyncio -async def test_prepare_spend_counter_increment_reseeds_from_db_on_counter_miss(): +async def test_prepare_spend_counter_increment_reseeds_from_db_on_counter_miss(monkeypatch): """When the Redis counter is missing, the reseed path reads the authoritative spend from the DB (not a stale cache), so the next increment continues from the correct base value.""" + import litellm.proxy.proxy_server as ps from litellm.caching.dual_cache import DualCache - counter_cache = DualCache() - recorded_increments: list = [] - - async def record_increment(key, value, ttl=None, **kwargs): - recorded_increments.append({"key": key, "value": value, "ttl": ttl}) - return value - - async def record_pipeline(increment_list, **kwargs): - results = [] - for op in increment_list: - await record_increment(key=op["key"], value=op["increment_value"], ttl=op["ttl"]) - results.append(op["increment_value"]) - return results - - fake_redis = AsyncMock() - fake_redis.async_increment = AsyncMock(side_effect=record_increment) - fake_redis.async_increment_pipeline = AsyncMock(side_effect=record_pipeline) - fake_redis.get_ttl = MagicMock(return_value=None) - fake_redis.async_get_cache = AsyncMock(return_value=None) # counter missing - fake_redis.async_set_cache = AsyncMock(return_value=True) # SET NX wins - counter_cache.redis_cache = fake_redis - - # Prisma returns spend=42.0 (authoritative) while the stale cached - # value (would be read only if prisma is None) is 10.0. The counter - # must seed from 42, not 10. + spend = FakeSpendCounters() db_row = MagicMock() db_row.spend = 42.0 fake_prisma = MagicMock() fake_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=db_row) - stale_cache = DualCache() stale_team = MagicMock() stale_team.spend = 10.0 stale_cache.in_memory_cache.set_cache(key="team_id:team-9", value=stale_team) + monkeypatch.setattr(ps, "user_api_key_cache", stale_cache) + monkeypatch.setattr(ps, "spend_counter_store", spend.store()) + monkeypatch.setattr(ps, "prisma_client", fake_prisma) - import litellm.proxy.proxy_server as ps - from litellm.proxy.proxy_server import ( - _apply_spend_counter_increments, - _prepare_spend_counter_increment, + pending = await ps._prepare_spend_counter_increment( + counter_key="spend:team:team-9", + source_cache_key="team_id:team-9", + increment=1.5, ) + await ps._apply_spend_counter_increments(pending=(pending,)) - orig_user, orig_counter, orig_prisma = ( - ps.user_api_key_cache, - ps.spend_counter_cache, - ps.prisma_client, - ) - ps.user_api_key_cache = stale_cache - ps.spend_counter_cache = counter_cache - ps.prisma_client = fake_prisma - try: - pending = await _prepare_spend_counter_increment( - counter_key="spend:team:team-9", - source_cache_key="team_id:team-9", - increment=1.5, - ) - await _apply_spend_counter_increments(pending=(pending,)) - - fake_prisma.db.litellm_teamtable.find_unique.assert_awaited_once_with(where={"team_id": "team-9"}) - # Seed uses SET NX with db_spend (42) — cross-pod safe, no INCR of 42. - # Only the per-request delta (1.5) goes through INCRBYFLOAT. - fake_redis.async_set_cache.assert_awaited_once_with(key="spend:team:team-9", value=42.0, nx=True) - writes = [(c["key"], c["value"]) for c in recorded_increments] - assert writes == [("spend:team:team-9", 1.5)] - finally: - ps.user_api_key_cache = orig_user - ps.spend_counter_cache = orig_counter - ps.prisma_client = orig_prisma + fake_prisma.db.litellm_teamtable.find_unique.assert_awaited_once_with(where={"team_id": "team-9"}) + # Seed uses SET NX with db_spend (42), cross-pod safe, no INCR of 42. + # Only the per-request delta (1.5) goes through INCRBYFLOAT. + assert spend.commands == [ + "GET spend:team:team-9", + "GET spend:team:team-9", + "SETNX spend:team:team-9 42.0", + "PIPELINE spend:team:team-9", + ] + assert spend.redis == {"spend:team:team-9": 43.5} @pytest.mark.asyncio async def test_primary_spend_counter_redis_concurrent_seed_does_not_double_seed(): """Two pods both observing a missing Redis counter must not both - INCRBYFLOAT the full DB spend. SpendCounterReseed.coalesced uses SET NX - so the loser reads the winner's value; final Redis = db_spend, not + INCRBYFLOAT the full DB spend. SpendCounterReseed.coalesced seeds with SET NX + so the loser answers with the winner's value; final Redis = db_spend, not 2 * db_spend. The per-counter asyncio.Lock is per-process, so it does NOT coordinate across pods. We simulate two pods by patching _get_lock to return a fresh lock per call (each "pod" has its own lock registry in real life). """ - from litellm.caching.dual_cache import DualCache from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed counter_key = "spend:team:team-concurrent-seed" - redis_store: dict = {} + shared_redis: dict[str, object] = {} + db_read_count = 0 - set_results: list = [] - get_after_set_count = 0 - set_completed_count = 0 - - async def redis_set_cache(key, value, nx=False, **_): - # Yield BEFORE the membership check so two concurrent callers - # interleave the way real atomic Redis SET NX does: the first - # to resume runs check + write atomically and wins; the second - # resumes after the key exists and loses. Yielding *after* the - # check would let both callers pass the empty-store check before - # either writes, so neither would ever lose. - await asyncio.sleep(0) - if nx and key in redis_store: - set_results.append(False) - return False - redis_store[key] = float(value) - set_results.append(True) - nonlocal set_completed_count - set_completed_count += 1 - return True - - async def redis_get_cache(key): - # Track reads that happen after at least one SET NX has completed - # — those are the loser-path fallback reads we want to verify. - if set_completed_count > 0: - nonlocal get_after_set_count - get_after_set_count += 1 - return redis_store.get(key) - - fake_redis = AsyncMock() - fake_redis.async_get_cache = AsyncMock(side_effect=redis_get_cache) - fake_redis.async_set_cache = AsyncMock(side_effect=redis_set_cache) async def slow_find_unique(**_): nonlocal db_read_count @@ -9158,48 +9071,30 @@ async def test_primary_spend_counter_redis_concurrent_seed_does_not_double_seed( fake_prisma = MagicMock() fake_prisma.db.litellm_teamtable.find_unique = AsyncMock(side_effect=slow_find_unique) + pod_a = FakeSpendCounters(redis=shared_redis) + pod_b = FakeSpendCounters(redis=shared_redis) - pod_a = DualCache() - pod_a.redis_cache = fake_redis - pod_b = DualCache() - pod_b.redis_cache = fake_redis - - # Each "pod" has its own per-process lock registry. Patch _get_lock to - # always return a fresh lock so the two coalesced calls do not serialize - # via one in-process lock (which is what would happen across pods). async def fresh_lock(_counter_key): return asyncio.Lock() with patch.object(SpendCounterReseed, "_get_lock", side_effect=fresh_lock): results = await asyncio.gather( SpendCounterReseed.coalesced( - prisma_client=fake_prisma, - spend_counter_cache=pod_a, - counter_key=counter_key, + prisma_client=fake_prisma, spend_counter_store=pod_a.store(), counter_key=counter_key ), SpendCounterReseed.coalesced( - prisma_client=fake_prisma, - spend_counter_cache=pod_b, - counter_key=counter_key, + prisma_client=fake_prisma, spend_counter_store=pod_b.store(), counter_key=counter_key ), ) - assert all(r == 506.0 for r in results), results - assert redis_store[counter_key] == pytest.approx(506.0), redis_store - # Both pods read the DB and both attempted SET NX; exactly one wrote - # (winner) and one was rejected (loser). + assert results == [506.0, 506.0] + assert shared_redis == {counter_key: 506.0} assert db_read_count == 2 - assert fake_redis.async_set_cache.await_count == 2 - nx_writes = [call for call in fake_redis.async_set_cache.await_args_list if call.kwargs.get("nx") is True] - assert len(nx_writes) == 2 - assert sorted(set_results) == [ - False, - True, - ], f"expected exactly one SET NX winner and one loser, got {set_results}" - # Loser path executed: after the winner's SET NX returned True, the - # losing coalesced() call falls back to async_get_cache to read the - # winner's value rather than re-seeding. - assert get_after_set_count >= 1, "loser branch (else: read back winner's value) was never exercised" + assert [command for command in pod_a.commands + pod_b.commands if command.startswith("SETNX")] == [ + f"SETNX {counter_key} 506.0", + f"SETNX {counter_key} 506.0", + ] + assert [command for command in pod_a.commands + pod_b.commands if command.startswith("INCRBYFLOAT")] == [] @pytest.mark.asyncio @@ -9261,13 +9156,12 @@ async def test_reseed_spend_from_db_skips_window_variant_keys(): @pytest.mark.asyncio async def test_window_spend_counter_reseeds_from_spend_logs_on_counter_miss(): - from litellm.caching.dual_cache import DualCache from litellm.proxy.proxy_server import ( _apply_spend_counter_increments, _prepare_window_spend_counter_increment, ) - counter_cache = DualCache() + counter_cache = SpendCounterStore(None) window_start = datetime.now(timezone.utc) - timedelta(hours=1) fake_prisma = MagicMock() fake_prisma.db.litellm_budgetwindowspend.find_unique = AsyncMock(return_value=None) @@ -9277,8 +9171,8 @@ async def test_window_spend_counter_reseeds_from_spend_logs_on_counter_miss(): import litellm.proxy.proxy_server as ps - orig_counter, orig_prisma = ps.spend_counter_cache, ps.prisma_client - ps.spend_counter_cache = counter_cache + orig_counter, orig_prisma = ps.spend_counter_store, ps.prisma_client + ps.spend_counter_store = counter_cache ps.prisma_client = fake_prisma try: pending = await _prepare_window_spend_counter_increment( @@ -9296,246 +9190,127 @@ async def test_window_spend_counter_reseeds_from_spend_logs_on_counter_miss(): where={"api_key": "key-window", "startTime": {"gte": window_start}}, sum={"spend": True}, ) - assert counter_cache.in_memory_cache.get_cache(key="spend:key:key-window:window:1h") == pytest.approx(2.75) + assert counter_cache.get_local("spend:key:key-window:window:1h") == pytest.approx(2.75) finally: - ps.spend_counter_cache = orig_counter + ps.spend_counter_store = orig_counter ps.prisma_client = orig_prisma @pytest.mark.asyncio -async def test_init_spend_counter_redis_clean_miss_skips_stale_in_memory(): +async def test_init_spend_counter_redis_clean_miss_skips_stale_in_memory(monkeypatch): + import litellm.proxy.proxy_server as ps from litellm.caching.dual_cache import DualCache - from litellm.proxy.proxy_server import ( - _apply_spend_counter_increments, - _prepare_spend_counter_increment, - ) - counter_cache = DualCache() counter_key = "spend:team:team-stale-local" - counter_cache.in_memory_cache.set_cache(key=counter_key, value=10.0) - - redis_store: dict = {} - - async def redis_increment(key, value, **_): - redis_store[key] = (redis_store.get(key) or 0.0) + value - return redis_store[key] - - async def redis_set_cache(key, value, nx=False, **_): - if nx and key in redis_store: - return False - redis_store[key] = float(value) - return True - - fake_redis = AsyncMock() - fake_redis.async_get_cache = AsyncMock(return_value=None) - fake_redis.async_increment = AsyncMock(side_effect=redis_increment) - fake_redis.async_set_cache = AsyncMock(side_effect=redis_set_cache) - - async def redis_increment_pipeline(increment_list, **_): - results = [] - for op in increment_list: - results.append(await redis_increment(key=op["key"], value=op["increment_value"])) - return results - - fake_redis.async_increment_pipeline = AsyncMock(side_effect=redis_increment_pipeline) - fake_redis.get_ttl = MagicMock(return_value=None) - counter_cache.redis_cache = fake_redis - + spend = FakeSpendCounters(memory={counter_key: 10.0}) db_row = MagicMock() db_row.spend = 42.0 fake_prisma = MagicMock() fake_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=db_row) + monkeypatch.setattr(ps, "spend_counter_store", spend.store()) + monkeypatch.setattr(ps, "prisma_client", fake_prisma) + monkeypatch.setattr(ps, "user_api_key_cache", DualCache()) - import litellm.proxy.proxy_server as ps - - orig_counter, orig_prisma, orig_user = ( - ps.spend_counter_cache, - ps.prisma_client, - ps.user_api_key_cache, + pending = await ps._prepare_spend_counter_increment( + counter_key=counter_key, + source_cache_key="team_id:team-stale-local", + increment=1.5, ) - ps.spend_counter_cache = counter_cache - ps.prisma_client = fake_prisma - ps.user_api_key_cache = DualCache() - try: - pending = await _prepare_spend_counter_increment( - counter_key=counter_key, - source_cache_key="team_id:team-stale-local", - increment=1.5, - ) - await _apply_spend_counter_increments(pending=(pending,)) + await ps._apply_spend_counter_increments(pending=(pending,)) - fake_prisma.db.litellm_teamtable.find_unique.assert_awaited_once_with(where={"team_id": "team-stale-local"}) - # Seed via SET NX (42) + delta via INCRBYFLOAT (1.5) = 43.5. - assert redis_store[counter_key] == pytest.approx(43.5) - assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(43.5) - finally: - ps.spend_counter_cache = orig_counter - ps.prisma_client = orig_prisma - ps.user_api_key_cache = orig_user + fake_prisma.db.litellm_teamtable.find_unique.assert_awaited_once_with(where={"team_id": "team-stale-local"}) + # Seed via SET NX (42) + delta via INCRBYFLOAT (1.5) = 43.5, ignoring the stale 10.0 in memory. + assert spend.redis == {counter_key: 43.5} + assert spend.memory == {counter_key: 43.5} @pytest.mark.asyncio -async def test_window_spend_counter_redis_clean_miss_skips_stale_in_memory(): - from litellm.caching.dual_cache import DualCache - from litellm.proxy.proxy_server import ( - _apply_spend_counter_increments, - _prepare_window_spend_counter_increment, - ) +async def test_window_spend_counter_redis_clean_miss_skips_stale_in_memory(monkeypatch): + import litellm.proxy.proxy_server as ps - counter_cache = DualCache() counter_key = "spend:key:key-window-stale-local:window:1h" - counter_cache.in_memory_cache.set_cache(key=counter_key, value=100.0) + spend = FakeSpendCounters(memory={counter_key: 100.0}) window_start = datetime.now(timezone.utc) - timedelta(hours=1) - - redis_store: dict = {} - - async def redis_increment(key, value, **_): - redis_store[key] = (redis_store.get(key) or 0.0) + value - return redis_store[key] - - async def redis_set_cache(key, value, **_): - if key in redis_store: - return False - redis_store[key] = value - return True - - fake_redis = AsyncMock() - fake_redis.async_get_cache = AsyncMock(return_value=None) - fake_redis.async_set_cache = AsyncMock(side_effect=redis_set_cache) - fake_redis.async_increment = AsyncMock(side_effect=redis_increment) - - async def redis_increment_pipeline(increment_list, **_): - results = [] - for op in increment_list: - results.append(await redis_increment(key=op["key"], value=op["increment_value"])) - return results - - fake_redis.async_increment_pipeline = AsyncMock(side_effect=redis_increment_pipeline) - fake_redis.get_ttl = MagicMock(return_value=None) - counter_cache.redis_cache = fake_redis - fake_prisma = MagicMock() fake_prisma.db.litellm_budgetwindowspend.find_unique = AsyncMock(return_value=None) fake_prisma.db.litellm_spendlogs.group_by = AsyncMock( return_value=[{"api_key": "key-window-stale-local", "_sum": {"spend": 2.25}}] ) + monkeypatch.setattr(ps, "spend_counter_store", spend.store()) + monkeypatch.setattr(ps, "prisma_client", fake_prisma) - import litellm.proxy.proxy_server as ps + pending = await ps._prepare_window_spend_counter_increment( + counter_key=counter_key, + entity_type="Key", + entity_id="key-window-stale-local", + window_duration="1h", + window_start=window_start, + increment=0.5, + ) + await ps._apply_spend_counter_increments(pending=(pending,) if pending is not None else ()) - orig_counter, orig_prisma = ps.spend_counter_cache, ps.prisma_client - ps.spend_counter_cache = counter_cache - ps.prisma_client = fake_prisma - try: - pending = await _prepare_window_spend_counter_increment( - counter_key=counter_key, - entity_type="Key", - entity_id="key-window-stale-local", - window_duration="1h", - window_start=window_start, - increment=0.5, - ) - await _apply_spend_counter_increments(pending=(pending,) if pending is not None else ()) - - fake_prisma.db.litellm_spendlogs.group_by.assert_awaited_once_with( - by=["api_key"], - where={ - "api_key": "key-window-stale-local", - "startTime": {"gte": window_start}, - }, - sum={"spend": True}, - ) - assert redis_store[counter_key] == pytest.approx(2.75) - assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(2.75) - finally: - ps.spend_counter_cache = orig_counter - ps.prisma_client = orig_prisma + fake_prisma.db.litellm_spendlogs.group_by.assert_awaited_once_with( + by=["api_key"], + where={ + "api_key": "key-window-stale-local", + "startTime": {"gte": window_start}, + }, + sum={"spend": True}, + ) + assert spend.redis == {counter_key: pytest.approx(2.75)} + assert spend.memory == {counter_key: pytest.approx(2.75)} @pytest.mark.asyncio -async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed(): - from litellm.caching.dual_cache import DualCache - from litellm.proxy.proxy_server import ( - _apply_spend_counter_increments, - _prepare_window_spend_counter_increment, - ) +async def test_window_spend_counter_seed_asks_the_store_to_increment_a_lost_seed(monkeypatch): + """A window seed that loses SET NX to a counter which then vanishes is added with INCRBYFLOAT, so a racing + pod's seed is never dropped; the primary counters answer with the seed instead.""" + import litellm.proxy.proxy_server as ps - counter_cache = DualCache() counter_key = "spend:key:key-window-concurrent-seed:window:1h" - window_start = datetime.now(timezone.utc) - timedelta(hours=1) - redis_store = {counter_key: 2.75} - redis_reads = 0 + seeds: list[tuple[str, float, bool]] = [] - async def redis_get_cache(key): - nonlocal redis_reads - redis_reads += 1 - if redis_reads <= 2: - return None - return redis_store.get(key) - - async def redis_increment(key, value, **_): - redis_store[key] = (redis_store.get(key) or 0.0) + value - return redis_store[key] - - fake_redis = AsyncMock() - fake_redis.async_get_cache = AsyncMock(side_effect=redis_get_cache) - fake_redis.async_set_cache = AsyncMock(return_value=False) - fake_redis.async_increment = AsyncMock(side_effect=redis_increment) - - async def redis_increment_pipeline(increment_list, **_): - results = [] - for op in increment_list: - results.append(await redis_increment(key=op["key"], value=op["increment_value"])) - return results - - fake_redis.async_increment_pipeline = AsyncMock(side_effect=redis_increment_pipeline) - fake_redis.get_ttl = MagicMock(return_value=None) - counter_cache.redis_cache = fake_redis + class RecordingSeeds(FakeSpendCounters): + def seed(self, counter: str, spend: float, *, increment_if_lost: bool) -> asyncio.Future[float]: + seeds.append((counter, spend, increment_if_lost)) + return super().seed(counter, spend, increment_if_lost=increment_if_lost) + spend = RecordingSeeds() fake_prisma = MagicMock() fake_prisma.db.litellm_budgetwindowspend.find_unique = AsyncMock(return_value=None) fake_prisma.db.litellm_spendlogs.group_by = AsyncMock( return_value=[{"api_key": "key-window-concurrent-seed", "_sum": {"spend": 2.25}}] ) + fake_prisma.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=MagicMock(spend=7.0)) + monkeypatch.setattr(ps, "spend_counter_store", spend.store()) + monkeypatch.setattr(ps, "prisma_client", fake_prisma) - import litellm.proxy.proxy_server as ps + await ps._prepare_window_spend_counter_increment( + counter_key=counter_key, + entity_type="Key", + entity_id="key-window-concurrent-seed", + window_duration="1h", + window_start=datetime.now(timezone.utc) - timedelta(hours=1), + increment=0.5, + ) + await ps._is_spend_counter_cache_warm(counter_key="spend:key:primary") + await ps.SpendCounterReseed.coalesced( + prisma_client=fake_prisma, spend_counter_store=spend.store(), counter_key="spend:key:primary" + ) - orig_counter, orig_prisma = ps.spend_counter_cache, ps.prisma_client - ps.spend_counter_cache = counter_cache - ps.prisma_client = fake_prisma - try: - pending = await _prepare_window_spend_counter_increment( - counter_key=counter_key, - entity_type="Key", - entity_id="key-window-concurrent-seed", - window_duration="1h", - window_start=window_start, - increment=0.5, - ) - await _apply_spend_counter_increments(pending=(pending,) if pending is not None else ()) - - fake_redis.async_set_cache.assert_awaited_once_with( - key=counter_key, - value=2.25, - nx=True, - ) - assert redis_store[counter_key] == pytest.approx(3.25) - assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(3.25) - finally: - ps.spend_counter_cache = orig_counter - ps.prisma_client = orig_prisma + assert seeds == [(counter_key, 2.25, True), ("spend:key:primary", 7.0, False)] @pytest.mark.asyncio async def test_window_spend_counter_skips_invalid_window_start(): - from litellm.caching.dual_cache import DualCache from litellm.proxy.proxy_server import _prepare_window_spend_counter_increment - counter_cache = DualCache() + counter_cache = SpendCounterStore(None) import litellm.proxy.proxy_server as ps - orig_counter = ps.spend_counter_cache - ps.spend_counter_cache = counter_cache + orig_counter = ps.spend_counter_store + ps.spend_counter_store = counter_cache try: pending = await _prepare_window_spend_counter_increment( counter_key="spend:key:key-invalid-window:window:not-a-duration", @@ -9547,23 +9322,22 @@ async def test_window_spend_counter_skips_invalid_window_start(): ) assert pending is None - assert counter_cache.in_memory_cache.get_cache(key="spend:key:key-invalid-window:window:not-a-duration") is None + assert counter_cache.get_local("spend:key:key-invalid-window:window:not-a-duration") is None finally: - ps.spend_counter_cache = orig_counter + ps.spend_counter_store = orig_counter @pytest.mark.asyncio async def test_window_spend_counter_does_not_seed_zero_when_db_unavailable(): - from litellm.caching.dual_cache import DualCache from litellm.proxy.proxy_server import ensure_window_spend_counter_initialized - counter_cache = DualCache() + counter_cache = SpendCounterStore(None) counter_key = "spend:key:key-window-db-unavailable:window:1h" import litellm.proxy.proxy_server as ps - orig_counter, orig_prisma = ps.spend_counter_cache, ps.prisma_client - ps.spend_counter_cache = counter_cache + orig_counter, orig_prisma = ps.spend_counter_store, ps.prisma_client + ps.spend_counter_store = counter_cache ps.prisma_client = None try: initialized = await ensure_window_spend_counter_initialized( @@ -9575,22 +9349,19 @@ async def test_window_spend_counter_does_not_seed_zero_when_db_unavailable(): ) assert initialized is False - assert counter_cache.in_memory_cache.get_cache(key=counter_key) is None + assert counter_cache.get_local(counter_key) is None finally: - ps.spend_counter_cache = orig_counter + ps.spend_counter_store = orig_counter ps.prisma_client = orig_prisma @pytest.mark.asyncio -async def test_increment_spend_counters_finalizes_after_unreserved_increments(): +async def test_increment_spend_counters_finalizes_after_unreserved_increments(monkeypatch): + import litellm.proxy.proxy_server as ps from litellm.caching.dual_cache import DualCache - from litellm.proxy.proxy_server import increment_spend_counters - counter_cache = DualCache() - counter_cache.in_memory_cache.set_cache( - key="spend:key:key-finalize-after-increments", - value=0.5, - ) + counter_cache = SpendCounterStore(None) + counter_cache.set_local("spend:key:key-finalize-after-increments", 0.5) budget_reservation = { "reserved_cost": 0.5, "entries": [ @@ -9611,44 +9382,31 @@ async def test_increment_spend_counters_finalizes_after_unreserved_increments(): incremented_counters.append(kwargs["counter_key"]) return ps.PendingSpendIncrement(counter_key=kwargs["counter_key"], increment=kwargs["increment"]) - import litellm.proxy.proxy_server as ps - - orig_counter, orig_user = ps.spend_counter_cache, ps.user_api_key_cache - ps.spend_counter_cache = counter_cache - ps.user_api_key_cache = DualCache() - try: - with patch( - "litellm.proxy.proxy_server._prepare_spend_counter_increment", - new=AsyncMock(side_effect=assert_reservation_not_finalized_yet), - ): - await increment_spend_counters( - token="key-finalize-after-increments", - team_id="team-finalize-after-increments", - user_id=None, - response_cost=0.25, - budget_reservation=budget_reservation, - ) - - assert incremented_counters == ["spend:team:team-finalize-after-increments"] - assert budget_reservation["finalized"] is True - assert counter_cache.in_memory_cache.get_cache(key="spend:key:key-finalize-after-increments") == pytest.approx( - 0.25 + monkeypatch.setattr(ps, "spend_counter_store", counter_cache) + monkeypatch.setattr(ps, "user_api_key_cache", DualCache()) + with patch( + "litellm.proxy.proxy_server._prepare_spend_counter_increment", + new=AsyncMock(side_effect=assert_reservation_not_finalized_yet), + ): + await ps.increment_spend_counters( + token="key-finalize-after-increments", + team_id="team-finalize-after-increments", + user_id=None, + response_cost=0.25, + budget_reservation=budget_reservation, ) - finally: - ps.spend_counter_cache = orig_counter - ps.user_api_key_cache = orig_user + + assert incremented_counters == ["spend:team:team-finalize-after-increments"] + assert budget_reservation["finalized"] is True + assert counter_cache.get_local("spend:key:key-finalize-after-increments") == pytest.approx(0.25) @pytest.mark.asyncio -async def test_increment_spend_counters_finalizes_none_cost_reservation(): - from litellm.caching.dual_cache import DualCache - from litellm.proxy.proxy_server import increment_spend_counters +async def test_increment_spend_counters_finalizes_none_cost_reservation(monkeypatch): + import litellm.proxy.proxy_server as ps - counter_cache = DualCache() - counter_cache.in_memory_cache.set_cache( - key="spend:key:key-finalize-none-cost", - value=0.5, - ) + counter_cache = SpendCounterStore(None) + counter_cache.set_local("spend:key:key-finalize-none-cost", 0.5) budget_reservation = { "reserved_cost": 0.5, "entries": [ @@ -9662,24 +9420,18 @@ async def test_increment_spend_counters_finalizes_none_cost_reservation(): ], "finalized": False, } + monkeypatch.setattr(ps, "spend_counter_store", counter_cache) - import litellm.proxy.proxy_server as ps + await ps.increment_spend_counters( + token="key-finalize-none-cost", + team_id=None, + user_id=None, + response_cost=None, + budget_reservation=budget_reservation, + ) - orig_counter = ps.spend_counter_cache - ps.spend_counter_cache = counter_cache - try: - await increment_spend_counters( - token="key-finalize-none-cost", - team_id=None, - user_id=None, - response_cost=None, - budget_reservation=budget_reservation, - ) - - assert budget_reservation["finalized"] is True - assert counter_cache.in_memory_cache.get_cache(key="spend:key:key-finalize-none-cost") == pytest.approx(0.0) - finally: - ps.spend_counter_cache = orig_counter + assert budget_reservation["finalized"] is True + assert counter_cache.get_local("spend:key:key-finalize-none-cost") == pytest.approx(0.0) @pytest.mark.asyncio @@ -9689,11 +9441,10 @@ async def test_increment_spend_counters_reseeds_from_db_on_bad_reserved_counter( behavior, which left the counter unenforced after a Redis reload). It reseeds from the authoritative DB and adds this request's settled cost, which the async spend flush has not written yet, so budget gating continues.""" - from litellm.caching.dual_cache import DualCache from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed from litellm.proxy.proxy_server import increment_spend_counters - counter_cache = DualCache() + counter_cache = SpendCounterStore(None) budget_reservation = { "reserved_cost": 0.5, "entries": [ @@ -9710,9 +9461,9 @@ async def test_increment_spend_counters_reseeds_from_db_on_bad_reserved_counter( import litellm.proxy.proxy_server as ps - orig_counter = ps.spend_counter_cache + orig_counter = ps.spend_counter_store orig_prisma = ps.prisma_client - ps.spend_counter_cache = counter_cache + ps.spend_counter_store = counter_cache ps.prisma_client = MagicMock() # truthy so reseed reaches from_db try: with patch.object(SpendCounterReseed, "from_db", AsyncMock(return_value=0.6)): @@ -9725,92 +9476,51 @@ async def test_increment_spend_counters_reseeds_from_db_on_bad_reserved_counter( ) assert budget_reservation["finalized"] is True - assert counter_cache.in_memory_cache.get_cache(key="spend:key:key-bad-reserved-counter") == pytest.approx(0.85) + assert counter_cache.get_local("spend:key:key-bad-reserved-counter") == pytest.approx(0.85) finally: - ps.spend_counter_cache = orig_counter + ps.spend_counter_store = orig_counter ps.prisma_client = orig_prisma @pytest.mark.asyncio -async def test_increment_spend_counter_invalidates_stale_cache_on_redis_failure(): - from litellm.caching.dual_cache import DualCache - from litellm.proxy.proxy_server import increment_spend_counter_cache - - counter_cache = DualCache() - counter_cache.in_memory_cache.set_cache(key="spend:team:redis-fail", value=4.0) - fake_redis = AsyncMock() - fake_redis.async_increment = AsyncMock(side_effect=RuntimeError("redis down")) - fake_redis.async_delete_cache = AsyncMock() - counter_cache.redis_cache = fake_redis - +async def test_increment_spend_counter_invalidates_stale_cache_on_redis_failure(monkeypatch): import litellm.proxy.proxy_server as ps - orig_counter = ps.spend_counter_cache - ps.spend_counter_cache = counter_cache - try: - with pytest.raises(RuntimeError): - await increment_spend_counter_cache( - counter_key="spend:team:redis-fail", - increment=0.5, - ) + spend = FakeSpendCounters(memory={"spend:team:redis-fail": 4.0}, fail=True) + monkeypatch.setattr(ps, "spend_counter_store", spend.store()) - assert counter_cache.in_memory_cache.get_cache(key="spend:team:redis-fail") is None - fake_redis.async_delete_cache.assert_awaited_once_with(key="spend:team:redis-fail") - finally: - ps.spend_counter_cache = orig_counter + with pytest.raises(NativeSpendCounterError): + await ps.increment_spend_counter_cache(counter_key="spend:team:redis-fail", increment=0.5) + + assert spend.memory == {} @pytest.mark.asyncio -async def test_get_current_spend_reseeds_from_db_when_counter_missing(): +async def test_get_current_spend_reseeds_from_db_when_counter_missing(monkeypatch): """ When both the Redis and in-memory counters are missing, the enforcement read path must reseed from the authoritative DB, not fall back to the caller-supplied stale value. Otherwise, every Redis TTL expiry lets a request through against a stale in-process `team_membership.spend`. """ - from litellm.caching.dual_cache import DualCache - from litellm.proxy.proxy_server import get_current_spend + import litellm.proxy.proxy_server as ps - counter_cache = DualCache() - recorded_seeds: list = [] - - async def record_set_cache(key, value, nx=False, **kwargs): - recorded_seeds.append({"key": key, "value": value, "nx": nx}) - return True - - fake_redis = AsyncMock() - fake_redis.async_set_cache = AsyncMock(side_effect=record_set_cache) - fake_redis.async_get_cache = AsyncMock(return_value=None) - counter_cache.redis_cache = fake_redis - - # DB has authoritative spend=362.0; caller hands us stale fallback=30.0 - # (the in-process team_membership.spend that hasn't caught up to DB). + spend = FakeSpendCounters() db_row = MagicMock() db_row.spend = 362.0 fake_prisma = MagicMock() fake_prisma.db.litellm_teammembership.find_unique = AsyncMock(return_value=db_row) + monkeypatch.setattr(ps, "spend_counter_store", spend.store()) + monkeypatch.setattr(ps, "prisma_client", fake_prisma) - import litellm.proxy.proxy_server as ps + result = await ps.get_current_spend(counter_key="spend:team_member:user-1:team-1", fallback_spend=30.0) - orig_counter, orig_prisma = ps.spend_counter_cache, ps.prisma_client - ps.spend_counter_cache = counter_cache - ps.prisma_client = fake_prisma - try: - spend = await get_current_spend( - counter_key="spend:team_member:user-1:team-1", - fallback_spend=30.0, - ) - assert spend == 362.0, ( - f"expected DB reseed to return 362.0, got {spend} (fallback would have returned 30.0 and caused bypass)" - ) - # Counter warmed via SET NX so subsequent reads are fast. - assert ("spend:team_member:user-1:team-1", 362.0, True) in [ - (s["key"], s["value"], s["nx"]) for s in recorded_seeds - ] - assert counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-1:team-1") == pytest.approx(362.0) - finally: - ps.spend_counter_cache = orig_counter - ps.prisma_client = orig_prisma + assert result == 362.0, ( + f"expected DB reseed to return 362.0, got {result} (fallback would have returned 30.0 and caused bypass)" + ) + # Counter warmed via SET NX so subsequent reads are fast. + assert "SETNX spend:team_member:user-1:team-1 362.0" in spend.commands + assert spend.memory == {"spend:team_member:user-1:team-1": 362.0} @pytest.mark.asyncio @@ -9819,18 +9529,17 @@ async def test_get_current_spend_uses_fallback_when_db_unavailable(): If prisma is unavailable and both counters are missing, the read path must degrade to the caller-supplied fallback rather than raising. """ - from litellm.caching.dual_cache import DualCache from litellm.proxy.proxy_server import get_current_spend - counter_cache = DualCache() + counter_cache = SpendCounterStore(None) fake_redis = AsyncMock() fake_redis.async_get_cache = AsyncMock(return_value=None) counter_cache.redis_cache = fake_redis import litellm.proxy.proxy_server as ps - orig_counter, orig_prisma = ps.spend_counter_cache, ps.prisma_client - ps.spend_counter_cache = counter_cache + orig_counter, orig_prisma = ps.spend_counter_store, ps.prisma_client + ps.spend_counter_store = counter_cache ps.prisma_client = None # simulate prisma unavailable try: spend = await get_current_spend( @@ -9839,7 +9548,7 @@ async def test_get_current_spend_uses_fallback_when_db_unavailable(): ) assert spend == 15.5 finally: - ps.spend_counter_cache = orig_counter + ps.spend_counter_store = orig_counter ps.prisma_client = orig_prisma @@ -9852,10 +9561,9 @@ async def test_get_current_spend_coalesces_concurrent_reseeds(): """ import asyncio as _asyncio - from litellm.caching.dual_cache import DualCache from litellm.proxy.proxy_server import get_current_spend - counter_cache = DualCache() + counter_cache = SpendCounterStore(None) counter_key = "spend:team_member:user-1:team-coalesce" # Track DB query calls and inject a small delay so the concurrent @@ -9896,8 +9604,8 @@ async def test_get_current_spend_coalesces_concurrent_reseeds(): import litellm.proxy.proxy_server as ps - orig_counter, orig_prisma = ps.spend_counter_cache, ps.prisma_client - ps.spend_counter_cache = counter_cache + orig_counter, orig_prisma = ps.spend_counter_store, ps.prisma_client + ps.spend_counter_store = counter_cache ps.prisma_client = fake_prisma try: results = await _asyncio.gather( @@ -9906,7 +9614,7 @@ async def test_get_current_spend_coalesces_concurrent_reseeds(): assert results == [100.0] * 5, f"all callers should see DB value, got {results}" assert db_call_count == 1, f"expected exactly 1 DB query for 5 concurrent reseeds, got {db_call_count}" finally: - ps.spend_counter_cache = orig_counter + ps.spend_counter_store = orig_counter ps.prisma_client = orig_prisma @@ -9918,10 +9626,9 @@ async def test_get_current_spend_uses_db_zero_over_stale_fallback(): fallback in production is the in-process team_membership.spend, which can still hold the pre-reset value across pods. """ - from litellm.caching.dual_cache import DualCache from litellm.proxy.proxy_server import get_current_spend - counter_cache = DualCache() + counter_cache = SpendCounterStore(None) fake_redis = AsyncMock() fake_redis.async_get_cache = AsyncMock(return_value=None) counter_cache.redis_cache = fake_redis @@ -9933,8 +9640,8 @@ async def test_get_current_spend_uses_db_zero_over_stale_fallback(): import litellm.proxy.proxy_server as ps - orig_counter, orig_prisma = ps.spend_counter_cache, ps.prisma_client - ps.spend_counter_cache = counter_cache + orig_counter, orig_prisma = ps.spend_counter_store, ps.prisma_client + ps.spend_counter_store = counter_cache ps.prisma_client = fake_prisma try: spend = await get_current_spend( @@ -9943,7 +9650,7 @@ async def test_get_current_spend_uses_db_zero_over_stale_fallback(): ) assert spend == 0.0, f"DB authoritative 0 must override stale fallback 42, got {spend}" finally: - ps.spend_counter_cache = orig_counter + ps.spend_counter_store = orig_counter ps.prisma_client = orig_prisma @@ -9964,7 +9671,7 @@ async def test_concurrent_read_and_write_paths_share_one_db_query(): get_current_spend, ) - counter_cache = DualCache() + counter_cache = SpendCounterStore(None) counter_key = "spend:team_member:user-1:team-cross-path" db_call_count = 0 @@ -10004,11 +9711,11 @@ async def test_concurrent_read_and_write_paths_share_one_db_query(): import litellm.proxy.proxy_server as ps orig_counter, orig_prisma, orig_user = ( - ps.spend_counter_cache, + ps.spend_counter_store, ps.prisma_client, ps.user_api_key_cache, ) - ps.spend_counter_cache = counter_cache + ps.spend_counter_store = counter_cache ps.prisma_client = fake_prisma ps.user_api_key_cache = DualCache() try: @@ -10028,7 +9735,7 @@ async def test_concurrent_read_and_write_paths_share_one_db_query(): assert results[0] in (50.0, 51.5), f"got {results[0]}" assert results[2] in (50.0, 51.5), f"got {results[2]}" finally: - ps.spend_counter_cache = orig_counter + ps.spend_counter_store = orig_counter ps.prisma_client = orig_prisma ps.user_api_key_cache = orig_user @@ -10071,39 +9778,17 @@ async def test_reseed_locks_dict_is_bounded(): @pytest.mark.asyncio -async def test_reseed_warms_cache_even_on_zero_db_spend(): +async def test_reseed_warms_cache_even_on_zero_db_spend(monkeypatch): """ When DB returns 0.0 (fresh entity / just after reset), the cache must still be warmed so subsequent reads hit the cache instead of issuing another DB query. Skipping the warm causes O(requests) DB load on zero-spend entities. """ - from litellm.caching.dual_cache import DualCache - from litellm.proxy.proxy_server import get_current_spend + import litellm.proxy.proxy_server as ps - counter_cache = DualCache() counter_key = "spend:team_member:user-1:team-zero-warm" - redis_store: dict = {} - - async def redis_get(key, **_): - return redis_store.get(key) - - async def redis_increment(key, value, **_): - redis_store[key] = (redis_store.get(key) or 0.0) + value - return redis_store[key] - - async def redis_set_cache(key, value, nx=False, **_): - if nx and key in redis_store: - return False - redis_store[key] = float(value) - return True - - fake_redis = AsyncMock() - fake_redis.async_get_cache = AsyncMock(side_effect=redis_get) - fake_redis.async_increment = AsyncMock(side_effect=redis_increment) - fake_redis.async_set_cache = AsyncMock(side_effect=redis_set_cache) - counter_cache.redis_cache = fake_redis - + spend = FakeSpendCounters() db_call_count = 0 async def find_unique(**kwargs): @@ -10115,23 +9800,15 @@ async def test_reseed_warms_cache_even_on_zero_db_spend(): fake_prisma = MagicMock() fake_prisma.db.litellm_teammembership.find_unique = AsyncMock(side_effect=find_unique) + monkeypatch.setattr(ps, "spend_counter_store", spend.store()) + monkeypatch.setattr(ps, "prisma_client", fake_prisma) - import litellm.proxy.proxy_server as ps + spend1 = await ps.get_current_spend(counter_key=counter_key, fallback_spend=0.0) + spend2 = await ps.get_current_spend(counter_key=counter_key, fallback_spend=0.0) - orig_counter, orig_prisma = ps.spend_counter_cache, ps.prisma_client - ps.spend_counter_cache = counter_cache - ps.prisma_client = fake_prisma - try: - # First call: cold cache, hits DB, returns 0. - spend1 = await get_current_spend(counter_key=counter_key, fallback_spend=0.0) - # Second call: cache should be warmed at 0, no second DB query. - spend2 = await get_current_spend(counter_key=counter_key, fallback_spend=0.0) - assert spend1 == 0.0 and spend2 == 0.0 - assert db_call_count == 1, f"second read should hit warmed cache, got {db_call_count} DB queries" - assert redis_store.get(counter_key) == 0.0, "cache must be warmed at 0" - finally: - ps.spend_counter_cache = orig_counter - ps.prisma_client = orig_prisma + assert spend1 == 0.0 and spend2 == 0.0 + assert db_call_count == 1, f"second read should hit warmed cache, got {db_call_count} DB queries" + assert spend.redis == {counter_key: 0.0}, "cache must be warmed at 0" # ----------------------------------------------------------------------------- @@ -11117,40 +10794,25 @@ async def test_get_current_spend_redis_clean_miss_skips_stale_in_memory(): @pytest.mark.asyncio -async def test_get_current_spend_redis_error_falls_back_to_in_memory(): +async def test_get_current_spend_redis_error_falls_back_to_in_memory(monkeypatch): """When Redis raises, the read should still degrade to in-memory rather than going straight to DB - in-memory is at least same-pod-fresh and cheaper than a DB query during a Redis outage.""" - from litellm.caching.dual_cache import DualCache - from litellm.proxy.proxy_server import get_current_spend - - counter_cache = DualCache() - counter_key = "spend:team_member:user-1:team-1" - - counter_cache.in_memory_cache.set_cache(key=counter_key, value=42.0) - - fake_redis = AsyncMock() - fake_redis.async_get_cache = AsyncMock(side_effect=ConnectionError("redis down")) - counter_cache.redis_cache = fake_redis - - fake_prisma = MagicMock() - fake_prisma.db.litellm_teammembership.find_unique = AsyncMock(return_value=MagicMock(spend=999.0)) - import litellm.proxy.proxy_server as ps - orig_counter, orig_prisma = ps.spend_counter_cache, ps.prisma_client - ps.spend_counter_cache = counter_cache - ps.prisma_client = fake_prisma - try: - spend = await get_current_spend(counter_key=counter_key, fallback_spend=0.0) - assert spend == 42.0, ( - f"expected in-memory fallback 42.0 on Redis error, got {spend} (should not have hit DB when Redis errored)" - ) - # DB query should NOT have fired - in-memory short-circuits. - fake_prisma.db.litellm_teammembership.find_unique.assert_not_awaited() - finally: - ps.spend_counter_cache = orig_counter - ps.prisma_client = orig_prisma + counter_key = "spend:team_member:user-1:team-1" + spend = FakeSpendCounters(memory={counter_key: 42.0}, fail=True) + fake_prisma = MagicMock() + fake_prisma.db.litellm_teammembership.find_unique = AsyncMock(return_value=MagicMock(spend=999.0)) + monkeypatch.setattr(ps, "spend_counter_store", spend.store()) + monkeypatch.setattr(ps, "prisma_client", fake_prisma) + + result = await ps.get_current_spend(counter_key=counter_key, fallback_spend=0.0) + + assert result == 42.0, ( + f"expected in-memory fallback 42.0 on Redis error, got {result} (should not have hit DB when Redis errored)" + ) + fake_prisma.db.litellm_teammembership.find_unique.assert_not_awaited() def test_realtime_websocket_route_aliases_registered(): @@ -13519,6 +13181,7 @@ class _EnvBuiltRedisCache(RedisCache): def __init__(self, **kwargs): self.init_kwargs = kwargs + self.default_ttl = 60 async def ping(self) -> bool: return True @@ -13553,6 +13216,13 @@ def _patched_coordination_redis_module_state( with ( patch.object(proxy_server_module, "redis_usage_cache", None), patch.object(proxy_server_module, "spend_counter_cache", spend_cache), + patch.object( + proxy_server_module, + "spend_counter_store", + SpendCounterStore( + None, handle=lambda client, _: FakeSpendCounters(redis=None if client is None else {}) + ), + ), patch.object(proxy_server_module, "user_api_key_cache", DualCache()), patch.object(proxy_server_module, "cli_sso_session_cache", DualCache()), patch.object(proxy_server_module, "llm_router", None), @@ -13639,6 +13309,7 @@ class _EnvBuiltClusterCache(RedisClusterCache): def __init__(self, **kwargs): self.init_kwargs = kwargs + self.default_ttl = 60 def _run_init_coordination_redis(config, env=None): @@ -14163,11 +13834,13 @@ def _window_spend_enqueue_env( originals = ( ps.user_api_key_cache, ps.spend_counter_cache, + ps.spend_counter_store, ps.prisma_client, ps.proxy_logging_obj, ) ps.user_api_key_cache = user_api_key_cache ps.spend_counter_cache = DualCache() + ps.spend_counter_store = FakeSpendCounters(redis=None).store() ps.prisma_client = prisma_client ps.proxy_logging_obj = proxy_logging_obj try: @@ -14176,6 +13849,7 @@ def _window_spend_enqueue_env( ( ps.user_api_key_cache, ps.spend_counter_cache, + ps.spend_counter_store, ps.prisma_client, ps.proxy_logging_obj, ) = originals @@ -14528,7 +14202,7 @@ async def test_a_key_update_that_lands_during_the_row_read_does_not_hide_the_win prisma_client.db.litellm_verificationtoken.find_unique = read_the_row_from_before_the_update with _window_spend_enqueue_env({}, prisma_client=prisma_client) as queue: - ps.spend_counter_cache.set_cache(key="spend:key:hashed-token", value=0.0) + ps.spend_counter_store.set_local("spend:key:hashed-token", 0.0) await increment_spend_counters(token="hashed-token", team_id=None, user_id=None, response_cost=0.25) _ = await _drain(queue) await increment_spend_counters(token="hashed-token", team_id=None, user_id=None, response_cost=0.5) @@ -14559,7 +14233,7 @@ async def test_late_requests_from_many_keys_with_windows_do_not_push_out_a_known reads_after_the_first_request: Final = no_window_key_reads() for token in windowed_tokens: await increment_spend_counters(token=token, team_id=None, user_id=None, response_cost=0.1) - ps.spend_counter_cache.set_cache(key="spend:key:no-windows", value=0.1) + ps.spend_counter_store.set_local("spend:key:no-windows", 0.1) await increment_spend_counters(token="no-windows", team_id=None, user_id=None, response_cost=0.1) assert no_window_key_reads() == reads_after_the_first_request @@ -14582,16 +14256,16 @@ async def test_a_slow_window_row_read_does_not_hold_back_the_key_spend_counter() prisma_client.db.litellm_verificationtoken.find_unique = slow_row_read with _window_spend_enqueue_env({}, prisma_client=prisma_client) as queue: - ps.spend_counter_cache.set_cache(key="spend:key:hashed-token", value=0.0) - ps.spend_counter_cache.set_cache(key="spend:key:hashed-token:window:7d", value=1.0) + ps.spend_counter_store.set_local("spend:key:hashed-token", 0.0) + ps.spend_counter_store.set_local("spend:key:hashed-token:window:7d", 1.0) charge: Final = asyncio.create_task( increment_spend_counters(token="hashed-token", team_id=None, user_id=None, response_cost=0.25) ) assert await asyncio.wait_for(read_started.wait(), timeout=5) - key_spend_while_the_row_read_waits: Final = ps.spend_counter_cache.get_cache(key="spend:key:hashed-token") + key_spend_while_the_row_read_waits: Final = ps.spend_counter_store.get_local("spend:key:hashed-token") finish_read.set() await charge - window_spend: Final = ps.spend_counter_cache.get_cache(key="spend:key:hashed-token:window:7d") + window_spend: Final = ps.spend_counter_store.get_local("spend:key:hashed-token:window:7d") charged: Final = await _drain(queue) assert key_spend_while_the_row_read_waits == pytest.approx(0.25) diff --git a/tests/unit/proxy/test_redis_auth_cache_flag.py b/tests/unit/proxy/test_redis_auth_cache_flag.py index d414e17b078..f0af8dc9220 100644 --- a/tests/unit/proxy/test_redis_auth_cache_flag.py +++ b/tests/unit/proxy/test_redis_auth_cache_flag.py @@ -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) diff --git a/tests/unit/redis/stores/spend_counters_double.py b/tests/unit/redis/stores/spend_counters_double.py new file mode 100644 index 00000000000..a13076541ba --- /dev/null +++ b/tests/unit/redis/stores/spend_counters_double.py @@ -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] diff --git a/tests/unit/redis/stores/test_spend_counters.py b/tests/unit/redis/stores/test_spend_counters.py new file mode 100644 index 00000000000..332af1c92d6 --- /dev/null +++ b/tests/unit/redis/stores/test_spend_counters.py @@ -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)