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

@@ -2476,7 +2476,6 @@ dependencies = [
  "reqwest 0.12.28",
  "serde_json",
  "tokio",
- "url",
  "wiremock",
 ]
 

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

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 {

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;

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))

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