From ca31149040170297c1a00044d6bb9650862dd895 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:28:27 +0000 Subject: [PATCH 1/3] feat(rust): add native GCS object-store cache backend Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 17 + litellm-rust/Cargo.toml | 2 + litellm-rust/crates/auth-gcp/src/lib.rs | 8 + litellm-rust/crates/cache-gcs/Cargo.toml | 21 ++ litellm-rust/crates/cache-gcs/src/cache.rs | 285 +++++++++++++++++ litellm-rust/crates/cache-gcs/src/lib.rs | 5 + litellm-rust/crates/cache-gcs/src/token.rs | 44 +++ litellm-rust/crates/cache-gcs/tests/cache.rs | 290 ++++++++++++++++++ litellm-rust/crates/cache/src/error.rs | 2 + litellm-rust/crates/python-bridge/Cargo.toml | 1 + .../crates/python-bridge/src/cache/config.rs | 143 ++++++++- .../crates/python-bridge/src/cache/facade.rs | 4 + .../crates/python-bridge/src/cache/handle.rs | 27 ++ .../crates/python-bridge/src/cache/mod.rs | 3 +- .../crates/python-bridge/src/cache/native.rs | 35 +++ 15 files changed, 876 insertions(+), 11 deletions(-) create mode 100644 litellm-rust/crates/cache-gcs/Cargo.toml create mode 100644 litellm-rust/crates/cache-gcs/src/cache.rs create mode 100644 litellm-rust/crates/cache-gcs/src/lib.rs create mode 100644 litellm-rust/crates/cache-gcs/src/token.rs create mode 100644 litellm-rust/crates/cache-gcs/tests/cache.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index ed4ae4e3353..99a6133b4f1 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2464,6 +2464,22 @@ dependencies = [ "tokio", ] +[[package]] +name = "litellm-cache-gcs" +version = "0.1.0" +dependencies = [ + "futures-util", + "litellm-auth-gcp", + "litellm-auth-types", + "litellm-cache", + "percent-encoding", + "reqwest 0.12.28", + "serde_json", + "tokio", + "url", + "wiremock", +] + [[package]] name = "litellm-cache-memory" version = "0.1.0" @@ -2666,6 +2682,7 @@ dependencies = [ "litellm-auth", "litellm-auth-gcp", "litellm-cache", + "litellm-cache-gcs", "litellm-cache-memory", "litellm-cache-redis", "litellm-cache-response", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 570d0dd3568..80425d567c0 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -29,6 +29,7 @@ litellm-core-utils = { path = "crates/core-utils" } litellm-cache = { path = "crates/cache" } litellm-cache-memory = { path = "crates/cache-memory" } litellm-cache-redis = { path = "crates/cache-redis" } +litellm-cache-gcs = { path = "crates/cache-gcs" } litellm-cache-response = { path = "crates/cache-response" } litellm-token-counter = { path = "crates/token-counter" } litellm-token-counter-fast = { path = "crates/token-counter-fast" } @@ -66,6 +67,7 @@ base64 = "0.22" moka = { version = "0.12.16", features = ["future"] } strum = { version = "0.28.0", features = ["derive"] } url = "2.5.8" +percent-encoding = "2.3" webpki-roots = "1" time = { version = "0.3.53", features = ["parsing"] } criterion = "0.8.2" diff --git a/litellm-rust/crates/auth-gcp/src/lib.rs b/litellm-rust/crates/auth-gcp/src/lib.rs index 8aeddae9efc..534d85acdb0 100644 --- a/litellm-rust/crates/auth-gcp/src/lib.rs +++ b/litellm-rust/crates/auth-gcp/src/lib.rs @@ -128,6 +128,14 @@ impl VertexAuth { } } + pub async fn access_token( + &self, + config: &VertexConfig, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + ) -> Result { + self.load_provider(config, env_lookup).await?.token().await + } + pub async fn validate_environment( &self, headers: Vec<(String, String)>, diff --git a/litellm-rust/crates/cache-gcs/Cargo.toml b/litellm-rust/crates/cache-gcs/Cargo.toml new file mode 100644 index 00000000000..3890e60843c --- /dev/null +++ b/litellm-rust/crates/cache-gcs/Cargo.toml @@ -0,0 +1,21 @@ +[package] +name = "litellm-cache-gcs" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +futures-util.workspace = true +litellm-auth-gcp.workspace = true +litellm-auth-types.workspace = true +litellm-cache.workspace = true +percent-encoding.workspace = true +reqwest.workspace = true +tokio.workspace = true +url.workspace = true + +[dev-dependencies] +serde_json.workspace = true +tokio.workspace = true +wiremock = "0.6.5" diff --git a/litellm-rust/crates/cache-gcs/src/cache.rs b/litellm-rust/crates/cache-gcs/src/cache.rs new file mode 100644 index 00000000000..1097f3252b1 --- /dev/null +++ b/litellm-rust/crates/cache-gcs/src/cache.rs @@ -0,0 +1,285 @@ +use std::{future::Future, sync::Arc, time::Duration}; + +use futures_util::future::try_join_all; +use litellm_cache::{ + BaseCache, BatchCache, BatchEntry, CacheCodec, CacheConnectionResult, Error, ExactCacheContext, + FlushCache, +}; +use percent_encoding::{AsciiSet, CONTROLS, percent_encode}; +use reqwest::Client; + +use crate::{GcpTokenSource, TokenSource}; + +pub const DEFAULT_ENDPOINT: &str = "https://storage.googleapis.com"; + +const OBJECT_NAME_ENCODE_SET: &AsciiSet = &CONTROLS + .add(b' ') + .add(b'!') + .add(b'"') + .add(b'#') + .add(b'$') + .add(b'%') + .add(b'&') + .add(b'\'') + .add(b'(') + .add(b')') + .add(b'*') + .add(b'+') + .add(b',') + .add(b'/') + .add(b':') + .add(b';') + .add(b'<') + .add(b'=') + .add(b'>') + .add(b'?') + .add(b'@') + .add(b'[') + .add(b'\\') + .add(b']') + .add(b'^') + .add(b'`') + .add(b'{') + .add(b'|') + .add(b'}'); + +pub fn key_prefix(gcs_path: Option<&str>) -> String { + match gcs_path { + Some(path) if !path.is_empty() => format!("{}/", path.trim_end_matches('/')), + _ => String::new(), + } +} + +#[derive(Clone, Debug)] +pub struct GcsConfig { + pub bucket_name: String, + pub gcs_path: Option, + pub path_service_account: Option, + pub endpoint: String, +} + +impl GcsConfig { + pub fn new(bucket_name: impl Into) -> Self { + Self { + bucket_name: bucket_name.into(), + gcs_path: None, + path_service_account: None, + endpoint: DEFAULT_ENDPOINT.to_string(), + } + } +} + +pub struct GcsCache { + config: GcsConfig, + key_prefix: String, + client: Client, + token: Arc, + codec: S, +} + +impl GcsCache { + pub fn new(config: GcsConfig, codec: S) -> Result { + let token = Arc::new(GcpTokenSource::new(config.path_service_account.clone())); + Self::with_token_source(config, codec, token) + } + + pub fn with_token_source( + config: GcsConfig, + codec: S, + token: Arc, + ) -> Result { + let client = Client::builder().build().map_err(|_| Error::Unavailable)?; + let key_prefix = key_prefix(config.gcs_path.as_deref()); + Ok(Self { + config, + key_prefix, + client, + token, + codec, + }) + } + + pub fn bucket_name(&self) -> &str { + &self.config.bucket_name + } + + pub fn key_prefix(&self) -> &str { + &self.key_prefix + } + + pub fn path_service_account(&self) -> Option<&str> { + self.config.path_service_account.as_deref() + } + + pub fn object_name(&self, key: &str) -> String { + format!("{}{}", self.key_prefix, key) + } + + fn encoded_object_name(&self, key: &str) -> String { + percent_encode(self.object_name(key).as_bytes(), OBJECT_NAME_ENCODE_SET).to_string() + } + + fn endpoint(&self, path: &str) -> String { + format!("{}{}", self.config.endpoint.trim_end_matches('/'), path) + } + + async fn async_set(&self, key: &str, value: S::Value) -> Result<(), Error> { + let token = self.token.bearer_token().await?; + let payload = self.codec.encode(&value)?; + let url = self.endpoint(&format!( + "/upload/storage/v1/b/{}/o?uploadType=media&name={}", + self.config.bucket_name, + self.encoded_object_name(key) + )); + let response = self + .client + .post(url) + .bearer_auth(token) + .header(reqwest::header::CONTENT_TYPE, "application/json") + .body(payload) + .send() + .await + .map_err(|_| Error::Unavailable)?; + if !response.status().is_success() { + return Err(Error::Unavailable); + } + Ok(()) + } + + async fn async_get(&self, key: &str) -> Result, Error> { + let token = self.token.bearer_token().await?; + let url = self.endpoint(&format!( + "/storage/v1/b/{}/o/{}?alt=media", + self.config.bucket_name, + self.encoded_object_name(key) + )); + let response = self + .client + .get(url) + .bearer_auth(token) + .header(reqwest::header::CONTENT_TYPE, "application/json") + .send() + .await + .map_err(|_| Error::Unavailable)?; + if response.status() == reqwest::StatusCode::NOT_FOUND { + return Ok(None); + } + if !response.status().is_success() { + return Err(Error::Unavailable); + } + let body = response.bytes().await.map_err(|_| Error::Unavailable)?; + self.codec + .decode(&body) + .map(Some) + .map_err(|_| Error::InvalidEntry) + } + + fn run_sync(future: F) -> Result + where + F: Future> + Send, + T: Send, + { + let run = || { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .map_err(|_| Error::Unavailable) + .and_then(|runtime| runtime.block_on(future)) + }; + if let Ok(handle) = tokio::runtime::Handle::try_current() { + if handle.runtime_flavor() == tokio::runtime::RuntimeFlavor::MultiThread { + return tokio::task::block_in_place(run); + } + return std::thread::scope(|scope| { + scope + .spawn(run) + .join() + .map_err(|_| Error::Unavailable) + .and_then(|result| result) + }); + } + run() + } +} + +impl BaseCache for GcsCache { + type Value = S::Value; + type Context = ExactCacheContext; + + fn get_ttl(&self, _: &Self::Context) -> Option { + None + } + + fn set_cache(&self, key: &str, value: Self::Value, _: &Self::Context) -> Result<(), Error> { + Self::run_sync(self.async_set(key, value)) + } + + fn get_cache(&self, key: &str, _: &Self::Context) -> Result, Error> { + Self::run_sync(self.async_get(key)) + } + + async fn async_set_cache( + &self, + key: &str, + value: Self::Value, + _: Self::Context, + ) -> Result<(), Error> { + self.async_set(key, value).await + } + + async fn async_get_cache( + &self, + key: &str, + _: &Self::Context, + ) -> Result, Error> { + self.async_get(key).await + } + + async fn async_set_cache_pipeline( + &self, + entries: Vec<(String, Self::Value)>, + context: Self::Context, + ) -> Result<(), Error> { + try_join_all(entries.into_iter().map(|(key, value)| { + let context = context.clone(); + async move { self.async_set_cache(&key, value, context).await } + })) + .await + .map(|_| ()) + } + + async fn disconnect(&self) -> Result<(), Error> { + Ok(()) + } + + async fn test_connection(&self) -> Result { + Err(Error::UnsupportedOperation) + } +} + +impl BatchCache for GcsCache { + async fn async_batch_get_cache( + &self, + keys: Vec, + context: Self::Context, + ) -> Result>, Error> { + try_join_all(keys.into_iter().map(|key| { + let context = context.clone(); + async move { + match self.async_get_cache(&key, &context).await { + Ok(Some(value)) => Ok(BatchEntry::Hit(value)), + Ok(None) => Ok(BatchEntry::Miss), + Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid), + Err(error) => Err(error), + } + } + })) + .await + } +} + +impl FlushCache for GcsCache { + fn flush_cache(&self) -> Result<(), Error> { + Ok(()) + } +} diff --git a/litellm-rust/crates/cache-gcs/src/lib.rs b/litellm-rust/crates/cache-gcs/src/lib.rs new file mode 100644 index 00000000000..cbb61cf0685 --- /dev/null +++ b/litellm-rust/crates/cache-gcs/src/lib.rs @@ -0,0 +1,5 @@ +mod cache; +mod token; + +pub use cache::{DEFAULT_ENDPOINT, GcsCache, GcsConfig, key_prefix}; +pub use token::{GcpTokenSource, StaticTokenSource, TokenSource}; diff --git a/litellm-rust/crates/cache-gcs/src/token.rs b/litellm-rust/crates/cache-gcs/src/token.rs new file mode 100644 index 00000000000..adb601c276c --- /dev/null +++ b/litellm-rust/crates/cache-gcs/src/token.rs @@ -0,0 +1,44 @@ +use std::{future::Future, pin::Pin}; + +use litellm_auth_gcp::{VertexAuth, VertexConfig}; +use litellm_auth_types::{InputSource, SecretValue, Sourced}; +use litellm_cache::Error; + +pub trait TokenSource: Send + Sync + 'static { + fn bearer_token(&self) -> Pin> + Send + '_>>; +} + +pub struct GcpTokenSource { + auth: VertexAuth, + config: VertexConfig, +} + +impl GcpTokenSource { + pub fn new(path_service_account: Option) -> Self { + let credentials = path_service_account + .map(|path| Sourced::new(SecretValue::new(path), InputSource::Deployment)); + Self { + auth: VertexAuth::default(), + config: VertexConfig::new(credentials, None, None), + } + } +} + +impl TokenSource for GcpTokenSource { + fn bearer_token(&self) -> Pin> + Send + '_>> { + Box::pin(async move { + self.auth + .access_token(&self.config, &|name| std::env::var(name).ok()) + .await + .map_err(|_| Error::Unavailable) + }) + } +} + +pub struct StaticTokenSource(pub String); + +impl TokenSource for StaticTokenSource { + fn bearer_token(&self) -> Pin> + Send + '_>> { + Box::pin(async move { Ok(self.0.clone()) }) + } +} diff --git a/litellm-rust/crates/cache-gcs/tests/cache.rs b/litellm-rust/crates/cache-gcs/tests/cache.rs new file mode 100644 index 00000000000..79e4ed100f8 --- /dev/null +++ b/litellm-rust/crates/cache-gcs/tests/cache.rs @@ -0,0 +1,290 @@ +use std::{sync::Arc, time::Duration}; + +use litellm_cache::{ + BaseCache, BatchCache, BatchEntry, CacheContext, Error, ExactCacheContext, FlushCache, + JsonCodec, +}; +use litellm_cache_gcs::{GcsCache, GcsConfig, StaticTokenSource, TokenSource, key_prefix}; +use serde_json::json; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{body_bytes, header, method, path, query_param}, +}; + +fn config(server: &MockServer, gcs_path: Option<&str>) -> GcsConfig { + GcsConfig { + bucket_name: "bucket".into(), + gcs_path: gcs_path.map(str::to_string), + path_service_account: None, + endpoint: server.uri(), + } +} + +fn cache(server: &MockServer, gcs_path: Option<&str>) -> GcsCache> { + GcsCache::with_token_source( + config(server, gcs_path), + JsonCodec::new(), + Arc::new(StaticTokenSource("tok".into())), + ) + .unwrap() +} + +#[tokio::test] +async fn set_writes_encoded_object_and_headers() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/upload/storage/v1/b/bucket/o")) + .and(query_param("uploadType", "media")) + .and(header("authorization", "Bearer tok")) + .and(header("content-type", "application/json")) + .and(body_bytes(br#"{"value":"entry"}"#)) + .respond_with(ResponseTemplate::new(200)) + .expect(1) + .mount(&server) + .await; + cache(&server, Some("cache/")) + .set_cache( + "team:a b/c", + json!({"value": "entry"}), + &ExactCacheContext::default(), + ) + .unwrap(); + let requests = server.received_requests().await.unwrap(); + assert_eq!(requests.len(), 1); + assert_eq!( + requests[0].url.query(), + Some("uploadType=media&name=cache%2Fteam%3Aa%20b%2Fc") + ); +} + +#[tokio::test] +async fn get_maps_statuses_and_decode_failures() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/storage/v1/b/bucket/o/hit")) + .and(query_param("alt", "media")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"value": "entry"}))) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/storage/v1/b/bucket/o/missing")) + .respond_with(ResponseTemplate::new(404)) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/storage/v1/b/bucket/o/server-error")) + .respond_with(ResponseTemplate::new(500)) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/storage/v1/b/bucket/o/invalid")) + .respond_with(ResponseTemplate::new(200).set_body_string("not json")) + .mount(&server) + .await; + + let cache = cache(&server, None); + assert_eq!( + cache + .get_cache("hit", &ExactCacheContext::default()) + .unwrap(), + Some(json!({"value": "entry"})) + ); + assert_eq!( + cache + .get_cache("missing", &ExactCacheContext::default()) + .unwrap(), + None + ); + assert_eq!( + cache + .get_cache("server-error", &ExactCacheContext::default()) + .unwrap_err(), + Error::Unavailable + ); + assert_eq!( + cache + .get_cache("invalid", &ExactCacheContext::default()) + .unwrap_err(), + Error::InvalidEntry + ); +} + +#[test] +fn key_prefix_normalizes_paths() { + assert_eq!(key_prefix(None), ""); + assert_eq!(key_prefix(Some("a/b/")), "a/b/"); + assert_eq!(key_prefix(Some("a/b")), "a/b/"); + assert_eq!(key_prefix(Some("")), ""); +} + +#[tokio::test] +async fn ignores_ttl_and_writes_pipeline_concurrently() { + let server = MockServer::start().await; + for key in ["one", "two", "three"] { + Mock::given(method("POST")) + .and(path("/upload/storage/v1/b/bucket/o")) + .and(query_param("uploadType", "media")) + .and(query_param("name", key)) + .respond_with(ResponseTemplate::new(200)) + .expect(1) + .mount(&server) + .await; + } + let cache = cache(&server, None); + assert_eq!(cache.get_ttl(&ExactCacheContext::default()), None); + assert_eq!( + cache.get_ttl(&ExactCacheContext::default().with_ttl(Some(Duration::from_secs(5)))), + None + ); + cache + .async_set_cache_pipeline( + vec![ + ("one".into(), json!({"key": "one"})), + ("two".into(), json!({"key": "two"})), + ("three".into(), json!({"key": "three"})), + ], + ExactCacheContext::default().with_ttl(Some(Duration::from_secs(5))), + ) + .await + .unwrap(); +} + +#[tokio::test] +async fn async_batch_get_preserves_hits_misses_and_invalid_entries() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/storage/v1/b/bucket/o/hit")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"value": "entry"}))) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/storage/v1/b/bucket/o/missing")) + .respond_with(ResponseTemplate::new(404)) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/storage/v1/b/bucket/o/invalid")) + .respond_with(ResponseTemplate::new(200).set_body_string("not json")) + .mount(&server) + .await; + + assert_eq!( + cache(&server, None) + .async_batch_get_cache( + vec!["hit".into(), "missing".into(), "invalid".into()], + ExactCacheContext::default(), + ) + .await + .unwrap(), + vec![ + BatchEntry::Hit(json!({"value": "entry"})), + BatchEntry::Miss, + BatchEntry::Invalid, + ] + ); +} + +#[tokio::test] +async fn lifecycle_operations_are_noops_and_connection_test_is_unsupported() { + let server = MockServer::start().await; + let cache = cache(&server, None); + assert_eq!(cache.flush_cache(), Ok(())); + assert_eq!(cache.disconnect().await, Ok(())); + assert_eq!( + cache.test_connection().await, + Err(Error::UnsupportedOperation) + ); +} + +#[test] +fn sync_operations_work_without_an_active_runtime() { + let runtime = tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .unwrap(); + let server = runtime.block_on(MockServer::start()); + runtime.block_on( + Mock::given(method("POST")) + .and(path("/upload/storage/v1/b/bucket/o")) + .respond_with(ResponseTemplate::new(200)) + .mount(&server), + ); + runtime.block_on( + Mock::given(method("GET")) + .and(path("/storage/v1/b/bucket/o/key")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"value": "entry"}))) + .mount(&server), + ); + let cache = cache(&server, None); + cache + .set_cache( + "key", + json!({"value": "entry"}), + &ExactCacheContext::default(), + ) + .unwrap(); + assert_eq!( + cache + .get_cache("key", &ExactCacheContext::default()) + .unwrap(), + Some(json!({"value": "entry"})) + ); +} + +#[tokio::test(flavor = "multi_thread")] +async fn sync_operations_work_inside_a_multi_thread_runtime() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/upload/storage/v1/b/bucket/o")) + .respond_with(ResponseTemplate::new(200)) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/storage/v1/b/bucket/o/key")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"value": "entry"}))) + .mount(&server) + .await; + let cache = cache(&server, None); + cache + .set_cache( + "key", + json!({"value": "entry"}), + &ExactCacheContext::default(), + ) + .unwrap(); + assert_eq!( + cache + .get_cache("key", &ExactCacheContext::default()) + .unwrap(), + Some(json!({"value": "entry"})) + ); +} + +struct FailingTokenSource; + +impl TokenSource for FailingTokenSource { + fn bearer_token( + &self, + ) -> std::pin::Pin> + Send + '_>> + { + Box::pin(async { Err(Error::Unavailable) }) + } +} + +#[tokio::test] +async fn token_source_failure_skips_http() { + let server = MockServer::start().await; + let cache = GcsCache::with_token_source( + config(&server, None), + JsonCodec::::new(), + Arc::new(FailingTokenSource), + ) + .unwrap(); + assert_eq!( + cache + .get_cache("key", &ExactCacheContext::default()) + .unwrap_err(), + Error::Unavailable + ); + assert_eq!(server.received_requests().await.unwrap().len(), 0); +} diff --git a/litellm-rust/crates/cache/src/error.rs b/litellm-rust/crates/cache/src/error.rs index ff3ff6572d4..1a381d0afd8 100644 --- a/litellm-rust/crates/cache/src/error.rs +++ b/litellm-rust/crates/cache/src/error.rs @@ -6,4 +6,6 @@ pub enum Error { InvalidEntry, #[error("flushing Redis requires an explicit namespace")] UnscopedFlush, + #[error("operation is not supported by this cache")] + UnsupportedOperation, } diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 1eb2ec28036..0f060c801a5 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -23,6 +23,7 @@ bytes.workspace = true litellm-cache.workspace = true litellm-cache-memory.workspace = true litellm-cache-redis.workspace = true +litellm-cache-gcs.workspace = true litellm-cache-response.workspace = true serde.workspace = true litellm-auth.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs index 0e7d6aee11d..66f493dd760 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -73,9 +73,17 @@ pub(super) struct RedisCacheConfig { pub(super) connection: RedisConnectionConfig, } +#[derive(Debug, PartialEq)] +pub(super) struct GcsCacheConfig { + pub(super) bucket_name: String, + pub(super) key_prefix: String, + pub(super) path_service_account: Option, +} + pub(super) enum CacheBackendConfig { Memory(MemoryCacheConfig), Redis(Box), + Gcs(GcsCacheConfig), } #[allow(dead_code, reason = "consumed by the cache activation follow-up")] @@ -90,6 +98,7 @@ pub(super) enum UnsupportedCacheConfig { RedisCredentials, RedisConnection, RedisOption, + GcsBucket, } impl UnsupportedCacheConfig { @@ -100,6 +109,7 @@ impl UnsupportedCacheConfig { Self::RedisCredentials => "native Redis credentials require Python", Self::RedisConnection => "native Redis connection type is not implemented", Self::RedisOption => "native Redis configuration requires Python", + Self::GcsBucket => "native GCS cache requires a configured bucket name", } } } @@ -142,14 +152,20 @@ impl NativeCacheConfig { }))), Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)), }, + Some(CacheType::Gcs) => match project_gcs(&backend)? { + Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self { + policy, + backend: CacheBackendConfig::Gcs(backend), + }))), + Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)), + }, Some( CacheType::RedisSemantic | CacheType::ValkeySemantic | CacheType::S3 | CacheType::Disk | CacheType::QdrantSemantic - | CacheType::AzureBlob - | CacheType::Gcs, + | CacheType::AzureBlob, ) | None => Ok(CacheConfigProjection::Unsupported( UnsupportedCacheConfig::Backend, @@ -158,12 +174,12 @@ impl NativeCacheConfig { } pub(super) fn service_mismatch(&self, service: &NativeResponseCache) -> Option<&'static str> { - if service.default_ttl() - != Some(match &self.backend { - CacheBackendConfig::Memory(config) => config.default_ttl, - CacheBackendConfig::Redis(config) => config.default_ttl, - }) - { + let expected = match &self.backend { + CacheBackendConfig::Memory(config) => Some(config.default_ttl), + CacheBackendConfig::Redis(config) => Some(config.default_ttl), + CacheBackendConfig::Gcs(_) => None, + }; + if service.default_ttl() != expected { return Some("facade and native backend default TTLs must match"); } match &self.backend { @@ -185,6 +201,31 @@ impl NativeCacheConfig { CacheBackendConfig::Redis(config) => (service.namespace() != config.namespace.as_deref()) .then_some("facade and native backend namespaces must match"), + CacheBackendConfig::Gcs(_) if service.kind() != "gcs" => { + Some("facade and native backend types must match") + } + CacheBackendConfig::Gcs(config) + if service + .gcs_backend() + .is_none_or(|backend| backend.bucket_name() != config.bucket_name) => + { + Some("facade and native backend buckets must match") + } + CacheBackendConfig::Gcs(config) + if service + .gcs_backend() + .is_none_or(|backend| backend.key_prefix() != config.key_prefix) => + { + Some("facade and native backend key prefixes must match") + } + CacheBackendConfig::Gcs(config) + if service.gcs_backend().is_none_or(|backend| { + backend.path_service_account() != config.path_service_account.as_deref() + }) => + { + Some("facade and native backend credentials must match") + } + CacheBackendConfig::Gcs(_) => None, } } } @@ -201,6 +242,23 @@ fn project_memory(backend: &Bound<'_, PyAny>) -> PyResult { }) } +#[inline(never)] +fn project_gcs( + backend: &Bound<'_, PyAny>, +) -> PyResult> { + let bucket_name = match backend.getattr("bucket_name")?.extract::>() { + Ok(Some(bucket_name)) if !bucket_name.is_empty() => bucket_name, + _ => return Ok(Err(UnsupportedCacheConfig::GcsBucket)), + }; + Ok(Ok(GcsCacheConfig { + bucket_name, + key_prefix: backend.getattr("key_prefix")?.extract::()?, + path_service_account: backend + .getattr("path_service_account")? + .extract::>()?, + })) +} + #[inline(never)] fn project_redis( backend: &Bound<'_, PyAny>, @@ -468,8 +526,8 @@ mod tests { use pyo3::{prelude::*, types::PyDict}; use super::{ - CacheBackendConfig, CacheConfigProjection, CertificateRequirement, NativeCacheConfig, - RedisProtocol, + CacheBackendConfig, CacheConfigProjection, CertificateRequirement, GcsCacheConfig, + NativeCacheConfig, RedisProtocol, UnsupportedCacheConfig, }; use crate::cache::native::NativeResponseCache; @@ -531,6 +589,71 @@ mod tests { }); } + #[test] + fn projects_gcs_configuration() { + Python::initialize(); + Python::attach(|py| { + let facade = facade( + py, + "backend = SimpleNamespace(bucket_name='bucket', key_prefix='cache/', path_service_account='credentials.json')\n\ + facade = SimpleNamespace(type='gcs', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)", + ); + let CacheConfigProjection::Native(config) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("GCS cache should be supported"); + }; + let CacheBackendConfig::Gcs(gcs) = config.backend else { + panic!("expected GCS configuration"); + }; + assert_eq!( + gcs, + GcsCacheConfig { + bucket_name: "bucket".into(), + key_prefix: "cache/".into(), + path_service_account: Some("credentials.json".into()), + } + ); + let matching = NativeResponseCache::gcs( + litellm_cache_gcs::GcsConfig { + bucket_name: "bucket".into(), + gcs_path: Some("cache/".into()), + path_service_account: Some("credentials.json".into()), + endpoint: litellm_cache_gcs::DEFAULT_ENDPOINT.into(), + }, + Some("token".into()), + ) + .unwrap(); + let matching_config = NativeCacheConfig { + policy: config.policy, + backend: CacheBackendConfig::Gcs(gcs), + }; + assert_eq!(matching_config.service_mismatch(&matching), None); + }); + } + + #[test] + fn rejects_gcs_without_a_bucket_name() { + Python::initialize(); + Python::attach(|py| { + let facade = facade( + py, + "backend = SimpleNamespace(bucket_name=None, key_prefix='', path_service_account=None)\n\ + facade = SimpleNamespace(type='gcs', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)", + ); + let CacheConfigProjection::Unsupported(reason) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("GCS cache without a bucket should be unsupported"); + }; + assert!(matches!(&reason, UnsupportedCacheConfig::GcsBucket)); + assert_eq!( + reason.message(), + "native GCS cache requires a configured bucket name" + ); + }); + } + #[test] fn projects_resolved_redis_tls_configuration() { Python::initialize(); diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index f2f86c14b37..6aa94754e2d 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -192,6 +192,7 @@ impl FacadeGuard { let (module, name, cache_kind) = match kind { "memory" => ("litellm.caching.in_memory_cache", "InMemoryCache", "local"), "redis" => ("litellm.caching.redis_cache", "RedisCache", "redis"), + "gcs" => ("litellm.caching.gcs_cache", "GCSCache", "gcs"), _ => unreachable!(), }; let backend = facade.getattr("cache")?; @@ -235,6 +236,9 @@ impl FacadeGuard { "max_size_per_item", "redis_kwargs", "redis_flush_size", + "bucket_name", + "key_prefix", + "path_service_account", ], )?, redis_pool: (kind == "redis") diff --git a/litellm-rust/crates/python-bridge/src/cache/handle.rs b/litellm-rust/crates/python-bridge/src/cache/handle.rs index 8251b3df06c..87d9300b08b 100644 --- a/litellm-rust/crates/python-bridge/src/cache/handle.rs +++ b/litellm-rust/crates/python-bridge/src/cache/handle.rs @@ -1,6 +1,8 @@ use litellm_host_python::release_gil; use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*}; +use litellm_cache_gcs::{DEFAULT_ENDPOINT, GcsConfig}; + use super::{cache_error, facade::FacadeGuard, native::NativeResponseCache, request::duration}; #[pyclass(frozen, name = "_CacheTestHandle")] @@ -51,6 +53,31 @@ impl CacheTestHandle { }) } + #[staticmethod] + #[pyo3(signature = (bucket_name, *, gcs_path=None, path_service_account=None, endpoint=None, token=None))] + fn gcs( + py: Python<'_>, + bucket_name: String, + gcs_path: Option, + path_service_account: Option, + endpoint: Option, + token: Option, + ) -> PyResult { + let config = GcsConfig { + bucket_name, + gcs_path, + path_service_account, + endpoint: endpoint.unwrap_or_else(|| DEFAULT_ENDPOINT.to_string()), + }; + let service = release_gil(py, move || NativeResponseCache::gcs(config, token)) + .map_err(cache_error)?; + Ok(Self { + service, + guard: None, + pid: std::process::id(), + }) + } + #[getter] fn backend(&self) -> &'static str { self.service.kind() diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index aec08610f6e..ac4494150d9 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -10,7 +10,7 @@ mod resolver; use litellm_cache::Error; use pyo3::{ - exceptions::{PyRuntimeError, PyValueError}, + exceptions::{PyNotImplementedError, PyRuntimeError, PyValueError}, prelude::*, }; @@ -21,6 +21,7 @@ pub(crate) use self::{ fn cache_error(error: Error) -> PyErr { match error { Error::InvalidEntry => PyValueError::new_err(error.to_string()), + Error::UnsupportedOperation => PyNotImplementedError::new_err(error.to_string()), _ => PyRuntimeError::new_err(error.to_string()), } } diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index a9475429e45..9a030ac0434 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -1,6 +1,7 @@ use std::{sync::Arc, time::Duration}; use litellm_cache::{CacheCodec, CacheConnectionResult, Error}; +use litellm_cache_gcs::{GcsCache, GcsConfig, StaticTokenSource}; use litellm_cache_memory::InMemoryCache; use litellm_cache_redis::RedisCache; use litellm_cache_response::{ @@ -15,6 +16,7 @@ pub(super) enum NativeResponseCache { cache: Arc>>, buffer: Option>, }, + Gcs(Arc>>), } impl NativeResponseCache { @@ -43,6 +45,18 @@ impl NativeResponseCache { buffer: None, }) } + + pub fn gcs(config: GcsConfig, token: Option) -> Result { + let backend = match token { + Some(token) => GcsCache::with_token_source( + config, + ResponseCacheCodec, + Arc::new(StaticTokenSource(token)), + )?, + None => GcsCache::new(config, ResponseCacheCodec)?, + }; + Ok(Self::Gcs(Arc::new(ResponseCache::new(Arc::new(backend))))) + } } impl NativeResponseCache { @@ -50,6 +64,7 @@ impl NativeResponseCache { match self { Self::Memory(_) => "memory", Self::Redis { .. } => "redis", + Self::Gcs(_) => "gcs", } } @@ -57,6 +72,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.default_ttl(), Self::Redis { cache, .. } => cache.default_ttl(), + Self::Gcs(cache) => cache.default_ttl(), } } @@ -64,6 +80,7 @@ impl NativeResponseCache { match self { Self::Memory(_) => None, Self::Redis { cache, .. } => cache.backend().namespace(), + Self::Gcs(_) => None, } } @@ -71,6 +88,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => Some(cache.backend().max_size_in_memory()), Self::Redis { .. } => None, + Self::Gcs(_) => None, } } @@ -78,6 +96,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.backend().max_entry_bytes(), Self::Redis { .. } => None, + Self::Gcs(_) => None, } } @@ -99,6 +118,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.lookup(request, now), Self::Redis { cache, .. } => cache.lookup(request, now), + Self::Gcs(cache) => cache.lookup(request, now), } } @@ -111,6 +131,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.store(request, response, now), Self::Redis { cache, .. } => cache.store(request, response, now), + Self::Gcs(cache) => cache.store(request, response, now), } } @@ -122,6 +143,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.lookup_batch(requests, now), Self::Redis { cache, .. } => cache.lookup_batch(requests, now), + Self::Gcs(cache) => cache.lookup_batch(requests, now), } } @@ -133,6 +155,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.async_lookup(request, now).await, Self::Redis { cache, .. } => cache.async_lookup(request, now).await, + Self::Gcs(cache) => cache.async_lookup(request, now).await, } } @@ -152,6 +175,7 @@ impl NativeResponseCache { cache, buffer: Some(buffer), } => buffer.async_store(cache, request, response, now).await, + Self::Gcs(cache) => cache.async_store(request, response, now).await, } } @@ -163,6 +187,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.async_lookup_batch(requests, now).await, Self::Redis { cache, .. } => cache.async_lookup_batch(requests, now).await, + Self::Gcs(cache) => cache.async_lookup_batch(requests, now).await, } } @@ -174,6 +199,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.async_store_batch(entries, now).await, Self::Redis { cache, .. } => cache.async_store_batch(entries, now).await, + Self::Gcs(cache) => cache.async_store_batch(entries, now).await, } } @@ -186,6 +212,7 @@ impl NativeResponseCache { } cache.async_flush().await } + Self::Gcs(cache) => cache.async_flush().await, } } @@ -193,6 +220,14 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.test_connection().await, Self::Redis { cache, .. } => cache.test_connection().await, + Self::Gcs(cache) => cache.test_connection().await, + } + } + + pub fn gcs_backend(&self) -> Option<&GcsCache> { + match self { + Self::Gcs(cache) => Some(cache.backend()), + _ => None, } } } From a6d1497932e257a809fdd4647d85a70039376816 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:38:20 +0000 Subject: [PATCH 2/3] test(rust): add GCS native cache parity fixtures Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 1 - litellm-rust/crates/cache-gcs/Cargo.toml | 1 - litellm-rust/crates/cache-gcs/src/cache.rs | 37 +-- litellm-rust/crates/cache-gcs/tests/cache.rs | 34 +++ tests/test_litellm_rust/support/fake_gcs.py | 152 +++++++++++++ tests/test_litellm_rust/test_cache.py | 228 +++++++++++++++++++ 6 files changed, 420 insertions(+), 33 deletions(-) create mode 100644 tests/test_litellm_rust/support/fake_gcs.py diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 99a6133b4f1..ae28553801d 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2476,7 +2476,6 @@ dependencies = [ "reqwest 0.12.28", "serde_json", "tokio", - "url", "wiremock", ] diff --git a/litellm-rust/crates/cache-gcs/Cargo.toml b/litellm-rust/crates/cache-gcs/Cargo.toml index 3890e60843c..4ec60bcfa3b 100644 --- a/litellm-rust/crates/cache-gcs/Cargo.toml +++ b/litellm-rust/crates/cache-gcs/Cargo.toml @@ -13,7 +13,6 @@ litellm-cache.workspace = true percent-encoding.workspace = true reqwest.workspace = true tokio.workspace = true -url.workspace = true [dev-dependencies] serde_json.workspace = true diff --git a/litellm-rust/crates/cache-gcs/src/cache.rs b/litellm-rust/crates/cache-gcs/src/cache.rs index 1097f3252b1..65282ac99d5 100644 --- a/litellm-rust/crates/cache-gcs/src/cache.rs +++ b/litellm-rust/crates/cache-gcs/src/cache.rs @@ -5,43 +5,18 @@ use litellm_cache::{ BaseCache, BatchCache, BatchEntry, CacheCodec, CacheConnectionResult, Error, ExactCacheContext, FlushCache, }; -use percent_encoding::{AsciiSet, CONTROLS, percent_encode}; +use percent_encoding::{AsciiSet, NON_ALPHANUMERIC, percent_encode}; use reqwest::Client; use crate::{GcpTokenSource, TokenSource}; pub const DEFAULT_ENDPOINT: &str = "https://storage.googleapis.com"; -const OBJECT_NAME_ENCODE_SET: &AsciiSet = &CONTROLS - .add(b' ') - .add(b'!') - .add(b'"') - .add(b'#') - .add(b'$') - .add(b'%') - .add(b'&') - .add(b'\'') - .add(b'(') - .add(b')') - .add(b'*') - .add(b'+') - .add(b',') - .add(b'/') - .add(b':') - .add(b';') - .add(b'<') - .add(b'=') - .add(b'>') - .add(b'?') - .add(b'@') - .add(b'[') - .add(b'\\') - .add(b']') - .add(b'^') - .add(b'`') - .add(b'{') - .add(b'|') - .add(b'}'); +const OBJECT_NAME_ENCODE_SET: &AsciiSet = &NON_ALPHANUMERIC + .remove(b'-') + .remove(b'_') + .remove(b'.') + .remove(b'~'); pub fn key_prefix(gcs_path: Option<&str>) -> String { match gcs_path { diff --git a/litellm-rust/crates/cache-gcs/tests/cache.rs b/litellm-rust/crates/cache-gcs/tests/cache.rs index 79e4ed100f8..45eecf01cec 100644 --- a/litellm-rust/crates/cache-gcs/tests/cache.rs +++ b/litellm-rust/crates/cache-gcs/tests/cache.rs @@ -117,6 +117,40 @@ fn key_prefix_normalizes_paths() { assert_eq!(key_prefix(Some("")), ""); } +#[tokio::test] +async fn object_names_use_python_quote_encoding() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/upload/storage/v1/b/bucket/o")) + .and(query_param("uploadType", "media")) + .respond_with(ResponseTemplate::new(200)) + .expect(2) + .mount(&server) + .await; + let cache = cache(&server, Some("p/")); + cache + .set_cache( + "a~b-c_d.e/f g%h", + json!({"value": "punctuation"}), + &ExactCacheContext::default(), + ) + .unwrap(); + cache + .set_cache( + "ключ", + json!({"value": "utf8"}), + &ExactCacheContext::default(), + ) + .unwrap(); + let requests = server.received_requests().await.unwrap(); + let queries: Vec<_> = requests + .iter() + .filter_map(|request| request.url.query()) + .collect(); + assert!(queries.contains(&"uploadType=media&name=p%2Fa~b-c_d.e%2Ff%20g%25h")); + assert!(queries.contains(&"uploadType=media&name=p%2F%D0%BA%D0%BB%D1%8E%D1%87")); +} + #[tokio::test] async fn ignores_ttl_and_writes_pipeline_concurrently() { let server = MockServer::start().await; diff --git a/tests/test_litellm_rust/support/fake_gcs.py b/tests/test_litellm_rust/support/fake_gcs.py new file mode 100644 index 00000000000..67eb61798b9 --- /dev/null +++ b/tests/test_litellm_rust/support/fake_gcs.py @@ -0,0 +1,152 @@ +from __future__ import annotations + +import json +import threading +from collections.abc import Mapping +from dataclasses import dataclass +from functools import partial +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from socket import socket +from types import MappingProxyType +from typing import Final, cast +from urllib.parse import unquote, urlsplit + + +@dataclass(frozen=True, slots=True) +class RecordedRequest: + method: str + path: str + query: str + headers: Mapping[str, str] + body: bytes + + +class _FakeGcsHandler(BaseHTTPRequestHandler): + def __init__( + self, + request: socket | tuple[bytes, socket], + client_address: tuple[str, int], + server: ThreadingHTTPServer, + *, + fake: FakeGcs, + ) -> None: + self._fake: Final = fake + super().__init__(request, client_address, server) + + def _handle(self) -> None: + parsed: Final = urlsplit(self.path) + content_length: Final = int(self.headers.get("Content-Length", "0")) + body: Final = self.rfile.read(content_length) if content_length else b"" + headers: Final = MappingProxyType( + {name.title(): value for name, value in self.headers.items()} + ) + self._fake.record( + RecordedRequest( + method=self.command, + path=parsed.path, + query=parsed.query, + headers=headers, + body=body, + ) + ) + if self.headers.get("Authorization") != f"Bearer {self._fake.token}": + self._send_json(401, {"error": "unauthorized"}) + return + + upload_prefix: Final = "/upload/storage/v1/b/" + download_prefix: Final = "/storage/v1/b/" + if parsed.path.startswith(upload_prefix) and parsed.path.endswith("/o"): + self._upload(parsed.path[len(upload_prefix) : -2], parsed.query, body) + return + if parsed.path.startswith(download_prefix): + self._download(parsed.path[len(download_prefix) :], parsed.query) + return + self._send_json(404, {"error": "not found"}) + + def _upload(self, path: str, query: str, body: bytes) -> None: + values: Final = { + unquote(pair.partition("=")[0]): unquote(pair.partition("=")[2]) + for pair in query.split("&") + if pair + } + if not path or values.get("uploadType") != "media" or "name" not in values: + self._send_json(404, {"error": "not found"}) + return + self._fake.put_object(path, values["name"], body) + self._send_json(200, {"name": values["name"], "bucket": path}) + + def _download(self, path: str, query: str) -> None: + bucket, separator, encoded_name = path.partition("/o/") + if not separator or query != "alt=media": + self._send_json(404, {"error": "not found"}) + return + name: Final = unquote(encoded_name) + if name.endswith("/server-error") or name == "server-error": + self._send_json(500, {"error": "server error"}) + return + body: Final = self._fake.get_object(bucket, name) + if body is None: + self._send_json(404, {"error": "not found"}) + return + self._send(200, body, "application/octet-stream") + + def _send_json(self, status: int, value: object) -> None: + payload: Final = json.dumps(value).encode() + self._send(status, payload, "application/json") + + def _send(self, status: int, body: bytes, content_type: str) -> None: + self.send_response(status) + self.send_header("Content-Type", content_type) + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, format: str, *args: object) -> None: + pass + + do_GET = _handle + do_POST = _handle + + +class FakeGcs: + def __init__(self) -> None: + self._objects: dict[tuple[str, str], bytes] = {} # mutable-ok: fake object store + self._requests: list[RecordedRequest] = [] # mutable-ok: recorded request history + self._server = ThreadingHTTPServer( + ("127.0.0.1", 0), + partial(_FakeGcsHandler, fake=self), + ) + self._worker = threading.Thread(target=self._server.serve_forever, daemon=True) + self._worker.start() + self.token: Final = "test-token" + + @property + def url(self) -> str: + address: Final = cast(tuple[str, int], self._server.server_address) + host, port = address + return f"http://{host}:{port}" + + @property + def objects(self) -> Mapping[tuple[str, str], bytes]: + return MappingProxyType(self._objects) + + @property + def requests(self) -> tuple[RecordedRequest, ...]: + return tuple(self._requests) + + def put(self, bucket: str, name: str, body: bytes) -> None: + self.put_object(bucket, name, body) + + def close(self) -> None: + self._server.shutdown() + self._server.server_close() + self._worker.join(timeout=5) + + def record(self, request: RecordedRequest) -> None: + self._requests.append(request) + + def put_object(self, bucket: str, name: str, body: bytes) -> None: + self._objects[(bucket, name)] = body + + def get_object(self, bucket: str, name: str) -> bytes | None: + return self._objects.get((bucket, name)) diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index c35cb1a20fb..cc409ebe49c 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -16,9 +16,11 @@ import redis import litellm from litellm.caching.caching import Cache, disable_cache, enable_cache, update_cache +from litellm.caching.gcs_cache import GCSCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.rust_bridge import _native from litellm.types.caching import LiteLLMCacheType +from tests.test_litellm_rust.support.fake_gcs import FakeGcs from tests.test_litellm_rust.support.isolation import rebound pytestmark: Final = pytest.mark.requires_rust_extension @@ -26,6 +28,7 @@ pytestmark: Final = pytest.mark.requires_rust_extension class CacheLookup(Protocol): def get_cache(self, **kwargs: object) -> object: ... + def flush_cache(self) -> object: ... def request(key: str = "key") -> dict[str, object]: @@ -45,6 +48,15 @@ def redis_url() -> Generator[str]: worker.join(timeout=5) +@pytest.fixture +def fake_gcs() -> Generator[FakeGcs]: + server: Final = FakeGcs() + try: + yield server + finally: + server.close() + + def test_existing_constructor_and_global_are_unchanged() -> None: facade: Final = Cache(type=LiteLLMCacheType.LOCAL) assert type(facade.cache) is InMemoryCache @@ -393,3 +405,219 @@ async def test_redis_facade_buffers_native_async_writes(redis_url: str) -> None: assert client.get("second") is not None await facade.cache.disconnect() client.close() + + +async def test_gcs_reads_python_entries_and_writes_python_compatible_objects( + fake_gcs: FakeGcs, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("GCS_PATH_SERVICE_ACCOUNT", raising=False) + monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) + response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}, "flag": True, "empty": None} + fake_gcs.put( + "bucket", + "cache/sync", + json.dumps({"timestamp": time.time(), "response": json.dumps(response)}).encode(), + ) + fake_gcs.put("bucket", "cache/async", json.dumps({"timestamp": time.time(), "response": response}).encode()) + fake_gcs.put("bucket", "cache/raw", json.dumps(response).encode()) + fake_gcs.put("bucket", "cache/invalid", b"not a cache entry") + binding: Final = _native._CacheTestResolver( + SimpleNamespace( + cache=_native._CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + ) + ).resolve() + + assert binding.lookup(request("sync")) == response + assert await binding.async_lookup(request("async")) == response + assert binding.lookup(request("raw")) == response + assert await binding.async_lookup(request("invalid")) is None + assert binding.lookup(request("missing")) is None + + await binding.async_store({**request("native"), "ttl_seconds": 12.0}, response) + stored: Final = fake_gcs.objects[("bucket", "cache/native")] + stored_value: Final = cast(dict[str, object], json.loads(stored)) + assert stored_value["response"] == response + assert isinstance(stored_value["timestamp"], float) + upload: Final = next(item for item in fake_gcs.requests if item.method == "POST") + assert upload.path == "/upload/storage/v1/b/bucket/o" + assert upload.query == "uploadType=media&name=cache%2Fnative" + assert upload.headers["Authorization"] == f"Bearer {fake_gcs.token}" + assert upload.headers["Content-Type"] == "application/json" + upload_text: Final = f"{upload.path}?{upload.query}{upload.headers}" + assert "ttl" not in upload_text.lower() + assert "expiry" not in upload_text.lower() + download: Final = next(item for item in fake_gcs.requests if item.path.endswith("/cache%2Fsync")) + assert download.path == "/storage/v1/b/bucket/o/cache%2Fsync" + assert download.query == "alt=media" + + binding.store(request("sync2"), response) + assert binding.lookup(request("sync2")) == response + assert GCSCache(bucket_name="bucket", gcs_path="cache").key_prefix == "cache/" + assert GCSCache(bucket_name="bucket", gcs_path="cache/").key_prefix == "cache/" + assert GCSCache(bucket_name="bucket").key_prefix == "" + + +async def test_gcs_batch_lookup_preserves_order_and_treats_malformed_entries_as_misses(fake_gcs: FakeGcs) -> None: + fake_gcs.put("bucket", "cache/hit", json.dumps({"timestamp": time.time(), "response": {"value": 1}}).encode()) + fake_gcs.put("bucket", "cache/invalid", b"not a cache entry") + binding: Final = _native._CacheTestResolver( + SimpleNamespace( + cache=_native._CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + ) + ).resolve() + requests: Final = [request("hit"), request("missing"), request("invalid")] + expected: Final = {"values": [{"value": 1}, None, None], "missing_indices": [1, 2]} + + assert await binding.async_lookup_batch(requests) == expected + assert binding.lookup_batch(requests) == expected + await binding.async_store_batch([request("first"), request("second")], [{"value": 1}, {"value": 2}]) + assert ("bucket", "cache/first") in fake_gcs.objects + assert ("bucket", "cache/second") in fake_gcs.objects + + +async def test_gcs_facade_binds_only_exact_matching_configuration( + fake_gcs: FakeGcs, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("GCS_PATH_SERVICE_ACCOUNT", raising=False) + monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) + monkeypatch.setenv("GOOGLE_APPLICATION_CREDENTIALS", "/nonexistent") + facade: Final = Cache(type=LiteLLMCacheType.GCS, gcs_bucket_name="bucket", gcs_path="cache/") + assert type(facade.cache) is GCSCache + + mismatched_bucket: Final = _native._CacheTestHandle.gcs( + "other", + gcs_path="cache", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + with pytest.raises(TypeError, match="buckets must match"): + mismatched_bucket._bind_facade(facade) + mismatched_prefix: Final = _native._CacheTestHandle.gcs( + "bucket", + gcs_path="x", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + with pytest.raises(TypeError, match="key prefixes must match"): + mismatched_prefix._bind_facade(facade) + mismatched_credentials: Final = _native._CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + path_service_account="sa.json", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + with pytest.raises(TypeError, match="credentials must match"): + mismatched_credentials._bind_facade(facade) + with pytest.raises(TypeError, match="types must match"): + _native._CacheTestHandle.memory()._bind_facade(facade) + + matching: Final = _native._CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + matching._bind_facade(facade) + resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)) + binding: Final = resolver.resolve() + assert binding.kind == "native" + await binding.async_store(request("native"), {"value": "native"}) + assert await binding.async_lookup(request("native")) == {"value": "native"} + assert cast(CacheLookup, facade).get_cache(cache_key="native") is None + + with rebound(facade.cache, "bucket_name", "other"): + assert resolver.resolve().kind == "python_callback" + with rebound(facade.cache, "key_prefix", "x/"): + assert resolver.resolve().kind == "python_callback" + with rebound(facade.cache, "path_service_account", "sa.json"): + assert resolver.resolve().kind == "python_callback" + def no_get_cache(*args: object, **kwargs: object) -> None: + return None + + with rebound(facade.cache, "get_cache", no_get_cache): + assert resolver.resolve().kind == "python_callback" + with rebound(facade, "ttl", 12): + assert resolver.resolve().kind == "python_callback" + + class CustomGcs(GCSCache): + pass + + with rebound(facade, "cache", CustomGcs(bucket_name="bucket", gcs_path="cache/")): + assert resolver.resolve().kind == "python_callback" + custom_facade: Final = Cache(type=LiteLLMCacheType.GCS, gcs_bucket_name="bucket", gcs_path="cache/") + with rebound(custom_facade, "cache", CustomGcs(bucket_name="bucket", gcs_path="cache/")): + with pytest.raises(TypeError, match="types must match"): + matching._bind_facade(custom_facade) + + missing_bucket: Final = Cache(type=LiteLLMCacheType.GCS) + with pytest.raises(TypeError, match="requires a configured bucket name"): + matching._bind_facade(missing_bucket) + + +async def test_gcs_flush_is_a_no_op_and_ping_is_not_implemented( + fake_gcs: FakeGcs, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("GCS_PATH_SERVICE_ACCOUNT", raising=False) + monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) + binding: Final = _native._CacheTestResolver( + SimpleNamespace( + cache=_native._CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + ) + ).resolve() + await binding.async_store(request("key"), {"value": "stored"}) + await binding.async_flush() + assert ("bucket", "cache/key") in fake_gcs.objects + assert await binding.async_lookup(request("key")) == {"value": "stored"} + with pytest.raises(NotImplementedError): + await binding.ping() + + facade: Final = Cache(type=LiteLLMCacheType.GCS, gcs_bucket_name="bucket", gcs_path="cache/") + with pytest.raises(AttributeError): + await facade.ping() + assert cast(CacheLookup, facade.cache).flush_cache() is None + + +async def test_gcs_unauthorized_and_server_errors_surface_as_runtime_errors(fake_gcs: FakeGcs) -> None: + wrong_token: Final = _native._CacheTestResolver( + SimpleNamespace( + cache=_native._CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + endpoint=fake_gcs.url, + token="wrong-token", + ) + ) + ).resolve() + with pytest.raises(RuntimeError): + wrong_token.lookup(request("missing")) + assert not fake_gcs.objects + + binding: Final = _native._CacheTestResolver( + SimpleNamespace( + cache=_native._CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + ) + ).resolve() + with pytest.raises(RuntimeError): + binding.lookup(request("server-error")) + assert binding.lookup(request("missing")) is None From 7c8f7d3f2c7d8356ac3662e153c40fd0a64e82fd Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:18:16 +0000 Subject: [PATCH 3/3] fix(rust): reduce native wheel size Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 80425d567c0..3e9d2d2fbb5 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -76,7 +76,7 @@ veil = "0.3.0" [profile.release] opt-level = 3 -lto = "thin" +lto = "fat" codegen-units = 1 panic = "unwind" debug = false