From 68fa7e5aaa609bafcfd13677c0f8d0a6632beb03 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 3 Nov 2026 17:05:11 +0900 Subject: [PATCH] feat(redis): move the auth cache keyspace onto the Rust store With enable_redis_auth_cache on, everything user_api_key_cache keeps (user-key objects under bare hashed tokens, team/user/alias/org/end-user/project and the other management objects, the global proxy spend scalar) now lives in litellm-redis-store's auth_cache module: both memory partitions plus Redis, behind litellm/redis/stores/auth_cache.py. redis-py never reads or writes the keyspace and UserApiKeyCache.redis_cache stays None; features that only borrow the connection (MCP refresh lock, user provider credential cache) get it from redis_connection. With the flag off the cache stays Python memory as before. Values cross pyo3 as plain data and are written as main's json.dumps (str() for the sync set_cache path), read back with json.loads then literal_eval, with main's TTL rules, pinned by matching inline cases in the Rust and Python tests. The identity and owner prefetches read on the request's native pre-call batch and arm the cooldown read right before, so both leave in one pipeline; write-backs and pre-call SET/DEL ride the same batch. The invalidation publisher and subscriber stay in Python and forward evictions to the native memory tier. The auth cache sends the commands UserApiKeyCache does. Prefetch reads are one MGET, set_many writes the key-object partition and then the rest as two pipelines, delete runs both partitions as concurrent multi-key DELs in chunks of 1000, and increment_if_present is a plain EVAL with the amount written the way redis-py encodes it. --- litellm-rust/crates/python-bridge/src/lib.rs | 8 + .../src/redis/store/auth_cache.rs | 338 +++++++++++ .../python-bridge/src/redis/store/mod.rs | 1 + .../crates/redis-store/src/auth_cache.rs | 463 +++++++++++++++ litellm-rust/crates/redis-store/src/lib.rs | 1 + .../crates/redis-store/tests/auth_cache.rs | 552 ++++++++++++++++++ litellm/caching/dual_cache.py | 4 +- litellm/caching/redis_batch.py | 12 + litellm/proxy/auth/auth_checks.py | 2 +- litellm/proxy/auth/auth_object_prefetch.py | 111 ++-- litellm/proxy/auth/user_api_key_auth.py | 20 +- .../auth_cache_invalidation_pubsub.py | 2 +- .../proxy/common_utils/user_api_key_cache.py | 286 +++++++-- .../user_provider_credentials.py | 18 +- .../per_user_oauth_store.py | 2 +- .../runtime_refresh_coordinator.py | 2 +- litellm/proxy/proxy_server.py | 18 +- litellm/redis/stores/auth_cache.py | 192 ++++++ litellm/router_utils/routing_read_batch.py | 7 +- litellm/rust_bridge/_native.pyi | 49 ++ .../test_request_redis_batch_pre_call.py | 184 +++--- tests/unit/proxy/auth/test_auth_checks.py | 20 +- ...st_auth_checks_object_access_and_lookup.py | 31 +- .../proxy/auth/test_auth_object_prefetch.py | 142 ++--- tests/unit/proxy/auth/test_handle_jwt.py | 2 +- .../test_auth_cache_invalidation_pubsub.py | 24 + .../common_utils/test_user_api_key_cache.py | 326 +++++------ .../credential_endpoints/test_endpoints.py | 18 +- .../test_user_provider_credentials.py | 45 +- .../test_team_endpoints.py | 73 +-- .../test_team_roster_sync.py | 71 +-- .../proxy/proxy_server/test_spend_counters.py | 133 +---- .../unit/proxy/test_litellm_pre_call_utils.py | 13 +- tests/unit/proxy/test_proxy_server.py | 24 +- .../unit/proxy/test_redis_auth_cache_flag.py | 48 +- tests/unit/redis/stores/auth_cache_double.py | 170 ++++++ tests/unit/redis/stores/test_auth_cache.py | 165 ++++++ 37 files changed, 2770 insertions(+), 807 deletions(-) create mode 100644 litellm-rust/crates/python-bridge/src/redis/store/auth_cache.rs create mode 100644 litellm-rust/crates/redis-store/src/auth_cache.rs create mode 100644 litellm-rust/crates/redis-store/tests/auth_cache.rs create mode 100644 litellm/redis/stores/auth_cache.py create mode 100644 tests/unit/redis/stores/auth_cache_double.py create mode 100644 tests/unit/redis/stores/test_auth_cache.py 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)