From 6edbd27f74734c77dd77ea49f31a5b23e29b32ca Mon Sep 17 00:00:00 2001 From: yujonglee Date: Mon, 26 Oct 2026 18:41:46 -0700 Subject: [PATCH] feat(rust): add redis-client crate for standalone Redis (#48269) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 48 ++- litellm-rust/Cargo.toml | 5 + .../crates/cache-redis-semantic/Cargo.toml | 4 +- litellm-rust/crates/cache-redis/Cargo.toml | 4 +- litellm-rust/crates/cache-response/Cargo.toml | 4 +- .../crates/cache-valkey-semantic/Cargo.toml | 4 +- litellm-rust/crates/db-testing/Cargo.toml | 2 +- litellm-rust/crates/python-bridge/Cargo.toml | 2 +- litellm-rust/crates/redis-client/AGENTS.md | 26 ++ litellm-rust/crates/redis-client/Cargo.toml | 15 + .../crates/redis-client/src/client.rs | 404 ++++++++++++++++++ .../crates/redis-client/src/config.rs | 60 +++ litellm-rust/crates/redis-client/src/error.rs | 67 +++ litellm-rust/crates/redis-client/src/lib.rs | 9 + .../crates/redis-client/src/namespace.rs | 53 +++ litellm-rust/crates/redis-testing/AGENTS.md | 5 + litellm-rust/crates/redis-testing/Cargo.toml | 11 + .../crates/redis-testing/src/error.rs | 11 + litellm-rust/crates/redis-testing/src/lib.rs | 6 + .../crates/redis-testing/src/ports.rs | 66 +++ .../crates/redis-testing/src/standalone.rs | 55 +++ .../crates/spend-clickhouse/Cargo.toml | 2 +- .../crates/storage-clickhouse/Cargo.toml | 2 +- ...test_proxy_server_endpoints_and_startup.py | 1 - 24 files changed, 851 insertions(+), 15 deletions(-) create mode 100644 litellm-rust/crates/redis-client/AGENTS.md create mode 100644 litellm-rust/crates/redis-client/Cargo.toml create mode 100644 litellm-rust/crates/redis-client/src/client.rs create mode 100644 litellm-rust/crates/redis-client/src/config.rs create mode 100644 litellm-rust/crates/redis-client/src/error.rs create mode 100644 litellm-rust/crates/redis-client/src/lib.rs create mode 100644 litellm-rust/crates/redis-client/src/namespace.rs create mode 100644 litellm-rust/crates/redis-testing/AGENTS.md create mode 100644 litellm-rust/crates/redis-testing/Cargo.toml create mode 100644 litellm-rust/crates/redis-testing/src/error.rs create mode 100644 litellm-rust/crates/redis-testing/src/lib.rs create mode 100644 litellm-rust/crates/redis-testing/src/ports.rs create mode 100644 litellm-rust/crates/redis-testing/src/standalone.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 32ae431f863..8b1be389343 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -876,6 +876,15 @@ dependencies = [ "time", ] +[[package]] +name = "backon" +version = "1.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cffb0e931875b666fc4fcb20fee52e9bbd1ef836fd9e9e04ec21555f9f85f7ef" +dependencies = [ + "fastrand", +] + [[package]] name = "base16ct" version = "0.2.0" @@ -1248,7 +1257,11 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cfc320937d09e6de266b31b9afb480f197d7a861be86be7cb2ea7e5d1bfffc5e" dependencies = [ "bytes", + "futures-core", "memchr", + "pin-project-lite", + "tokio", + "tokio-util", ] [[package]] @@ -4282,6 +4295,26 @@ dependencies = [ "thiserror 2.0.19", ] +[[package]] +name = "litellm-redis-client" +version = "0.1.0" +dependencies = [ + "litellm-redis-testing", + "redis", + "rstest", + "thiserror 2.0.19", + "tokio", +] + +[[package]] +name = "litellm-redis-testing" +version = "0.1.0" +dependencies = [ + "testcontainers-modules", + "thiserror 2.0.19", + "tokio", +] + [[package]] name = "litellm-router" version = "0.1.0" @@ -5757,22 +5790,33 @@ dependencies = [ [[package]] name = "redis" -version = "1.7.0" +version = "1.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2acbc41a996f7652b2ddd9dfd98cc4ff602cfd742ae35382f07f608405ab50ed" +checksum = "7acd809e80afa754c63f586955699d8990efe1c4b6613183ee5f027a11fbc8b5" dependencies = [ + "arc-swap", "arcstr", + "async-lock", + "backon", + "bytes", + "cfg-if", "combine", "crc16", + "futures-channel", + "futures-util", "itoa", "num-bigint 0.5.1", "percent-encoding", + "pin-project-lite", "rand 0.10.2", "rustls 0.23.42", "rustls-native-certs", "ryu", "sha1_smol", "socket2 0.6.5", + "tokio", + "tokio-rustls 0.26.4", + "tokio-util", "url", "xxhash-rust", ] diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index fc65355c047..8994fdd5b7d 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -69,6 +69,8 @@ litellm-token-counter-huggingface = { path = "crates/token-counter-huggingface" litellm-token-counter-tiktoken = { path = "crates/token-counter-tiktoken" } litellm-host-python = { path = "crates/host-python" } litellm-python-compat = { path = "crates/python-compat" } +litellm-redis-client = { path = "crates/redis-client" } +litellm-redis-testing = { path = "crates/redis-testing" } askama = { version = "0.16.1", default-features = false, features = ["derive", "std"] } tracing = "0.1" @@ -88,6 +90,9 @@ rand = "0.8" macro_rules_attribute = "0.2.3" schemars = "1" reqwest = { version = "0.12", default-features = false, features = ["json", "multipart", "rustls-tls", "http2", "stream"] } +redis = "1.8.0" +redis-test = "1.0.4" +testcontainers-modules = "0.15.0" qdrant-client = { version = "1.19.0", default-features = false } uuid = { version = "1", features = ["v4"] } rstest = "0.26.1" diff --git a/litellm-rust/crates/cache-redis-semantic/Cargo.toml b/litellm-rust/crates/cache-redis-semantic/Cargo.toml index fe5a317cec0..6a5c794bde7 100644 --- a/litellm-rust/crates/cache-redis-semantic/Cargo.toml +++ b/litellm-rust/crates/cache-redis-semantic/Cargo.toml @@ -8,12 +8,12 @@ repository.workspace = true [dependencies] litellm-cache.workspace = true litellm-cache-redis.workspace = true -redis = { version = "1.7.0", features = ["tls-rustls"] } +redis = { workspace = true, features = ["tls-rustls"] } sha2.workspace = true [dev-dependencies] litellm-cache-testing.workspace = true -redis-test = "1.0.4" +redis-test.workspace = true rstest.workspace = true serde_json.workspace = true tokio.workspace = true diff --git a/litellm-rust/crates/cache-redis/Cargo.toml b/litellm-rust/crates/cache-redis/Cargo.toml index a234286a338..bebfcafd98a 100644 --- a/litellm-rust/crates/cache-redis/Cargo.toml +++ b/litellm-rust/crates/cache-redis/Cargo.toml @@ -7,12 +7,12 @@ repository.workspace = true [dependencies] litellm-cache.workspace = true -redis = { version = "1.7.0", features = ["cluster", "tls-rustls"] } +redis = { workspace = true, features = ["cluster", "tls-rustls"] } r2d2 = "0.8.10" tokio.workspace = true [dev-dependencies] litellm-cache-testing.workspace = true -redis-test = "1.0.4" +redis-test.workspace = true rstest.workspace = true serde_json.workspace = true diff --git a/litellm-rust/crates/cache-response/Cargo.toml b/litellm-rust/crates/cache-response/Cargo.toml index 869c40a12ab..8f4570edc6f 100644 --- a/litellm-rust/crates/cache-response/Cargo.toml +++ b/litellm-rust/crates/cache-response/Cargo.toml @@ -17,8 +17,8 @@ litellm-cache-gcs.workspace = true litellm-http = { workspace = true, features = ["test-support"] } litellm-cache-memory.workspace = true litellm-cache-redis.workspace = true -redis = "1.7.0" -redis-test = "1.0.4" +redis.workspace = true +redis-test.workspace = true rstest.workspace = true tokio.workspace = true wiremock.workspace = true diff --git a/litellm-rust/crates/cache-valkey-semantic/Cargo.toml b/litellm-rust/crates/cache-valkey-semantic/Cargo.toml index da9bb60c853..ec755e61b8c 100644 --- a/litellm-rust/crates/cache-valkey-semantic/Cargo.toml +++ b/litellm-rust/crates/cache-valkey-semantic/Cargo.toml @@ -8,13 +8,13 @@ repository.workspace = true [dependencies] litellm-cache.workspace = true litellm-cache-redis.workspace = true -redis = { version = "1.7.0", features = ["tls-rustls"] } +redis = { workspace = true, features = ["tls-rustls"] } sha2.workspace = true uuid = { version = "1", features = ["v4"] } [dev-dependencies] litellm-cache-testing.workspace = true -redis-test = "1.0.4" +redis-test.workspace = true rstest.workspace = true serde_json.workspace = true tokio.workspace = true diff --git a/litellm-rust/crates/db-testing/Cargo.toml b/litellm-rust/crates/db-testing/Cargo.toml index 3ba71fbf644..457c0f2ad87 100644 --- a/litellm-rust/crates/db-testing/Cargo.toml +++ b/litellm-rust/crates/db-testing/Cargo.toml @@ -7,7 +7,7 @@ repository.workspace = true [dependencies] sqlx = { workspace = true, features = ["migrate"] } -testcontainers-modules = { version = "0.15.0", features = ["postgres"] } +testcontainers-modules = { workspace = true, features = ["postgres"] } thiserror.workspace = true tokio.workspace = true diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 6fc05c43146..8afde8362f6 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -53,7 +53,7 @@ pyo3.workspace = true prost.workspace = true pyo3-async-runtimes.workspace = true reqwest.workspace = true -redis = { version = "1.7.0", features = ["tls-rustls"] } +redis = { workspace = true, features = ["tls-rustls"] } serde_json.workspace = true strum.workspace = true veil.workspace = true diff --git a/litellm-rust/crates/redis-client/AGENTS.md b/litellm-rust/crates/redis-client/AGENTS.md new file mode 100644 index 00000000000..4625eafc780 --- /dev/null +++ b/litellm-rust/crates/redis-client/AGENTS.md @@ -0,0 +1,26 @@ +- The one async Redis client for the gateway and the Rust cache crates. It owns connecting, timeouts, error classes and the key namespace +- It knows nothing about domains, so key builders, TTL policy, value formats and fallback decisions stay in the crate that owns the state. Every write takes its TTL from the caller, and there is no default +- Async only. Build one `RedisClient` per Redis instance at the process boundary and clone it. Never add a pool, a sync API, `spawn_blocking`, or a client built inside a route +- Set every redis-rs timeout and retry explicitly. Its defaults (500ms response, 1s connect, 6 reconnect retries) are not the Python proxy's (5s socket timeout from `REDIS_SOCKET_TIMEOUT`) +- Resolving config from `config.yaml` / `REDIS_*` env the way `litellm/_redis.py` does belongs to the consumer that reads that config, not here. `RedisConfig` takes already-resolved values + +## Errors +- `Transport`, `Server`, `Decode`, each wrapping the redis-rs error as `#[source]`. `ConnectError` is separate because only `connect` returns it +- Classify in `From for Error` by `ErrorKind`. LOADING, CLUSTERDOWN, MASTERDOWN, TRYAGAIN and READONLY are typed server replies, but they mean the node can't serve right now, so they are `Transport`. Codes redis-rs doesn't type (WRONGTYPE) come back as `ErrorKind::Extension` and are `Server`. Each code gets its own rstest case built from a parsed RESP error +- A `Server` error never drops the connection +- `Transport` includes timeouts, where the command may still have applied. Callers of counters and locks treat it as an unknown outcome +- Never swallow an error or turn it into a miss. The caller picks the fallback +- Reads return `Option`: a missing key is `None`, and a value that doesn't decode to `T` (`"15.0"` as `i64`) is `Decode` + +## Namespace +- Applied once, here, to every key. A key that already starts with `{namespace}:` is left alone, matching `RedisCache.check_and_fix_namespace` +- The rule never changes, because Python and Rust pods share keys in the same Redis + +## Commands +- One typed method per command, built with `redis::cmd`. Empty batches (MGET, DEL) return without a round trip, since Redis rejects an empty MGET +- TTLs are `Duration` and are sent in milliseconds (`PX`, `PEXPIRE`, `PTTL`) + +## Tests +- Against real Redis from `litellm-redis-testing`, never a mocked wire or asserted command bytes +- Assert values and Redis state. Outages come from the container (`restart`, `freeze` / `thaw`), not `DEBUG`, which the official image disables +- Use a short client timeout in outage tests so they stay fast. Never sleep to wait for a state; poll with a deadline diff --git a/litellm-rust/crates/redis-client/Cargo.toml b/litellm-rust/crates/redis-client/Cargo.toml new file mode 100644 index 00000000000..ae9ae651456 --- /dev/null +++ b/litellm-rust/crates/redis-client/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "litellm-redis-client" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +redis = { workspace = true, features = ["connection-manager", "tokio-rustls-comp"] } +thiserror.workspace = true + +[dev-dependencies] +litellm-redis-testing.workspace = true +rstest.workspace = true +tokio.workspace = true diff --git a/litellm-rust/crates/redis-client/src/client.rs b/litellm-rust/crates/redis-client/src/client.rs new file mode 100644 index 00000000000..1bc86ff110a --- /dev/null +++ b/litellm-rust/crates/redis-client/src/client.rs @@ -0,0 +1,404 @@ +use std::time::Duration; + +use redis::{ + Cmd, FromRedisValue, ToRedisArgs, ToSingleRedisArg, + aio::{ConnectionManager, ConnectionManagerConfig}, +}; + +use crate::{ConnectError, Error, Namespace, RedisConfig}; + +const RECONNECT_RETRIES: usize = 3; + +#[derive(Clone)] +pub struct RedisClient { + connection: ConnectionManager, + namespace: Namespace, +} + +#[derive(Debug, PartialEq, Eq)] +pub enum Ttl { + Missing, + Persistent, + Expires(Duration), +} + +impl RedisClient { + pub async fn connect(config: RedisConfig) -> Result { + let client = redis::Client::open(config.endpoint.connection_info()?)?; + let manager_config = ConnectionManagerConfig::new() + .set_connection_timeout(Some(config.timeout)) + .set_response_timeout(Some(config.timeout)) + .set_number_of_retries(RECONNECT_RETRIES); + Ok(Self { + connection: client + .get_connection_manager_with_config(manager_config) + .await?, + namespace: config.namespace, + }) + } + + pub async fn get(&self, key: &str) -> Result, Error> { + self.query(redis::cmd("GET").arg(self.key(key))).await + } + + pub async fn mget(&self, keys: &[&str]) -> Result>, Error> { + if keys.is_empty() { + return Ok(Vec::new()); + } + self.query(redis::cmd("MGET").arg(self.keys(keys))).await + } + + pub async fn set( + &self, + key: &str, + value: impl ToSingleRedisArg, + ttl: Duration, + ) -> Result<(), Error> { + self.query( + redis::cmd("SET") + .arg(self.key(key)) + .arg(value) + .arg("PX") + .arg(millis(ttl)), + ) + .await + } + + pub async fn set_nx( + &self, + key: &str, + value: impl ToSingleRedisArg, + ttl: Duration, + ) -> Result { + let reply: Option = self + .query( + redis::cmd("SET") + .arg(self.key(key)) + .arg(value) + .arg("NX") + .arg("PX") + .arg(millis(ttl)), + ) + .await?; + Ok(reply.is_some()) + } + + pub async fn del(&self, keys: &[&str]) -> Result { + if keys.is_empty() { + return Ok(0); + } + self.query(redis::cmd("DEL").arg(self.keys(keys))).await + } + + pub async fn expire(&self, key: &str, ttl: Duration) -> Result { + self.query(redis::cmd("PEXPIRE").arg(self.key(key)).arg(millis(ttl))) + .await + } + + pub async fn ttl(&self, key: &str) -> Result { + let reply: i64 = self.query(redis::cmd("PTTL").arg(self.key(key))).await?; + Ok(match reply { + -2 => Ttl::Missing, + -1 => Ttl::Persistent, + millis => Ttl::Expires(Duration::from_millis(millis.unsigned_abs())), + }) + } + + pub async fn incr_by(&self, key: &str, delta: i64) -> Result { + self.query(redis::cmd("INCRBY").arg(self.key(key)).arg(delta)) + .await + } + + pub async fn incr_by_float(&self, key: &str, delta: f64) -> Result { + self.query(redis::cmd("INCRBYFLOAT").arg(self.key(key)).arg(delta)) + .await + } + + pub async fn rpush(&self, key: &str, values: impl ToRedisArgs) -> Result { + self.query(redis::cmd("RPUSH").arg(self.key(key)).arg(values)) + .await + } + + fn key(&self, key: &str) -> String { + self.namespace.apply(key) + } + + fn keys(&self, keys: &[&str]) -> Vec { + keys.iter().map(|key| self.key(key)).collect() + } + + async fn query(&self, command: &Cmd) -> Result { + Ok(command.query_async(&mut self.connection.clone()).await?) + } +} + +fn millis(duration: Duration) -> u64 { + u64::try_from(duration.as_millis()).unwrap_or(u64::MAX) +} + +#[cfg(test)] +mod tests { + use std::time::{Duration, Instant}; + + use litellm_redis_testing::RedisServer; + use rstest::rstest; + + use super::{RedisClient, Ttl}; + use crate::{DEFAULT_TIMEOUT, Endpoint, Error, HostEndpoint, Namespace, RedisConfig}; + + async fn client(server: &RedisServer, namespace: &str, timeout: Duration) -> RedisClient { + RedisClient::connect(RedisConfig { + endpoint: Endpoint::Url(server.url().to_owned()), + namespace: Namespace::new(namespace), + timeout, + }) + .await + .unwrap() + } + + async fn client_id(client: &RedisClient) -> i64 { + redis::cmd("CLIENT") + .arg("ID") + .query_async(&mut client.connection.clone()) + .await + .unwrap() + } + + #[rstest] + #[tokio::test] + async fn namespaces_isolate_keys_on_the_server() { + let server = RedisServer::start().await.unwrap(); + let a = client(&server, "a", DEFAULT_TIMEOUT).await; + let b = client(&server, "b", DEFAULT_TIMEOUT).await; + let bare = client(&server, "", DEFAULT_TIMEOUT).await; + + a.set("x", "1", Duration::from_secs(60)).await.unwrap(); + + assert_eq!(a.get::("x").await.unwrap().as_deref(), Some("1")); + assert_eq!(a.get::("a:x").await.unwrap().as_deref(), Some("1")); + assert_eq!( + bare.get::("a:x").await.unwrap().as_deref(), + Some("1") + ); + assert_eq!(b.get::("x").await.unwrap(), None); + } + + #[rstest] + #[tokio::test] + async fn a_rejected_command_keeps_the_connection() { + let server = RedisServer::start().await.unwrap(); + let client = client(&server, "", DEFAULT_TIMEOUT).await; + client.rpush("list", "item").await.unwrap(); + let before = client_id(&client).await; + + let rejected = client.incr_by("list", 1).await; + + assert!(matches!(rejected, Err(Error::Server(_))), "{rejected:?}"); + assert_eq!(client.incr_by("counter", 2).await.unwrap(), 2); + assert_eq!(client_id(&client).await, before); + } + + #[rstest] + #[tokio::test] + async fn a_reply_of_another_type_is_a_decode_error_and_a_missing_key_is_none() { + let server = RedisServer::start().await.unwrap(); + let client = client(&server, "", DEFAULT_TIMEOUT).await; + client + .set("float", "15.0", Duration::from_secs(60)) + .await + .unwrap(); + + let decoded = client.get::("float").await; + + assert!(matches!(decoded, Err(Error::Decode(_))), "{decoded:?}"); + assert_eq!(client.get::("missing").await.unwrap(), None); + } + + #[rstest] + #[tokio::test] + async fn a_timed_out_command_does_not_leak_its_reply_into_the_next_one() { + let server = RedisServer::start().await.unwrap(); + let client = client(&server, "", Duration::from_millis(500)).await; + client + .set("first", "1", Duration::from_secs(60)) + .await + .unwrap(); + client + .set("second", "2", Duration::from_secs(60)) + .await + .unwrap(); + + server.freeze().await.unwrap(); + let timed_out = client.get::("first").await; + server.thaw().await.unwrap(); + + assert!( + matches!(timed_out, Err(Error::Transport(_))), + "{timed_out:?}" + ); + assert_eq!( + client.get::("second").await.unwrap().as_deref(), + Some("2") + ); + } + + #[rstest] + #[tokio::test] + async fn the_same_client_recovers_after_redis_restarts() { + let server = RedisServer::start().await.unwrap(); + let client = client(&server, "", Duration::from_millis(500)).await; + + server.restart().await.unwrap(); + + let deadline = Instant::now() + Duration::from_secs(30); + let recovered = loop { + match client.incr_by("after_restart", 1).await { + Ok(value) => break Some(value), + Err(_) if Instant::now() < deadline => { + tokio::time::sleep(Duration::from_millis(200)).await + } + Err(_) => break None, + } + }; + assert_eq!(recovered, Some(1)); + } + + #[rstest] + #[tokio::test] + async fn empty_batches_succeed_without_asking_redis() { + let server = RedisServer::start().await.unwrap(); + let client = client(&server, "", DEFAULT_TIMEOUT).await; + + assert_eq!( + client.mget::(&[]).await.unwrap(), + Vec::>::new() + ); + assert_eq!(client.del(&[]).await.unwrap(), 0); + } + + #[rstest] + #[tokio::test] + async fn mget_returns_values_in_input_order_with_gaps() { + let server = RedisServer::start().await.unwrap(); + let client = client(&server, "ns", DEFAULT_TIMEOUT).await; + client.set("a", "1", Duration::from_secs(60)).await.unwrap(); + client.set("c", "3", Duration::from_secs(60)).await.unwrap(); + + let values = client.mget::(&["c", "b", "a"]).await.unwrap(); + + assert_eq!( + values, + vec![Some("3".to_owned()), None, Some("1".to_owned())] + ); + } + + #[rstest] + #[tokio::test] + async fn set_nx_only_writes_a_missing_key() { + let server = RedisServer::start().await.unwrap(); + let client = client(&server, "", DEFAULT_TIMEOUT).await; + + let first = client + .set_nx("lock", "owner-1", Duration::from_secs(60)) + .await + .unwrap(); + let second = client + .set_nx("lock", "owner-2", Duration::from_secs(60)) + .await + .unwrap(); + + assert!(first); + assert!(!second); + assert_eq!( + client.get::("lock").await.unwrap().as_deref(), + Some("owner-1") + ); + } + + #[rstest] + #[tokio::test] + async fn ttl_reports_missing_persistent_and_remaining_time() { + let server = RedisServer::start().await.unwrap(); + let client = client(&server, "", DEFAULT_TIMEOUT).await; + client.incr_by("persistent", 1).await.unwrap(); + client + .set("expiring", "1", Duration::from_secs(60)) + .await + .unwrap(); + client.incr_by_float("expired_later", 0.5).await.unwrap(); + assert!( + client + .expire("expired_later", Duration::from_secs(30)) + .await + .unwrap() + ); + + assert_eq!(client.ttl("missing").await.unwrap(), Ttl::Missing); + assert_eq!(client.ttl("persistent").await.unwrap(), Ttl::Persistent); + assert!(matches!( + client.ttl("expiring").await.unwrap(), + Ttl::Expires(left) if left > Duration::from_secs(50) && left <= Duration::from_secs(60) + )); + assert!(matches!( + client.ttl("expired_later").await.unwrap(), + Ttl::Expires(left) if left > Duration::from_secs(20) && left <= Duration::from_secs(30) + )); + assert!( + !client + .expire("missing", Duration::from_secs(30)) + .await + .unwrap() + ); + } + + #[rstest] + #[tokio::test] + async fn incr_by_float_accumulates_fractions() { + let server = RedisServer::start().await.unwrap(); + let client = client(&server, "", DEFAULT_TIMEOUT).await; + + client.incr_by_float("spend", 0.25).await.unwrap(); + + assert_eq!(client.incr_by_float("spend", 0.5).await.unwrap(), 0.75); + } + + #[rstest] + #[tokio::test] + async fn a_host_endpoint_connects_to_its_database() { + let server = RedisServer::start().await.unwrap(); + let port = server.url().rsplit(':').next().unwrap().parse().unwrap(); + let in_db_3 = RedisClient::connect(RedisConfig { + endpoint: Endpoint::Host(HostEndpoint { + host: "127.0.0.1".to_owned(), + port, + username: None, + password: None, + tls: false, + db: 3, + }), + namespace: Namespace::default(), + timeout: DEFAULT_TIMEOUT, + }) + .await + .unwrap(); + let in_db_0 = client(&server, "", DEFAULT_TIMEOUT).await; + + in_db_3 + .set("k", "3", Duration::from_secs(60)) + .await + .unwrap(); + + assert_eq!(in_db_0.get::("k").await.unwrap(), None); + let via_url = RedisClient::connect(RedisConfig { + endpoint: Endpoint::Url(format!("{}/3", server.url())), + namespace: Namespace::default(), + timeout: DEFAULT_TIMEOUT, + }) + .await + .unwrap(); + assert_eq!( + via_url.get::("k").await.unwrap().as_deref(), + Some("3") + ); + } +} diff --git a/litellm-rust/crates/redis-client/src/config.rs b/litellm-rust/crates/redis-client/src/config.rs new file mode 100644 index 00000000000..bb6e8a8eae6 --- /dev/null +++ b/litellm-rust/crates/redis-client/src/config.rs @@ -0,0 +1,60 @@ +use std::time::Duration; + +use redis::{ConnectionAddr, ConnectionInfo, IntoConnectionInfo, RedisConnectionInfo, RedisResult}; + +use crate::Namespace; + +pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(5); + +pub struct RedisConfig { + pub endpoint: Endpoint, + pub namespace: Namespace, + pub timeout: Duration, +} + +pub enum Endpoint { + Url(String), + Host(HostEndpoint), +} + +pub struct HostEndpoint { + pub host: String, + pub port: u16, + pub username: Option, + pub password: Option, + pub tls: bool, + pub db: i64, +} + +impl Endpoint { + pub(crate) fn connection_info(&self) -> RedisResult { + match self { + Self::Url(url) => url.as_str().into_connection_info(), + Self::Host(host) => host.connection_info(), + } + } +} + +impl HostEndpoint { + fn connection_info(&self) -> RedisResult { + let addr = match self.tls { + true => ConnectionAddr::TcpTls { + host: self.host.clone(), + port: self.port, + insecure: false, + tls_params: None, + }, + false => ConnectionAddr::Tcp(self.host.clone(), self.port), + }; + let settings = RedisConnectionInfo::default().set_db(self.db); + let settings = match &self.username { + Some(username) => settings.set_username(username), + None => settings, + }; + let settings = match &self.password { + Some(password) => settings.set_password(password), + None => settings, + }; + Ok(addr.into_connection_info()?.set_redis_settings(settings)) + } +} diff --git a/litellm-rust/crates/redis-client/src/error.rs b/litellm-rust/crates/redis-client/src/error.rs new file mode 100644 index 00000000000..61b9b2b89c0 --- /dev/null +++ b/litellm-rust/crates/redis-client/src/error.rs @@ -0,0 +1,67 @@ +use redis::{ErrorKind, RedisError, ServerErrorKind}; + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("Redis did not answer, so the command may or may not have applied")] + Transport(#[source] RedisError), + #[error("Redis rejected the command")] + Server(#[source] RedisError), + #[error("Redis replied with a value of a different type")] + Decode(#[source] RedisError), +} + +#[derive(Debug, thiserror::Error)] +#[error("connecting to Redis")] +pub struct ConnectError(#[from] RedisError); + +impl From for Error { + fn from(error: RedisError) -> Self { + match error.kind() { + ErrorKind::Server( + ServerErrorKind::BusyLoading + | ServerErrorKind::ClusterDown + | ServerErrorKind::MasterDown + | ServerErrorKind::TryAgain + | ServerErrorKind::ReadOnly, + ) => Self::Transport(error), + ErrorKind::Server(_) | ErrorKind::Extension => Self::Server(error), + ErrorKind::Parse | ErrorKind::UnexpectedReturnType => Self::Decode(error), + _ => Self::Transport(error), + } + } +} + +#[cfg(test)] +mod tests { + use redis::{RedisError, Value}; + use rstest::rstest; + + use super::Error; + + fn reply(resp: &str) -> RedisError { + match redis::parse_redis_value(resp.as_bytes()).unwrap() { + Value::ServerError(error) => error.into(), + other => panic!("{resp:?} parsed as {other:?}"), + } + } + + #[rstest] + #[case::wrong_type("-WRONGTYPE Operation against a key holding the wrong kind of value\r\n")] + #[case::no_script("-NOSCRIPT No matching script\r\n")] + #[case::cross_slot("-CROSSSLOT Keys in request don't hash to the same slot\r\n")] + #[case::generic("-ERR value is not an integer or out of range\r\n")] + #[case::no_perm("-NOPERM this user has no permissions to access one of the keys\r\n")] + fn rejections_are_server_errors(#[case] resp: &str) { + assert!(matches!(Error::from(reply(resp)), Error::Server(_))); + } + + #[rstest] + #[case::loading("-LOADING Redis is loading the dataset in memory\r\n")] + #[case::cluster_down("-CLUSTERDOWN The cluster is down\r\n")] + #[case::master_down("-MASTERDOWN Link with MASTER is down\r\n")] + #[case::try_again("-TRYAGAIN Multiple keys request during rehashing of slot\r\n")] + #[case::read_only("-READONLY You can't write against a read only replica.\r\n")] + fn unavailable_nodes_are_transport_errors(#[case] resp: &str) { + assert!(matches!(Error::from(reply(resp)), Error::Transport(_))); + } +} diff --git a/litellm-rust/crates/redis-client/src/lib.rs b/litellm-rust/crates/redis-client/src/lib.rs new file mode 100644 index 00000000000..d86d31c3f5b --- /dev/null +++ b/litellm-rust/crates/redis-client/src/lib.rs @@ -0,0 +1,9 @@ +mod client; +mod config; +mod error; +mod namespace; + +pub use client::{RedisClient, Ttl}; +pub use config::{DEFAULT_TIMEOUT, Endpoint, HostEndpoint, RedisConfig}; +pub use error::{ConnectError, Error}; +pub use namespace::Namespace; diff --git a/litellm-rust/crates/redis-client/src/namespace.rs b/litellm-rust/crates/redis-client/src/namespace.rs new file mode 100644 index 00000000000..b64293db421 --- /dev/null +++ b/litellm-rust/crates/redis-client/src/namespace.rs @@ -0,0 +1,53 @@ +#[derive(Clone, Default)] +pub struct Namespace(Option); + +impl Namespace { + pub fn new(namespace: impl Into) -> Self { + let namespace = namespace.into(); + Self((!namespace.is_empty()).then_some(namespace)) + } + + pub(crate) fn apply(&self, key: &str) -> String { + match &self.0 { + Some(namespace) if !is_prefixed(key, namespace) => { + format!("{namespace}:{key}") + } + _ => key.to_owned(), + } + } +} + +fn is_prefixed(key: &str, namespace: &str) -> bool { + key.strip_prefix(namespace) + .is_some_and(|rest| rest.starts_with(':')) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::Namespace; + + #[rstest] + #[case::plain(Namespace::new("litellm"), "spend:key:abc", "litellm:spend:key:abc")] + #[case::no_namespace(Namespace::default(), "spend:key:abc", "spend:key:abc")] + #[case::empty_namespace(Namespace::new(""), "spend:key:abc", "spend:key:abc")] + #[case::already_prefixed(Namespace::new("litellm"), "litellm:abc", "litellm:abc")] + #[case::shares_prefix_without_delimiter( + Namespace::new("litellm"), + "litellm_spend_update_buffer", + "litellm:litellm_spend_update_buffer" + )] + #[case::hash_tagged_key( + Namespace::new("litellm"), + "{api_key:sk-1}:requests", + "litellm:{api_key:sk-1}:requests" + )] + fn prefixes_keys_like_python_check_and_fix_namespace( + #[case] namespace: Namespace, + #[case] key: &str, + #[case] expected: &str, + ) { + assert_eq!(namespace.apply(key), expected); + } +} diff --git a/litellm-rust/crates/redis-testing/AGENTS.md b/litellm-rust/crates/redis-testing/AGENTS.md new file mode 100644 index 00000000000..b116ba365c9 --- /dev/null +++ b/litellm-rust/crates/redis-testing/AGENTS.md @@ -0,0 +1,5 @@ +- Dev-only crate that starts real Redis for tests. Never a dependency outside `[dev-dependencies]` +- Same shape as `db-testing`: one container per test, a handle that removes it on drop +- The host port is fixed at start so `restart` keeps the same URL and a client can reconnect to it +- Outages come from the container (`restart`, `freeze` / `thaw`), never `DEBUG`, which the official image disables +- No env gate that makes a test pass without Redis. A test that can't start its container fails diff --git a/litellm-rust/crates/redis-testing/Cargo.toml b/litellm-rust/crates/redis-testing/Cargo.toml new file mode 100644 index 00000000000..c9bad227acc --- /dev/null +++ b/litellm-rust/crates/redis-testing/Cargo.toml @@ -0,0 +1,11 @@ +[package] +name = "litellm-redis-testing" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +testcontainers-modules = { workspace = true, features = ["redis"] } +thiserror.workspace = true +tokio = { workspace = true, features = ["io-util"] } diff --git a/litellm-rust/crates/redis-testing/src/error.rs b/litellm-rust/crates/redis-testing/src/error.rs new file mode 100644 index 00000000000..f21cc959812 --- /dev/null +++ b/litellm-rust/crates/redis-testing/src/error.rs @@ -0,0 +1,11 @@ +use std::io; + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("reserving a host port for the Redis container")] + Port(#[source] io::Error), + #[error("Redis on host port {port} did not answer PING")] + Unreachable { port: u16 }, + #[error("controlling the Redis container")] + Container(#[from] testcontainers_modules::testcontainers::TestcontainersError), +} diff --git a/litellm-rust/crates/redis-testing/src/lib.rs b/litellm-rust/crates/redis-testing/src/lib.rs new file mode 100644 index 00000000000..ac9389fedc0 --- /dev/null +++ b/litellm-rust/crates/redis-testing/src/lib.rs @@ -0,0 +1,6 @@ +mod error; +mod ports; +mod standalone; + +pub use error::Error; +pub use standalone::RedisServer; diff --git a/litellm-rust/crates/redis-testing/src/ports.rs b/litellm-rust/crates/redis-testing/src/ports.rs new file mode 100644 index 00000000000..e7b94988c10 --- /dev/null +++ b/litellm-rust/crates/redis-testing/src/ports.rs @@ -0,0 +1,66 @@ +use std::{ + future::Future, + io, + net::TcpListener, + time::{Duration, Instant}, +}; + +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::TcpStream, +}; + +use crate::Error; + +pub(crate) const HOST: &str = "127.0.0.1"; +const READY_TIMEOUT: Duration = Duration::from_secs(10); +const START_ATTEMPTS: usize = 3; + +fn free_ports() -> Result<[u16; N], Error> { + let listeners = (0..N) + .map(|_| TcpListener::bind((HOST, 0))) + .collect::>>() + .map_err(Error::Port)?; + let ports = listeners + .iter() + .map(|listener| listener.local_addr().map(|address| address.port())) + .collect::>>() + .map_err(Error::Port)?; + Ok(ports + .try_into() + .unwrap_or_else(|ports: Vec| unreachable!("{} ports", ports.len()))) +} + +pub(crate) async fn wait_for_pong(port: u16) -> Result<(), Error> { + let deadline = Instant::now() + READY_TIMEOUT; + while Instant::now() < deadline { + if ping(port).await.unwrap_or(false) { + return Ok(()); + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + Err(Error::Unreachable { port }) +} + +async fn ping(port: u16) -> io::Result { + let mut stream = TcpStream::connect((HOST, port)).await?; + stream.write_all(b"PING\r\n").await?; + let mut reply = [0; 7]; + stream.read_exact(&mut reply).await?; + Ok(&reply == b"+PONG\r\n") +} + +pub(crate) async fn start_on_free_ports( + start: impl Fn([u16; N]) -> F, +) -> Result +where + F: Future>, +{ + let mut attempt = 1; + loop { + match start(free_ports()?).await { + Err(_) if attempt < START_ATTEMPTS => attempt += 1, + started => return started, + } + } +} diff --git a/litellm-rust/crates/redis-testing/src/standalone.rs b/litellm-rust/crates/redis-testing/src/standalone.rs new file mode 100644 index 00000000000..01cfa82acdf --- /dev/null +++ b/litellm-rust/crates/redis-testing/src/standalone.rs @@ -0,0 +1,55 @@ +use testcontainers_modules::{ + redis::{REDIS_PORT, Redis}, + testcontainers::core::IntoContainerPort, + testcontainers::{ContainerAsync, ImageExt, runners::AsyncRunner}, +}; + +use crate::{ + Error, + ports::{HOST, start_on_free_ports, wait_for_pong}, +}; + +const REDIS_TAG: &str = + "7.4@sha256:44960b3100c48159a66d28c18b0588c193492e25f7551d60e51dd464fc3eca5b"; + +pub struct RedisServer { + container: ContainerAsync, + url: String, +} + +impl RedisServer { + pub async fn start() -> Result { + start_on_free_ports(Self::start_on).await + } + + async fn start_on([host_port]: [u16; 1]) -> Result { + let container = Redis::default() + .with_tag(REDIS_TAG) + .with_mapped_port(host_port, REDIS_PORT.tcp()) + .start() + .await?; + wait_for_pong(host_port).await?; + Ok(Self { + container, + url: format!("redis://{HOST}:{host_port}"), + }) + } + + pub fn url(&self) -> &str { + &self.url + } + + pub async fn restart(&self) -> Result<(), Error> { + self.container.stop().await?; + self.container.start().await?; + Ok(()) + } + + pub async fn freeze(&self) -> Result<(), Error> { + Ok(self.container.pause().await?) + } + + pub async fn thaw(&self) -> Result<(), Error> { + Ok(self.container.unpause().await?) + } +} diff --git a/litellm-rust/crates/spend-clickhouse/Cargo.toml b/litellm-rust/crates/spend-clickhouse/Cargo.toml index 661a75ac02e..893c5daded3 100644 --- a/litellm-rust/crates/spend-clickhouse/Cargo.toml +++ b/litellm-rust/crates/spend-clickhouse/Cargo.toml @@ -19,7 +19,7 @@ time = { workspace = true, features = ["formatting"] } [dev-dependencies] litellm-http = { workspace = true, features = ["test-support"] } rstest.workspace = true -testcontainers-modules = { version = "0.15.0", features = ["clickhouse"] } +testcontainers-modules = { workspace = true, features = ["clickhouse"] } tokio.workspace = true uuid.workspace = true wiremock.workspace = true diff --git a/litellm-rust/crates/storage-clickhouse/Cargo.toml b/litellm-rust/crates/storage-clickhouse/Cargo.toml index 4e72e823dc2..39d04f2de5d 100644 --- a/litellm-rust/crates/storage-clickhouse/Cargo.toml +++ b/litellm-rust/crates/storage-clickhouse/Cargo.toml @@ -19,6 +19,6 @@ url.workspace = true [dev-dependencies] litellm-http = { workspace = true, features = ["test-support"] } rstest.workspace = true -testcontainers-modules = { version = "0.15.0", features = ["clickhouse"] } +testcontainers-modules = { workspace = true, features = ["clickhouse"] } tokio.workspace = true wiremock.workspace = true diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index a859ce96c70..745488bf93f 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -15587,4 +15587,3 @@ async def test_spend_capture_rate_check_job_clears_the_gauge_once_the_setting_is call(api_provider="openai", capture_rate=0.97), call(api_provider="openai", capture_rate=None), ] -