diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index faf121c6a56..469b994559d 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -41,6 +41,10 @@ mod _native { #[pymodule_export] use crate::redis::client::native_redis_breaker_states; #[pymodule_export] + use crate::redis::store::auth_cache::{ + NativeAuthCache, NativeAuthCacheError, NativeDeclaredAuthReads, NativeDeclaredAuthWrite, + }; + #[pymodule_export] use crate::redis::store::request::NativeRequestRedis; #[pymodule_export] use crate::redis::store::router_usage::NativeRouterUsage; @@ -125,6 +129,10 @@ mod tests { "NativeDiagnosticProcessor", "NativeRedisClient", "NativeRequestRedis", + "NativeAuthCache", + "NativeAuthCacheError", + "NativeDeclaredAuthReads", + "NativeDeclaredAuthWrite", "NativePending", "NativeRouterCooldowns", "NativeRouterUsage", diff --git a/litellm-rust/crates/python-bridge/src/redis/store/auth_cache.rs b/litellm-rust/crates/python-bridge/src/redis/store/auth_cache.rs new file mode 100644 index 00000000000..ac65d97020e --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/store/auth_cache.rs @@ -0,0 +1,338 @@ +use std::{ + sync::{Arc, Mutex, PoisonError}, + time::Duration, +}; + +use litellm_python_compat::Value; +use litellm_redis_store::auth_cache::{AuthCache, DeclaredGets, DeclaredWrite, Error, Settings}; +use pyo3::{exceptions::PyRuntimeError, prelude::*}; + +use super::{PyNumber, call, request::NativeRequestRedis, store_error, value::PyValue}; +use crate::execution::{run_async_value, run_sync_value}; +use crate::redis::client::NativeRedisClient; + +pyo3::create_exception!( + _native, + NativeAuthCacheError, + pyo3::exceptions::PyException, + "An auth cache operation failed. Args are (kind, message), kind one of \"breaker_open\", \"timeout\", \"transport\", \"rejected\", \"decode\", \"encode\" or \"memory\"." +); + +fn auth_cache_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::Encode(_) => "encode", + Error::Memory(_) => "memory", + }; + NativeAuthCacheError::new_err((kind, litellm_error::Report::new(&error).to_string())) +} + +fn ttl(seconds: Option) -> Option { + seconds.map(|seconds| Duration::from_secs_f64(seconds.max(0.0))) +} + +fn values(values: Vec>) -> Vec> { + values.into_iter().map(|value| value.map(PyValue)).collect() +} + +/// `UserApiKeyCache` with `enable_redis_auth_cache` on. The sync methods back `DualCache`'s sync +/// API and block with the GIL released, like redis-py's sync client. +#[pyclass(frozen, name = "NativeAuthCache")] +pub(crate) struct NativeAuthCache(Arc); + +#[pymethods] +impl NativeAuthCache { + #[new] + #[pyo3(signature = (client, *, memory_entries = 200, memory_ttl_seconds = None, redis_ttl_seconds = 60))] + fn new( + client: Option>, + memory_entries: usize, + memory_ttl_seconds: Option, + redis_ttl_seconds: u64, + ) -> Self { + Self(Arc::new(AuthCache::new( + client.map(|client| client.0.clone()), + Settings { + memory_entries, + memory_ttl: ttl(memory_ttl_seconds), + redis_ttl: Duration::from_secs(redis_ttl_seconds), + }, + ))) + } + + #[getter] + fn has_redis(&self) -> bool { + self.0.has_redis() + } + + fn get<'py>(&self, py: Python<'py>, key: String) -> PyResult> { + let store = self.0.clone(); + let scope = call(py, "async_get_cache")?; + run_async_value(py, async move { + let value = scope.scope(store.get(&key)).await.map_err(store_error)?; + Ok(value.map(PyValue)) + }) + } + + fn get_sync(&self, py: Python<'_>, key: String) -> PyResult> { + let store = self.0.clone(); + let scope = call(py, "get_cache")?; + run_sync_value(py, async move { + let value = scope.scope(store.get(&key)).await.map_err(store_error)?; + Ok(value.map(PyValue)) + }) + } + + #[pyo3(signature = (key, value, ttl_seconds = None))] + fn set<'py>( + &self, + py: Python<'py>, + key: String, + value: PyValue, + 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(&key, value.0, ttl(ttl_seconds))) + .await + .map_err(store_error) + }) + } + + #[pyo3(signature = (key, value, ttl_seconds = None))] + fn set_sync( + &self, + py: Python<'_>, + key: String, + value: PyValue, + ttl_seconds: Option, + ) -> PyResult<()> { + let store = self.0.clone(); + let scope = call(py, "set_cache")?; + run_sync_value(py, async move { + scope + .scope(store.set_str(&key, value.0, ttl(ttl_seconds))) + .await + .map_err(store_error) + }) + } + + /// Memory now and the `SET` on the request's pre-call batch; `None` without Redis. + #[pyo3(signature = (request, key, value, ttl_seconds = None))] + fn set_pre_call( + &self, + request: &NativeRequestRedis, + key: String, + value: PyValue, + ttl_seconds: Option, + ) -> PyResult> { + Ok(self + .0 + .set_pre_call(&request.0, &key, value.0, ttl(ttl_seconds)) + .map_err(store_error)? + .map(NativeDeclaredAuthWrite::new)) + } + + fn delete_pre_call( + &self, + request: &NativeRequestRedis, + key: String, + ) -> PyResult> { + Ok(self + .0 + .delete_pre_call(&request.0, &key) + .map_err(store_error)? + .map(NativeDeclaredAuthWrite::new)) + } + + #[pyo3(signature = (entries, ttl_seconds = None))] + fn set_many<'py>( + &self, + py: Python<'py>, + entries: Vec<(String, PyValue)>, + ttl_seconds: Option, + ) -> PyResult> { + let store = self.0.clone(); + let entries = entries + .into_iter() + .map(|(key, value)| (key, value.0)) + .collect(); + let scope = call(py, "async_set_cache_pipeline")?; + run_async_value(py, async move { + scope + .scope(store.set_many(entries, ttl(ttl_seconds))) + .await + .map_err(store_error) + }) + } + + /// On the request's pre-call batch when `request` is given, returning before anything is sent. + fn write_back<'py>( + &self, + py: Python<'py>, + request: Option<&NativeRequestRedis>, + entries: Vec<(String, PyValue, Option)>, + ) -> PyResult> { + let store = self.0.clone(); + let request = request.map(|request| request.0.clone()); + let entries = entries + .into_iter() + .map(|(key, value, seconds)| (key, value.0, ttl(seconds))) + .collect(); + let scope = call(py, "async_set_cache_pipeline_with_ttls")?; + run_async_value(py, async move { + scope + .scope(store.write_back(request.as_ref(), entries)) + .await + .map_err(auth_cache_error) + }) + } + + fn delete<'py>(&self, py: Python<'py>, keys: Vec) -> PyResult> { + let store = self.0.clone(); + let scope = call(py, "async_delete_cache")?; + run_async_value(py, async move { + scope + .scope(store.delete(&keys)) + .await + .map_err(auth_cache_error) + }) + } + + fn delete_sync(&self, py: Python<'_>, keys: Vec) -> PyResult<()> { + let store = self.0.clone(); + let scope = call(py, "delete_cache")?; + run_sync_value(py, async move { + scope + .scope(store.delete(&keys)) + .await + .map_err(auth_cache_error) + }) + } + + /// Reads of keys the caller found missing in memory, on the request's pre-call batch; `None` + /// without Redis. + fn declare_get_many( + &self, + request: &NativeRequestRedis, + keys: Vec, + ) -> Option { + self.0 + .declare_get_many(&request.0, &keys) + .map(|declared| NativeDeclaredAuthReads(Mutex::new(Some(declared)))) + } + + /// One round trip for keys missing in memory; resolves to `None` without Redis. + fn get_many<'py>(&self, py: Python<'py>, keys: Vec) -> PyResult> { + let store = self.0.clone(); + let scope = call(py, "async_batch_get_cache")?; + run_async_value(py, async move { + scope + .scope(store.get_many(&keys)) + .await + .transpose() + .map(|read| read.map(values)) + .map_err(auth_cache_error) + }) + } + + fn increment_if_present<'py>( + &self, + py: Python<'py>, + key: String, + amount: PyNumber, + ) -> PyResult> { + let store = self.0.clone(); + let scope = call(py, "async_increment_if_present")?; + run_async_value(py, async move { + scope + .scope(store.increment_if_present(&key, amount.0)) + .await + .map_err(auth_cache_error) + }) + } + + fn get_local(&self, key: &str) -> PyResult> { + Ok(self.0.get_local(key).map_err(store_error)?.map(PyValue)) + } + + #[pyo3(signature = (key, value, ttl_seconds = None))] + fn set_local(&self, key: &str, value: PyValue, ttl_seconds: Option) -> PyResult<()> { + self.0 + .set_local(key, value.0, ttl(ttl_seconds)) + .map_err(store_error) + } + + fn delete_local(&self, key: &str) -> PyResult<()> { + self.0.delete_local(key).map_err(store_error) + } + + fn flush_local(&self) -> PyResult<()> { + self.0.flush_local().map_err(store_error) + } + + /// `(management objects, user-key objects)` held in memory. + fn len_local(&self) -> PyResult<(usize, usize)> { + self.0.len_local().map_err(store_error) + } +} + +/// A `declare_get_many` on a request's pre-call batch. `resolve` answers it once. +#[pyclass(frozen, name = "NativeDeclaredAuthReads")] +pub(crate) struct NativeDeclaredAuthReads(Mutex>); + +#[pymethods] +impl NativeDeclaredAuthReads { + 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 auth cache read was already resolved"))?; + let scope = call(py, "async_batch_get_cache")?; + run_async_value(py, async move { + scope + .scope(declared.resolve()) + .await + .map(values) + .map_err(auth_cache_error) + }) + } +} + +/// A `set_pre_call` or `delete_pre_call`. `resolve` settles it once, raising when Redis refused it. +#[pyclass(frozen, name = "NativeDeclaredAuthWrite")] +pub(crate) struct NativeDeclaredAuthWrite(Mutex>); + +impl NativeDeclaredAuthWrite { + fn new(declared: DeclaredWrite) -> Self { + Self(Mutex::new(Some(declared))) + } +} + +#[pymethods] +impl NativeDeclaredAuthWrite { + 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 auth cache write was already resolved"))?; + let scope = call(py, "async_set_cache")?; + run_async_value(py, async move { + scope + .scope(declared.resolve()) + .await + .map_err(auth_cache_error) + }) + } +} 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 ee38003d22f..42dcc91b5f9 100644 --- a/litellm-rust/crates/python-bridge/src/redis/store/mod.rs +++ b/litellm-rust/crates/python-bridge/src/redis/store/mod.rs @@ -1,3 +1,4 @@ +pub(crate) mod auth_cache; mod policy; pub(crate) mod request; pub(crate) mod router_cooldowns; diff --git a/litellm-rust/crates/redis-store/src/auth_cache.rs b/litellm-rust/crates/redis-store/src/auth_cache.rs new file mode 100644 index 00000000000..ae333e539a7 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/auth_cache.rs @@ -0,0 +1,463 @@ +use std::time::Duration; + +use litellm_cache::Error as MemoryError; +use litellm_cache_memory::InMemoryCache; +use litellm_python_compat::{Value, json, repr::to_str}; +use litellm_redis_client::{Declared, RedisBatch, RedisClient}; + +use crate::tier::{Clock, system_clock}; +use crate::{DecodeError, Number, RequestRedis, codec::python_value}; + +/// `InMemoryCache`'s TTL for a write that names none and has no `default_in_memory_ttl`. +pub const MEMORY_TTL: Duration = Duration::from_secs(600); +/// `BaseCache.default_ttl`, which `RedisCache` keeps when `litellm.default_redis_ttl` is unset. +pub const REDIS_TTL: Duration = Duration::from_secs(60); +pub const MEMORY_ENTRIES: usize = 200; + +/// `RedisCache.async_increment_if_present`, byte for byte, sent as `EVAL` like it. +const INCREMENT_IF_PRESENT: &str = "if redis.call('EXISTS', KEYS[1]) == 0 then return false end return redis.call('INCRBYFLOAT', KEYS[1], ARGV[1])"; +/// `DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE`, the chunk `DualCache.async_delete_cache_keys` deletes at once. +const DELETE_CHUNK: usize = 1000; + +/// `UserApiKeyCache`'s configuration: each memory partition's size, `default_in_memory_ttl`, and the +/// `RedisCache.default_ttl` of the Redis it was attached to. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Settings { + pub memory_entries: usize, + pub memory_ttl: Option, + pub redis_ttl: Duration, +} + +impl Default for Settings { + fn default() -> Self { + Self { + memory_entries: MEMORY_ENTRIES, + memory_ttl: None, + redis_ttl: REDIS_TTL, + } + } +} + +#[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), + #[error("the value is not JSON serializable: {0}")] + Encode(#[from] litellm_python_compat::Error), +} + +/// `is_user_key_cache_key`: user-key objects are cached under a bare `hash_token` digest, every +/// other object under a prefixed key. +pub fn is_key_object(key: &str) -> bool { + key.len() == 64 + && key + .bytes() + .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte)) +} + +/// `BaseCache.get_ttl` sends `int(ttl)`, so Redis sees whole seconds. +fn whole_seconds(ttl: Duration) -> Duration { + Duration::from_secs(ttl.as_secs()) +} + +fn decode(raw: Option>) -> Result, DecodeError> { + raw.as_deref().map(python_value).transpose() +} + +/// The proxy's auth object cache with `enable_redis_auth_cache` on: two memory partitions, so churn +/// in management objects cannot evict user-key objects, in front of one Redis shared by every pod. +/// Keys are the logical cache keys `UserApiKeyCache` callers already use; values are what +/// `json.loads` returns, written as `json.dumps`, like `RedisCache.async_set_cache`. +pub struct AuthCache { + key_objects: InMemoryCache, + objects: InMemoryCache, + remote: Option, + settings: Settings, +} + +impl AuthCache { + pub fn new(client: Option, settings: Settings) -> Self { + Self::with_clock(client, settings, system_clock()) + } + + pub fn with_clock(client: Option, settings: Settings, now: Clock) -> Self { + let memory = || { + let now = now.clone(); + InMemoryCache::with_clock(Some(settings.memory_entries), Some(MEMORY_TTL), move || { + now() + }) + }; + Self { + key_objects: memory(), + objects: memory(), + remote: client, + settings, + } + } + + pub fn has_redis(&self) -> bool { + self.remote.is_some() + } + + pub fn settings(&self) -> Settings { + self.settings + } + + fn memory(&self, key: &str) -> &InMemoryCache { + if is_key_object(key) { + &self.key_objects + } else { + &self.objects + } + } + + /// `DualCache`'s rule: a write without a TTL takes `default_in_memory_ttl` in memory and Redis. + fn memory_ttl(&self, ttl: Option) -> Option { + ttl.or(self.settings.memory_ttl) + } + + fn redis_ttl(&self, ttl: Option) -> Duration { + whole_seconds(self.memory_ttl(ttl).unwrap_or(self.settings.redis_ttl)) + } + + pub fn get_local(&self, key: &str) -> Result, MemoryError> { + self.memory(key).get_cache(key) + } + + /// A memory write with `ttl`, or `default_in_memory_ttl` when it names none. + pub fn set_local( + &self, + key: &str, + value: Value, + ttl: Option, + ) -> Result<(), MemoryError> { + self.memory(key) + .set_cache(key, value, self.memory_ttl(ttl)) + .map(drop) + } + + pub fn delete_local(&self, key: &str) -> Result<(), MemoryError> { + self.memory(key).delete_cache(key) + } + + pub fn flush_local(&self) -> Result<(), MemoryError> { + self.key_objects.flush_cache()?; + self.objects.flush_cache() + } + + /// Entries held by the management-object and the user-key partitions. + pub fn len_local(&self) -> Result<(usize, usize), MemoryError> { + Ok((self.objects.len()?, self.key_objects.len()?)) + } + + /// `DualCache.async_get_cache`: memory, then Redis, backfilling memory. A failed Redis read is a + /// miss, as `DualCache` logs and swallows it. + pub async fn get(&self, key: &str) -> Result, MemoryError> { + if let Some(value) = self.get_local(key)? { + return Ok(Some(value)); + } + let Some(remote) = &self.remote else { + return Ok(None); + }; + let read = remote + .get::>(key) + .await + .map_err(Error::from) + .and_then(|raw| Ok(decode(raw)?)); + match read { + Ok(Some(value)) => { + self.set_local(key, value.clone(), None)?; + Ok(Some(value)) + } + Ok(None) => Ok(None), + Err(error) => { + tracing::warn!(%error, "auth cache Redis read failed, treating it as a miss"); + Ok(None) + } + } + } + + /// `DualCache.async_set_cache`: memory, then `SET` of `json.dumps(value)`. Redis failures and + /// values `json.dumps` refuses only log, as there. + pub async fn set( + &self, + key: &str, + value: Value, + ttl: Option, + ) -> Result<(), MemoryError> { + let encoded = json::dumps(&value); + self.set_local(key, value, ttl)?; + let Some(remote) = &self.remote else { + return Ok(()); + }; + let written = match encoded { + Ok(bytes) => remote + .set(key, bytes, self.redis_ttl(ttl)) + .await + .map_err(Error::from), + Err(error) => Err(error.into()), + }; + if let Err(error) = written { + tracing::warn!(%error, "writing an auth cache entry to Redis failed"); + } + Ok(()) + } + + /// `DualCache.set_cache`, whose `RedisCache.set_cache` writes `str(value)`. + pub async fn set_str( + &self, + key: &str, + value: Value, + ttl: Option, + ) -> Result<(), MemoryError> { + let encoded = to_str(&value); + self.set_local(key, value, ttl)?; + let Some(remote) = &self.remote else { + return Ok(()); + }; + if let Err(error) = remote.set(key, encoded, self.redis_ttl(ttl)).await { + tracing::warn!(%error, "writing an auth cache entry to Redis failed"); + } + Ok(()) + } + + /// `DualCache.async_set_cache_pre_call`: memory now, the `SET` on the request's pre-call batch, + /// sent with whatever read goes out next. `None` without Redis. + pub fn set_pre_call( + &self, + request: &RequestRedis, + key: &str, + value: Value, + ttl: Option, + ) -> Result, MemoryError> { + let encoded = json::dumps(&value); + let Some(remote) = &self.remote else { + return Ok(None); + }; + self.set_local(key, value, ttl)?; + Ok(Some(match encoded { + Ok(bytes) => { + DeclaredWrite::Set(request.pre(remote).set(key, bytes, self.redis_ttl(ttl))) + } + Err(error) => DeclaredWrite::Unencodable(error), + })) + } + + /// `DualCache.async_delete_cache_pre_call`. `None` without Redis. + pub fn delete_pre_call( + &self, + request: &RequestRedis, + key: &str, + ) -> Result, MemoryError> { + let Some(remote) = &self.remote else { + return Ok(None); + }; + self.delete_local(key)?; + Ok(Some(DeclaredWrite::Delete(request.pre(remote).del(key)))) + } + + /// `UserApiKeyCache.async_set_cache_pipeline`: the key-object entries in one pipeline, then + /// the rest in another, each with one TTL. + pub async fn set_many( + &self, + entries: Vec<(String, Value)>, + ttl: Option, + ) -> Result<(), MemoryError> { + let (key_objects, others): (Vec<_>, Vec<_>) = + entries.into_iter().partition(|(key, _)| is_key_object(key)); + for partition in [key_objects, others] { + if !partition.is_empty() { + self.set_partition(partition, ttl).await?; + } + } + Ok(()) + } + + async fn set_partition( + &self, + entries: Vec<(String, Value)>, + ttl: Option, + ) -> Result<(), MemoryError> { + let mut pipeline = self.remote.as_ref().map(RedisClient::pipeline); + for (key, value) in entries { + if let Some(writes) = pipeline.take() { + pipeline = Some(match json::dumps(&value) { + Ok(bytes) => writes.set(&key, bytes, self.redis_ttl(ttl)), + Err(error) => { + tracing::warn!(%error, "an auth cache value is not JSON serializable, kept in memory only"); + writes + } + }); + } + self.set_local(&key, value, ttl)?; + } + if let (Some(remote), Some(pipeline)) = (&self.remote, pipeline) { + for reply in remote.run(&pipeline).await { + if let Err(error) = reply { + tracing::warn!(%error, "writing an auth cache entry to Redis failed"); + } + } + } + Ok(()) + } + + /// The auth prefetch write-back, `async_set_cache_pipeline_with_ttls`: an entry without a TTL + /// takes `default_in_memory_ttl` in memory but the connection's default in Redis. On the + /// request's pre-call batch when there is one, else in a round trip of its own. + pub async fn write_back( + &self, + request: Option<&RequestRedis>, + entries: Vec<(String, Value, Option)>, + ) -> Result<(), Error> { + let Some(remote) = &self.remote else { + for (key, value, ttl) in entries { + self.set_local(&key, value, ttl)?; + } + return Ok(()); + }; + if let Some(request) = request { + let batch = request.pre(remote); + for (key, value, ttl) in entries { + let redis_ttl = whole_seconds(ttl.unwrap_or(self.settings.redis_ttl)); + drop(batch.set(&key, json::dumps(&value)?, redis_ttl)); + self.set_local(&key, value, ttl)?; + } + return Ok(()); + } + let mut pipeline = remote.pipeline(); + for (key, value, ttl) in entries { + let redis_ttl = whole_seconds(ttl.unwrap_or(self.settings.redis_ttl)); + pipeline = pipeline.set(&key, json::dumps(&value)?, redis_ttl); + self.set_local(&key, value, ttl)?; + } + for reply in remote.run(&pipeline).await { + reply?; + } + Ok(()) + } + + /// `UserApiKeyCache.async_delete_cache_keys`: memory, then the key-object partition and the + /// rest as two concurrent `DualCache.async_delete_cache_keys`, each a multi-key `DEL` per + /// chunk of 1000. Both run even when one fails. + pub async fn delete(&self, keys: &[String]) -> Result<(), Error> { + for key in keys { + self.delete_local(key)?; + } + let Some(remote) = &self.remote else { + return Ok(()); + }; + let (key_objects, others): (Vec<&str>, Vec<&str>) = keys + .iter() + .map(String::as_str) + .partition(|key| is_key_object(key)); + let (key_objects, others) = tokio::join!( + delete_chunks(remote, &key_objects), + delete_chunks(remote, &others) + ); + key_objects.and(others) + } + + /// Reads for keys the caller found missing in memory, on the request's pre-call batch. `None` + /// without Redis. + pub fn declare_get_many( + &self, + request: &RequestRedis, + keys: &[String], + ) -> Option { + Some(self.declare_gets(&request.pre(self.remote.as_ref()?), keys)) + } + + /// `RedisCache.async_batch_get_cache` for keys missing in memory, in one round trip. `None` + /// without Redis. + pub async fn get_many(&self, keys: &[String]) -> Option>, Error>> { + Some( + self.declare_gets(&self.remote.as_ref()?.batch(), keys) + .resolve() + .await, + ) + } + + fn declare_gets(&self, batch: &RedisBatch, keys: &[String]) -> DeclaredGets { + let keys: Vec<&str> = keys.iter().map(String::as_str).collect(); + DeclaredGets(batch.mget::>(&keys)) + } + + /// `RedisCache.async_increment_if_present`: `INCRBYFLOAT` only an entry Redis still holds. + /// `None` without Redis or for an absent entry. + pub async fn increment_if_present( + &self, + key: &str, + amount: Number, + ) -> Result, Error> { + let Some(remote) = &self.remote else { + return Ok(None); + }; + let total: Option> = remote + .eval_source( + INCREMENT_IF_PRESENT, + &[key], + litellm_redis_client::Amount::from(amount), + ) + .await?; + total + .as_deref() + .map(|raw| { + std::str::from_utf8(raw) + .ok() + .and_then(|text| text.trim().parse::().ok()) + .ok_or(Error::Decode(DecodeError::Shape( + "INCRBYFLOAT returned a non-number", + ))) + }) + .transpose() + } +} + +/// A write declared on a request's pre-call batch, settled once the batch goes out. +pub enum DeclaredWrite { + Set(Declared<()>), + Delete(Declared), + /// Kept in memory only, as `json.dumps` refused the value. + Unencodable(litellm_python_compat::Error), +} + +impl DeclaredWrite { + /// Flushes the batch if nobody has yet. + pub async fn resolve(self) -> Result<(), Error> { + match self { + Self::Set(written) => written.resolve().await?, + Self::Delete(deleted) => drop(deleted.resolve().await?), + Self::Unencodable(error) => { + tracing::warn!(%error, "an auth cache value is not JSON serializable, kept in memory only"); + } + } + Ok(()) + } +} + +/// Reads declared on a batch, answered once it goes out. +pub struct DeclaredGets(Declared>>>); + +impl DeclaredGets { + /// Flushes the batch if nobody has yet. Any failed read or undecodable value fails the whole + /// read, as a failed `MGET` does. + pub async fn resolve(self) -> Result>, Error> { + self.0 + .resolve() + .await? + .into_iter() + .map(|reply| Ok(decode(reply)?)) + .collect() + } +} + +async fn delete_chunks(remote: &RedisClient, keys: &[&str]) -> Result<(), Error> { + for chunk in keys.chunks(DELETE_CHUNK) { + remote.del(chunk).await?; + } + Ok(()) +} diff --git a/litellm-rust/crates/redis-store/src/lib.rs b/litellm-rust/crates/redis-store/src/lib.rs index 2f46dcd541b..b53df6d46fd 100644 --- a/litellm-rust/crates/redis-store/src/lib.rs +++ b/litellm-rust/crates/redis-store/src/lib.rs @@ -1,3 +1,4 @@ +pub mod auth_cache; mod codec; mod error; mod number; diff --git a/litellm-rust/crates/redis-store/tests/auth_cache.rs b/litellm-rust/crates/redis-store/tests/auth_cache.rs new file mode 100644 index 00000000000..d3f5de9ae62 --- /dev/null +++ b/litellm-rust/crates/redis-store/tests/auth_cache.rs @@ -0,0 +1,552 @@ +mod support; + +use std::time::Duration; + +use litellm_python_compat::Value; +use litellm_redis_client::{Operation, RedisClient, Ttl}; +use litellm_redis_store::{ + Number, RequestRedis, + auth_cache::{self, AuthCache, Settings}, + router_cooldowns::{self, RouterCooldowns}, +}; +use litellm_redis_testing::{RedisDeployment, Topology}; +use rstest::rstest; +use support::{connect, connect_observed, within}; +use tokio::runtime::Handle; + +const LONG_DEADLINE: Duration = Duration::from_secs(60); +const KEY_OBJECT: &str = "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"; +const TEAM_JSON: &str = concat!( + r#"{"team_id": "team-\u00e9", "models": ["gpt-5", "claude-opus-5"], "max_budget": null, "#, + r#""spend": 0.30000000000000004, "blocked": false, "tpm_limit": 10000000000, "#, + r#""metadata": {"tags": [], "quoted": "it's \"x\""}}"#, +); +const TEAM_STR: &str = concat!( + "{'team_id': 'team-é', 'models': ['gpt-5', 'claude-opus-5'], 'max_budget': None, ", + "'spend': 0.30000000000000004, 'blocked': False, 'tpm_limit': 10000000000, ", + r#"'metadata': {'tags': [], 'quoted': 'it\'s "x"'}}"#, +); + +fn dict(entries: [(&str, Value); N]) -> Value { + Value::Dict( + entries + .into_iter() + .map(|(key, value)| (Value::from(key), value)) + .collect(), + ) +} + +fn strings(items: &[&str]) -> Value { + Value::List(items.iter().copied().map(Value::from).collect()) +} + +fn team_record() -> Value { + dict([ + ("team_id", Value::from("team-é")), + ("models", strings(&["gpt-5", "claude-opus-5"])), + ("max_budget", Value::None), + ("spend", Value::Float(0.30000000000000004)), + ("blocked", Value::Bool(false)), + ("tpm_limit", Value::from(10000000000)), + ( + "metadata", + dict([ + ("tags", strings(&[])), + ("quoted", Value::from("it's \"x\"")), + ]), + ), + ]) +} + +/// The writes `tests/unit/redis/stores/test_auth_cache.py` checks against Python's own +/// `json.dumps` and `str`. +fn written() -> [(Value, &'static str, &'static str); 7] { + [ + (team_record(), TEAM_JSON, TEAM_STR), + ( + dict([ + ("token", Value::from(KEY_OBJECT)), + ("spend", Value::Float(1e-05)), + ("rpm_limit", Value::from(5)), + ("expires", Value::from("2026-11-03T00:00:00Z")), + ]), + concat!( + r#"{"token": "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b", "spend": 1e-05, "#, + r#""rpm_limit": 5, "expires": "2026-11-03T00:00:00Z"}"#, + ), + concat!( + "{'token': '88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b', 'spend': 1e-05, ", + "'rpm_limit': 5, 'expires': '2026-11-03T00:00:00Z'}", + ), + ), + ( + Value::from("__no_team_membership__"), + "\"__no_team_membership__\"", + "__no_team_membership__", + ), + ( + strings(&["mcp-a", "mcp-b"]), + r#"["mcp-a", "mcp-b"]"#, + "['mcp-a', 'mcp-b']", + ), + (Value::Float(12.5), "12.5", "12.5"), + (Value::from(7), "7", "7"), + (Value::Bool(true), "true", "True"), + ] +} + +/// The reads `tests/unit/redis/stores/test_auth_cache.py` checks against Python's own +/// `RedisCache._get_cache_logic`. +fn read_back() -> [(&'static str, Value); 5] { + [ + (TEAM_JSON, team_record()), + (TEAM_STR, team_record()), + ("['mcp-a', 'mcp-b']", strings(&["mcp-a", "mcp-b"])), + ( + "{'user_id': 'u1', 'teams': ['t1'], 'active': True, 'budget': None}", + dict([ + ("user_id", Value::from("u1")), + ("teams", strings(&["t1"])), + ("active", Value::Bool(true)), + ("budget", Value::None), + ]), + ), + ("\"__NO_MAPPING__\"", Value::from("__NO_MAPPING__")), + ] +} + +fn settings(memory_ttl: Option) -> Settings { + Settings { + memory_ttl: memory_ttl.map(Duration::from_secs), + ..Settings::default() + } +} + +fn store(client: Option) -> AuthCache { + AuthCache::new(client, settings(Some(60))) +} + +fn team(id: &str) -> Value { + Value::Dict(vec![(Value::from("team_id"), Value::from(id))]) +} + +async fn redis_ttl(raw: &RedisClient, key: &str) -> Duration { + let Ttl::Expires(remaining) = raw.ttl(key).await.unwrap() else { + panic!("{key} has no TTL"); + }; + remaining +} + +#[rstest] +#[case(KEY_OBJECT, true)] +#[case( + "88DC28D0F030C55ED4AB77ED8FAF098196CB1C05DF778539800C9F1243FE6B4B", + false +)] +#[case( + "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4", + false +)] +#[case( + "team_id:88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b", + false +)] +#[case("team_membership:user-1:team-1", false)] +#[case("end_user_restricted_registry", false)] +fn only_a_bare_lowercase_sha256_digest_is_a_user_key_object( + #[case] key: &str, + #[case] key_object: bool, +) { + assert_eq!(auth_cache::is_key_object(key), key_object, "{key}"); +} + +#[rstest] +#[tokio::test] +async fn writes_mains_json_and_str_bytes( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let store = store(Some(connect(&deployment))); + + for (index, (value, json, printed)) in written().into_iter().enumerate() { + let key = format!("team_id:t{index}"); + store.set(&key, value.clone(), None).await.unwrap(); + assert_eq!( + raw.get::(&key).await.unwrap().as_deref(), + Some(json) + ); + store.set_str(&key, value, None).await.unwrap(); + assert_eq!( + raw.get::(&key).await.unwrap().as_deref(), + Some(printed) + ); + } +} + +#[rstest] +#[tokio::test] +async fn reads_every_encoding_main_writes_and_backfills_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))); + + for (index, (bytes, expected)) in read_back().into_iter().enumerate() { + let key = format!("user_id:u{index}"); + raw.set(&key, bytes, Duration::from_secs(60)).await.unwrap(); + + assert_eq!(store.get(&key).await.unwrap(), Some(expected.clone())); + raw.del(&[&key]).await.unwrap(); + assert_eq!(store.get(&key).await.unwrap(), Some(expected), "{key}"); + } + assert_eq!(store.get("user_id:absent").await.unwrap(), None); +} + +#[rstest] +#[tokio::test] +async fn ttls_follow_dual_cache_and_the_write_back_rule( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let store = AuthCache::new( + Some(connect(&deployment)), + Settings { + memory_ttl: Some(Duration::from_secs(20)), + redis_ttl: Duration::from_secs(45), + ..Settings::default() + }, + ); + + store + .set( + "team_id:given", + team("a"), + Some(Duration::from_secs_f64(30.9)), + ) + .await + .unwrap(); + store.set("team_id:default", team("b"), None).await.unwrap(); + store + .write_back( + None, + vec![ + ("user-1".to_owned(), team("c"), None), + ("user-2".to_owned(), team("d"), Some(Duration::from_secs(5))), + ], + ) + .await + .unwrap(); + + assert!(within(redis_ttl(&raw, "team_id:given").await, 30)); + assert!(within(redis_ttl(&raw, "team_id:default").await, 20)); + assert!(within(redis_ttl(&raw, "user-1").await, 45)); + assert!(within(redis_ttl(&raw, "user-2").await, 5)); + assert_eq!(store.get_local("user-1").unwrap(), Some(team("c"))); +} + +#[tokio::test] +async fn user_key_objects_survive_churn_in_the_management_object_partition() { + let store = AuthCache::new( + None, + Settings { + memory_entries: 2, + ..Settings::default() + }, + ); + store.set(KEY_OBJECT, team("key"), None).await.unwrap(); + + for index in 0..5 { + store + .set(&format!("team_id:t{index}"), team("t"), None) + .await + .unwrap(); + } + + assert_eq!(store.get(KEY_OBJECT).await.unwrap(), Some(team("key"))); + assert_eq!(store.len_local().unwrap(), (2, 1)); +} + +#[rstest] +#[tokio::test] +async fn a_value_json_refuses_stays_in_memory_and_never_reaches_redis( + #[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 members = Value::Set(vec![Value::from("a")]); + + store + .set("team_id:set", members.clone(), None) + .await + .unwrap(); + + assert_eq!(store.get_local("team_id:set").unwrap(), Some(members)); + assert_eq!(raw.get::("team_id:set").await.unwrap(), None); +} + +#[rstest] +#[tokio::test] +async fn pre_call_writes_and_reads_ride_one_pipeline_with_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 store = store(Some(client)); + let request = RequestRedis::new(Handle::current(), LONG_DEADLINE); + raw.set(KEY_OBJECT, "{\"token\": \"t\"}", Duration::from_secs(60)) + .await + .unwrap(); + raw.set("team_alias:alpha", "\"stale\"", Duration::from_secs(60)) + .await + .unwrap(); + + let reads = store + .declare_get_many( + &request, + &[KEY_OBJECT.to_owned(), "end_user_id:absent".to_owned()], + ) + .unwrap(); + let active = cooldowns + .declare_active(&request, &["dep-a".to_owned()]) + .unwrap(); + let written = store + .set_pre_call(&request, "team_id:t1", team("t1"), None) + .unwrap() + .unwrap(); + let deleted = store + .delete_pre_call(&request, "team_alias:alpha") + .unwrap() + .unwrap(); + + assert!(trips.seen().is_empty()); + assert_eq!(store.get_local("team_id:t1").unwrap(), Some(team("t1"))); + request.flush_pre().await; + + let mut seen = trips.seen(); + seen.sort_by_key(|operation| format!("{operation:?}")); + let expected = match topology { + Topology::Standalone => vec![Operation::Pipeline { ops: 4 }], + Topology::Cluster => vec![ + Operation::Command("DEL"), + Operation::Command("MGET"), + Operation::Command("MGET"), + Operation::Command("SET"), + ], + }; + assert_eq!(seen, expected); + assert_eq!( + reads.resolve().await.unwrap(), + vec![ + Some(Value::Dict(vec![(Value::from("token"), Value::from("t"))])), + None + ] + ); + assert_eq!(cooldowns.resolve_active(active).await.unwrap(), vec![]); + written.resolve().await.unwrap(); + deleted.resolve().await.unwrap(); + assert_eq!( + raw.get::("team_id:t1").await.unwrap(), + Some("{\"team_id\": \"t1\"}".to_owned()) + ); + assert!(within(redis_ttl(&raw, "team_id:t1").await, 60)); + assert_eq!(raw.get::("team_alias:alpha").await.unwrap(), None); + assert_eq!(trips.seen().len(), expected.len()); +} + +#[rstest] +#[tokio::test] +async fn a_write_back_on_the_request_waits_for_the_next_flush( + #[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); + + store + .write_back( + Some(&request), + vec![( + "user-1".to_owned(), + team("u"), + Some(Duration::from_secs(42)), + )], + ) + .await + .unwrap(); + + assert!(trips.seen().is_empty()); + assert_eq!(store.get_local("user-1").unwrap(), Some(team("u"))); + request.flush_pre().await; + let write = match topology { + Topology::Standalone => Operation::Pipeline { ops: 1 }, + Topology::Cluster => Operation::Command("SET"), + }; + assert_eq!(trips.seen(), [write]); + assert!(within(redis_ttl(&raw, "user-1").await, 42)); +} + +#[rstest] +#[tokio::test] +async fn an_eviction_broadcast_drops_only_this_pods_copy_and_the_next_read_refills_from_redis( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let this_pod = store(Some(connect(&deployment))); + let other_pod = store(Some(connect(&deployment))); + this_pod.set("team_id:t1", team("old"), None).await.unwrap(); + other_pod + .set("team_id:t1", team("new"), None) + .await + .unwrap(); + + assert_eq!(this_pod.get("team_id:t1").await.unwrap(), Some(team("old"))); + this_pod.delete_local("team_id:t1").unwrap(); + + assert_eq!(this_pod.get("team_id:t1").await.unwrap(), Some(team("new"))); + assert!(raw.get::("team_id:t1").await.unwrap().is_some()); +} + +#[rstest] +#[tokio::test] +async fn delete_drops_memory_and_redis( + #[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 keys = [KEY_OBJECT.to_owned(), "team_id:t1".to_owned()]; + store + .set_many( + keys.iter().map(|key| (key.clone(), team("t"))).collect(), + None, + ) + .await + .unwrap(); + assert!(raw.get::("team_id:t1").await.unwrap().is_some()); + + store.delete(&keys).await.unwrap(); + + for key in &keys { + assert_eq!(store.get_local(key).unwrap(), None); + assert_eq!(raw.get::(key).await.unwrap(), None); + } +} + +#[rstest] +#[tokio::test] +async fn increment_if_present_only_raises_an_entry_redis_holds( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let store = store(Some(connect(&deployment))); + + assert_eq!( + store + .increment_if_present("litellm-proxy-budget:spend", Number::Float(2.5)) + .await + .unwrap(), + None + ); + assert_eq!( + raw.get::("litellm-proxy-budget:spend") + .await + .unwrap(), + None + ); + store + .set("litellm-proxy-budget:spend", Value::Float(1.5), None) + .await + .unwrap(); + + assert_eq!( + store + .increment_if_present("litellm-proxy-budget:spend", Number::Float(2.5)) + .await + .unwrap(), + Some(4.0) + ); +} + +#[tokio::test] +async fn without_redis_every_operation_is_memory_only() { + let store = store(None); + let request = RequestRedis::new(Handle::current(), LONG_DEADLINE); + + store.set("team_id:t1", team("t"), None).await.unwrap(); + + assert!( + store + .set_pre_call(&request, "team_id:t2", team("t"), None) + .unwrap() + .is_none() + ); + assert!( + store + .declare_get_many(&request, &["team_id:t1".to_owned()]) + .is_none() + ); + assert!(store.get_many(&["team_id:t1".to_owned()]).await.is_none()); + assert_eq!( + store + .increment_if_present("k", Number::Float(1.0)) + .await + .unwrap(), + None + ); + assert_eq!(store.get("team_id:t1").await.unwrap(), Some(team("t"))); + assert_eq!(store.get_local("team_id:t2").unwrap(), None); +} + +// UserApiKeyCache writes the key-object partition, then the rest, and deletes both partitions +// concurrently, while async_increment_if_present sends its script as a plain EVAL +#[rstest] +#[tokio::test] +async fn the_auth_cache_sends_the_commands_user_api_key_cache_sends( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let (client, trips) = connect_observed(&deployment); + let store = store(Some(client)); + deployment + .cli_on_every_node(&["CONFIG", "RESETSTAT"]) + .await + .unwrap(); + + store + .set_many( + vec![ + (KEY_OBJECT.to_owned(), team("t0")), + ("team_id:t1".to_owned(), team("t1")), + ], + None, + ) + .await + .unwrap(); + let written = trips.seen(); + store + .delete(&[KEY_OBJECT.to_owned(), "team_id:t1".to_owned()]) + .await + .unwrap(); + store + .increment_if_present("litellm-proxy-budget:spend", Number::Int(5)) + .await + .unwrap(); + + assert_eq!(written, vec![Operation::Pipeline { ops: 1 }; 2]); + assert_eq!( + trips.seen()[written.len()..written.len() + 2], + [Operation::Command("DEL"), Operation::Command("DEL")] + ); + assert_eq!(support::calls(&deployment, "eval").await, 1); + assert_eq!(support::calls(&deployment, "evalsha").await, 0); +} diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index c82d25c9425..815599252b8 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -12,7 +12,7 @@ import asyncio import itertools import logging import time -from collections.abc import Mapping, Sequence +from collections.abc import Awaitable, Mapping, Sequence from dataclasses import dataclass from threading import Lock from typing import TYPE_CHECKING, Any, Final @@ -497,7 +497,7 @@ class DualCache(BaseCache): batch: Final = None if self.redis_cache is None else active_request_redis_batch(self.redis_cache) return None if batch is None else await self._set_on_batch(batch, key, value, ttl) - async def async_delete_cache_pre_call(self, key: str) -> BatchResult[None] | None: + async def async_delete_cache_pre_call(self, key: str) -> Awaitable[None] | None: """Memory now, the Redis DEL on the request's pipeline; None when no pipeline is open, so the caller takes its direct path.""" batch: Final = None if self.redis_cache is None else active_request_redis_batch(self.redis_cache) diff --git a/litellm/caching/redis_batch.py b/litellm/caching/redis_batch.py index dd886b09067..60330a4275e 100644 --- a/litellm/caching/redis_batch.py +++ b/litellm/caching/redis_batch.py @@ -498,6 +498,7 @@ class RequestRedisBatches: "_native", "_native_factory", "_native_flush_hooks", + "_native_misses", "_post_call", "post_call_deadline", "prefetched", @@ -514,6 +515,7 @@ class RequestRedisBatches: self._native_factory: Final = native_factory self._native: NativeRequestRedis | None = None self._native_flush_hooks: Final[list[Callable[[], None]]] = [] # mutable-ok: append-only registry + self._native_misses: frozenset[str] = frozenset() self._deadline: asyncio.TimerHandle | None = None self._deadline_flush: asyncio.Task[None] | None = None # Reads declared early for a consumer that runs later in the request, keyed by consumer name. @@ -558,6 +560,16 @@ class RequestRedisBatches: """Called before every native pre-call flush so late readers can declare into the same trip.""" self._native_flush_hooks.append(hook) + def note_native_misses(self, keys: Sequence[str]) -> None: + """Keys a native pre-call read found absent (or a native DEL removed), like ``RedisBatch.read_as_missing``.""" + self._native_misses = self._native_misses | frozenset(keys) + + def forget_native_miss(self, key: str) -> None: + self._native_misses = self._native_misses - frozenset((key,)) + + def native_read_as_missing(self, key: str) -> bool: + return key in self._native_misses + async def flush_native_pre(self) -> None: for hook in self._native_flush_hooks: hook() diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 4db340a0d48..11ca3d9083c 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -3198,7 +3198,7 @@ async def cache_team_object( usage_cache: Final = None if proxy_logging_obj is None else proxy_logging_obj.internal_usage_cache.dual_cache # On a shared Redis the write below replaces the team entry and the alias DEL below removes the alias entry # for both caches, so the usage cache only has its own memory to clear. - redis_shared: Final = usage_cache is not None and usage_cache.redis_cache is user_api_key_cache.redis_cache + redis_shared: Final = usage_cache is not None and usage_cache.redis_cache is user_api_key_cache.redis_connection await _invalidate_usage_cache_entry(usage_cache, key, redis_shared=redis_shared, stale="team object") diff --git a/litellm/proxy/auth/auth_object_prefetch.py b/litellm/proxy/auth/auth_object_prefetch.py index fd5b99951f4..db2a2b30e0c 100644 --- a/litellm/proxy/auth/auth_object_prefetch.py +++ b/litellm/proxy/auth/auth_object_prefetch.py @@ -5,17 +5,15 @@ the readers and the fallback, so enforcement never depends on this running.""" from __future__ import annotations import time -from collections.abc import Iterator, Mapping, Sequence +from collections.abc import Callable, Iterator, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType -from typing import Final, Literal, Protocol, TypeAlias +from typing import Final, Literal, TypeAlias from pydantic import BaseModel, TypeAdapter, ValidationError from litellm._internal_context import service_target from litellm._logging import verbose_proxy_logger -from litellm.caching.redis_batch import active_request_redis_batch -from litellm.caching.redis_cache import RedisCache from litellm.constants import DEFAULT_IN_MEMORY_TTL, REGISTRY_ERROR_NEGATIVE_CACHE_TTL from litellm.models.organization import LiteLLM_OrganizationTable from litellm.models.team import LiteLLM_TeamTableCachedObj @@ -32,6 +30,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( ) from litellm.proxy.db.db_span import db_span from litellm.proxy.utils import PrismaClient +from litellm.redis.stores.auth_cache import AuthCacheStore _RowKind: TypeAlias = Literal["user_row", "team_row", "membership_row", "organization_row", "project_row"] @@ -162,11 +161,6 @@ class AuthObjectRefs: ) -class _InMemoryCache(Protocol): - def get_cache(self, key: str) -> object: ... - def set_cache(self, key: str, value: object, *, ttl: float | None = ...) -> None: ... - - @dataclass(frozen=True, slots=True) class _CacheEntry: cache_key: str @@ -211,39 +205,32 @@ def _entries(refs: AuthObjectRefs, cache: UserApiKeyCache) -> tuple[_CacheEntry, return tuple(_iter_entries(refs, get_management_object_ttl(cache))) -def _missing_in_memory(entries: Sequence[_CacheEntry], memory: _InMemoryCache) -> tuple[_CacheEntry, ...]: - return tuple(entry for entry in entries if memory.get_cache(key=entry.cache_key) is None) +def _missing_in_memory(entries: Sequence[_CacheEntry], cache: UserApiKeyCache) -> tuple[_CacheEntry, ...]: + return tuple(entry for entry in entries if cache.get_local(entry.cache_key) is None) -def _set_in_memory(memory: _InMemoryCache, cache_key: str, value: object, ttl: float | None) -> None: - if ttl is None: - memory.set_cache(key=cache_key, value=value) - else: - memory.set_cache(key=cache_key, value=value, ttl=ttl) +def _nothing_to_arm() -> None: + return None -async def _read_redis_rows(keys: list[str], redis_cache: RedisCache) -> Mapping[str, object]: - """On the request pipeline when one is open; a failed pipeline reads as a miss, like ``async_batch_get_cache``.""" - batch: Final = active_request_redis_batch(redis_cache) - with service_target(AUTH_OBJECTS_TARGET): - if batch is None: - return await redis_cache.async_batch_get_cache(key_list=keys) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # untyped cache API - try: - return await batch.mget(keys) - except Exception as e: # noqa: BLE001 # the DB fill below takes over, as it does after a failed MGET today - verbose_proxy_logger.debug("auth prefetch Redis read failed, filling from the database: %s", e) - return MappingProxyType({}) - - -async def _fill_from_redis(entries: Sequence[_CacheEntry], redis_cache: RedisCache, memory: _InMemoryCache) -> None: +async def _fill_from_redis( + entries: Sequence[_CacheEntry], shared: AuthCacheStore, cache: UserApiKeyCache, before_read: Callable[[], None] +) -> None: + """On the request's native pipeline when one is open; a failed read is a miss, like ``async_batch_get_cache``.""" if not entries: return - found: Final = _RowValues.validate_python( - await _read_redis_rows(sorted(entry.cache_key for entry in entries), redis_cache) - ) + before_read() + try: + with service_target(AUTH_OBJECTS_TARGET): + found: Final = await shared.read_missing(sorted(entry.cache_key for entry in entries)) + except Exception as e: # noqa: BLE001 # the DB fill below takes over, as it does after a failed MGET + verbose_proxy_logger.debug("auth prefetch Redis read failed, filling from the database: %s", e) + return for entry, value in ((entry, found.get(entry.cache_key)) for entry in entries): if value is not None: - _set_in_memory(memory, entry.cache_key, value, entry.ttl) + shared.set_local( + entry.cache_key, value, cache.in_memory_cache.default_ttl if entry.ttl is None else entry.ttl + ) def _validate_row( @@ -282,18 +269,13 @@ async def _write_back(entries: Sequence[tuple[_CacheEntry, BaseModel]], cache: U (entry.cache_key, CacheCodec.serialize(value, model_type=entry.model_type), entry.ttl) for entry, value in entries ) - memory: Final[_InMemoryCache] = cache.in_memory_cache - for cache_key, payload, ttl in payloads: - _set_in_memory(memory, cache_key, payload, cache.default_in_memory_ttl if ttl is None else ttl) - if cache.redis_cache is None: + shared: Final = cache.shared + if shared is not None: + with service_target(AUTH_OBJECTS_TARGET): + await shared.write_back(payloads) return - batch: Final = active_request_redis_batch(cache.redis_cache) - with service_target(AUTH_OBJECTS_TARGET): - if batch is None: - await cache.redis_cache.async_set_cache_pipeline_with_ttls(payloads) - return - for cache_key, payload, ttl in payloads: # rides the request's next round trip; the scope drains leftovers - batch.set(cache_key, payload, ttl) + for cache_key, payload, ttl in payloads: + cache.set_local(cache_key, payload, ttl) async def _fill_from_db( @@ -318,16 +300,18 @@ async def prefetch_auth_objects( refs: AuthObjectRefs, user_api_key_cache: UserApiKeyCache, prisma_client: PrismaClient | None, + before_read: Callable[[], None] = _nothing_to_arm, ) -> None: - """Best effort: any failure leaves the per-object getters to fetch as before.""" + """Best effort: any failure leaves the per-object getters to fetch as before. ``before_read`` runs right before + a Redis read goes out, so a caller can declare other reads into the same round trip.""" try: - memory: Final[_InMemoryCache] = user_api_key_cache.in_memory_cache - missing: Final = _missing_in_memory(_entries(refs, user_api_key_cache), memory) - if user_api_key_cache.redis_cache is not None: - await _fill_from_redis(missing, user_api_key_cache.redis_cache, memory) + missing: Final = _missing_in_memory(_entries(refs, user_api_key_cache), user_api_key_cache) + shared: Final = user_api_key_cache.shared + if shared is not None: + await _fill_from_redis(missing, shared, user_api_key_cache, before_read) if prisma_client is None: return - await _fill_from_db(refs, _missing_in_memory(missing, memory), user_api_key_cache, prisma_client) + await _fill_from_db(refs, _missing_in_memory(missing, user_api_key_cache), user_api_key_cache, prisma_client) except Exception as e: # noqa: BLE001 # warm-up only; the getters enforce and fail closed on their own verbose_proxy_logger.warning("auth prefetch skipped, falling back to per-object lookups: %s", e) @@ -337,28 +321,29 @@ def _identity_memory_ttl(value: object, management_ttl: float) -> float: return min(REGISTRY_ERROR_NEGATIVE_CACHE_TTL, management_ttl) if isinstance(value, str) else management_ttl -async def prefetch_identity_keys(cache_keys: Sequence[str], user_api_key_cache: UserApiKeyCache) -> None: +async def prefetch_identity_keys( + cache_keys: Sequence[str], + user_api_key_cache: UserApiKeyCache, + before_read: Callable[[], None] = _nothing_to_arm, +) -> None: """Warm the entries auth reads before it knows the key's owners (the key object, the end user and the two - registries) in one MGET on the request pipeline. Keys the MGET finds absent stay noted on the pipeline, so the + registries) in one read on the request's native pipeline. Keys it finds absent stay noted on the request, so the per-key getters that follow go to the database without a GET of their own. Best effort, like the owner prefetch: the getters read and enforce on their own.""" try: - redis_cache: Final = user_api_key_cache.redis_cache - if redis_cache is None: + shared: Final = user_api_key_cache.shared + if shared is None: return - missing: Final = tuple( - key - for key in dict.fromkeys(cache_keys) - if user_api_key_cache.in_memory_cache_for(key).get_cache(key=key) is None - ) + missing: Final = tuple(key for key in dict.fromkeys(cache_keys) if shared.get_local(key) is None) if not missing: return - found: Final = _RowValues.validate_python(await _read_redis_rows(sorted(missing), redis_cache)) + before_read() + with service_target(AUTH_OBJECTS_TARGET): + found: Final = await shared.read_missing(sorted(missing)) management_ttl: Final = get_management_object_ttl(user_api_key_cache) except Exception as e: # noqa: BLE001 # warm-up only; the getters read Redis and the database on their own verbose_proxy_logger.warning("auth identity prefetch skipped, falling back to per-key lookups: %s", e) return for key, value in ((key, found.get(key)) for key in missing): if value is not None: - memory: _InMemoryCache = user_api_key_cache.in_memory_cache_for(key) - _set_in_memory(memory, key, value, _identity_memory_ttl(value, management_ttl)) + shared.set_local(key, value, _identity_memory_ttl(value, management_ttl)) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 0a15b9aa62c..05fcf7cd3f3 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -188,7 +188,13 @@ from litellm.redis.stores.spend_counters import ( from litellm.repositories.table_repositories import JWTKeyMappingRepository, TeamMembershipRepository from litellm.repositories.user_repository import UserRepository from litellm.repositories.verification_token_repository import VerificationTokenRepository -from litellm.router_api import fallbacks_disabled, references_provider_scoped_resource, request_fallback_chain +from litellm.router_api import ( + RoutingContext, + fallbacks_disabled, + references_provider_scoped_resource, + request_fallback_chain, +) +from litellm.router_utils.routing_read_batch import RoutingPrefetch from litellm.secret_managers.main import get_secret_bool from litellm.types.services import ServiceTypes @@ -2651,6 +2657,7 @@ async def user_api_key_auth_builder( await prefetch_identity_keys( _identity_cache_keys(api_key, end_user_id=end_user_id, key_is_resolved=valid_token is not None), user_api_key_cache=user_api_key_cache, + before_read=arm_routing_cooldowns, ) if end_user_id: try: @@ -4303,9 +4310,20 @@ async def _prefetch_referenced_auth_objects( refs=AuthObjectRefs.from_token(valid_token), user_api_key_cache=user_api_key_cache, prisma_client=prisma_client, + before_read=arm_routing_cooldowns, ) +def arm_routing_cooldowns() -> None: + """Declares the cooldown read routing takes later into the auth prefetch's native round trip, so a request whose + auth objects come from Redis pays one round trip for both. Armed only when a prefetch reads Redis, so a request + served from memory, or denied, reads no cooldowns early.""" + from litellm.proxy.proxy_server import llm_router # noqa: PLC0415 # runtime global + + if llm_router is not None: + RoutingPrefetch.arm_cooldowns(RoutingContext(llm_router)) + + def _seed_request_destinations(user_api_key_dict: UserAPIKeyAuth, request: Request | None = None) -> None: """Anchor the OTLP destinations this key or team overrides its traces to. diff --git a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py index 39d300bb898..36650629300 100644 --- a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py +++ b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py @@ -190,7 +190,7 @@ def evict_local(cache_keys: Sequence[str], user_api_key_cache: "UserApiKeyCache" def _drop_from_memory(cache_key: str, user_api_key_cache: "UserApiKeyCache") -> None: - user_api_key_cache.delete_from_memory(cache_key) + user_api_key_cache.delete_local(cache_key) @with_service_target(AUTH_OBJECTS_TARGET) diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index 011b4e744ea..4c3931bc5ed 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -2,10 +2,10 @@ from __future__ import annotations import asyncio import re -from collections.abc import Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, TypeVar, cast, overload -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.caching.dual_cache import DualCache @@ -13,12 +13,14 @@ from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec, CacheValue +from litellm.redis.native_settings import NativeRedisSettings +from litellm.redis.stores._connection import native_client +from litellm.redis.stores.auth_cache import AuthCacheSettings, AuthCacheStore, HandleFactory, native_handle +from litellm.rust_bridge._native import NativeRedisClient if TYPE_CHECKING: from opentelemetry.trace import Span - from litellm.caching.redis_batch import BatchResult - T = TypeVar("T", bound=BaseModel) from litellm.constants import AUTH_OBJECTS_TARGET as AUTH_OBJECTS_TARGET # noqa: E402, PLC0414 # public re-export @@ -33,13 +35,18 @@ def is_user_key_cache_key(key: str) -> bool: _PIPELINED_SET_OPTIONS: Final = frozenset(("ttl",)) +_TTL: Final[TypeAdapter[float | None]] = TypeAdapter(float | None) + + +def _ttl(kwargs: Mapping[str, object]) -> float | None: + return _TTL.validate_python(kwargs.get("ttl")) class UserApiKeyCache(DualCache): """ - DualCache wrapper for UserAPIKeyAuth-like payloads. + DualCache for UserAPIKeyAuth-like payloads. - Stores a Redis-safe JSON payload in BOTH in-memory and Redis to avoid + Stores a Redis-safe JSON payload in memory and, with ``enable_redis_auth_cache``, in Redis, to avoid "memory returns BaseModel, Redis returns dict" format drift. When ``model_type`` is provided: @@ -53,12 +60,12 @@ class UserApiKeyCache(DualCache): When ``model_type`` is omitted, the interface behaves like ``DualCache``: raw cached payload is returned (dict/str/etc.). - ``async_set_cache_pipeline`` applies the same untyped Codec pass as omitting - ``model_type`` on ``async_set_cache`` (so ``BaseModel`` rows are dumped before Redis). - User-key objects (see ``is_user_key_cache_key``) live in their own in-memory partition, - ``key_object_cache``, so churn in the other management objects cannot evict them. Both - partitions share the same Redis backend and TTL settings. + ``key_object_cache``, so churn in the other management objects cannot evict them. + + Without Redis both partitions are this class's Python memory. ``attach_redis_cache`` moves the whole + keyspace, both memory partitions included, onto the native ``AuthCacheStore`` (``shared``); redis-py never + reads or writes it, and ``redis_cache`` stays None. ``get_cache`` / ``async_get_cache`` overloads and implementations must be contiguous (no other methods in between) so mypy resolves ``@overload`` + implementation correctly. @@ -72,16 +79,16 @@ class UserApiKeyCache(DualCache): default_redis_ttl: float | None = None, key_object_in_memory_cache: InMemoryCache | None = None, budget_windows_absent_in_memory_cache: InMemoryCache | None = None, + auth_cache_handle: HandleFactory = native_handle, + redis_client: Callable[[NativeRedisSettings], NativeRedisClient | None] = native_client, ) -> None: super().__init__( in_memory_cache=in_memory_cache, - redis_cache=redis_cache, default_in_memory_ttl=default_in_memory_ttl, default_redis_ttl=default_redis_ttl, ) self.key_object_cache: Final = DualCache( in_memory_cache=key_object_in_memory_cache or InMemoryCache(), - redis_cache=redis_cache, default_in_memory_ttl=default_in_memory_ttl, default_redis_ttl=default_redis_ttl, ) @@ -90,14 +97,34 @@ class UserApiKeyCache(DualCache): or InMemoryCache(default_ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL), default_in_memory_ttl=default_in_memory_ttl, ) + self._auth_cache_handle: Final = auth_cache_handle + self._redis_client: Final = redis_client + self._redis_connection: RedisCache | None = None + self._shared: AuthCacheStore | None = None + self.attach_redis_cache(redis_cache, default_redis_ttl=default_redis_ttl) - def in_memory_cache_for(self, key: str) -> InMemoryCache: - return self.key_object_cache.in_memory_cache if is_user_key_cache_key(key) else self.in_memory_cache + @property + def shared(self) -> AuthCacheStore | None: + """The native store holding the keyspace once ``enable_redis_auth_cache`` attached Redis.""" + return self._shared - def delete_from_memory(self, key: str) -> None: - self.in_memory_cache_for(key).delete_cache(key) - if key.startswith(_BUDGET_WINDOWS_ABSENT_PREFIX): - self.budget_windows_absent_cache.delete_cache(key) + @property + def has_redis(self) -> bool: + return self._shared is not None + + @property + def redis_connection(self) -> RedisCache | None: + """The Redis the auth cache was attached to, for features that share it under keys of their own.""" + return self._redis_connection + + def _settings(self) -> AuthCacheSettings: + return AuthCacheSettings( + memory_entries=self.in_memory_cache.max_size_in_memory, memory_ttl=self.default_in_memory_ttl + ) + + def _python_only(self, method: str) -> None: + if self._shared is not None: + raise NotImplementedError(f"UserApiKeyCache.{method} is not part of the shared auth cache keyspace") def update_cache_ttl(self, default_in_memory_ttl: float | None, default_redis_ttl: float | None) -> None: super().update_cache_ttl(default_in_memory_ttl=default_in_memory_ttl, default_redis_ttl=default_redis_ttl) @@ -107,17 +134,62 @@ class UserApiKeyCache(DualCache): self.budget_windows_absent_cache.update_cache_ttl( default_in_memory_ttl=default_in_memory_ttl, default_redis_ttl=None ) + if self._shared is not None: + self._shared.reconfigure(self._settings()) def update_in_memory_max_size(self, max_size: int | None) -> None: super().update_in_memory_max_size(max_size) self.key_object_cache.update_in_memory_max_size(max_size) self.budget_windows_absent_cache.update_in_memory_max_size(max_size) + if self._shared is not None: + self._shared.reconfigure(self._settings()) def attach_redis_cache( self, redis_cache: RedisCache | None = None, *, default_redis_ttl: float | None = None ) -> None: - super().attach_redis_cache(redis_cache, default_redis_ttl=default_redis_ttl) - self.key_object_cache.attach_redis_cache(redis_cache, default_redis_ttl=default_redis_ttl) + if redis_cache is None or self._shared is not None: + return + if default_redis_ttl is not None: + self.update_cache_ttl(default_in_memory_ttl=None, default_redis_ttl=default_redis_ttl) + self._redis_connection = redis_cache + native_settings: Final = redis_cache.native_settings + if native_settings is None: + return + self._shared = AuthCacheStore( + self._redis_client(native_settings), + self._settings(), + redis_ttl=native_settings.default_ttl, + handle=self._auth_cache_handle, + ) + + def get_local(self, key: str) -> object: + """This worker's memory copy alone.""" + if self._shared is not None: + return self._shared.get_local(key) + return self._python_memory(key).get_cache(key=key) + + def set_local(self, key: str, value: object, ttl: float | None) -> None: + """A memory write alone, with ``default_in_memory_ttl`` when ``ttl`` is None.""" + if self._shared is not None: + self._shared.set_local(key, value, ttl) + return + effective_ttl: Final = self.default_in_memory_ttl if ttl is None else ttl + if effective_ttl is None: + self._python_memory(key).set_cache(key=key, value=value) + return + self._python_memory(key).set_cache(key=key, value=value, ttl=effective_ttl) + + def delete_local(self, key: str) -> None: + """Drops this worker's memory copy alone, as a cross-worker eviction broadcast does.""" + if key.startswith(_BUDGET_WINDOWS_ABSENT_PREFIX): + self.budget_windows_absent_cache.delete_cache(key) + if self._shared is not None: + self._shared.delete_local(key) + return + self._python_memory(key).delete_cache(key) + + def _python_memory(self, key: str) -> InMemoryCache: + return self.key_object_cache.in_memory_cache if is_user_key_cache_key(key) else self.in_memory_cache @overload def get_cache( @@ -150,24 +222,17 @@ class UserApiKeyCache(DualCache): ) -> object: if model_type is None and "model_type" in kwargs: model_type = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) + shared: Final = self._shared cached: Final = ( - self.key_object_cache.get_cache(key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs) + self._get_shared_sync(shared, key, local_only) + if shared is not None + else self.key_object_cache.get_cache( + key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs + ) if is_user_key_cache_key(key) else super().get_cache(key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs) ) - if model_type is None: - return cached - if cached is None: - return None - decoded: Final = CacheCodec.deserialize(cached, model_type=model_type) - if decoded is None: - verbose_proxy_logger.error( - "UserApiKeyCache.get_cache failed to deserialize cached value for key=%r model_type=%s", - key, - getattr(model_type, "__name__", str(model_type)), - ) - return None - return decoded + return _decoded(cached, key, model_type, "get_cache") @overload async def async_get_cache( @@ -200,8 +265,11 @@ class UserApiKeyCache(DualCache): ) -> object: if model_type is None and "model_type" in kwargs: model_type = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) + shared: Final = self._shared cached: Final = ( - await self.key_object_cache.async_get_cache( + await self._get_shared(shared, key, local_only) + if shared is not None + else await self.key_object_cache.async_get_cache( key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs ) if is_user_key_cache_key(key) @@ -209,23 +277,39 @@ class UserApiKeyCache(DualCache): key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs ) ) - if model_type is None: - return cached - if cached is None: + return _decoded(cached, key, model_type, "async_get_cache") + + @staticmethod + def _get_shared_sync(shared: AuthCacheStore, key: str, local_only: bool) -> object: + try: + return shared.get_local(key) if local_only else shared.get_sync(key) + except Exception as e: # noqa: BLE001 # DualCache.get_cache reads a failure as a miss + verbose_proxy_logger.error("UserApiKeyCache.get_cache failed for key=%r: %s", key, e) return None - decoded: Final = CacheCodec.deserialize(cached, model_type=model_type) - if decoded is None: - verbose_proxy_logger.error( - "UserApiKeyCache.async_get_cache failed to deserialize cached value for key=%r model_type=%s", - key, - getattr(model_type, "__name__", str(model_type)), - ) + + @staticmethod + async def _get_shared(shared: AuthCacheStore, key: str, local_only: bool) -> object: + try: + return shared.get_local(key) if local_only else await shared.get(key) + except Exception as e: # noqa: BLE001 # DualCache.async_get_cache reads a failure as a miss + verbose_proxy_logger.error("UserApiKeyCache.async_get_cache failed for key=%r: %s", key, e) return None - return decoded def set_cache(self, key: str | None, value: CacheValue, local_only: bool = False, **kwargs: object): model_type: Final = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) payload: Final[object] = CacheCodec.serialize(value, model_type=model_type) + shared: Final = self._shared + if shared is not None: + if key is None: + return None + try: + if local_only: + shared.set_local(key, payload, _ttl(kwargs)) + else: + shared.set_sync(key, payload, _ttl(kwargs)) + except Exception as e: # noqa: BLE001 # DualCache.set_cache swallows a failed write + verbose_proxy_logger.error("UserApiKeyCache.set_cache failed for key=%r: %s", key, e) + return None if key is not None and is_user_key_cache_key(key): return self.key_object_cache.set_cache(key=key, value=payload, local_only=local_only, **kwargs) return super().set_cache(key=key, value=payload, local_only=local_only, **kwargs) @@ -242,6 +326,11 @@ class UserApiKeyCache(DualCache): and kwargs.keys() <= _PIPELINED_SET_OPTIONS and (ttl is None or isinstance(ttl, (int, float))) ) + shared: Final = self._shared + if shared is not None: + if key is not None: + await self._set_shared(shared, key, payload, local_only=local_only, pipelined=pipelined, kwargs=kwargs) + return None if key is not None and is_user_key_cache_key(key): if pipelined and await self.key_object_cache.async_set_cache_pre_call(key, payload, ttl) is not None: return None @@ -250,19 +339,46 @@ class UserApiKeyCache(DualCache): return None return await super().async_set_cache(key=key, value=payload, local_only=local_only, **kwargs) + @staticmethod + async def _set_shared( + shared: AuthCacheStore, + key: str, + payload: object, + *, + local_only: bool, + pipelined: bool, + kwargs: Mapping[str, object], + ) -> None: + try: + ttl: Final = _ttl(kwargs) + if local_only: + shared.set_local(key, payload, ttl) + elif not pipelined or shared.set_pre_call(key, payload, ttl) is None: + await shared.set(key, payload, ttl) + except Exception as e: # noqa: BLE001 # DualCache.async_set_cache swallows a failed write + verbose_proxy_logger.error("UserApiKeyCache.async_set_cache failed for key=%r: %s", key, e) + def delete_cache(self, key: str) -> None: + if self._shared is not None: + self._shared.delete_sync((key,)) + return if is_user_key_cache_key(key): self.key_object_cache.delete_cache(key) return super().delete_cache(key) async def async_delete_cache(self, key: str) -> None: + if self._shared is not None: + await self._shared.delete((key,)) + return if is_user_key_cache_key(key): await self.key_object_cache.async_delete_cache(key) return await super().async_delete_cache(key) - async def async_delete_cache_pre_call(self, key: str) -> BatchResult[None] | None: + async def async_delete_cache_pre_call(self, key: str) -> Awaitable[None] | None: + if self._shared is not None: + return self._shared.delete_pre_call(key) if is_user_key_cache_key(key): return await self.key_object_cache.async_delete_cache_pre_call(key) return await super().async_delete_cache_pre_call(key) @@ -274,6 +390,10 @@ class UserApiKeyCache(DualCache): Both partitions are cleared even when one raises, because a caller batching these has already committed the rows they cache. """ + if self._shared is not None: + if keys: + await self._shared.delete(keys) + return key_object_keys: Final = tuple(key for key in keys if is_user_key_cache_key(key)) other_keys: Final = tuple(key for key in keys if not is_user_key_cache_key(key)) outcomes: Final = await asyncio.gather( @@ -286,9 +406,12 @@ class UserApiKeyCache(DualCache): raise failed[0] def flush_cache(self) -> None: + self.budget_windows_absent_cache.flush_cache() + if self._shared is not None: + self._shared.flush_local() + return super().flush_cache() self.key_object_cache.in_memory_cache.flush_cache() - self.budget_windows_absent_cache.flush_cache() async def async_set_cache_pipeline( self, cache_list: Sequence[tuple[str, CacheValue]], local_only: bool = False, **kwargs: object @@ -298,6 +421,10 @@ class UserApiKeyCache(DualCache): ``model_type``: ``BaseModel`` values become JSON-safe dicts; dicts/scalars unchanged. """ normalized: Final = tuple((key, CacheCodec.serialize(value, model_type=None)) for key, value in cache_list) + shared: Final = self._shared + if shared is not None: + await self._set_many_shared(shared, normalized, local_only=local_only, ttl=_ttl(kwargs)) + return key_object_entries: Final = tuple(entry for entry in normalized if is_user_key_cache_key(entry[0])) other_entries: Final = tuple(entry for entry in normalized if not is_user_key_cache_key(entry[0])) if key_object_entries: @@ -307,6 +434,67 @@ class UserApiKeyCache(DualCache): if other_entries: await super().async_set_cache_pipeline(cache_list=other_entries, local_only=local_only, **kwargs) + @staticmethod + async def _set_many_shared( + shared: AuthCacheStore, entries: Sequence[tuple[str, object]], *, local_only: bool, ttl: float | None + ) -> None: + try: + if not local_only: + await shared.set_many(entries, ttl) + return + for key, value in entries: + shared.set_local(key, value, ttl) + except Exception as e: # noqa: BLE001 # DualCache.async_set_cache_pipeline swallows a failed write + verbose_proxy_logger.error("UserApiKeyCache.async_set_cache_pipeline failed: %s", e) + + async def increment_if_present(self, key: str, amount: float) -> float | None: + """``INCRBYFLOAT`` of an entry Redis still holds; None without Redis or when Redis lacks the entry. Raises + when Redis failed.""" + if self._shared is None: + return None + return await self._shared.increment_if_present(key, amount) + + async def async_batch_get_cache( + self, + keys: list[str], + parent_otel_span: Span | None = None, + local_only: bool = False, + throttle_redis: bool = True, + **kwargs: object, + ) -> list[object | None] | None: + self._python_only("async_batch_get_cache") + return await super().async_batch_get_cache( + keys, parent_otel_span=parent_otel_span, local_only=local_only, throttle_redis=throttle_redis, **kwargs + ) + + async def async_increment_cache( + self, + key: str, + value: float, + parent_otel_span: Span | None = None, + local_only: bool = False, + refresh_ttl: bool = False, + **kwargs: object, + ) -> float | None: + self._python_only("async_increment_cache") + return await super().async_increment_cache( + key, value, parent_otel_span=parent_otel_span, local_only=local_only, refresh_ttl=refresh_ttl, **kwargs + ) + + +def _decoded(cached: object, key: str, model_type: type[BaseModel] | None, method: str) -> object: + if model_type is None or cached is None: + return cached + decoded: Final = CacheCodec.deserialize(cached, model_type=model_type) + if decoded is None: + verbose_proxy_logger.error( + "UserApiKeyCache.%s failed to deserialize cached value for key=%r model_type=%s", + method, + key, + getattr(model_type, "__name__", str(model_type)), + ) + return decoded + #: Value cached under ``user_object_permission_id_cache_key`` when the user links no permission row, #: so a human without an entitlement costs no DB read per request. Lives beside the key builder diff --git a/litellm/proxy/credential_endpoints/user_provider_credentials.py b/litellm/proxy/credential_endpoints/user_provider_credentials.py index 22ee271410e..077f0a43663 100644 --- a/litellm/proxy/credential_endpoints/user_provider_credentials.py +++ b/litellm/proxy/credential_endpoints/user_provider_credentials.py @@ -20,10 +20,10 @@ from pydantic import BaseModel, ConfigDict, Field, ValidationError from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger -from litellm.caching.dual_cache import DualCache if TYPE_CHECKING: from litellm.caching.redis_cache import RedisCache + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.constants import ( GITHUB_COPILOT_DEFAULT_HOST, GITHUB_COPILOT_USER_CREDENTIAL_CACHE_TTL_SECONDS, @@ -203,7 +203,7 @@ async def _try_cache_set(token_cache: "RedisCache", key: str, value: str, nx: bo @with_service_target("user_provider_connections") async def invalidate_user_provider_credential_cache( - cache: DualCache, + cache: "UserApiKeyCache", user_id: str, credential_name: str, ) -> bool: @@ -211,7 +211,7 @@ async def invalidate_user_provider_credential_cache( before the disconnect must not fill the token back in behind it (fills are set-if-absent). Returns False when Redis is attached and the write fails, so the disconnect can refuse to drop the row while a stale token might linger.""" - token_cache: Final = cache.redis_cache + token_cache: Final = cache.redis_connection if token_cache is None: return True try: @@ -235,13 +235,13 @@ async def invalidate_user_provider_credential_cache( @with_service_target("user_provider_connections") async def drop_user_provider_credential_cache( - cache: DualCache, + cache: "UserApiKeyCache", user_id: str, credential_name: str, ) -> None: """Connect path: delete the key outright so the next read fills fresh (a tombstone would linger for the TTL and hide the new connection).""" - token_cache: Final = cache.redis_cache + token_cache: Final = cache.redis_connection if token_cache is None: return try: @@ -252,7 +252,7 @@ async def drop_user_provider_credential_cache( @with_service_target("user_provider_connections") async def set_user_provider_credential_cache( - cache: DualCache, + cache: "UserApiKeyCache", user_id: str, credential_name: str, payload: GithubCopilotUserConnectionPayload, @@ -262,7 +262,7 @@ async def set_user_provider_credential_cache( not-connected marker. The overwrite is verified by decoding the key back (async_set_cache swallows write errors); on a mismatch the key is deleted, and False means a stale entry survived both attempts.""" - token_cache: Final = cache.redis_cache + token_cache: Final = cache.redis_connection if token_cache is None: return True key: Final = _cache_key(user_id, credential_name) @@ -296,7 +296,7 @@ async def set_user_provider_credential_cache( @with_service_target("user_provider_connections") async def aget_user_provider_credentials( prisma_client: "PrismaClient", - cache: DualCache, + cache: "UserApiKeyCache", user_id: str, credential_names: Sequence[str], ) -> Mapping[str, Mapping[str, str]]: @@ -306,7 +306,7 @@ async def aget_user_provider_credentials( plaintext connection data only ever lives in the returned mapping. Redis is the only cache layer used: without it there is no shared invalidation, so a worker serving a stale local entry after a disconnect is worse than a DB read.""" - token_cache: Final = cache.redis_cache + token_cache: Final = cache.redis_connection names: Final = tuple(dict.fromkeys(credential_names)) if not names: return {} diff --git a/litellm/proxy/mcp_server/outbound_credentials/per_user_oauth_store.py b/litellm/proxy/mcp_server/outbound_credentials/per_user_oauth_store.py index f3748f99ba4..f03d42ad2fa 100644 --- a/litellm/proxy/mcp_server/outbound_credentials/per_user_oauth_store.py +++ b/litellm/proxy/mcp_server/outbound_credentials/per_user_oauth_store.py @@ -121,7 +121,7 @@ async def _post_token_endpoint(url: str, form: dict[str, str], headers: dict[str def _redis_cache_is_available() -> bool: from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415 - return user_api_key_cache.redis_cache is not None + return user_api_key_cache.has_redis def _runtime_backend_and_coordinator() -> tuple[TokenCacheBackend | None, RefreshCoordinator | None, bool]: diff --git a/litellm/proxy/mcp_server/outbound_credentials/runtime_refresh_coordinator.py b/litellm/proxy/mcp_server/outbound_credentials/runtime_refresh_coordinator.py index 9376d1bda50..63ab22d0a9d 100644 --- a/litellm/proxy/mcp_server/outbound_credentials/runtime_refresh_coordinator.py +++ b/litellm/proxy/mcp_server/outbound_credentials/runtime_refresh_coordinator.py @@ -28,7 +28,7 @@ from litellm.proxy.mcp_server.outbound_credentials.redis_refresh_coordinator imp def runtime_refresh_coordinator() -> RefreshCoordinator | None: from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415 # runtime global - redis_cache: Final = user_api_key_cache.redis_cache + redis_cache: Final = user_api_key_cache.redis_connection if redis_cache is None: return None # The Redis client from init_async_client() is only partially typed; the lock validates every diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 08902bce1d6..0166b25d050 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -4435,15 +4435,9 @@ def _add_spend_to_locally_cached_object( ) -async def _add_spend_to_shared_global_proxy_spend( - cache: UserApiKeyCache, response_cost: float, parent_otel_span: Span | None -) -> float | None: - if cache.redis_cache is None: - return None +async def _add_spend_to_shared_global_proxy_spend(cache: UserApiKeyCache, response_cost: float) -> float | None: try: - return await cache.redis_cache.async_increment_if_present( - key=GLOBAL_PROXY_SPEND_CACHE_KEY, value=response_cost, parent_otel_span=parent_otel_span - ) + return await cache.increment_if_present(GLOBAL_PROXY_SPEND_CACHE_KEY, response_cost) except Exception as e: verbose_proxy_logger.warning( "Spend tracking - failed to add spend to the global proxy spend in Redis. " @@ -4454,11 +4448,9 @@ async def _add_spend_to_shared_global_proxy_spend( return None -async def _add_spend_to_global_proxy_spend( - cache: UserApiKeyCache, response_cost: float, parent_otel_span: Span | None -) -> None: +async def _add_spend_to_global_proxy_spend(cache: UserApiKeyCache, response_cost: float) -> None: local_before: Final = cache.get_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY, local_only=True) - shared_total: Final = await _add_spend_to_shared_global_proxy_spend(cache, response_cost, parent_otel_span) + shared_total: Final = await _add_spend_to_shared_global_proxy_spend(cache, response_cost) local_total: Final = cache.get_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY, local_only=True) if not isinstance(local_total, (int, float)): return @@ -4574,7 +4566,7 @@ async def update_cache( if litellm.max_budget <= 0: return with service_target(SPEND_COUNTERS_TARGET): - await _add_spend_to_global_proxy_spend(target, response_cost, parent_otel_span) + await _add_spend_to_global_proxy_spend(target, response_cost) def run_ollama_serve(): diff --git a/litellm/redis/stores/auth_cache.py b/litellm/redis/stores/auth_cache.py new file mode 100644 index 00000000000..534b33e2e46 --- /dev/null +++ b/litellm/redis/stores/auth_cache.py @@ -0,0 +1,192 @@ +from __future__ import annotations + +from collections.abc import Awaitable, Callable, Sequence +from dataclasses import dataclass +from typing import Final, Protocol, TypeAlias + +from litellm.caching.redis_batch import ( + NativeBatchResult, + NativePending, + active_request_redis_batches, +) +from litellm.rust_bridge._native import NativeAuthCache, NativeRedisClient, NativeRequestRedis + +WriteBack: TypeAlias = tuple[str, object, float | None] + + +@dataclass(frozen=True, slots=True) +class AuthCacheSettings: + """``UserApiKeyCache``'s memory configuration: each partition's size and ``default_in_memory_ttl``.""" + + memory_entries: int + memory_ttl: float | None + + +class AuthCacheHandle(Protocol): + def get(self, key: str) -> Awaitable[object]: ... + def get_sync(self, key: str) -> object: ... + def set(self, key: str, value: object, ttl_seconds: float | None = None) -> Awaitable[None]: ... + def set_sync(self, key: str, value: object, ttl_seconds: float | None = None) -> None: ... + def set_pre_call( + self, request: NativeRequestRedis, key: str, value: object, ttl_seconds: float | None = None + ) -> NativePending[None] | None: ... + def delete_pre_call(self, request: NativeRequestRedis, key: str) -> NativePending[None] | None: ... + def set_many(self, entries: Sequence[tuple[str, object]], ttl_seconds: float | None = None) -> Awaitable[None]: ... + def write_back(self, request: NativeRequestRedis | None, entries: Sequence[WriteBack]) -> Awaitable[None]: ... + def delete(self, keys: Sequence[str]) -> Awaitable[None]: ... + def delete_sync(self, keys: Sequence[str]) -> None: ... + def declare_get_many( + self, request: NativeRequestRedis, keys: Sequence[str] + ) -> NativePending[Sequence[object]] | None: ... + def get_many(self, keys: Sequence[str]) -> Awaitable[Sequence[object] | None]: ... + def increment_if_present(self, key: str, amount: float) -> Awaitable[float | None]: ... + def get_local(self, key: str) -> object: ... + def set_local(self, key: str, value: object, ttl_seconds: float | None = None) -> None: ... + def delete_local(self, key: str) -> None: ... + def flush_local(self) -> None: ... + def len_local(self) -> tuple[int, int]: ... + + +HandleFactory: TypeAlias = Callable[[NativeRedisClient | None, AuthCacheSettings, int], AuthCacheHandle] + + +def native_handle(client: NativeRedisClient | None, settings: AuthCacheSettings, redis_ttl: int) -> AuthCacheHandle: + return NativeAuthCache( + client, + memory_entries=settings.memory_entries, + memory_ttl_seconds=settings.memory_ttl, + redis_ttl_seconds=redis_ttl, + ) + + +class AuthCacheStore: + """The auth object cache ``enable_redis_auth_cache`` shares across workers: memory partitions in front of Redis, + all native. Keys are ``UserApiKeyCache``'s logical cache keys; values are plain data (what ``json.loads`` + returns), so callers validate models themselves. Redis calls are reported under the caller's service target. + + Inside a request, a key a prefetch read found absent in Redis is answered as absent by ``get`` until something + writes it, so per-key getters after a prefetch cost no round trip.""" + + def __init__( + self, + client: NativeRedisClient | None, + settings: AuthCacheSettings, + *, + redis_ttl: int, + handle: HandleFactory = native_handle, + ) -> None: + self._client: Final = client + self._redis_ttl: Final = redis_ttl + self._build: Final = handle + self._settings = settings + self._handle = handle(client, settings, redis_ttl) + + @property + def settings(self) -> AuthCacheSettings: + return self._settings + + def reconfigure(self, settings: AuthCacheSettings) -> None: + """A changed memory configuration starts from an empty memory tier, as a resized partition would.""" + if settings == self._settings: + return + self._settings = settings + self._handle = self._build(self._client, settings, self._redis_ttl) + + async def get(self, key: str) -> object: + request: Final = active_request_redis_batches() + if request is not None and request.native_read_as_missing(key) and self._handle.get_local(key) is None: + return None + return await self._handle.get(key) + + def get_sync(self, key: str) -> object: + return self._handle.get_sync(key) + + async def set(self, key: str, value: object, ttl: float | None) -> None: + self._forget_miss(key) + await self._handle.set(key, value, ttl) + + def set_sync(self, key: str, value: object, ttl: float | None) -> None: + self._forget_miss(key) + self._handle.set_sync(key, value, ttl) + + def set_pre_call(self, key: str, value: object, ttl: float | None) -> NativeBatchResult[None] | None: + """Memory now and the SET on the request's native pre-call batch, sent with its next flush; None outside a + request, so the caller writes through.""" + request: Final = active_request_redis_batches() + if request is None: + return None + pending: Final = self._handle.set_pre_call(request.native(), key, value, ttl) + if pending is None: + return None + request.forget_native_miss(key) + return request.native_result(pending) + + def delete_pre_call(self, key: str) -> NativeBatchResult[None] | None: + request: Final = active_request_redis_batches() + if request is None: + return None + pending: Final = self._handle.delete_pre_call(request.native(), key) + if pending is None: + return None + request.note_native_misses((key,)) + return request.native_result(pending) + + async def set_many(self, entries: Sequence[tuple[str, object]], ttl: float | None) -> None: + for key, _value in entries: + self._forget_miss(key) + await self._handle.set_many(entries, ttl) + + async def write_back(self, entries: Sequence[WriteBack]) -> None: + """Memory now; inside a request the SETs ride its native pre-call batch, sent with the next flush, else they + go out in a round trip of their own. An entry without a TTL keeps ``default_in_memory_ttl`` in memory and the + connection's default TTL in Redis.""" + request: Final = active_request_redis_batches() + for key, _value, _ttl in entries: + self._forget_miss(key) + await self._handle.write_back(None if request is None else request.native(), entries) + + async def delete(self, keys: Sequence[str]) -> None: + """Memory, then Redis in one round trip; raises when Redis failed, with memory already dropped.""" + await self._handle.delete(keys) + + def delete_sync(self, keys: Sequence[str]) -> None: + self._handle.delete_sync(keys) + + async def read_missing(self, keys: Sequence[str]) -> dict[str, object]: + """One read of ``keys`` (already missing in memory): on the request's native pre-call batch inside a request, + noting the keys Redis lacks, else a round trip of its own. Absent keys map to None; raises when Redis failed.""" + request: Final = active_request_redis_batches() + if request is None: + values: Final = await self._handle.get_many(keys) + return {} if values is None else dict(zip(keys, values, strict=True)) + pending: Final = self._handle.declare_get_many(request.native(), keys) + if pending is None: + return {} + found: Final = dict(zip(keys, await request.native_result(pending), strict=True)) + request.note_native_misses(tuple(key for key, value in found.items() if value is None)) + return found + + async def increment_if_present(self, key: str, amount: float) -> float | None: + return await self._handle.increment_if_present(key, amount) + + def get_local(self, key: str) -> object: + return self._handle.get_local(key) + + def set_local(self, key: str, value: object, ttl: float | None) -> None: + self._handle.set_local(key, value, ttl) + + def delete_local(self, key: str) -> None: + self._handle.delete_local(key) + + def flush_local(self) -> None: + self._handle.flush_local() + + def len_local(self) -> tuple[int, int]: + """(management objects, user-key objects) held in memory.""" + return self._handle.len_local() + + @staticmethod + def _forget_miss(key: str) -> None: + request: Final = active_request_redis_batches() + if request is not None: + request.forget_native_miss(key) diff --git a/litellm/router_utils/routing_read_batch.py b/litellm/router_utils/routing_read_batch.py index 7622991a78d..c7cee04c946 100644 --- a/litellm/router_utils/routing_read_batch.py +++ b/litellm/router_utils/routing_read_batch.py @@ -35,11 +35,16 @@ class RoutingPrefetch: @staticmethod @renamed_router_param("litellm_router_instance") def arm(ctx: "RoutingContext", usage_selector: LowestTPMLoggingHandler_v2 | None, deployments: list) -> None: - routing_prefetch.arm_cooldowns(ctx.health().cooldowns.cooldown_store, ctx.registry().names(kind="ids")) + RoutingPrefetch.arm_cooldowns(ctx) read: Final = _selector_read(usage_selector, deployments) if read is not None: routing_prefetch.arm(read) + @staticmethod + def arm_cooldowns(ctx: "RoutingContext") -> None: + """The cooldown read routing takes later, declared on the request's native pre-call batch now.""" + routing_prefetch.arm_cooldowns(ctx.health().cooldowns.cooldown_store, ctx.registry().names(kind="ids")) + @staticmethod def armed() -> bool: return routing_prefetch.armed() diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index a7bd580d56d..fb24f97b4b1 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -15,6 +15,7 @@ class RustUpstreamError(Exception): ... class ForkedAfterNativeRuntimeStarted(RuntimeError): ... class ProcessReservedForForking(RuntimeError): ... class NativeRedisStoreError(Exception): ... +class NativeAuthCacheError(Exception): ... def trace_encode_error(message: str) -> bytes: ... def bundled_model_cost_map() -> bytes: ... @@ -92,6 +93,50 @@ class NativeRouterUsage: def flush_local_counters(self) -> None: ... def snapshot(self) -> dict[str, float]: ... +@final +class NativeDeclaredAuthReads: + def resolve(self) -> Future[list[object]]: ... + +@final +class NativeDeclaredAuthWrite: + def resolve(self) -> Future[None]: ... + +@final +class NativeAuthCache: + def __new__( + cls, + client: NativeRedisClient | None, + /, + *, + memory_entries: int = 200, + memory_ttl_seconds: float | None = None, + redis_ttl_seconds: int = 60, + ) -> NativeAuthCache: ... + @property + def has_redis(self) -> bool: ... + def get(self, key: str) -> Future[object]: ... + def get_sync(self, key: str) -> object: ... + def set(self, key: str, value: object, ttl_seconds: float | None = None) -> Future[None]: ... + def set_sync(self, key: str, value: object, ttl_seconds: float | None = None) -> None: ... + def set_pre_call( + self, request: NativeRequestRedis, key: str, value: object, ttl_seconds: float | None = None + ) -> NativeDeclaredAuthWrite | None: ... + def delete_pre_call(self, request: NativeRequestRedis, key: str) -> NativeDeclaredAuthWrite | None: ... + def set_many(self, entries: Sequence[tuple[str, object]], ttl_seconds: float | None = None) -> Future[None]: ... + def write_back( + self, request: NativeRequestRedis | None, entries: Sequence[tuple[str, object, float | None]] + ) -> Future[None]: ... + def delete(self, keys: Sequence[str]) -> Future[None]: ... + def delete_sync(self, keys: Sequence[str]) -> None: ... + def declare_get_many(self, request: NativeRequestRedis, keys: Sequence[str]) -> NativeDeclaredAuthReads | None: ... + def get_many(self, keys: Sequence[str]) -> Future[list[object] | None]: ... + def increment_if_present(self, key: str, amount: float) -> Future[float | None]: ... + def get_local(self, key: str) -> object: ... + def set_local(self, key: str, value: object, ttl_seconds: float | None = None) -> None: ... + def delete_local(self, key: str) -> None: ... + def flush_local(self) -> None: ... + def len_local(self) -> tuple[int, int]: ... + @final class NativeSpendCounters: def __new__(cls, client: NativeRedisClient | None, redis_ttl_seconds: int = 60, /) -> NativeSpendCounters: ... @@ -303,9 +348,13 @@ def reserve_process_for_forking() -> None: ... __all__ = [ "ForkedAfterNativeRuntimeStarted", "HuggingFaceEncoding", + "NativeAuthCache", + "NativeAuthCacheError", "NativeClickHouseSpendConfig", "NativeClickHouseSpendStorage", "NativeDatabase", + "NativeDeclaredAuthReads", + "NativeDeclaredAuthWrite", "NativeDiagnosticProcessor", "NativePending", "NativeRedisClient", diff --git a/tests/unit/caching/test_request_redis_batch_pre_call.py b/tests/unit/caching/test_request_redis_batch_pre_call.py index 31878658049..eded068cf33 100644 --- a/tests/unit/caching/test_request_redis_batch_pre_call.py +++ b/tests/unit/caching/test_request_redis_batch_pre_call.py @@ -6,7 +6,8 @@ from __future__ import annotations import asyncio import hashlib import json -from collections.abc import Callable +from collections.abc import Callable, Sequence +from dataclasses import dataclass from datetime import datetime, timezone from typing import Any, Final from unittest.mock import AsyncMock, MagicMock @@ -21,7 +22,6 @@ from litellm.caching.redis_batch import active_request_redis_batches, request_re from litellm.proxy._types import LiteLLM_TeamTableCachedObj, LiteLLM_UserTable from litellm.proxy.auth.auth_checks import cache_team_object from litellm.proxy.auth.auth_object_prefetch import _CacheEntry, _write_back, prefetch_identity_keys -from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.hooks.parallel_request_limiter_v3 import ( CHECK_AND_INCREMENT_BY_N_SCRIPT, PROXY_MaxParallelRequestsHandler_v3, @@ -29,14 +29,22 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import ( RateLimitUnverifiableError, ) from litellm.proxy.utils import InternalUsageCache +from litellm.redis.stores.router_cooldowns import ActiveCooldowns, RouterCooldownStore from litellm.router_api import RoutingContext -from litellm.router_utils.routing_prefetch import ArmedRead, TierRead +from litellm.router_utils.routing_prefetch import ArmedRead, TierRead, arm_cooldowns, read_cooldowns from litellm.router_utils.routing_read_batch import ( ROUTER_USAGE_TARGET, RoutingPrefetch, ) +from tests.unit.redis.stores.auth_cache_double import FakeAuthCache, shared_cache -from .test_redis_batch import FakeClient, FakeNativePending, FakeNatives, FakeRedisCache, replies +from .test_redis_batch import ( + FakeClient, + FakeNativePending, + FakeNativeRequest, + FakeNatives, + FakeRedisCache, +) def _cooldown_key(model_id: str) -> str: @@ -75,6 +83,11 @@ def _refunds(limiter: PROXY_MaxParallelRequestsHandler_v3) -> list[tuple[str, fl return [(call.kwargs["keys"][1], call.kwargs["args"][1]) for call in refund_script.await_args_list] +def _declared_before_the_first_flush(natives: FakeNatives) -> list[str]: + events: Final = natives.built[0].events + return events[: events.index("flush_pre")] + + def _lua_ok_replies(command: tuple[Any, ...]) -> Any: if command[0] == "EVALSHA": return [0, 1, 1700000000] # OK: one counter, new_counter=1, window_start @@ -764,125 +777,144 @@ def _user_entry() -> tuple[_CacheEntry, LiteLLM_UserTable]: @pytest.mark.asyncio -async def test_auth_write_back_rides_the_next_round_trip_and_the_scope_drains_what_nobody_awaited(): - client = FakeClient(_lua_ok_replies) - redis_cache = FakeRedisCache(client) - cache = UserApiKeyCache(redis_cache=redis_cache) - with request_redis_batch_scope() as request: +async def test_auth_write_back_rides_the_native_pre_call_batch_and_the_scope_drains_what_nobody_awaited(): + cache, fake = shared_cache() + natives: Final = FakeNatives() + with request_redis_batch_scope(native_factory=natives) as request: # pyright: ignore[reportArgumentType] # test double await _write_back([_user_entry()], cache) - assert client.pipelines == [] # not sent yet: the SET waits for the next round trip - await request.batch(redis_cache).mget(["spend:key:k1"]) - assert len(client.pipelines) == 1 - kinds = [c[0] for c in client.pipelines[0].commands] - assert kinds == ["MGET", "SET"] or kinds == ["SET", "MGET"] - set_command = next(c for c in client.pipelines[0].commands if c[0] == "SET") - assert set_command[1] == "user-1" and set_command[3] == 42 - assert json.loads(set_command[2])["user_id"] == "user-1" - assert cache.in_memory_cache.get_cache("user-1") is not None - - await _write_back([_user_entry()], cache) - assert len(client.pipelines) == 1 + assert cache.get_local("user-1")["user_id"] == "user-1" + assert natives.built[0].events == ["declare write_back (('user-1', 42),)"] await request.flush_all() - assert len(client.pipelines) == 2 - assert [c[0] for c in client.pipelines[1].commands] == ["SET"] + assert natives.built[0].events[-1] == "flush_pre" + assert fake.commands == [] @pytest.mark.asyncio async def test_auth_write_back_outside_a_scope_writes_through_as_before(): - redis_cache = FakeRedisCache(FakeClient(_lua_ok_replies)) - cache = UserApiKeyCache(redis_cache=redis_cache) + cache, fake = shared_cache() await _write_back([_user_entry()], cache) - assert [(op[0], [(key, ttl) for key, _value, ttl in op[1]]) for op in redis_cache.alone] == [ - ("SET_PIPELINE", [("user-1", 42)]) - ] + assert fake.commands == [("WRITE_BACK", (("user-1", 42),))] @pytest.mark.asyncio -async def test_a_key_the_request_mget_read_as_absent_is_not_read_again_by_a_per_key_get(): - client = FakeClient(_lua_ok_replies) - redis_cache = FakeRedisCache(client) - cache = UserApiKeyCache(redis_cache=redis_cache) - with request_redis_batch_scope() as request: - assert await request.batch(redis_cache).mget(["absent-key"]) == {"absent-key": None} +async def test_a_key_the_request_read_as_absent_is_not_read_again_by_a_per_key_get(): + cache, fake = shared_cache() + with request_redis_batch_scope(native_factory=FakeNatives()) as request: # pyright: ignore[reportArgumentType] # test double + await prefetch_identity_keys(["absent-key"], cache) assert await cache.async_get_cache("absent-key") is None - assert redis_cache.alone == [] and len(client.pipelines) == 1 await cache.async_set_cache("absent-key", {"v": 1}, ttl=5) + assert await cache.async_get_cache("absent-key") == {"v": 1} await request.flush_all() - assert [c[:2] for c in client.pipelines[1].commands] == [("SET", "absent-key")] + assert fake.commands == [] @pytest.mark.asyncio async def test_management_object_writes_inside_a_request_ride_its_pipeline_and_write_through_outside(): - client = FakeClient(_lua_ok_replies) - redis_cache = FakeRedisCache(client) - cache = UserApiKeyCache(redis_cache=redis_cache) - with request_redis_batch_scope() as request: + cache, fake = shared_cache() + natives: Final = FakeNatives() + with request_redis_batch_scope(native_factory=natives) as request: # pyright: ignore[reportArgumentType] # test double await cache.async_set_cache("team_id:t1", {"team_id": "t1"}, ttl=60) await cache.async_set_cache("hashed-key-object", {"token": "hashed-key-object"}, ttl=60) - assert client.pipelines == [] - assert cache.in_memory_cache.get_cache("team_id:t1") == {"team_id": "t1"} + assert cache.get_local("team_id:t1") == {"team_id": "t1"} assert await cache.async_get_cache("hashed-key-object") == {"token": "hashed-key-object"} await request.flush_all() - assert sorted((c[0], c[1], c[3]) for c in client.pipelines[0].commands) == [ - ("SET", "hashed-key-object", 60), - ("SET", "team_id:t1", 60), + assert natives.built[0].events == [ + "declare set team_id:t1 ttl=60.0", + "declare set hashed-key-object ttl=60.0", + "flush_pre", ] await cache.async_set_cache("team_id:t2", {"team_id": "t2"}, ttl=60) - assert len(client.pipelines) == 1 - assert redis_cache.alone == [("SET", "team_id:t2", {"team_id": "t2"})] + assert fake.commands == [("SET", "team_id:t2", 60.0)] @pytest.mark.asyncio async def test_a_team_refresh_inside_a_request_sends_its_set_and_alias_del_in_one_pipeline_before_returning(): - client = FakeClient(_lua_ok_replies) - redis_cache = FakeRedisCache(client) - cache = UserApiKeyCache(redis_cache=redis_cache) + redis_cache = FakeRedisCache(FakeClient(_lua_ok_replies)) + cache, fake = shared_cache(connection=redis_cache) usage_cache = DualCache(redis_cache=redis_cache) usage_cache.in_memory_cache.set_cache("team_id:t1", "stale team") usage_cache.in_memory_cache.set_cache("team_alias:alpha", "stale alias") - cache.in_memory_cache.set_cache("team_alias:alpha", "stale alias") + cache.set_local("team_alias:alpha", "stale alias", None) proxy_logging_obj = MagicMock() proxy_logging_obj.internal_usage_cache = InternalUsageCache(dual_cache=usage_cache) team = LiteLLM_TeamTableCachedObj(team_id="t1", team_alias="alpha") - with request_redis_batch_scope() as request: + natives: Final = FakeNatives() + with request_redis_batch_scope(native_factory=natives) as request: # pyright: ignore[reportArgumentType] # test double await cache_team_object("t1", team, cache, proxy_logging_obj) - assert [c[:2] for c in client.pipelines[0].commands] == [("SET", "team_id:t1"), ("DEL", "team_alias:alpha")], ( - "the alias DEL must reach Redis before the refresh returns, or another request can refill memory from it" - ) - assert redis_cache.alone == [] + assert natives.built[0].events == [ + "declare set team_id:t1 ttl=60.0", + "declare del team_alias:alpha", + "flush_pre", + ], "the alias DEL must reach Redis before the refresh returns, or another request can refill memory from it" assert usage_cache.in_memory_cache.get_cache("team_id:t1") is None assert usage_cache.in_memory_cache.get_cache("team_alias:alpha") is None - assert cache.in_memory_cache.get_cache("team_alias:alpha") is None - assert cache.in_memory_cache.get_cache("team_id:t1")["team_id"] == "t1" + assert cache.get_local("team_alias:alpha") is None + assert cache.get_local("team_id:t1")["team_id"] == "t1" await request.flush_all() - assert len(client.pipelines) == 1 and redis_cache.alone == [] + assert fake.commands == [] and redis_cache.alone == [] @pytest.mark.asyncio -async def test_a_pipelined_management_write_without_a_ttl_expires_in_redis_like_the_direct_path(): - client = FakeClient(_lua_ok_replies) - redis_cache = FakeRedisCache(client) - cache = UserApiKeyCache(redis_cache=redis_cache) +async def test_a_changed_default_in_memory_ttl_reaches_the_native_store(): + cache, fake = shared_cache() cache.update_cache_ttl(default_in_memory_ttl=5, default_redis_ttl=None) - with request_redis_batch_scope() as request: - await cache.async_set_cache("team_id:t1", {"team_id": "t1"}) - await request.flush_all() - assert [(c[0], c[1], c[3]) for c in client.pipelines[0].commands] == [("SET", "team_id:t1", 5)] + assert [settings.memory_ttl for settings in fake.settings] == [None, 5] @pytest.mark.asyncio -async def test_identity_prefetch_is_one_mget_after_which_hits_and_misses_alike_cost_no_read(): - client = FakeClient(replies) - redis_cache = FakeRedisCache(client) - cache = UserApiKeyCache(redis_cache=redis_cache) - with request_redis_batch_scope(): +async def test_identity_prefetch_is_one_read_after_which_hits_and_misses_alike_cost_no_read(): + cache, fake = shared_cache(FakeAuthCache(redis={"key-hit": {"k": "key-hit"}})) + natives: Final = FakeNatives() + with request_redis_batch_scope(native_factory=natives): # pyright: ignore[reportArgumentType] # test double await prefetch_identity_keys(["key-hit", "end_user_id:eu-miss", "key-hit"], cache) - assert [c[0] for c in client.pipelines[0].commands] == ["MGET"] - assert sorted(client.pipelines[0].commands[0][1:]) == ["end_user_id:eu-miss", "key-hit"] assert await cache.async_get_cache("key-hit") == {"k": "key-hit"} assert await cache.async_get_cache("end_user_id:eu-miss") is None - assert len(client.pipelines) == 1 and redis_cache.alone == [] - assert cache.in_memory_cache.get_cache("end_user_id:eu-miss") is None + assert natives.built[0].events == ["declare auth_cache end_user_id:eu-miss,key-hit", "flush_pre"] + assert fake.commands == [] + assert cache.get_local("end_user_id:eu-miss") is None + + +@dataclass +class _FakeCooldownsHandle: + """A connected cooldown store whose ``declare_active`` lands on the request's fake native batch.""" + + def declare_active( + self, request: FakeNativeRequest, model_ids: Sequence[str] + ) -> FakeNativePending[ActiveCooldowns]: + request.events.append("declare cooldowns") + return FakeNativePending(request, lambda: ()) + + +@pytest.mark.asyncio +async def test_the_auth_prefetch_and_the_armed_cooldown_read_go_out_in_one_native_pipeline(): + cache, _fake = shared_cache(FakeAuthCache(redis={"key-hit": {"k": "key-hit"}})) + cooldowns: Final = RouterCooldownStore( + FakeRedisCache(FakeClient(_lua_ok_replies)), handle=lambda _c: _FakeCooldownsHandle() + ) # pyright: ignore[reportArgumentType] # test double + natives: Final = FakeNatives() + + with request_redis_batch_scope(native_factory=natives) as request: # pyright: ignore[reportArgumentType] # test double + await prefetch_identity_keys(["key-hit"], cache, before_read=lambda: arm_cooldowns(cooldowns, ["dep-a"])) + arm_cooldowns(cooldowns, ["dep-a"]) + active = await read_cooldowns(cooldowns, ["dep-a"], None) + await request.flush_all() + + assert active == () + assert len(natives.built) == 1 + assert _declared_before_the_first_flush(natives) == ["declare cooldowns", "declare auth_cache key-hit"] + assert not [event for event in natives.built[0].events[2:] if event.startswith("declare")] + + +@pytest.mark.asyncio +async def test_a_prefetch_served_from_memory_arms_no_early_cooldown_read(): + cache, _fake = shared_cache() + cache.set_local("key-hit", {"k": "key-hit"}, None) + armed: Final[list[str]] = [] # mutable-ok: records the arming callback + + with request_redis_batch_scope(native_factory=FakeNatives()): # pyright: ignore[reportArgumentType] # test double + await prefetch_identity_keys(["key-hit"], cache, before_read=lambda: armed.append("armed")) + + assert armed == [] @pytest.mark.asyncio diff --git a/tests/unit/proxy/auth/test_auth_checks.py b/tests/unit/proxy/auth/test_auth_checks.py index 9d4fe67614e..b191f7835a8 100644 --- a/tests/unit/proxy/auth/test_auth_checks.py +++ b/tests/unit/proxy/auth/test_auth_checks.py @@ -16,6 +16,8 @@ from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch import pytest, litellm + +from tests.unit.redis.stores.auth_cache_double import FakeAuthCache, shared_cache import httpx from fastapi import Request, status from prisma import Prisma @@ -2342,10 +2344,8 @@ async def test_bulk_eviction_clears_this_worker_before_waiting_on_a_wedged_redis for index in range(600) ) for hashed_token in hashed_tokens: - cache.in_memory_cache_for(hashed_token).set_cache(key=hashed_token, value={"token": hashed_token}) - seeded: Final = tuple( - token for token in hashed_tokens if cache.in_memory_cache_for(token).get_cache(token) is not None - ) + cache.set_local(hashed_token, {"token": hashed_token}, None) + seeded: Final = tuple(token for token in hashed_tokens if cache.get_local(token) is not None) with patch( "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache", @@ -2355,9 +2355,7 @@ async def test_bulk_eviction_clears_this_worker_before_waiting_on_a_wedged_redis delete_cache_key_objects(hashed_tokens=hashed_tokens, user_api_key_cache=cache, proxy_logging_obj=None) ) await asyncio.sleep(0.05) - still_cached: Final = tuple( - token for token in seeded if cache.in_memory_cache_for(token).get_cache(token) is not None - ) + still_cached: Final = tuple(token for token in seeded if cache.get_local(token) is not None) client.release.set() await eviction await _drain_publishes() @@ -2404,11 +2402,9 @@ async def test_bulk_eviction_clears_the_redis_copies_before_waiting_on_a_wedged_ hashlib.sha256(f"token-{index}".encode()).hexdigest() for index in range(600) ) - redis = _AuthRedisCache({token: {"token": token} for token in hashed_tokens}) - cache: Final = UserApiKeyCache(key_object_in_memory_cache=InMemoryCache(max_size_in_memory=2_000)) - cache.attach_redis_cache(redis) # pyright: ignore[reportArgumentType] # a dict-backed fake of the Redis layer + cache, redis = shared_cache(FakeAuthCache(redis={token: {"token": token} for token in hashed_tokens})) for hashed_token in hashed_tokens: - cache.in_memory_cache_for(hashed_token).set_cache(key=hashed_token, value={"token": hashed_token}) + cache.set_local(hashed_token, {"token": hashed_token}, None) with patch( "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache", @@ -2424,7 +2420,7 @@ async def test_bulk_eviction_clears_the_redis_copies_before_waiting_on_a_wedged_ await _drain_publishes() assert refilled == [] - assert redis.entries == {} + assert redis.redis == {} class _SignallingPublishRedisClient(_SlowPublishRedisClient): 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 ddcfe36deb6..79e33ed10c6 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 @@ -106,6 +106,7 @@ from litellm.llms.base_llm.managed_resources.utils import encode_unified_id, gen 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.auth_cache_double import shared_cache from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters @@ -6975,7 +6976,7 @@ async def test_team_update_not_shadowed_by_internal_usage_cache_lit_4391(): team_id = "team-lit-4391" shared_redis = _SharedFakeRedis() - user_api_key_cache = UserApiKeyCache(redis_cache=shared_redis) + user_api_key_cache, auth_store = shared_cache(connection=shared_redis) proxy_logging_obj = MagicMock() proxy_logging_obj.internal_usage_cache.dual_cache = DualCache( redis_cache=shared_redis, @@ -7016,12 +7017,13 @@ async def test_team_update_not_shadowed_by_internal_usage_cache_lit_4391(): f"refreshed it. Got models={refreshed.models if refreshed else None}" ) - redis_copy = await shared_redis.async_get_cache(f"team_id:{team_id}") - assert redis_copy is not None and redis_copy["models"] == ["model-a", "model-b"], ( + redis_copy = auth_store.redis.get(f"team_id:{team_id}") + assert isinstance(redis_copy, dict) and redis_copy["models"] == ["model-a", "model-b"], ( "The shared Redis lost the refreshed team object — the internal-cache " "invalidation must run BEFORE the fresh write, not after. " f"Got: {redis_copy}" ) + assert f"team_id:{team_id}" not in shared_redis._store, "redis-py never touches the auth cache keyspace" class _CountingFakeRedis(_SharedFakeRedis): @@ -8195,7 +8197,6 @@ async def test_get_team_member_default_budget_caches_json_safe_payload(): must be dumped to ISO strings before the write, and the read side must give back a model. """ from litellm.proxy.auth.auth_checks import get_team_member_default_budget - from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache budget_row = MagicMock() budget_row.dict = lambda: { @@ -8208,20 +8209,7 @@ async def test_get_team_member_default_budget_caches_json_safe_payload(): mock_prisma_client = MagicMock() mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=budget_row) - class _JsonOnlyRedis: - """Stands in for RedisCache, which serializes with a bare json.dumps().""" - - def __init__(self): - self.writes = [] - - async def async_set_cache(self, key, value, **kwargs): - self.writes.append((key, json.dumps(value))) - - async def async_get_cache(self, key, **kwargs): - return None - - redis_cache = _JsonOnlyRedis() - cache = UserApiKeyCache(redis_cache=redis_cache) + cache, auth_store = shared_cache() budget = await get_team_member_default_budget( budget_id="tm-budget-1", @@ -8231,9 +8219,10 @@ async def test_get_team_member_default_budget_caches_json_safe_payload(): assert isinstance(budget, LiteLLM_BudgetTable) assert budget.max_budget == 25.0 - assert len(redis_cache.writes) == 1 - written_key, written_payload = redis_cache.writes[0] - assert written_key == "team_member_default_budget:tm-budget-1" + assert [command[:2] for command in auth_store.commands if command[0] == "SET"] == [ + ("SET", "team_member_default_budget:tm-budget-1") + ] + written_payload = json.dumps(auth_store.redis["team_member_default_budget:tm-budget-1"]) assert json.loads(written_payload)["max_budget"] == 25.0 cached = await get_team_member_default_budget( diff --git a/tests/unit/proxy/auth/test_auth_object_prefetch.py b/tests/unit/proxy/auth/test_auth_object_prefetch.py index ffac95d6815..a80c637d3df 100644 --- a/tests/unit/proxy/auth/test_auth_object_prefetch.py +++ b/tests/unit/proxy/auth/test_auth_object_prefetch.py @@ -1,15 +1,12 @@ """Counts the Redis round trips and DB queries auth object reads cost per cache regime, and checks that the per-object getters still enforce on their own when the prefetch cannot help.""" -import json -from collections.abc import Sequence from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import HTTPException from litellm.caching.in_memory_cache import InMemoryCache -from litellm.caching.redis_cache import RedisCache from litellm.proxy._types import ( LiteLLM_OrganizationTable, LiteLLM_TeamMembership, @@ -30,6 +27,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( end_user_cache_key, end_user_restricted_registry_cache_key, ) +from tests.unit.redis.stores.auth_cache_double import FakeAuthCache, shared_cache USER_ID = "prefetch-user" TEAM_ID = "prefetch-team" @@ -77,49 +75,6 @@ ALL_ROWS = { } -class CountingRedis(RedisCache): - """Redis fake that counts commands and round trips (an MGET or a pipeline is one round trip).""" - - def __init__(self, store: dict[str, str] | None = None, fail: bool = False) -> None: - self.store: dict[str, str] = dict(store or {}) - self.fail = fail - self.round_trips = 0 - self.commands: list[str] = [] - - def _trip(self, *commands: str) -> None: - if self.fail: - raise ConnectionError("redis down") - self.round_trips += 1 - self.commands.extend(commands) - - async def async_get_cache(self, key: str, **kwargs: object) -> object: - self._trip(f"GET {key}") - raw = self.store.get(key) - return json.loads(raw) if raw is not None else None - - async def async_batch_get_cache(self, key_list: Sequence[str], **kwargs: object) -> dict[str, object]: - self._trip(f"MGET {' '.join(key_list)}") - return {key: (json.loads(self.store[key]) if key in self.store else None) for key in key_list} - - async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None: - self._trip(f"SET {key}") - self.store[key] = json.dumps(value) - - async def async_set_cache_pipeline(self, cache_list: Sequence[tuple[str, object]], **kwargs: object) -> None: - self._trip(*(f"SET {key}" for key, _ in cache_list)) - for key, value in cache_list: - self.store[key] = json.dumps(value) - - async def async_set_cache_pipeline_with_ttls(self, cache_list: Sequence[tuple[str, object, float | None]]) -> None: - self._trip(*(f"SET {key} ttl={ttl}" for key, _, ttl in cache_list)) - for key, value, _ in cache_list: - self.store[key] = json.dumps(value) - - async def async_delete_cache(self, key: str) -> None: - self._trip(f"DEL {key}") - self.store.pop(key, None) - - def _prisma(rows: dict[str, object] | None = ALL_ROWS) -> MagicMock: prisma = MagicMock(name="prisma_client") prisma.db.query_first = AsyncMock(return_value=rows) @@ -130,8 +85,21 @@ def _non_prefetch_db_calls(prisma: MagicMock) -> list[str]: return [str(call) for call in prisma.db.mock_calls if not str(call).startswith("call.query_first(")] -def _cache(redis: RedisCache | None) -> UserApiKeyCache: - return UserApiKeyCache(in_memory_cache=InMemoryCache(), redis_cache=redis) +def _cache(redis: FakeAuthCache | None) -> UserApiKeyCache: + """A worker of its own (empty memory) on ``redis``'s Redis, or a worker without Redis.""" + if redis is None: + return UserApiKeyCache(in_memory_cache=InMemoryCache()) + return shared_cache(FakeAuthCache(redis=redis.redis, commands=redis.commands, fail=redis.fail))[0] + + +def _round_trips(redis: FakeAuthCache) -> int: + return len(redis.commands) + + +def _writes(redis: FakeAuthCache) -> list[str]: + return sorted( + f"SET {key} ttl={ttl}" for command in redis.commands if command[0] == "WRITE_BACK" for key, ttl in command[1] + ) def _refs() -> AuthObjectRefs: @@ -164,7 +132,7 @@ def test_refs_from_token_only_names_membership_when_both_ids_present(): @pytest.mark.asyncio async def test_cold_regime_is_one_mget_one_query_and_the_getters_never_touch_io_again(): - redis = CountingRedis() + redis = FakeAuthCache() prisma = _prisma() cache = _cache(redis) @@ -172,9 +140,9 @@ async def test_cold_regime_is_one_mget_one_query_and_the_getters_never_touch_io_ assert prisma.db.query_first.await_count == 1 assert prisma.db.query_first.await_args.args[1:] == (USER_ID, TEAM_ID, USER_ID, ORG_ID, None) - mgets = [c for c in redis.commands if c.startswith("MGET")] + mgets = [c for c in redis.commands if c[0] == "MGET"] assert len(mgets) == 1 - assert set(mgets[0].split()[1:]) == { + assert set(mgets[0][1]) == { USER_ID, f"team_id:{TEAM_ID}", f"{TEAM_ID}_{USER_ID}", @@ -182,8 +150,7 @@ async def test_cold_regime_is_one_mget_one_query_and_the_getters_never_touch_io_ f"org_id:{ORG_ID}", f"org_id:{ORG_ID}:with_budget", } - sets = sorted(c for c in redis.commands if c.startswith("SET")) - assert sets == sorted( + assert _writes(redis) == sorted( [ f"SET {TEAM_ID}_{USER_ID} ttl=5", f"SET org_id:{ORG_ID} ttl=5", @@ -193,11 +160,11 @@ async def test_cold_regime_is_one_mget_one_query_and_the_getters_never_touch_io_ f"SET team_membership:{USER_ID}:{TEAM_ID} ttl=None", ] ) - assert redis.round_trips == 2, "one MGET, one pipeline" + assert _round_trips(redis) == 2, "one MGET, one pipeline" - before = (redis.round_trips, prisma.db.query_first.await_count) + before = (_round_trips(redis), prisma.db.query_first.await_count) user, team, membership, org = await _read_all_through_getters(cache, prisma) - assert (redis.round_trips, prisma.db.query_first.await_count) == before + assert (_round_trips(redis), prisma.db.query_first.await_count) == before assert _non_prefetch_db_calls(prisma) == [] assert isinstance(user, LiteLLM_UserTable) and user.max_budget == 50.0 @@ -211,46 +178,46 @@ async def test_cold_regime_is_one_mget_one_query_and_the_getters_never_touch_io_ @pytest.mark.asyncio async def test_redis_warm_regime_is_exactly_one_mget_and_zero_queries(): - seeded = CountingRedis() + seeded = FakeAuthCache() await prefetch_auth_objects(refs=_refs(), user_api_key_cache=_cache(seeded), prisma_client=_prisma()) - redis = CountingRedis(store=seeded.store) + redis = FakeAuthCache(redis=seeded.redis) prisma = _prisma() cache = _cache(redis) await prefetch_auth_objects(refs=_refs(), user_api_key_cache=cache, prisma_client=prisma) - assert redis.round_trips == 1 - assert redis.commands[0].startswith("MGET") + assert _round_trips(redis) == 1 + assert redis.commands[0][0] == "MGET" assert prisma.db.query_first.await_count == 0 user, team, membership, org = await _read_all_through_getters(cache, prisma) - assert redis.round_trips == 1 + assert _round_trips(redis) == 1 assert prisma.db.mock_calls == [] assert (user.user_id, team.team_id, membership.team_id, org.organization_id) == (USER_ID, TEAM_ID, TEAM_ID, ORG_ID) @pytest.mark.asyncio async def test_hot_regime_costs_nothing(): - redis = CountingRedis() + redis = FakeAuthCache() prisma = _prisma() cache = _cache(redis) await prefetch_auth_objects(refs=_refs(), user_api_key_cache=cache, prisma_client=prisma) - redis.round_trips, redis.commands = 0, [] + redis.commands.clear() prisma.db.reset_mock() await prefetch_auth_objects(refs=_refs(), user_api_key_cache=cache, prisma_client=prisma) await _read_all_through_getters(cache, prisma) - assert redis.round_trips == 0 + assert _round_trips(redis) == 0 assert prisma.db.mock_calls == [] @pytest.mark.asyncio async def test_partial_redis_hit_queries_only_the_missing_objects(): - seeded = CountingRedis() + seeded = FakeAuthCache() await prefetch_auth_objects(refs=_refs(), user_api_key_cache=_cache(seeded), prisma_client=_prisma()) for key in (f"team_id:{TEAM_ID}", f"{TEAM_ID}_{USER_ID}", f"team_membership:{USER_ID}:{TEAM_ID}"): - del seeded.store[key] + del seeded.redis[key] prisma = _prisma() await prefetch_auth_objects(refs=_refs(), user_api_key_cache=_cache(seeded), prisma_client=prisma) @@ -258,14 +225,14 @@ async def test_partial_redis_hit_queries_only_the_missing_objects(): assert prisma.db.query_first.await_count == 1 assert prisma.db.query_first.await_args.args[1:] == (None, TEAM_ID, USER_ID, None, None) - del seeded.store[f"org_id:{ORG_ID}:with_budget"] + del seeded.redis[f"org_id:{ORG_ID}:with_budget"] await prefetch_auth_objects(refs=_refs(), user_api_key_cache=_cache(seeded), prisma_client=prisma) assert prisma.db.query_first.await_args.args[1:] == (None, None, None, ORG_ID, None) @pytest.mark.asyncio async def test_deleted_team_cache_entry_is_refetched_and_the_update_is_visible(): - redis = CountingRedis() + redis = FakeAuthCache() prisma = _prisma() cache = _cache(redis) await prefetch_auth_objects(refs=_refs(), user_api_key_cache=cache, prisma_client=prisma) @@ -284,7 +251,7 @@ async def test_deleted_team_cache_entry_is_refetched_and_the_update_is_visible() async def test_row_missing_a_required_column_is_not_cached_so_the_getter_still_fails_closed(): prisma = _prisma({**ALL_ROWS, "team_row": {"max_budget": 1.0}}) prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) - cache = _cache(CountingRedis()) + cache = _cache(FakeAuthCache()) await prefetch_auth_objects(refs=_refs(), user_api_key_cache=cache, prisma_client=prisma) @@ -292,44 +259,43 @@ async def test_row_missing_a_required_column_is_not_cached_so_the_getter_still_f await get_team_object(team_id=TEAM_ID, prisma_client=prisma, user_api_key_cache=cache) assert exc.value.status_code == 404 assert prisma.db.litellm_teamtable.find_unique.await_count == 1 - assert cache.in_memory_cache.get_cache(USER_ID) is not None + assert cache.get_local(USER_ID) is not None @pytest.mark.asyncio async def test_absent_rows_are_not_cached_as_present(): - redis = CountingRedis() + redis = FakeAuthCache() prisma = _prisma({key: None for key in ALL_ROWS}) cache = _cache(redis) await prefetch_auth_objects(refs=_refs(), user_api_key_cache=cache, prisma_client=prisma) - assert [c for c in redis.commands if c.startswith("SET")] == [] - assert cache.in_memory_cache.get_cache(f"team_id:{TEAM_ID}") is None + assert _writes(redis) == [] + assert cache.get_local(f"team_id:{TEAM_ID}") is None @pytest.mark.asyncio -async def test_redis_failure_is_swallowed_and_getters_fall_back_to_their_own_reads(): +async def test_a_failed_redis_read_reads_as_a_miss_and_the_database_fill_takes_over(): prisma = _prisma() - prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) - cache = _cache(CountingRedis(fail=True)) + cache = _cache(FakeAuthCache(fail=ConnectionError("redis down"))) await prefetch_auth_objects(refs=_refs(), user_api_key_cache=cache, prisma_client=prisma) - assert prisma.db.query_first.await_count == 0 - assert cache.in_memory_cache.get_cache(USER_ID) is None + assert prisma.db.query_first.await_count == 1 + assert cache.get_local(USER_ID)["user_id"] == USER_ID @pytest.mark.asyncio async def test_no_prisma_still_uses_redis_but_never_queries(): - seeded = CountingRedis() + seeded = FakeAuthCache() await prefetch_auth_objects(refs=_refs(), user_api_key_cache=_cache(seeded), prisma_client=_prisma()) - redis = CountingRedis(store=seeded.store) + redis = FakeAuthCache(redis=seeded.redis) cache = _cache(redis) await prefetch_auth_objects(refs=_refs(), user_api_key_cache=cache, prisma_client=None) - assert redis.round_trips == 1 - assert cache.in_memory_cache.get_cache(f"org_id:{ORG_ID}") is not None + assert _round_trips(redis) == 1 + assert cache.get_local(f"org_id:{ORG_ID}") is not None @pytest.mark.asyncio @@ -340,13 +306,13 @@ async def test_no_redis_goes_straight_to_one_query(): await prefetch_auth_objects(refs=_refs(), user_api_key_cache=cache, prisma_client=prisma) assert prisma.db.query_first.await_count == 1 - assert cache.in_memory_cache.get_cache(f"team_membership:{USER_ID}:{TEAM_ID}") is not None + assert cache.get_local(f"team_membership:{USER_ID}:{TEAM_ID}") is not None @pytest.mark.asyncio async def test_identity_prefetch_warms_the_end_user_so_its_getter_needs_neither_redis_nor_the_database(): end_user_key = end_user_cache_key("eu-1") - redis = CountingRedis({end_user_key: json.dumps({"user_id": "eu-1", "blocked": False, "spend": 0.0})}) + redis = FakeAuthCache(redis={end_user_key: {"user_id": "eu-1", "blocked": False, "spend": 0.0}}) cache = _cache(redis) prisma = _prisma() @@ -354,16 +320,16 @@ async def test_identity_prefetch_warms_the_end_user_so_its_getter_needs_neither_ end_user = await get_end_user_object(end_user_id="eu-1", prisma_client=prisma, user_api_key_cache=cache) assert end_user is not None and end_user.user_id == "eu-1" - assert redis.commands == [f"MGET {end_user_key} {end_user_restricted_registry_cache_key()}"] + assert redis.commands == [("MGET", (end_user_key, end_user_restricted_registry_cache_key()))] assert prisma.db.mock_calls == [] @pytest.mark.asyncio async def test_identity_prefetch_does_not_cache_an_absent_entry_as_present(): - redis = CountingRedis({}) + redis = FakeAuthCache() cache = _cache(redis) await prefetch_identity_keys([end_user_cache_key("eu-absent")], cache) - assert redis.round_trips == 1 - assert cache.in_memory_cache.get_cache(end_user_cache_key("eu-absent")) is None + assert _round_trips(redis) == 1 + assert cache.get_local(end_user_cache_key("eu-absent")) is None diff --git a/tests/unit/proxy/auth/test_handle_jwt.py b/tests/unit/proxy/auth/test_handle_jwt.py index 4c223edc20d..ae6bf95c13f 100644 --- a/tests/unit/proxy/auth/test_handle_jwt.py +++ b/tests/unit/proxy/auth/test_handle_jwt.py @@ -7744,7 +7744,7 @@ def _entra_signed_app_token(monkeypatch, azp: str, scope: str) -> tuple[JWTHandl monkeypatch.setenv("JWT_PUBLIC_KEY_URL", jwks_url) monkeypatch.delenv("JWT_AUDIENCE", raising=False) private_key, jwk = _get_rsa_key_and_jwk(kid="entra-kid") - cache = DualCache() + cache = UserApiKeyCache() cache.set_cache(key=f"litellm_jwt_auth_keys_{jwks_url}", value=[jwk]) jwt_handler = JWTHandler() jwt_handler.update_environment( 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 fa2f189703a..32aca696d0e 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 @@ -26,6 +26,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( budget_windows_absent_cache_key, ) from litellm.redis.stores.spend_counters import SpendCounterStore +from tests.unit.redis.stores.auth_cache_double import shared_cache class _RecordingRedisClient(Redis): @@ -271,6 +272,29 @@ async def test_subscriber_evicts_a_user_whose_id_looks_like_a_no_window_marker_a assert cache.budget_windows_absent_cache.get_cache(key=user_id) is None +@pytest.mark.asyncio +async def test_subscriber_drops_the_shared_auth_cache_copy_so_the_next_read_refills_from_redis() -> None: + """With enable_redis_auth_cache the keyspace's memory tier is native; the Python subscriber forwards each + eviction broadcast to it, and the next read takes the other worker's write from Redis.""" + cache, store = shared_cache() + await cache.async_set_cache("team_id:t-1", {"team_id": "t-1", "models": ["old"]}) + store.redis["team_id:t-1"] = {"team_id": "t-1", "models": ["new"]} + + pubsub = _QueuePubSub(initial_messages=[_invalidation_message("team_id:t-1")]) + subscriber = AuthCacheInvalidationSubscriber( + redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[pubsub])), + user_api_key_cache=cache, + ) + subscriber.start() + try: + assert await asyncio.wait_for(pubsub.drained.wait(), timeout=5) + finally: + await subscriber.stop() + + assert cache.get_local("team_id:t-1") is None + assert await cache.async_get_cache("team_id:t-1") == {"team_id": "t-1", "models": ["new"]} + + @pytest.mark.asyncio async def test_subscriber_deletes_additional_in_memory_cache_entry_on_message() -> None: """ diff --git a/tests/unit/proxy/common_utils/test_user_api_key_cache.py b/tests/unit/proxy/common_utils/test_user_api_key_cache.py index 21ae8655b97..0a13712e6ec 100644 --- a/tests/unit/proxy/common_utils/test_user_api_key_cache.py +++ b/tests/unit/proxy/common_utils/test_user_api_key_cache.py @@ -1,7 +1,7 @@ import datetime import hashlib -import json from typing import Any, Final +from unittest.mock import MagicMock import pytest @@ -18,6 +18,8 @@ from litellm.proxy.common_utils.user_api_key_cache import ( is_user_key_cache_key, ) from litellm.proxy.proxy_server import UserAPIKeyCacheTTLEnum +from litellm.rust_bridge._native import NativeAuthCache +from tests.unit.redis.stores.auth_cache_double import FakeAuthCache, shared_cache HASHED_TOKEN = hashlib.sha256(b"sk-lit7563-hot-key").hexdigest() @@ -45,70 +47,6 @@ class _ManualClock: self.seconds += seconds -class FakeRedisCache(RedisCache): - """ - In-memory fake that enforces the UserApiKeyCache Redis payload contract. - - For user_api_key_cache entries we expect Redis to store a JSON object (dict) - produced by `CacheCodec.serialize(..., model_type=...)`. - - This fake: - - raises TypeError if the value is not a dict - - raises TypeError if the dict is not JSON-serializable - - Records the ``ttl`` kwarg DualCache forwards on each Redis write for tests. - """ - - def __init__(self): # noqa: super().__init__ skipped intentionally - self._store: dict[str, str] = {} - self.last_ttl: Any = None - - def set_cache(self, key: str, value: Any, **kwargs): # type: ignore[override] - if not isinstance(value, dict): - raise TypeError("FakeRedisCache only accepts dict payloads") - self.last_ttl = kwargs.get("ttl") - self._store[key] = json.dumps(value) - return True - - def get_cache(self, key: str, **kwargs): # type: ignore[override] - raw = self._store.get(key) - if raw is None: - return None - return json.loads(raw) - - async def async_set_cache(self, key: str, value: Any, **kwargs): # type: ignore[override] - if not isinstance(value, dict): - raise TypeError("FakeRedisCache only accepts dict payloads") - self.last_ttl = kwargs.get("ttl") - self._store[key] = json.dumps(value) - return True - - async def async_get_cache(self, key: str, **kwargs): # type: ignore[override] - raw = self._store.get(key) - if raw is None: - return None - return json.loads(raw) - - def delete_cache(self, key: str): # type: ignore[override] - self._store.pop(key, None) - - async def async_delete_cache(self, key: str): # type: ignore[override] - self._store.pop(key, None) - - async def delete_cache_keys(self, keys): # type: ignore[override] - for key in keys: - self._store.pop(key, None) - - -class PartitionFailingRedisCache(FakeRedisCache): - """Fails the batch delete for the key-object partition and no other.""" - - async def delete_cache_keys(self, keys): # type: ignore[override] - if any(is_user_key_cache_key(key) for key in keys): - raise ConnectionError("redis unavailable") - await super().delete_cache_keys(keys) - - def _make_key_obj(token: str = "tok") -> UserAPIKeyAuth: # Minimal object (UserAPIKeyAuth inherits token from base view). return UserAPIKeyAuth(token=token) @@ -128,7 +66,6 @@ class TestUserApiKeyCache: mem = CapturingInMemoryCache() cache = UserApiKeyCache( in_memory_cache=mem, - redis_cache=FakeRedisCache(), default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value, ) await cache.async_set_cache( @@ -145,127 +82,138 @@ class TestUserApiKeyCache: mem = CapturingInMemoryCache() cache = UserApiKeyCache( in_memory_cache=mem, - redis_cache=FakeRedisCache(), default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value, ) cache.set_cache("sk", _make_key_obj("s"), model_type=UserAPIKeyAuth) assert mem.last_ttl == UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value @pytest.mark.asyncio - async def test_async_set_forwards_default_in_memory_ttl_to_redis_layer(self): - """ - DualCache injects missing ``ttl`` from ``default_in_memory_ttl`` into kwargs - before calling ``redis_cache.async_set_cache`` — Redis should receive the same - TTL as memory (matches proxy defaults: enum 60s). - """ - fake = FakeRedisCache() - cache = UserApiKeyCache( - redis_cache=fake, - default_in_memory_ttl=60, - ) + async def test_the_shared_store_gets_default_in_memory_ttl_and_writes_without_a_ttl(self): + """The native store applies ``default_in_memory_ttl`` to a write that names no TTL, in memory and Redis, + as DualCache does; the Rust side pins the resulting Redis TTL.""" + cache, fake = shared_cache(default_in_memory_ttl=60) - await cache.async_set_cache( - key="ttl-key", - value=_make_key_obj("ttl-tok"), - model_type=UserAPIKeyAuth, - ) + await cache.async_set_cache(key="ttl-key", value=_make_key_obj("ttl-tok"), model_type=UserAPIKeyAuth) - assert fake.last_ttl == 60 + assert [settings.memory_ttl for settings in fake.settings] == [60] + assert fake.commands == [("SET", "ttl-key", None)] @pytest.mark.asyncio async def test_async_set_explicit_ttl_override_reaches_redis(self): - fake = FakeRedisCache() - cache = UserApiKeyCache( - redis_cache=fake, - default_in_memory_ttl=60, - ) + cache, fake = shared_cache(default_in_memory_ttl=60) - await cache.async_set_cache( - key="k", - value=_make_key_obj("x"), - model_type=UserAPIKeyAuth, - ttl=900, - ) + await cache.async_set_cache(key="k", value=_make_key_obj("x"), model_type=UserAPIKeyAuth, ttl=900) - assert fake.last_ttl == 900 + assert fake.commands == [("SET", "k", 900)] - def test_sync_set_forwards_default_in_memory_ttl_to_redis_layer(self): - fake = FakeRedisCache() - cache = UserApiKeyCache( - redis_cache=fake, - default_in_memory_ttl=45, - ) - cache.set_cache( - "sk", - _make_key_obj("sync"), - model_type=UserAPIKeyAuth, - ) - assert fake.last_ttl == 45 + def test_sync_set_writes_through_the_shared_store_as_str(self): + cache, fake = shared_cache(default_in_memory_ttl=45) + + cache.set_cache("sk", _make_key_obj("sync"), model_type=UserAPIKeyAuth, ttl=45) + + assert fake.commands == [("SET_STR", "sk", 45)] + assert fake.redis["sk"]["token"] == "sync" @pytest.mark.asyncio async def test_async_set_typed_stores_serialized_payload_in_memory_and_redis(self): - cache = UserApiKeyCache(redis_cache=FakeRedisCache()) + cache, fake = shared_cache() obj = _make_key_obj("abc") await cache.async_set_cache("k", obj, model_type=UserAPIKeyAuth) - # In-memory hit should still be raw dict (not BaseModel) because wrapper - # stores the serialized payload into both layers. - raw = await cache.in_memory_cache.async_get_cache("k") # type: ignore[union-attr] + raw = cache.get_local("k") assert isinstance(raw, dict) assert raw["token"] == "abc" - - # Redis should also hold the same serialized dict - redis_raw = await cache.redis_cache.async_get_cache("k") # type: ignore[union-attr] - assert redis_raw == raw + assert fake.redis["k"] == raw + assert cache.in_memory_cache.get_cache("k") is None, "redis-py's DualCache tier is never written" @pytest.mark.asyncio async def test_async_get_typed_returns_model_on_valid_hit(self): - cache = UserApiKeyCache(redis_cache=FakeRedisCache()) - await cache.async_set_cache("k", {"token": "abc"}, model_type=UserAPIKeyAuth) + cache, fake = shared_cache(FakeAuthCache(redis={"k": {"token": "abc"}})) value = await cache.async_get_cache("k", model_type=UserAPIKeyAuth) assert value is not None assert isinstance(value, UserAPIKeyAuth) assert value.token == "abc" + assert fake.commands == [("GET", "k")] @pytest.mark.asyncio async def test_async_get_typed_returns_none_on_validation_failure_after_hit(self): - cache = UserApiKeyCache(redis_cache=FakeRedisCache()) - - # Bypass UserApiKeyCache.serialize: CacheCodec rejects non-dict cached values - # for dict-based models (deserialize returns None). - await cache.in_memory_cache.async_set_cache(key="k", value="invalid-payload-not-a-dict") + cache, _fake = shared_cache() + cache.set_local("k", "invalid-payload-not-a-dict", None) value = await cache.async_get_cache("k", model_type=UserAPIKeyAuth) assert value is None - def test_fake_redis_cache_rejects_non_json_serializable_values(self): - fake = FakeRedisCache() + @pytest.mark.asyncio + async def test_a_failed_shared_read_is_a_miss_and_a_failed_write_keeps_the_request_going(self): + cache, fake = shared_cache(FakeAuthCache(fail=ConnectionError("redis down"))) - class NotSerializable: - pass + assert await cache.async_get_cache("team_id:t1") is None + await cache.async_set_cache("team_id:t1", {"team_id": "t1"}, ttl=5) + assert cache.get_cache("team_id:t2") is None - with pytest.raises(TypeError): - fake.set_cache("k", NotSerializable()) + assert [command[0] for command in fake.commands] == ["GET", "SET", "GET"] + assert cache.get_local("team_id:t1") == {"team_id": "t1"} - with pytest.raises(TypeError): - fake.set_cache("k2", {"ok": NotSerializable()}) + @pytest.mark.asyncio + async def test_a_value_the_bridge_refuses_is_logged_not_raised(self): + cache = UserApiKeyCache( + redis_cache=MagicMock(spec=RedisCache), + auth_cache_handle=lambda _c, _s, _t: _refusing(), + redis_client=lambda _: None, + ) + + await cache.async_set_cache("team_id:t1", {"team_id": "t1"}) + cache.set_cache("team_id:t1", {"team_id": "t1"}) @pytest.mark.asyncio async def test_a_model_holding_a_datetime_is_cached_as_its_json_dump(self): - cache = UserApiKeyCache() + cache = UserApiKeyCache( + redis_cache=MagicMock(spec=RedisCache), + auth_cache_handle=lambda _c, _s, _t: NativeAuthCache(None), + redis_client=lambda _: None, + ) key_obj = UserAPIKeyAuth(token="abc", expires=datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc)) - await cache.async_set_cache("k", key_obj, model_type=UserAPIKeyAuth) + await cache.async_set_cache("k", key_obj, model_type=UserAPIKeyAuth, local_only=True) - stored = cache.get_cache("k", local_only=True) - assert stored == key_obj.model_dump(mode="json") - assert json.loads(json.dumps(stored)) == stored + assert cache.get_local("k") == key_obj.model_dump(mode="json") cached = await cache.async_get_cache("k", model_type=UserAPIKeyAuth) assert cached is not None assert cached.model_dump(mode="json") == key_obj.model_dump(mode="json") + @pytest.mark.asyncio + async def test_keyspace_operations_redis_py_cannot_reach_are_refused_once_shared(self): + cache, _fake = shared_cache() + + with pytest.raises(NotImplementedError): + await cache.async_increment_cache("team_id:t1", 1) + with pytest.raises(NotImplementedError): + await cache.async_batch_get_cache(["team_id:t1"]) + assert cache.redis_cache is None + + @pytest.mark.asyncio + async def test_global_spend_increments_only_an_entry_redis_still_holds(self): + cache, fake = shared_cache(FakeAuthCache(redis={"litellm-proxy-budget:spend": 1.5})) + memory_only = UserApiKeyCache() + + assert await cache.increment_if_present("litellm-proxy-budget:spend", 2.0) == 3.5 + assert await cache.increment_if_present("absent", 2.0) is None + assert await memory_only.increment_if_present("litellm-proxy-budget:spend", 2.0) is None + assert [command[:2] for command in fake.commands] == [ + ("INCRBYFLOAT_IF_PRESENT", "litellm-proxy-budget:spend"), + ("INCRBYFLOAT_IF_PRESENT", "absent"), + ] + + +def _refusing() -> FakeAuthCache: + class Refusing(FakeAuthCache): + def set_local(self, key: str, value: object, ttl_seconds: float | None = None) -> None: + raise TypeError("not a plain data value") + + return Refusing() + class TestUserKeyObjectPartition: """ @@ -311,19 +259,16 @@ class TestUserKeyObjectPartition: assert key_obj.token == HASHED_TOKEN @pytest.mark.asyncio - async def test_redis_hit_backfills_key_object_partition_with_configured_ttl(self): - redis = FakeRedisCache() - writer = UserApiKeyCache(redis_cache=redis, default_in_memory_ttl=30) + async def test_redis_hit_backfills_this_workers_shared_memory(self): + writer, fake = shared_cache(default_in_memory_ttl=30) await writer.async_set_cache(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), model_type=UserAPIKeyAuth) + fake.flush_local() - key_partition = CapturingInMemoryCache() - reader = UserApiKeyCache(redis_cache=redis, default_in_memory_ttl=30, key_object_in_memory_cache=key_partition) - key_obj = await reader.async_get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth) + key_obj = await writer.async_get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth) assert key_obj is not None assert key_obj.token == HASHED_TOKEN - assert key_partition.last_ttl == 30 - assert HASHED_TOKEN not in reader.in_memory_cache.cache_dict + assert writer.get_local(HASHED_TOKEN)["token"] == HASHED_TOKEN @pytest.mark.asyncio async def test_update_cache_ttl_applies_to_key_object_partition(self): @@ -336,80 +281,65 @@ class TestUserKeyObjectPartition: assert key_partition.last_ttl == 7 @pytest.mark.asyncio - async def test_attach_redis_cache_applies_to_key_object_partition(self): - redis = FakeRedisCache() - cache = UserApiKeyCache() - cache.attach_redis_cache(redis) + async def test_attach_redis_cache_moves_both_partitions_onto_the_shared_store(self): + fake = FakeAuthCache() + cache = UserApiKeyCache(auth_cache_handle=fake.factory, redis_client=lambda _: None) + cache.set_cache(end_user_cache_key("before"), {"user_id": "before"}) + assert not cache.has_redis + cache.attach_redis_cache(MagicMock(spec=RedisCache)) await cache.async_set_cache(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), model_type=UserAPIKeyAuth) + await cache.async_set_cache(end_user_cache_key("u1"), {"user_id": "u1"}) - other_worker = UserApiKeyCache(redis_cache=redis) - key_obj = await other_worker.async_get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth) - assert key_obj is not None - assert key_obj.token == HASHED_TOKEN + assert cache.has_redis + assert set(fake.redis) == {HASHED_TOKEN, end_user_cache_key("u1")} + assert HASHED_TOKEN not in cache.key_object_cache.in_memory_cache.cache_dict @pytest.mark.asyncio - async def test_delete_removes_key_object_from_partition_and_redis(self): - redis = FakeRedisCache() - cache = UserApiKeyCache(redis_cache=redis) + async def test_delete_removes_key_object_from_memory_and_redis(self): + cache, fake = shared_cache() await cache.async_set_cache(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), model_type=UserAPIKeyAuth) - assert await cache.async_get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth) is not None cache.delete_cache(HASHED_TOKEN) - assert await cache.async_get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth) is None - assert await redis.async_get_cache(HASHED_TOKEN) is None + assert cache.get_local(HASHED_TOKEN) is None + assert HASHED_TOKEN not in fake.redis @pytest.mark.asyncio - async def test_async_delete_removes_key_object_from_partition_and_redis(self): - redis = FakeRedisCache() - cache = UserApiKeyCache(redis_cache=redis) + async def test_async_delete_removes_key_object_from_memory_and_redis(self): + cache, fake = shared_cache() await cache.async_set_cache(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), model_type=UserAPIKeyAuth) await cache.async_delete_cache(HASHED_TOKEN) assert await cache.async_get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth) is None - assert await redis.async_get_cache(HASHED_TOKEN) is None + assert HASHED_TOKEN not in fake.redis @pytest.mark.asyncio - async def test_batch_delete_routes_each_key_to_its_partition(self): - """A batch delete has to clear the same partition the single delete does. - - ``DualCache``'s batch delete only knows about the main in-memory cache, so - inheriting it unchanged leaves a key object sitting in ``key_object_cache`` - with its pre-reset spend, and the next request is authorized against that - stale copy until the local entry expires. - """ - redis = FakeRedisCache() - cache = UserApiKeyCache(redis_cache=redis) + async def test_batch_delete_clears_every_key_in_one_round_trip(self): + cache, fake = shared_cache() await cache.async_set_cache(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), model_type=UserAPIKeyAuth) await cache.async_set_cache(end_user_cache_key("u1"), {"user_id": "u1"}) await cache.async_delete_cache_keys([HASHED_TOKEN, end_user_cache_key("u1")]) - assert await cache.async_get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth) is None - assert await cache.async_get_cache(end_user_cache_key("u1")) is None - assert await redis.async_get_cache(HASHED_TOKEN) is None - assert await redis.async_get_cache(end_user_cache_key("u1")) is None + assert cache.get_local(HASHED_TOKEN) is None + assert cache.get_local(end_user_cache_key("u1")) is None + assert fake.redis == {} + assert fake.commands[-1] == ("DEL", (HASHED_TOKEN, end_user_cache_key("u1"))) @pytest.mark.asyncio - async def test_batch_delete_clears_the_other_partition_when_one_fails(self): - """One partition failing must not cost the other its deletions. - - A caller batching these has already committed the rows they cache, so a - partition that is skipped keeps authorizing against pre-reset spend until - the entry expires. The failure is still raised for the caller to report. - """ - redis = PartitionFailingRedisCache() - cache = UserApiKeyCache(redis_cache=redis) - await cache.async_set_cache(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), model_type=UserAPIKeyAuth) + async def test_a_failed_batch_delete_still_clears_memory_and_raises(self): + """A caller batching these has already committed the rows they cache, so memory is dropped first and the + Redis failure is still raised for the caller to report.""" + cache, fake = shared_cache() await cache.async_set_cache(end_user_cache_key("u1"), {"user_id": "u1"}) + fake.fail = ConnectionError("redis unavailable") with pytest.raises(ConnectionError): await cache.async_delete_cache_keys([HASHED_TOKEN, end_user_cache_key("u1")]) - assert await cache.async_get_cache(end_user_cache_key("u1")) is None - assert await redis.async_get_cache(end_user_cache_key("u1")) is None + assert cache.get_local(end_user_cache_key("u1")) is None @pytest.mark.asyncio async def test_pipeline_write_routes_each_entry_to_its_partition(self): @@ -436,10 +366,21 @@ class TestUserKeyObjectPartition: assert cache.get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth) is None assert cache.get_cache(end_user_cache_key("u1")) is None - def test_in_memory_cache_for_routes_by_key(self): - cache = UserApiKeyCache() - assert cache.in_memory_cache_for(HASHED_TOKEN) is cache.key_object_cache.in_memory_cache - assert cache.in_memory_cache_for(end_user_cache_key("u1")) is cache.in_memory_cache + def test_local_access_routes_by_key_without_redis(self): + cache = UserApiKeyCache(default_in_memory_ttl=60) + cache.set_local(HASHED_TOKEN, {"token": HASHED_TOKEN}, None) + cache.set_local(end_user_cache_key("u1"), {"user_id": "u1"}, 5) + + assert cache.key_object_cache.in_memory_cache.get_cache(HASHED_TOKEN) == {"token": HASHED_TOKEN} + assert cache.in_memory_cache.get_cache(end_user_cache_key("u1")) == {"user_id": "u1"} + cache.delete_local(HASHED_TOKEN) + assert cache.get_local(HASHED_TOKEN) is None + + def test_a_resized_memory_tier_reaches_the_shared_store(self): + cache, fake = shared_cache() + cache.update_in_memory_max_size(500) + + assert [settings.memory_entries for settings in fake.settings] == [200, 500] def test_update_in_memory_max_size_applies_to_key_object_partition(self): cache = UserApiKeyCache( @@ -522,7 +463,6 @@ class TestManagementObjectTTL: mem = CapturingInMemoryCache() cache = UserApiKeyCache( in_memory_cache=mem, - redis_cache=FakeRedisCache(), default_in_memory_ttl=300, ) assert get_management_object_ttl(cache) != (DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL) diff --git a/tests/unit/proxy/credential_endpoints/test_endpoints.py b/tests/unit/proxy/credential_endpoints/test_endpoints.py index cdd18fd4eef..f9f8235361e 100644 --- a/tests/unit/proxy/credential_endpoints/test_endpoints.py +++ b/tests/unit/proxy/credential_endpoints/test_endpoints.py @@ -1810,7 +1810,7 @@ def test_update_credential_purges_connections_when_github_host_changes(monkeypat redis_cache: Final = MagicMock() redis_cache.async_set_cache = AsyncMock(side_effect=set_cache) redis_cache.async_get_cache = AsyncMock(side_effect=get_cache) - monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache(redis_cache=redis_cache)) + monkeypatch.setattr(proxy_server, "user_api_key_cache", shared_cache(connection=redis_cache)[0]) session_key: Final = _session_cache_key("user-a", "copilot", "github.com") _SESSION_CACHE.set_cache( @@ -2063,7 +2063,8 @@ def test_update_credential_still_accepts_a_body_without_credential_values(creden from contextlib import contextmanager as _ctx -from litellm.caching.dual_cache import DualCache +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from tests.unit.redis.stores.auth_cache_double import shared_cache from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy.credential_endpoints import user_provider_credentials as upc @@ -2159,7 +2160,7 @@ def _patch_user_connection_env(monkeypatch, credentials, rows=()): prisma_client.db.litellm_userprovidercredentials = table prisma_client.writer_db = prisma_client.db monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client) - cache = DualCache() + cache = UserApiKeyCache() monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache) return table @@ -2417,7 +2418,7 @@ def test_deleting_the_credential_purges_user_connections(monkeypatch, credential prisma_client.writer_db = prisma_client.db prisma_client.db.litellm_credentialstable.find_unique = AsyncMock(return_value=None) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client) - monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", DualCache()) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache()) response = _delete_credential("copilot-cred") assert response.status_code == 200, response.text @@ -2882,10 +2883,9 @@ def test_user_connection_poll_valid_handle_works_with_a_fresh_cache(monkeypatch) start = _call_as("POST", "/credentials/copilot-cred/user_connection/start", auth=_as_user()) assert start.status_code == 200, start.text # simulate a different worker: swap in a brand new empty cache - from litellm.caching.dual_cache import DualCache from litellm.proxy import proxy_server - monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache()) + monkeypatch.setattr(proxy_server, "user_api_key_cache", UserApiKeyCache()) poll = _call_as( "POST", "/credentials/copilot-cred/user_connection/poll", @@ -2910,7 +2910,7 @@ def test_delete_user_connection_fails_closed_when_the_tombstone_write_fails(monk async_get_cache=AsyncMock(return_value=None), async_delete_cache=AsyncMock(return_value=None), ) - monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache(redis_cache=failing_redis)) + monkeypatch.setattr(proxy_server, "user_api_key_cache", shared_cache(connection=failing_redis)[0]) response = _call_as("DELETE", "/credentials/copilot-cred/user_connection", auth=_as_user()) assert response.status_code == 503, response.text @@ -2932,7 +2932,7 @@ def test_delete_user_connection_fails_when_the_tombstone_write_silently_noops(mo async_get_cache=AsyncMock(side_effect=lambda key: store.get(key)), async_delete_cache=AsyncMock(return_value=None), ) - monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache(redis_cache=silent_redis)) + monkeypatch.setattr(proxy_server, "user_api_key_cache", shared_cache(connection=silent_redis)[0]) response = _call_as("DELETE", "/credentials/copilot-cred/user_connection", auth=_as_user()) assert response.status_code == 503, response.text @@ -2953,7 +2953,7 @@ def test_user_connection_poll_reports_503_when_the_cache_overwrite_cannot_land(m async_get_cache=AsyncMock(return_value="stale-value"), async_delete_cache=AsyncMock(return_value=None), ) - monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache(redis_cache=stale_redis)) + monkeypatch.setattr(proxy_server, "user_api_key_cache", shared_cache(connection=stale_redis)[0]) with _github_http( { diff --git a/tests/unit/proxy/credential_endpoints/test_user_provider_credentials.py b/tests/unit/proxy/credential_endpoints/test_user_provider_credentials.py index dae3a0459e2..1612b6d8a16 100644 --- a/tests/unit/proxy/credential_endpoints/test_user_provider_credentials.py +++ b/tests/unit/proxy/credential_endpoints/test_user_provider_credentials.py @@ -7,9 +7,10 @@ from unittest.mock import AsyncMock, MagicMock import pytest -from litellm.caching.dual_cache import DualCache +from litellm.caching.redis_cache import RedisCache from litellm.constants import GITHUB_COPILOT_DEFAULT_HOST from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.credential_endpoints.user_provider_credentials import ( GithubCopilotUserConnectionPayload, aget_user_provider_credentials, @@ -22,6 +23,12 @@ from litellm.proxy.credential_endpoints.user_provider_credentials import ( set_user_provider_credential_cache, upsert_user_provider_credential, ) +from tests.unit.redis.stores.auth_cache_double import shared_cache + + +def _auth_cache(redis: RedisCache) -> UserApiKeyCache: + """A worker's auth cache with ``enable_redis_auth_cache`` on, whose connection the credential cache shares.""" + return shared_cache(connection=redis)[0] def _prisma(table=None): @@ -122,7 +129,7 @@ async def test_aget_credentials_returns_plaintext_and_caches_ciphertext(monkeypa table.find_many = AsyncMock(return_value=[_row(credential_b64=ciphertext)]) prisma_client = _prisma(table) redis = _fake_redis(monkeypatch) - cache = DualCache(redis_cache=redis) + cache = _auth_cache(redis) connections = await aget_user_provider_credentials(prisma_client, cache, "user-a", ["copilot-cred"]) assert connections == {"copilot-cred": {"access_token": "gho_secret", "github_host": "github.com"}} @@ -149,7 +156,7 @@ async def test_aget_connections_preserves_github_host(monkeypatch): table: Final = MagicMock() table.find_many = AsyncMock(return_value=[_row(credential_b64=upc._encode(payload))]) prisma_client: Final = _prisma(table) - cache: Final = DualCache(redis_cache=_fake_redis(monkeypatch)) + cache: Final = _auth_cache(_fake_redis(monkeypatch)) connections: Final = await aget_user_provider_credentials( prisma_client, @@ -166,7 +173,7 @@ async def test_legacy_connection_without_github_host_defaults_to_github_com(monk table: Final = MagicMock() table.find_many = AsyncMock(return_value=[_row(credential_b64=ciphertext)]) prisma_client: Final = _prisma(table) - cache: Final = DualCache(redis_cache=_fake_redis(monkeypatch)) + cache: Final = _auth_cache(_fake_redis(monkeypatch)) connections: Final = await aget_user_provider_credentials( prisma_client, @@ -187,7 +194,7 @@ async def test_aget_credentials_negative_caches_missing_connection(monkeypatch): table = MagicMock() table.find_many = AsyncMock(return_value=[]) prisma_client = _prisma(table) - cache = DualCache(redis_cache=_fake_redis(monkeypatch)) + cache = _auth_cache(_fake_redis(monkeypatch)) assert await aget_user_provider_credentials(prisma_client, cache, "user-a", ["copilot-cred"]) == {} assert await aget_user_provider_credentials(prisma_client, cache, "user-a", ["copilot-cred"]) == {} @@ -202,7 +209,7 @@ async def test_invalidate_cache_forces_db_refetch(monkeypatch): table = MagicMock() table.find_many = AsyncMock(return_value=[_row(credential_b64=upc._encode(payload))]) prisma_client = _prisma(table) - cache = DualCache(redis_cache=_fake_redis(monkeypatch)) + cache = _auth_cache(_fake_redis(monkeypatch)) await aget_user_provider_credentials(prisma_client, cache, "user-a", ["copilot-cred"]) await drop_user_provider_credential_cache(cache, "user-a", "copilot-cred") @@ -222,8 +229,8 @@ async def test_disconnect_on_one_worker_invalidates_other_workers_via_redis(monk from litellm.proxy.credential_endpoints import user_provider_credentials as upc redis = _fake_redis(monkeypatch) - cache_a = DualCache(redis_cache=redis) - cache_b = DualCache(redis_cache=redis) + cache_a = _auth_cache(redis) + cache_b = _auth_cache(redis) payload = GithubCopilotUserConnectionPayload(access_token="gho_secret", github_login="octo") ciphertext = upc._encode(payload) @@ -255,7 +262,7 @@ async def test_stale_fill_cannot_overwrite_a_disconnect_tombstone(monkeypatch): from litellm.proxy.credential_endpoints import user_provider_credentials as upc redis = _fake_redis(monkeypatch) - cache = DualCache(redis_cache=redis) + cache = _auth_cache(redis) payload = GithubCopilotUserConnectionPayload(access_token="gho_secret", github_login="octo") ciphertext = upc._encode(payload) @@ -281,7 +288,7 @@ async def test_redis_read_failure_falls_back_to_the_database(monkeypatch): redis = _fake_redis(monkeypatch) monkeypatch.setattr(redis, "async_get_cache", AsyncMock(side_effect=RuntimeError("redis down"))) - cache = DualCache(redis_cache=redis) + cache = _auth_cache(redis) payload = GithubCopilotUserConnectionPayload(access_token="gho_secret", github_login="octo") table = MagicMock() @@ -298,7 +305,7 @@ async def test_redis_write_and_delete_failures_do_not_fail_the_request(monkeypat redis = _fake_redis(monkeypatch) monkeypatch.setattr(redis, "async_set_cache", AsyncMock(side_effect=RuntimeError("redis down"))) monkeypatch.setattr(redis, "async_delete_cache", AsyncMock(side_effect=RuntimeError("redis down"))) - cache = DualCache(redis_cache=redis) + cache = _auth_cache(redis) payload = GithubCopilotUserConnectionPayload(access_token="gho_secret", github_login="octo") from litellm.proxy.credential_endpoints import user_provider_credentials as upc @@ -324,8 +331,8 @@ async def test_without_redis_no_worker_local_entries_are_kept(): table = MagicMock() table.find_many = AsyncMock(side_effect=[[_row(credential_b64=ciphertext)], [_row(credential_b64=ciphertext)], []]) prisma_client = _prisma(table) - cache_a = DualCache() - cache_b = DualCache() + cache_a = UserApiKeyCache() + cache_b = UserApiKeyCache() assert await aget_user_provider_credentials(prisma_client, cache_a, "user-a", ["copilot-cred"]) == { "copilot-cred": {"access_token": "gho_secret", "github_host": "github.com"} @@ -345,7 +352,7 @@ async def test_connect_overwrites_a_not_connected_tombstone(monkeypatch): from litellm.proxy.credential_endpoints import user_provider_credentials as upc redis = _fake_redis(monkeypatch) - cache = DualCache(redis_cache=redis) + cache = _auth_cache(redis) await invalidate_user_provider_credential_cache(cache, "user-a", "copilot-cred") payload = GithubCopilotUserConnectionPayload(access_token="gho_secret", github_login="octo") @@ -368,7 +375,7 @@ async def test_stale_not_connected_fill_cannot_overwrite_a_connect_write(monkeyp from litellm.proxy.credential_endpoints import user_provider_credentials as upc redis = _fake_redis(monkeypatch) - cache = DualCache(redis_cache=redis) + cache = _auth_cache(redis) payload = GithubCopilotUserConnectionPayload(access_token="gho_secret", github_login="octo") @@ -391,7 +398,7 @@ async def test_set_reports_connected_without_redis(): """No Redis configured: there is no stale entry to defend against, so the cache refresh succeeds by definition.""" payload = GithubCopilotUserConnectionPayload(access_token="gho_secret", github_login="octo") - assert await set_user_provider_credential_cache(DualCache(), "user-a", "copilot-cred", payload) + assert await set_user_provider_credential_cache(UserApiKeyCache(), "user-a", "copilot-cred", payload) @pytest.mark.asyncio @@ -405,7 +412,7 @@ async def test_set_reports_connected_when_delete_clears_a_dropped_write(): async_get_cache=AsyncMock(side_effect=lambda key: store.get(key)), async_delete_cache=AsyncMock(side_effect=lambda key: store.pop(key, None)), ) - cache = DualCache(redis_cache=silent_redis) + cache = _auth_cache(silent_redis) payload = GithubCopilotUserConnectionPayload(access_token="gho_secret", github_login="octo") assert await set_user_provider_credential_cache(cache, "user-a", "copilot-cred", payload) @@ -421,7 +428,7 @@ async def test_set_reports_failure_when_a_stale_entry_survives_both_attempts(): async_get_cache=AsyncMock(return_value=stale), async_delete_cache=AsyncMock(return_value=None), ) - cache = DualCache(redis_cache=silent_redis) + cache = _auth_cache(silent_redis) payload = GithubCopilotUserConnectionPayload(access_token="gho_secret", github_login="octo") assert not await set_user_provider_credential_cache(cache, "user-a", "copilot-cred", payload) @@ -439,7 +446,7 @@ async def test_reads_hit_the_writer_engine_not_a_stale_replica(): prisma_client.db.litellm_userprovidercredentials = reader_table prisma_client.writer_db.litellm_userprovidercredentials = writer_table - await aget_user_provider_credentials(prisma_client, DualCache(), "user-a", ["copilot-cred"]) + await aget_user_provider_credentials(prisma_client, UserApiKeyCache(), "user-a", ["copilot-cred"]) writer_table.find_many.assert_awaited_once() reader_table.find_many.assert_not_awaited() diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index d0ea7e2391b..f463640d274 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -2522,85 +2522,32 @@ async def test_refresh_cached_team_broadcasts_team_and_alias_invalidations( @pytest.mark.asyncio async def test_refresh_cached_team_flushes_redis_set_before_broadcast(monkeypatch) -> None: - from datetime import timedelta - from types import SimpleNamespace - from unittest.mock import AsyncMock - - from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_batch import request_redis_batch_scope - from litellm.caching.redis_cache import RedisCache, RedisCircuitBreaker - from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.management_endpoints import team_endpoints + from tests.unit.caching.test_redis_batch import FakeNatives + from tests.unit.redis.stores.auth_cache_double import shared_cache - events: Final = MagicMock() - - class RecordingPipeline: - def __init__(self, calls: MagicMock) -> None: - self.calls: Final = MagicMock() - self.events: Final = calls - - async def __aenter__(self) -> "RecordingPipeline": - return self - - async def __aexit__(self, *exc: object) -> None: - return None - - def set(self, key: str, _value: str, ex: timedelta | None = None) -> None: - self.calls("SET", key) - - def delete(self, *keys: str) -> None: - for key in keys: - self.calls("DEL", key) - - async def execute(self, raise_on_error: bool = True) -> list[object]: - assert raise_on_error is False - for command in self.calls.call_args_list: - self.events(*command.args) - return [True for _ in self.calls.call_args_list] - - class RecordingRedisClient: - def __init__(self, calls: MagicMock) -> None: - self.events: Final = calls - - def pipeline(self, transaction: bool = True) -> RecordingPipeline: - assert transaction is False - return RecordingPipeline(self.events) - - class RecordingRedisCache(RedisCache): - def __init__(self) -> None: - self.redis_kwargs: Final[dict[str, object]] = {} - self.namespace: Final[str | None] = None - self.client: Final = RecordingRedisClient(events) - self._circuit_breaker: Final = RedisCircuitBreaker(failure_threshold=5, recovery_timeout=30) - self.service_logger_obj: Final = SimpleNamespace( - async_service_success_hook=AsyncMock(), async_service_failure_hook=AsyncMock() - ) - self.default_ttl: Final[int | None] = None - - def init_async_client( - self, - ) -> RecordingRedisClient: # pyright: ignore[reportIncompatibleMethodOverride] # in-memory Redis client - return self.client + natives: Final = FakeNatives() async def record_publish(cache_key: str) -> None: - events("PUBLISH", cache_key) + natives.built[0].events.append(f"PUBLISH {cache_key}") - redis_cache: Final = RecordingRedisCache() - cache: Final = UserApiKeyCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache) + cache, _auth_store = shared_cache() proxy_logging_obj: Final = MagicMock() proxy_logging_obj.internal_usage_cache.dual_cache = None monkeypatch.setattr(team_endpoints, "publish_auth_cache_invalidation", record_publish) - with request_redis_batch_scope(): + with request_redis_batch_scope(native_factory=natives): # pyright: ignore[reportArgumentType] # test double await team_endpoints.refresh_cached_team( team_row=LiteLLM_TeamTable(team_id="team-order", team_alias=None), user_api_key_cache=cache, proxy_logging_obj=proxy_logging_obj, ) - assert [(call.args[0], call.args[1]) for call in events.call_args_list] == [ - ("SET", "team_id:team-order"), - ("PUBLISH", "team_id:team-order"), + assert natives.built[0].events == [ + "declare set team_id:team-order ttl=60.0", + "flush_pre", + "PUBLISH team_id:team-order", ] diff --git a/tests/unit/proxy/management_helpers/test_team_roster_sync.py b/tests/unit/proxy/management_helpers/test_team_roster_sync.py index 74c3dee9d56..fcbd3c9a51c 100644 --- a/tests/unit/proxy/management_helpers/test_team_roster_sync.py +++ b/tests/unit/proxy/management_helpers/test_team_roster_sync.py @@ -29,6 +29,7 @@ from litellm.proxy.management_helpers.team_roster_sync import ( TeamGone, sync_team_roster, ) +from tests.unit.redis.stores.auth_cache_double import FakeAuthCache, shared_cache ADMIN: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin") TEAM: Final = "t1" @@ -666,8 +667,8 @@ async def test_removing_members_clears_this_worker_before_waiting_on_a_wedged_re key_object_in_memory_cache=InMemoryCache(max_size_in_memory=5_000), ) for key in member_keys: - cache.in_memory_cache_for(key).set_cache(key=key, value={"cache_key": key}) - seeded = [key for key in member_keys if cache.in_memory_cache_for(key).get_cache(key) is not None] + cache.set_local(key, {"cache_key": key}, None) + seeded = [key for key in member_keys if cache.get_local(key) is not None] prisma = _FakePrisma( users=[_user(u, TEAM) for u in removed], teams=[_team(*(_member(u) for u in removed))], @@ -681,7 +682,7 @@ async def test_removing_members_clears_this_worker_before_waiting_on_a_wedged_re ): sync = asyncio.create_task(_sync(prisma, RosterTarget(member_ids=frozenset()), cache)) await _until_reached_or_done(client.reached, sync) - still_cached = [key for key in seeded if cache.in_memory_cache_for(key).get_cache(key) is not None] + still_cached = [key for key in seeded if cache.get_local(key) is not None] client.release.set() outcome = await sync await _drain_publishes() @@ -691,35 +692,6 @@ async def test_removing_members_clears_this_worker_before_waiting_on_a_wedged_re assert still_cached == [] -class _AuthRedisCache: - namespace = None - - def __init__(self, entries: dict[str, object]) -> None: - self.entries = dict(entries) - self.round_trips = 0 - - async def async_get_cache(self, key: str, **kwargs: object) -> object | None: - return self.entries.get(key) - - def set_cache(self, key: str, value: object, **kwargs: object) -> None: - self.entries[key] = value - - async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None: - self.entries[key] = value - - def delete_cache(self, key: str) -> None: - self.entries.pop(key, None) - - async def async_delete_cache(self, key: str) -> None: - self.round_trips += 1 - self.entries.pop(key, None) - - async def delete_cache_keys(self, keys: Sequence[str]) -> None: - self.round_trips += 1 - for key in keys: - self.entries.pop(key, None) - - @pytest.mark.asyncio async def test_removing_members_clears_the_redis_copies_before_waiting_on_a_wedged_redis( monkeypatch: pytest.MonkeyPatch, @@ -740,14 +712,9 @@ async def test_removing_members_clears_the_redis_copies_before_waiting_on_a_wedg team_membership_reservation_cache_key(user_id=u, team_id=TEAM), ) ] # comprehension-ok: a test fixture listing each member's four cache keys - redis = _AuthRedisCache({key: {"cache_key": key} for key in member_keys}) - cache = UserApiKeyCache( - in_memory_cache=InMemoryCache(max_size_in_memory=5_000), - key_object_in_memory_cache=InMemoryCache(max_size_in_memory=5_000), - ) - cache.attach_redis_cache(redis) # pyright: ignore[reportArgumentType] # a dict-backed fake of the Redis layer + cache, redis = shared_cache(FakeAuthCache(redis={key: {"cache_key": key} for key in member_keys})) for key in member_keys: - cache.in_memory_cache_for(key).set_cache(key=key, value={"cache_key": key}) + cache.set_local(key, {"cache_key": key}, None) prisma = _FakePrisma( users=[_user(u, TEAM) for u in removed], teams=[_team(*(_member(u) for u in removed))], @@ -768,17 +735,17 @@ async def test_removing_members_clears_the_redis_copies_before_waiting_on_a_wedg assert isinstance(outcome, RosterSync) and outcome.removed == frozenset(removed) assert refilled == [] - assert not any(key in redis.entries for key in member_keys) + assert not any(key in redis.redis for key in member_keys) -class _SlowToAcknowledgeRedisCache(_AuthRedisCache): - def __init__(self, entries: dict[str, object]) -> None: - super().__init__(entries) +class _SlowToAcknowledgeRedis(FakeAuthCache): + def __init__(self, redis: dict[str, object]) -> None: + super().__init__(redis=redis) self.acknowledge = asyncio.Event() self.reached = asyncio.Event() - async def delete_cache_keys(self, keys: Sequence[str]) -> None: - await super().delete_cache_keys(keys) + async def delete(self, keys: Sequence[str]) -> None: + await super().delete(keys) self.reached.set() await self.acknowledge.wait() @@ -789,10 +756,8 @@ async def test_a_removed_members_team_key_cannot_refill_from_redis_while_the_evi ): monkeypatch.setattr(pubsub_module, "_BULK_EVICTION_WAIT_SECONDS", 0.5) token = hashlib.sha256(b"key-bob").hexdigest() - redis = _SlowToAcknowledgeRedisCache({token: {"token": token}}) - cache = UserApiKeyCache() - cache.attach_redis_cache(redis) # pyright: ignore[reportArgumentType] # a dict-backed fake of the Redis layer - cache.in_memory_cache_for(token).set_cache(key=token, value={"token": token}) + cache, redis = shared_cache(_SlowToAcknowledgeRedis({token: {"token": token}})) + cache.set_local(token, {"token": token}, None) prisma = _FakePrisma( users=[_user("bob", TEAM)], teams=[_team(_member("bob"))], @@ -808,7 +773,7 @@ async def test_a_removed_members_team_key_cannot_refill_from_redis_while_the_evi assert isinstance(outcome, RosterSync) and outcome.removed == frozenset({"bob"}) assert served_while_waiting is None - assert token not in redis.entries + assert token not in redis.redis def _removal_of(removed: Sequence[str]) -> tuple[_FakePrisma, set[str]]: @@ -863,9 +828,7 @@ async def test_the_redis_round_trips_do_not_grow_with_the_member_count(monkeypat for size in (5, 300): removed = [f"u{index:03d}" for index in range(size)] prisma, _ = _removal_of(removed) - redis = _AuthRedisCache({}) - cache = UserApiKeyCache() - cache.attach_redis_cache(redis) # pyright: ignore[reportArgumentType] # a dict-backed fake of the Redis layer + cache, redis = shared_cache() client = _SlowPublishRedisClient() with patch( @@ -876,7 +839,7 @@ async def test_the_redis_round_trips_do_not_grow_with_the_member_count(monkeypat await _drain_publishes() assert isinstance(outcome, RosterSync) and outcome.removed == frozenset(removed) - round_trips[size] = (redis.round_trips, client.round_trips) + round_trips[size] = (len(redis.commands), client.round_trips) assert round_trips[300] == round_trips[5] diff --git a/tests/unit/proxy/proxy_server/test_spend_counters.py b/tests/unit/proxy/proxy_server/test_spend_counters.py index e77284cb4bc..19db80b8930 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counters.py +++ b/tests/unit/proxy/proxy_server/test_spend_counters.py @@ -28,20 +28,12 @@ from typing import Final, Literal, cast from unittest.mock import AsyncMock, MagicMock, call import pytest -from fakeredis import FakeAsyncRedis, FakeServer -from opentelemetry.sdk.trace import TracerProvider -from opentelemetry.sdk.trace.export import SimpleSpanProcessor -from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter -from opentelemetry.trace import StatusCode import litellm import litellm.proxy.proxy_server as ps -from litellm._service_logger import ServiceLogging from litellm._internal_context import current_service_target from litellm.caching.in_memory_cache import InMemoryCache -from litellm.caching.redis_cache import RedisCache, RedisCircuitBreaker from litellm.constants import GLOBAL_PROXY_SPEND_CACHE_KEY -from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig from litellm.models.end_user import LiteLLM_EndUserTable from litellm.models.tag import LiteLLM_TagTable from litellm.models.team import LiteLLM_TeamTable, LiteLLM_TeamTableCachedObj @@ -53,6 +45,7 @@ from litellm.proxy.management_endpoints.team_endpoints import refresh_cached_tea from .conftest import normalize from litellm.rust_bridge._native import NativeRedisStoreError +from tests.unit.redis.stores.auth_cache_double import FakeAuthCache, shared_cache from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters @@ -1531,46 +1524,31 @@ async def test_update_cache_never_brings_back_a_global_proxy_spend_a_budget_rese assert cache.get_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY) is None -class _FakeRedisRunningIncrementIfPresent(FakeAsyncRedis): - def __init__(self, server: FakeServer, after_increment: Callable[[], object] = lambda: None) -> None: - super().__init__(server=server) +class _SharedSpend(FakeAuthCache): + """The shared auth cache's Redis as the global spend scalar sees it; ``after_increment`` runs once Redis applied + an increment, before its reply reaches the caller, and an increment equal to ``held`` waits for ``released``.""" + + def __init__( + self, + total: float | None, + after_increment: Callable[[], object] = lambda: None, + held: float | None = None, + applied: asyncio.Event | None = None, + released: asyncio.Event | None = None, + ) -> None: + super().__init__(redis={} if total is None else {GLOBAL_PROXY_SPEND_CACHE_KEY: total}) self._after_increment: Final = after_increment - - async def eval(self, script: str, numkeys: int, *keys_and_args: str | bytes | float) -> bytes | None: - key, increment = keys_and_args - if not await self.exists(key): - return None - total: Final = await self.incrbyfloat(key, increment) - self._after_increment() - return str(total).encode() - - -class _RedisCacheOnClient(RedisCache): - def __init__(self, client: FakeAsyncRedis) -> None: - self._injected_client: Final = client - self.namespace = None - self._circuit_breaker = RedisCircuitBreaker(failure_threshold=5, recovery_timeout=30) - self.service_logger_obj = ServiceLogging() - self.default_ttl = None - - def init_async_client(self) -> FakeAsyncRedis: - return self._injected_client - - -class _FakeRedisHoldingOneReply(FakeAsyncRedis): - def __init__(self, server: FakeServer, held_increment: float, applied: asyncio.Event, released: asyncio.Event) -> None: - super().__init__(server=server) - self._held_increment: Final = held_increment + self._held: Final = held self._applied: Final = applied self._released: Final = released - async def eval(self, script: str, numkeys: int, *keys_and_args: str | bytes | float) -> bytes | None: - key, increment = keys_and_args - total: Final = await self.incrbyfloat(key, increment) - if float(increment) == self._held_increment: + async def increment_if_present(self, key: str, amount: float) -> float | None: + total: Final = await super().increment_if_present(key, amount) + self._after_increment() + if amount == self._held and self._applied is not None and self._released is not None: self._applied.set() await self._released.wait() - return str(total).encode() + return total _SharedGlobalSpendOutcome = Literal["peers_added_spend", "reset_elsewhere", "reset_while_redis_answers", "redis_down"] @@ -1592,14 +1570,12 @@ async def test_update_cache_moves_this_workers_global_proxy_spend_to_the_shared_ ) -> None: def _reset_this_worker() -> None: if outcome == "reset_while_redis_answers": - cache.in_memory_cache.delete_cache(GLOBAL_PROXY_SPEND_CACHE_KEY) + cache.delete_local(GLOBAL_PROXY_SPEND_CACHE_KEY) - server: Final = FakeServer() - redis: Final = _FakeRedisRunningIncrementIfPresent(server, after_increment=_reset_this_worker) - if shared_total is not None: - await redis.set(GLOBAL_PROXY_SPEND_CACHE_KEY, shared_total) - server.connected = outcome != "redis_down" - cache: Final = UserApiKeyCache(in_memory_cache=InMemoryCache(), redis_cache=_RedisCacheOnClient(redis)) + redis: Final = _SharedSpend(shared_total, after_increment=_reset_this_worker) + if outcome == "redis_down": + redis.fail = ConnectionError("redis down") + cache, _ = shared_cache(redis) cache.set_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY, value=10.0, local_only=True) await ps.update_cache( @@ -1620,12 +1596,9 @@ async def test_update_cache_moves_this_workers_global_proxy_spend_to_the_shared_ async def test_update_cache_keeps_the_newest_global_proxy_spend_when_redis_replies_out_of_order() -> None: held_increment_applied: Final = asyncio.Event() held_reply_released: Final = asyncio.Event() - server: Final = FakeServer() - redis: Final = _FakeRedisHoldingOneReply( - server, held_increment=0.25, applied=held_increment_applied, released=held_reply_released + cache, _ = shared_cache( + _SharedSpend(10.0, held=0.25, applied=held_increment_applied, released=held_reply_released) ) - await redis.set(GLOBAL_PROXY_SPEND_CACHE_KEY, 10.0) - cache: Final = UserApiKeyCache(in_memory_cache=InMemoryCache(), redis_cache=_RedisCacheOnClient(redis)) cache.set_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY, value=10.0, local_only=True) async def _request(response_cost: float) -> None: @@ -1652,11 +1625,8 @@ async def test_update_cache_keeps_the_newest_global_proxy_spend_when_redis_repli @pytest.mark.usefixtures("proxy_max_budget") async def test_update_cache_adds_the_global_spend_scalar_in_redis_as_spend_counters() -> None: target_during_increment: Final = MagicMock() - redis: Final = _FakeRedisRunningIncrementIfPresent( - FakeServer(), after_increment=lambda: target_during_increment(current_service_target()) - ) - await redis.set(GLOBAL_PROXY_SPEND_CACHE_KEY, 10.0) - cache: Final = UserApiKeyCache(in_memory_cache=InMemoryCache(), redis_cache=_RedisCacheOnClient(redis)) + redis: Final = _SharedSpend(10.0, after_increment=lambda: target_during_increment(current_service_target())) + cache, _ = shared_cache(redis) await ps.update_cache( token=None, @@ -1669,53 +1639,10 @@ async def test_update_cache_adds_the_global_spend_scalar_in_redis_as_spend_count ) assert target_during_increment.call_args_list == [call("spend_counters")] - assert float(await redis.get(GLOBAL_PROXY_SPEND_CACHE_KEY)) == 12.0 + assert redis.redis[GLOBAL_PROXY_SPEND_CACHE_KEY] == 12.0 assert current_service_target() is None -@pytest.mark.asyncio -@pytest.mark.usefixtures("proxy_max_budget") -@pytest.mark.parametrize(("redis_up", "expected_status"), [(True, StatusCode.OK), (False, StatusCode.ERROR)]) -async def test_update_cache_traces_the_global_spend_redis_call_under_the_request_span( - monkeypatch: pytest.MonkeyPatch, redis_up: bool, expected_status: StatusCode -) -> None: - exporter: Final = InMemorySpanExporter() - provider: Final = TracerProvider() - provider.add_span_processor(SimpleSpanProcessor(exporter)) - monkeypatch.setattr( - litellm, - "service_callback", - [OpenTelemetry(config=OpenTelemetryConfig(exporter=exporter), tracer_provider=provider)], - ) - server: Final = FakeServer() - redis: Final = _FakeRedisRunningIncrementIfPresent(server) - await redis.set(GLOBAL_PROXY_SPEND_CACHE_KEY, 10.0) - server.connected = redis_up - cache: Final = UserApiKeyCache(in_memory_cache=InMemoryCache(), redis_cache=_RedisCacheOnClient(redis)) - request_span: Final = provider.get_tracer("test").start_span("request") - - await ps.update_cache( - token=None, - user_id=None, - end_user_id=None, - team_id=None, - response_cost=2.0, - parent_otel_span=request_span, - cache=cache, - ) - await _drain_spend_cache_writes() - request_span.end() - - redis_spans: Final = tuple( - (span.attributes.get("call_type"), span.status.status_code, span.parent.span_id) - for span in exporter.get_finished_spans() - if span.name == "redis" and span.attributes is not None and span.parent is not None - ) - assert redis_spans == ( - ("async_increment_if_present", expected_status, request_span.get_span_context().span_id), - ) - - class _UnwritableKeyUserApiKeyCache(UserApiKeyCache): def __init__(self, unwritable_key: str) -> None: super().__init__(in_memory_cache=InMemoryCache()) diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 91c87c63067..75ab0ae14f9 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -17,6 +17,7 @@ from pydantic import ValidationError as PydanticValidationError from starlette.datastructures import Headers import litellm +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from tests.unit.proxy.router_entry_point_stubs import route_entry_points_to_stubs from litellm.constants import ( CALLER_SESSION_ID_METADATA_KEY, @@ -10219,7 +10220,6 @@ class TestResolveUserProviderCredentials: ) def _env(self, monkeypatch): - from litellm.caching.dual_cache import DualCache monkeypatch.setattr( litellm, @@ -10239,7 +10239,7 @@ class TestResolveUserProviderCredentials: prisma_client.db.litellm_userprovidercredentials = table prisma_client.writer_db = prisma_client.db monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client) - monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", DualCache()) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache()) return table @pytest.mark.asyncio @@ -11057,7 +11057,6 @@ class TestResolveUserProviderCredentials: def _per_user_env(monkeypatch, credential_names=("copilot-cred",)): - from litellm.caching.dual_cache import DualCache monkeypatch.setattr( litellm, @@ -11077,7 +11076,7 @@ def _per_user_env(monkeypatch, credential_names=("copilot-cred",)): prisma_client.db.litellm_userprovidercredentials = table prisma_client.writer_db = prisma_client.db monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client) - monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", DualCache()) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache()) monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "sk-test-master") return table @@ -11221,7 +11220,6 @@ async def test_header_mapped_user_id_cannot_load_another_users_connection(monkey request header. Per-user token resolution is bound to the key-authenticated identity, so a header naming user B still loads only user A's connections.""" from litellm import Router - from litellm.caching.dual_cache import DualCache from litellm.proxy.credential_endpoints import user_provider_credentials as upc from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request @@ -11242,7 +11240,7 @@ async def test_header_mapped_user_id_cannot_load_another_users_connection(monkey ] ) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) - monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", DualCache()) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache()) request_mock = MagicMock(spec=Request) request_mock.url = MagicMock() @@ -11277,7 +11275,6 @@ async def test_header_mapped_user_id_cannot_load_another_users_connection(monkey async def test_per_user_token_resolution_uses_key_identity_without_header_mapping(monkeypatch): """Control: with no user_header_mappings configured the key's user_id is used as before.""" from litellm import Router - from litellm.caching.dual_cache import DualCache from litellm.proxy.credential_endpoints import user_provider_credentials as upc from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request @@ -11298,7 +11295,7 @@ async def test_per_user_token_resolution_uses_key_identity_without_header_mappin ] ) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) - monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", DualCache()) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache()) request_mock = MagicMock(spec=Request) request_mock.url = MagicMock() diff --git a/tests/unit/proxy/test_proxy_server.py b/tests/unit/proxy/test_proxy_server.py index dae9ffcf7fd..0bf7d199205 100644 --- a/tests/unit/proxy/test_proxy_server.py +++ b/tests/unit/proxy/test_proxy_server.py @@ -1392,29 +1392,23 @@ async def test_team_update_redis(): """ Tests if team update, updates the redis cache if set """ - from litellm.caching.caching import DualCache, RedisCache from litellm.proxy._types import LiteLLM_TeamTableCachedObj from litellm.proxy.auth.auth_checks import cache_team_object + from tests.unit.redis.stores.auth_cache_double import shared_cache proxy_logging_obj: ProxyLogging = getattr( litellm.proxy.proxy_server, "proxy_logging_obj" ) + cache, redis = shared_cache() - redis_cache = RedisCache(host="localhost") + await cache_team_object( + team_id="1234", + team_table=LiteLLM_TeamTableCachedObj(team_id="1234"), + user_api_key_cache=cache, + proxy_logging_obj=proxy_logging_obj, + ) - with patch.object( - redis_cache, - "async_set_cache", - new=AsyncMock(), - ) as mock_client: - await cache_team_object( - team_id="1234", - team_table=LiteLLM_TeamTableCachedObj(team_id="1234"), - user_api_key_cache=DualCache(redis_cache=redis_cache), - proxy_logging_obj=proxy_logging_obj, - ) - - mock_client.assert_called() + assert redis.redis["team_id:1234"]["team_id"] == "1234" @pytest.mark.asyncio diff --git a/tests/unit/proxy/test_redis_auth_cache_flag.py b/tests/unit/proxy/test_redis_auth_cache_flag.py index f0af8dc9220..3cac31c5f54 100644 --- a/tests/unit/proxy/test_redis_auth_cache_flag.py +++ b/tests/unit/proxy/test_redis_auth_cache_flag.py @@ -5,8 +5,8 @@ Verifies that _init_cache attaches Redis to user_api_key_cache only when the flag is explicitly set to True, and leaves it in-memory-only otherwise. """ -from contextlib import ExitStack, contextmanager import json +from contextlib import ExitStack, contextmanager from unittest.mock import MagicMock, patch import pytest @@ -15,7 +15,10 @@ import litellm import litellm.proxy.proxy_server as ps from litellm.caching.caching import RedisCache from litellm.caching.dual_cache import DualCache +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.redis.native_settings import NativeRedisSettings from litellm.redis.stores.spend_counters import SpendCounterStore +from tests.unit.redis.stores.auth_cache_double import FakeAuthCache from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters # --------------------------------------------------------------------------- @@ -33,6 +36,9 @@ class _FakeRedisCache(RedisCache): def __init__(self): # super().__init__ skipped intentionally self._store = {} self.default_ttl = 60 + self.native_settings = NativeRedisSettings( + redis_kwargs={"host": "127.0.0.1", "port": 1, "socket_timeout": 0.2}, namespace=None, default_ttl=60 + ) def set_cache(self, key, value, **kwargs): # type: ignore[override] # Enforce Redis JSON-serializable payload contract. @@ -51,7 +57,7 @@ def _spend_counter_store() -> SpendCounterStore: @contextmanager -def _patched_init_cache(litellm_settings: dict, cache_params: dict): +def _patched_init_cache(litellm_settings: dict, cache_params: dict, user_cache: DualCache | None = None): """ Context manager that: 1. Replaces the module-level globals with fresh DualCache instances. @@ -68,7 +74,7 @@ def _patched_init_cache(litellm_settings: dict, cache_params: dict): mock_litellm_cache = MagicMock() mock_litellm_cache.cache = fake_redis - fresh_user_cache = DualCache() + fresh_user_cache = DualCache() if user_cache is None else user_cache fresh_spend_store = _spend_counter_store() fresh_cli_sso_cache = DualCache() fresh_config_cache = DualCache() @@ -97,6 +103,42 @@ def _patched_init_cache(litellm_settings: dict, cache_params: dict): class TestRedisAuthCacheFlag: + @pytest.mark.asyncio + async def test_flag_true_moves_the_whole_auth_keyspace_onto_the_native_store(self): + store = FakeAuthCache() + with _patched_init_cache( + litellm_settings={"enable_redis_auth_cache": True}, + cache_params={"type": "redis", "host": "localhost", "port": 6379}, + user_cache=UserApiKeyCache( + default_in_memory_ttl=60, auth_cache_handle=store.factory, redis_client=lambda _: None + ), + ) as (user_cache, _, _cli_sso_cache): + assert isinstance(user_cache, UserApiKeyCache) + await user_cache.async_set_cache("team_id:t1", {"team_id": "t1"}) + assert await user_cache.async_get_cache("team_id:t1") == {"team_id": "t1"} + + assert user_cache.redis_cache is None, "redis-py must never read or write the auth keyspace" + assert isinstance(user_cache.redis_connection, _FakeRedisCache) + assert store.redis == {"team_id:t1": {"team_id": "t1"}} + assert [settings.memory_ttl for settings in store.settings] == [60] + + @pytest.mark.asyncio + async def test_flag_false_keeps_the_auth_cache_in_python_memory_only(self): + store = FakeAuthCache() + with _patched_init_cache( + litellm_settings={"enable_redis_auth_cache": False}, + cache_params={"type": "redis", "host": "localhost", "port": 6379}, + user_cache=UserApiKeyCache( + default_in_memory_ttl=60, auth_cache_handle=store.factory, redis_client=lambda _: None + ), + ) as (user_cache, _, _cli_sso_cache): + assert isinstance(user_cache, UserApiKeyCache) + await user_cache.async_set_cache("team_id:t1", {"team_id": "t1"}) + + assert not user_cache.has_redis and user_cache.redis_connection is None + assert user_cache.in_memory_cache.get_cache("team_id:t1") == {"team_id": "t1"} + assert store.settings == [] and store.redis == {} + def test_flag_true_attaches_redis_to_user_api_key_cache(self): """When enable_redis_auth_cache=True, user_api_key_cache.redis_cache must be set.""" with _patched_init_cache( diff --git a/tests/unit/redis/stores/auth_cache_double.py b/tests/unit/redis/stores/auth_cache_double.py new file mode 100644 index 00000000000..2f9eb2a7546 --- /dev/null +++ b/tests/unit/redis/stores/auth_cache_double.py @@ -0,0 +1,170 @@ +"""A stand-in for ``NativeAuthCache`` with the native store's semantics: memory in front of a Redis dict, every +direct round trip logged in ``commands`` and every pre-call declaration logged on the request's ``events``.""" + +from __future__ import annotations + +from collections.abc import Callable, Sequence +from dataclasses import dataclass, field +from typing import Final, TypeVar +from unittest.mock import MagicMock + +from litellm.caching.redis_cache import RedisCache +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.redis.stores.auth_cache import AuthCacheSettings, WriteBack +from litellm.rust_bridge._native import NativeRedisClient + +T = TypeVar("T") + + +@dataclass +class DeclaredOp: + """Settles only once its request's pre-call batch was flushed, like a native declaration.""" + + request: object + answer: Callable[[], object] + + async def resolve(self) -> object: + 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.answer() + + +@dataclass +class FakeAuthCache: + """``redis`` holds what each key decodes to; ``fail`` makes every direct Redis round trip raise.""" + + redis: dict[str, object] = field(default_factory=dict) + memory: dict[str, tuple[object, float | None]] = field(default_factory=dict) + commands: list[tuple[object, ...]] = field(default_factory=list) + settings: list[AuthCacheSettings] = field(default_factory=list) + fail: BaseException | None = None + + def _round_trip(self, *command: object) -> None: + self.commands.append(command) + if self.fail is not None: + raise self.fail + + @staticmethod + def _declare(request: object, event: str) -> None: + events: Final = getattr(request, "events", None) + if events is not None: + events.append(event) + + async def get(self, key: str) -> object: + return self.get_sync(key) + + def get_sync(self, key: str) -> object: + cached: Final = self.get_local(key) + if cached is not None: + return cached + self._round_trip("GET", key) + value: Final = self.redis.get(key) + if value is not None: + self.set_local(key, value) + return value + + async def set(self, key: str, value: object, ttl_seconds: float | None = None) -> None: + self.set_local(key, value, ttl_seconds) + self._round_trip("SET", key, ttl_seconds) + self.redis[key] = value + + def set_sync(self, key: str, value: object, ttl_seconds: float | None = None) -> None: + self.set_local(key, value, ttl_seconds) + self._round_trip("SET_STR", key, ttl_seconds) + self.redis[key] = value + + def set_pre_call(self, request: object, key: str, value: object, ttl_seconds: float | None = None) -> DeclaredOp: + self.set_local(key, value, ttl_seconds) + self._declare(request, f"declare set {key} ttl={ttl_seconds}") + self.redis[key] = value + return DeclaredOp(request, lambda: None) + + def delete_pre_call(self, request: object, key: str) -> DeclaredOp: + self.delete_local(key) + self._declare(request, f"declare del {key}") + self.redis.pop(key, None) + return DeclaredOp(request, lambda: None) + + async def set_many(self, entries: Sequence[tuple[str, object]], ttl_seconds: float | None = None) -> None: + for key, value in entries: + self.set_local(key, value, ttl_seconds) + self._round_trip("SET_MANY", tuple(key for key, _ in entries), ttl_seconds) + self.redis.update(entries) + + async def write_back(self, request: object | None, entries: Sequence[WriteBack]) -> None: + for key, value, ttl in entries: + self.set_local(key, value, ttl) + self.redis[key] = value + written: Final = tuple((key, ttl) for key, _value, ttl in entries) + if request is None: + self._round_trip("WRITE_BACK", written) + return + self._declare(request, f"declare write_back {written}") + + async def delete(self, keys: Sequence[str]) -> None: + self.delete_sync(keys) + + def delete_sync(self, keys: Sequence[str]) -> None: + for key in keys: + self.delete_local(key) + self._round_trip("DEL", tuple(keys)) + for key in keys: + self.redis.pop(key, None) + + def declare_get_many(self, request: object, keys: Sequence[str]) -> DeclaredOp: + self._declare(request, f"declare auth_cache {','.join(keys)}") + + def answer() -> list[object]: + if self.fail is not None: + raise self.fail + return [self.redis.get(key) for key in keys] + + return DeclaredOp(request, answer) + + async def get_many(self, keys: Sequence[str]) -> list[object]: + self._round_trip("MGET", tuple(keys)) + return [self.redis.get(key) for key in keys] + + async def increment_if_present(self, key: str, amount: float) -> float | None: + self._round_trip("INCRBYFLOAT_IF_PRESENT", key, amount) + current: Final = self.redis.get(key) + if not isinstance(current, (int, float)): + return None + total: Final = float(current) + amount + self.redis[key] = total + return total + + def get_local(self, key: str) -> object: + entry: Final = self.memory.get(key) + return None if entry is None else entry[0] + + def set_local(self, key: str, value: object, ttl_seconds: float | None = None) -> None: + self.memory[key] = (value, ttl_seconds) + + def delete_local(self, key: str) -> None: + self.memory.pop(key, None) + + def flush_local(self) -> None: + self.memory.clear() + + def len_local(self) -> tuple[int, int]: + return (len(self.memory), 0) + + def factory(self, client: NativeRedisClient | None, settings: AuthCacheSettings, redis_ttl: int) -> FakeAuthCache: + self.settings.append(settings) + return self + + +def shared_cache( + fake: FakeAuthCache | None = None, connection: RedisCache | None = None, **kwargs: float | None +) -> tuple[UserApiKeyCache, FakeAuthCache]: + """A ``UserApiKeyCache`` with ``enable_redis_auth_cache`` on, its native store replaced by ``fake``.""" + double: Final = fake or FakeAuthCache() + cache: Final = UserApiKeyCache( + redis_cache=connection or MagicMock(spec=RedisCache), + auth_cache_handle=double.factory, + redis_client=lambda _: None, + **kwargs, + ) + return cache, double diff --git a/tests/unit/redis/stores/test_auth_cache.py b/tests/unit/redis/stores/test_auth_cache.py new file mode 100644 index 00000000000..e11a05b9561 --- /dev/null +++ b/tests/unit/redis/stores/test_auth_cache.py @@ -0,0 +1,165 @@ +import datetime +import json +from typing import Final +from unittest.mock import MagicMock + +import pytest + +from litellm.caching.redis_batch import request_redis_batch_scope +from litellm.caching.redis_cache import RedisCache +from litellm.proxy.common_utils.user_api_key_cache import is_user_key_cache_key +from litellm.redis.stores.auth_cache import AuthCacheSettings, AuthCacheStore +from litellm.rust_bridge._native import NativeAuthCache +from tests.unit.caching.test_redis_batch import FakeNatives + +from .auth_cache_double import FakeAuthCache + +_KEY_OBJECT: Final = "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b" +_TEAM_RECORD: Final = { + "team_id": "team-é", + "models": ["gpt-5", "claude-opus-5"], + "max_budget": None, + "spend": 0.30000000000000004, + "blocked": False, + "tpm_limit": 10000000000, + "metadata": {"tags": [], "quoted": 'it\'s "x"'}, +} +_TEAM_JSON: Final = ( + r'{"team_id": "team-\u00e9", "models": ["gpt-5", "claude-opus-5"], "max_budget": null, ' + r'"spend": 0.30000000000000004, "blocked": false, "tpm_limit": 10000000000, ' + r""""metadata": {"tags": [], "quoted": "it's \"x\""}}""" +) +_TEAM_STR: Final = ( + "{'team_id': 'team-é', 'models': ['gpt-5', 'claude-opus-5'], 'max_budget': None, " + "'spend': 0.30000000000000004, 'blocked': False, 'tpm_limit': 10000000000, " + """'metadata': {'tags': [], 'quoted': 'it\\'s "x"'}}""" +) +_WRITTEN: Final[tuple[tuple[object, str, str], ...]] = ( + (_TEAM_RECORD, _TEAM_JSON, _TEAM_STR), + ( + {"token": _KEY_OBJECT, "spend": 1e-05, "rpm_limit": 5, "expires": "2026-11-03T00:00:00Z"}, + '{"token": "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b", "spend": 1e-05, ' + '"rpm_limit": 5, "expires": "2026-11-03T00:00:00Z"}', + "{'token': '88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b', 'spend': 1e-05, " + "'rpm_limit': 5, 'expires': '2026-11-03T00:00:00Z'}", + ), + ("__no_team_membership__", '"__no_team_membership__"', "__no_team_membership__"), + (["mcp-a", "mcp-b"], '["mcp-a", "mcp-b"]', "['mcp-a', 'mcp-b']"), + (12.5, "12.5", "12.5"), + (7, "7", "7"), + (True, "true", "True"), +) + + +@pytest.mark.parametrize( + ("key", "key_object"), + [ + (_KEY_OBJECT, True), + ("88DC28D0F030C55ED4AB77ED8FAF098196CB1C05DF778539800C9F1243FE6B4B", False), + ("88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4", False), + ("team_id:88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b", False), + ("team_membership:user-1:team-1", False), + ("end_user_restricted_registry", False), + ], +) +def test_the_native_partition_rule_is_is_user_key_cache_key(key: str, key_object: bool) -> None: + assert is_user_key_cache_key(key) is key_object + + +@pytest.mark.parametrize(("value", "encoded", "printed"), _WRITTEN, ids=[encoded for _, encoded, _ in _WRITTEN]) +def test_writes_are_what_redis_cache_writes(value: object, encoded: str, printed: str) -> None: + assert json.dumps(value) == encoded + assert str(value) == printed + + +@pytest.mark.parametrize( + ("stored", "value"), + [ + (_TEAM_JSON, _TEAM_RECORD), + (_TEAM_STR, _TEAM_RECORD), + ("['mcp-a', 'mcp-b']", ["mcp-a", "mcp-b"]), + ( + "{'user_id': 'u1', 'teams': ['t1'], 'active': True, 'budget': None}", + {"user_id": "u1", "teams": ["t1"], "active": True, "budget": None}, + ), + ('"__NO_MAPPING__"', "__NO_MAPPING__"), + ], + ids=str, +) +def test_reads_are_what_redis_cache_reads(stored: str, value: object) -> None: + assert RedisCache._get_cache_logic(MagicMock(spec=RedisCache), stored.encode()) == value # pyright: ignore[reportPrivateUsage] # main's own decoder is the contract + + +def test_plain_values_cross_the_bridge_unchanged(value: object) -> None: + store: Final = NativeAuthCache(None) + + store.set_local("team_id:t1", value) + + assert store.get_local("team_id:t1") == value + assert type(store.get_local("team_id:t1")) is type(value) + + +def test_a_live_object_is_refused_rather_than_cached_as_something_else() -> None: + store: Final = NativeAuthCache(None) + + with pytest.raises(TypeError): + store.set_local("team_id:t1", {"created_at": datetime.datetime(2026, 1, 1)}) + + assert store.get_local("team_id:t1") is None + + +def test_user_key_objects_live_in_their_own_memory_partition() -> None: + store: Final = NativeAuthCache(None, memory_entries=2) + store.set_local(_KEY_OBJECT, {"token": _KEY_OBJECT}) + + for index in range(5): + store.set_local(f"team_id:t{index}", {"team_id": index}) + + assert store.get_local(_KEY_OBJECT) == {"token": _KEY_OBJECT} + assert store.len_local() == (2, 1) + + +@pytest.mark.asyncio +async def test_a_key_the_request_read_found_absent_is_answered_without_another_read_until_written() -> None: + fake: Final = FakeAuthCache(redis={"team_id:hit": {"team_id": "hit"}}) + store: Final = AuthCacheStore(None, AuthCacheSettings(200, 60), redis_ttl=60, handle=fake.factory) + + with request_redis_batch_scope(native_factory=FakeNatives()) as request: # pyright: ignore[reportArgumentType] # test double + found = await store.read_missing(["team_id:absent", "team_id:hit"]) + absent = await store.get("team_id:absent") + await store.set("team_id:absent", {"team_id": "now"}, None) + store.delete_local("team_id:absent") + refilled = await store.get("team_id:absent") + store.delete_pre_call("team_id:hit") + deleted = await store.get("team_id:hit") + await request.flush_all() + + assert found == {"team_id:absent": None, "team_id:hit": {"team_id": "hit"}} + assert absent is None + assert refilled == {"team_id": "now"} + assert deleted is None + assert [command[:2] for command in fake.commands] == [("SET", "team_id:absent"), ("GET", "team_id:absent")] + + +@pytest.mark.asyncio +async def test_outside_a_request_a_read_is_a_round_trip_of_its_own() -> None: + fake: Final = FakeAuthCache(redis={"team_id:t1": {"team_id": "t1"}}) + store: Final = AuthCacheStore(None, AuthCacheSettings(200, 60), redis_ttl=60, handle=fake.factory) + + assert await store.read_missing(["team_id:t1", "team_id:t2"]) == { + "team_id:t1": {"team_id": "t1"}, + "team_id:t2": None, + } + assert store.set_pre_call("team_id:t3", {}, None) is None + assert fake.commands == [("MGET", ("team_id:t1", "team_id:t2"))] + + +def test_only_a_changed_memory_configuration_rebuilds_the_native_store() -> None: + fake: Final = FakeAuthCache() + store: Final = AuthCacheStore(None, AuthCacheSettings(200, 60), redis_ttl=60, handle=fake.factory) + + store.reconfigure(AuthCacheSettings(200, 60)) + store.reconfigure(AuthCacheSettings(500, 5)) + + assert fake.settings == [AuthCacheSettings(200, 60), AuthCacheSettings(500, 5)] + assert store.settings == AuthCacheSettings(500, 5)