Merge remote-tracking branch 'origin/main' into litellm_langfuse_sdk_v4

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

# Conflicts:
#	tests/logging_callback_tests/langfuse_expected_request_body/completion_with_vertex_call.json

.circleci/scripts/run_integration.sh

@@ -9,7 +9,6 @@ fi
 suite="${1:?integration suite required}"
 results="test-results/integration-${suite}"
 mkdir -p "$results"
-shard_timeout=11m
 integration_identity="$(.venv/bin/python -c 'import uuid; print(uuid.uuid4().hex)')"
 upstream_pid=""
 proxy_pid=""
@@ -131,6 +130,7 @@ start_proxy() {
   fi
   setsid env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" INTEGRATION_RUN_ID="$integration_identity" \
     DATABASE_URL="$DATABASE_URL" REDIS_HOST="$REDIS_HOST" REDIS_PORT="$REDIS_PORT" \
+    INTEGRATION_UPSTREAM_URL="$INTEGRATION_UPSTREAM_URL" \
     LITELLM_MASTER_KEY="$LITELLM_MASTER_KEY" LITELLM_SALT_KEY="$LITELLM_SALT_KEY" LITELLM_UI_PATH="$LITELLM_UI_PATH" PROXY_BASE_URL="http://127.0.0.1:$port" \
     LITELLM_MODE=PRODUCTION STORE_MODEL_IN_DB=True "${cost_map_env[@]}" \
     AWS_EC2_METADATA_DISABLED=true DO_NOT_TRACK=1 \
@@ -176,7 +176,7 @@ if [ "$suite" = browser ]; then
   exit 0
 fi
 
-timeout --signal=TERM --kill-after=20s "$shard_timeout" env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" \
+env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" \
   INTEGRATION_RUN_ID="$integration_identity" \
   DATABASE_URL="$DATABASE_URL" REDIS_HOST="$REDIS_HOST" REDIS_PORT="$REDIS_PORT" \
   INTEGRATION_PROXY_URL="$INTEGRATION_PROXY_URL" INTEGRATION_PEER_URL="$INTEGRATION_PEER_URL" \

.github/e2e-stack/select_tests.py

@@ -10,6 +10,8 @@ UNSUPPORTED: Final = re.compile(
     r"|^tests/e2e/batches/test_managed_files_enforcement_e2e\.py$"
     r"|^tests/e2e/guardrails/test_presidio_masking_e2e\.py$"
     r"|^tests/e2e/logging/test_otel_v2_langfuse_generation_output_e2e\.py$"
+    r"|^tests/e2e/logging/test_langsmith_batch_serialization_e2e\.py$"
+    r"|^tests/e2e/secret_manager/"
 )
 HARNESS: Final = re.compile(
     r"^tests/e2e/[A-Za-z0-9_.-]+\.(py|ini)$"

.github/e2e-stack/up.sh

@@ -24,6 +24,7 @@ DATABASE_USER="${E2E_DATABASE_USER:-litellm}"
 DATABASE_PASSWORD="${E2E_DATABASE_PASSWORD:-dbpassword9090}"
 DATABASE_NAME="${E2E_DATABASE_NAME:-litellm}"
 JAEGER_OTLP_PORT="${E2E_JAEGER_OTLP_PORT:-4318}"
+JAEGER_OTLP_TLS_PORT="${E2E_JAEGER_OTLP_TLS_PORT:-4319}"
 JAEGER_QUERY_PORT="${E2E_JAEGER_QUERY_PORT:-16686}"
 KEYCLOAK_PORT="${E2E_KEYCLOAK_PORT:-8081}"
 
@@ -122,7 +123,7 @@ SERVER_ENV=(
   "CONFIG_FILE_PATH=${CONFIG_PATH}"
   "STORE_MODEL_IN_DB=True"
   "OTEL_EXPORTER_OTLP_PROTOCOL=http/protobuf"
-  "OTEL_EXPORTER_OTLP_ENDPOINT=http://127.0.0.1:${JAEGER_OTLP_PORT}"
+  "OTEL_EXPORTER_OTLP_ENDPOINT=https://127.0.0.1:${JAEGER_OTLP_TLS_PORT}"
   "SSL_CERT_FILE=${CERTS_DIR}/ca-bundle.pem"
   "PYTHONPATH=${REPO_ROOT}"
   "JWT_PUBLIC_KEY_URL=http://127.0.0.1:${KEYCLOAK_PORT}/realms/litellm-e2e/protocol/openid-connect/certs"
@@ -147,16 +148,12 @@ start_server() {
   echo $! > "${PIDS_DIR}/${name}.pid"
 }
 
-start_server backend uv run --no-sync uvicorn backend.main:app --host 0.0.0.0 --port "${BACKEND_PORT}"
-start_server gateway-1 uv run --no-sync uvicorn gateway.main:app --workers 1 --host 0.0.0.0 --port "${GATEWAY_PORT_1}"
-start_server gateway-2 uv run --no-sync uvicorn gateway.main:app --workers 1 --host 0.0.0.0 --port "${GATEWAY_PORT_2}"
-
 if [[ "$(uname)" == "Linux" ]]; then
   NGINX_UPSTREAM_HOST=127.0.0.1
   NGINX_DOCKER_ARGS=(--network host)
 else
   NGINX_UPSTREAM_HOST=host.docker.internal
-  NGINX_DOCKER_ARGS=(-p "${LB_PORT}:${LB_PORT}")
+  NGINX_DOCKER_ARGS=(-p "${LB_PORT}:${LB_PORT}" -p "${JAEGER_OTLP_TLS_PORT}:${JAEGER_OTLP_TLS_PORT}")
 fi
 
 cat > "${STACK_DIR}/nginx.conf" <<EOF
@@ -186,12 +183,29 @@ http {
       proxy_send_timeout 600s;
     }
   }
+  server {
+    listen ${JAEGER_OTLP_TLS_PORT} ssl;
+    ssl_certificate /certs/server.crt;
+    ssl_certificate_key /certs/server.key;
+    client_max_body_size 100m;
+    location / {
+      proxy_pass http://${NGINX_UPSTREAM_HOST}:${JAEGER_OTLP_PORT};
+    }
+  }
 }
 EOF
 
 docker rm -f e2e-nginx >/dev/null 2>&1 || true
 docker run -d --name e2e-nginx "${NGINX_DOCKER_ARGS[@]}" \
-  -v "${STACK_DIR}/nginx.conf:/etc/nginx/nginx.conf:ro" "${NGINX_IMAGE}" >/dev/null
+  -v "${STACK_DIR}/nginx.conf:/etc/nginx/nginx.conf:ro" \
+  -v "${CERTS_DIR}:/certs:ro" "${NGINX_IMAGE}" >/dev/null
+
+wait_for "Jaeger OTLP TLS listener" \
+  "curl -sS --cacert ${CERTS_DIR}/ca.crt https://127.0.0.1:${JAEGER_OTLP_TLS_PORT}/ -o /dev/null -w '%{http_code}' | grep -qE '^[2345]'"
+
+start_server backend uv run --no-sync uvicorn backend.main:app --host 0.0.0.0 --port "${BACKEND_PORT}"
+start_server gateway-1 uv run --no-sync uvicorn gateway.main:app --workers 1 --host 0.0.0.0 --port "${GATEWAY_PORT_1}"
+start_server gateway-2 uv run --no-sync uvicorn gateway.main:app --workers 1 --host 0.0.0.0 --port "${GATEWAY_PORT_2}"
 
 wait_for "backend" "curl -fs http://127.0.0.1:${BACKEND_PORT}/health/liveliness >/dev/null" 300
 wait_for "gateway-1" "curl -fs http://127.0.0.1:${GATEWAY_PORT_1}/health/liveliness >/dev/null" 300
@@ -206,6 +220,7 @@ LITELLM_MASTER_KEY=${MASTER_KEY}
 REDIS_HOST=127.0.0.1
 REDIS_PORT=${REDIS_PORT}
 E2E_OTEL_QUERY_URL=http://127.0.0.1:${JAEGER_QUERY_PORT}
+E2E_OTEL_EXPORTER_ENDPOINT=https://127.0.0.1:${JAEGER_OTLP_TLS_PORT}
 E2E_KEYCLOAK_URL=http://127.0.0.1:${KEYCLOAK_PORT}
 E2E_KEYCLOAK_ADMIN_USER=admin
 E2E_KEYCLOAK_ADMIN_PASSWORD=e2e-ephemeral-idp-not-a-secret

.github/workflows/_test-unit-base.yml

@@ -4,7 +4,13 @@ on:
   workflow_call:
     inputs:
       test-path:
-        description: "Pytest path(s) to run"
+        description: >-
+          Space-separated pytest paths to run. A path that no longer exists is
+          dropped with a warning instead of being passed to pytest, because one
+          missing path makes pytest-xdist collect nothing and report exit 5, which
+          the step treats as a drained shard. Options are passed through as
+          written, so use the `--flag=value` form: a bare `--ignore path` would
+          have its path existence-checked like any other token.
         required: true
         type: string
       workers:
@@ -165,14 +171,22 @@ jobs:
           DIST: ${{ inputs.dist }}
           COVERAGE_CORE: sysmon
         run: |
-          found_path=false
-          for path in ${TEST_PATH}; do
-            if [ -e "${path%%::*}" ]; then
-              found_path=true
-              break
-            fi
+          pytest_args=()
+          existing_paths=0
+          for token in ${TEST_PATH:?}; do
+            case "${token}" in
+              -*) pytest_args+=("${token}") ;;
+              *)
+                if [ -e "${token%%::*}" ]; then
+                  pytest_args+=("${token}")
+                  existing_paths=$((existing_paths + 1))
+                else
+                  echo "::warning::${token} does not exist; drop it from this shard's test-path"
+                fi
+                ;;
+            esac
           done
-          if [ "$found_path" = false ]; then
+          if [ "${existing_paths}" -eq 0 ]; then
             echo "No path in TEST_PATH exists (${TEST_PATH}); nothing to run"
             exit 0
           fi
@@ -181,7 +195,7 @@ jobs:
             xdist_args=(-n "${WORKERS}" --dist="${DIST}")
           fi
           set +e
-          uv run --no-sync pytest ${TEST_PATH:?} \
+          uv run --no-sync pytest "${pytest_args[@]}" \
             --tb=short -vv \
             --maxfail="${MAX_FAILURES}" \
             "${xdist_args[@]}" \

.github/workflows/compat-matrix-image.yml

@@ -0,0 +1,33 @@
+name: Compat Matrix Image
+
+on:
+  pull_request:
+    paths:
+      - tests/e2e/claude_code/cron_vm/**
+      - .github/workflows/compat-matrix-image.yml
+  workflow_dispatch:
+
+permissions: {}
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  compat-matrix-image:
+    name: compat-matrix-image
+    runs-on: ubuntu-latest
+    timeout-minutes: 15
+    permissions:
+      contents: read
+    steps:
+      - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
+        with:
+          persist-credentials: false
+
+      - name: Build the Render cron image
+        run: docker build -f tests/e2e/claude_code/cron_vm/Dockerfile -t compat-matrix:${{ github.sha }} tests/e2e
+
+      - name: Run the pinned binaries as the cron user
+        run: |
+          docker run --rm compat-matrix:${{ github.sha }} bash -c 'set -e; whoami; claude --version; gh --version; uv --version'

.github/workflows/issue_fixed_comment.yml

@@ -6,8 +6,12 @@ on:
   workflow_dispatch:
     inputs:
       issue_number:
-        description: "Closed issue number to comment on manually."
-        required: true
+        description: "Closed issue number to comment on and close the superseded pull requests of. Ignored by a sweep."
+        required: false
+      sweep:
+        description: "Close every open pull request whose linked issues were all fixed on the default branch. Reads every open pull request, so run it at most once an hour."
+        type: boolean
+        default: false
   pull_request:
     paths:
       - .github/workflows/issue_fixed_comment.yml
@@ -39,16 +43,17 @@ jobs:
         with:
           bun-version: "1.4.0"
 
-      - name: Test the closer lookup, the release placement and the comment
+      - name: Test the closer lookup, the release placement, the comment and the superseded pull request close
         run: bun test scripts/comment-fixed-issue.test.ts
 
   comment-fixed-issue:
     if: github.event_name != 'pull_request' && github.repository == 'BerriAI/litellm'
     runs-on: ubuntu-latest
-    timeout-minutes: 5
+    timeout-minutes: 15
     permissions:
       contents: read
       issues: write
+      pull-requests: write
     steps:
       - name: Checkout scripts
         uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
@@ -59,13 +64,16 @@ jobs:
       - name: Setup Bun
         uses: oven-sh/setup-bun@0c5077e51419868618aeaa5fe8019c62421857d6 # v2.2.0
         with:
-          # Exact version, never latest: the next step holds an issues: write token
+          # Exact version, never latest: the next step holds issues: write and pull-requests: write tokens
           bun-version: "1.4.0"
 
-      - name: Name the release that carries the fix
+      - name: Name the release that carries the fix and close the pull requests it supersedes
+        shell: bash
         run: bun run scripts/comment-fixed-issue.ts | tee -a "${GITHUB_STEP_SUMMARY}"
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
           ISSUE_NUMBER: ${{ github.event.issue.number || github.event.inputs.issue_number }}
+          SWEEP: ${{ github.event.inputs.sweep }}
           DEFAULT_BRANCH: ${{ github.event.repository.default_branch }}
           DRY_RUN: ${{ vars.ISSUE_FIXED_COMMENT_ENABLED != 'true' }}
+          CLOSE_PRS_DRY_RUN: ${{ vars.ISSUE_FIXED_CLOSE_PRS_ENABLED != 'true' }}

.github/workflows/test-rust.yml

@@ -105,6 +105,16 @@ jobs:
         with:
           python-version: "3.12"
 
+      - uses: ./.github/actions/setup-uv-with-retries
+        with:
+          version: "0.10.9"
+
+      - name: Install Python dependencies for the bridge tests
+        working-directory: .
+        run: |
+          uv sync --frozen --no-install-project
+          echo "PYTHONPATH=$PWD/.venv/lib/$(ls .venv/lib)/site-packages" >> "$GITHUB_ENV"
+
       - run: rustup toolchain install --no-self-update
 
       - uses: taiki-e/install-action@d438492cf8a250514fa2d34b30bc3c0dc37c65ff # v2.87.8

.github/workflows/test-unit.yml

@@ -107,26 +107,18 @@ jobs:
               tests/test_litellm/batches
               tests/test_litellm/secret_managers
               tests/test_litellm/a2a_protocol
-              tests/test_litellm/anthropic_interface
               tests/test_litellm/chat_completions
               tests/test_litellm/completion_extras
-              tests/test_litellm/compression
               tests/test_litellm/containers
               tests/test_litellm/endpoints
-              tests/test_litellm/models
-              tests/test_litellm/repositories
               tests/test_litellm/images
               tests/test_litellm/interactions
               tests/test_litellm/messages
               tests/test_litellm/ocr
               tests/test_litellm/passthrough
               tests/test_litellm/rag
-              tests/test_litellm/realtime_api
               tests/test_litellm/rerank_api
               tests/test_litellm/rust_bridge
-              tests/test_litellm/sandbox
-              tests/test_litellm/skills
-              tests/test_litellm/test_router
               tests/test_litellm/vector_stores
               tests/test_litellm/videos
               tests/test_litellm/test_*.py

Dockerfile

@@ -1,10 +1,10 @@
 # syntax=docker/dockerfile:1.7
 
 # Base image for building
-ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
+ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
 
 # Runtime image
-ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
+ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
 # Pinned by digest like the other base images; bump explicitly on Node upgrades.
 ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43

backend/Dockerfile

@@ -1,5 +1,5 @@
-ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
-ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
+ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
+ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
 
 FROM $UV_IMAGE AS uvbin

docker/Dockerfile.database

@@ -1,10 +1,10 @@
 # syntax=docker/dockerfile:1.7
 
 # Base image for building
-ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
+ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
 
 # Runtime image
-ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
+ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
 # Pinned by digest like the other base images; bump explicitly on Node upgrades.
 ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43

docker/Dockerfile.non_root

@@ -1,8 +1,8 @@
 # syntax=docker/dockerfile:1.7
 
 # Base images
-ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
-ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
+ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
+ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
 ARG PROXY_EXTRAS_SOURCE=published
 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
 # Pinned by digest like the other base images; bump explicitly on Node upgrades.

enterprise/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "litellm-enterprise"
-version = "0.1.69"
+version = "0.1.70"
 description = "Package for LiteLLM Enterprise features"
 readme = "README.md"
 requires-python = ">=3.9"
@@ -26,7 +26,7 @@ required-version = ">=0.10.9"
 module-root = ""
 
 [tool.commitizen]
-version = "0.1.69"
+version = "0.1.70"
 version_files = [
     "pyproject.toml:^version",
     "../pyproject.toml:litellm-enterprise==",

gateway/Dockerfile

@@ -1,5 +1,5 @@
-ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
-ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
+ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
+ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
 # Checksum from https://www.pgbouncer.org/downloads/ (the Wolfi repo only carries 1.24.x)
 ARG PGBOUNCER_VERSION=1.25.2

gateway/routes/allowlist.py

@@ -131,6 +131,7 @@ GATEWAY_EXACT_PATHS: frozenset[str] = frozenset(
         "/redoc",
         "/test",
         "/debug/memory/summary",
+        "/api/event_logging/batch",
     }
 )
 

helm/litellm/templates/ingress.yaml

@@ -61,7 +61,7 @@
     "/v1/fine-tuning" "/fine-tuning" "/v1/responses" "/responses" "/v1/threads" "/threads"
     "/v1/assistants" "/assistants" "/v1/vector_stores" "/vector_stores" "/v1/indexes"
     "/v1/models" "/models" "/openai" "/engines"
-    "/v1/messages" "/messages" "/v1/skills" "/v1/a2a" "/a2a"
+    "/v1/messages" "/messages" "/v1/skills" "/v1/a2a" "/a2a" "/api/event_logging"
     "/v1/rerank" "/v2/rerank" "/rerank" "/v1/ocr" "/ocr" "/v1/rag" "/rag"
     "/v1/video" "/v1/videos" "/video" "/videos" "/v1/search" "/search"
     "/v1/containers" "/containers" "/v1/evals" "/v1/memory" "/queue/chat"

litellm-proxy-extras/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "litellm-proxy-extras"
-version = "0.4.100"
+version = "0.4.101"
 description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
 readme = "README.md"
 requires-python = ">=3.9"
@@ -26,7 +26,7 @@ required-version = ">=0.10.9"
 module-root = ""
 
 [tool.commitizen]
-version = "0.4.100"
+version = "0.4.101"
 version_files = [
     "pyproject.toml:^version",
     "../pyproject.toml:litellm-proxy-extras==",

litellm-rust/Cargo.lock

@@ -2727,9 +2727,13 @@ dependencies = [
  "litellm-auth-types",
  "litellm-cache",
  "litellm-cache-response",
+ "litellm-cache-testing",
+ "reqwest 0.12.28",
+ "rstest",
  "serde_json",
  "tokio",
  "url",
+ "wiremock",
 ]
 
 [[package]]
@@ -2737,6 +2741,7 @@ name = "litellm-cache-disk"
 version = "0.1.0"
 dependencies = [
  "litellm-cache",
+ "litellm-cache-testing",
  "py_literal",
  "rand 0.8.7",
  "rstest",
@@ -2755,8 +2760,10 @@ dependencies = [
  "litellm-auth-gcp",
  "litellm-auth-types",
  "litellm-cache",
+ "litellm-cache-testing",
  "percent-encoding",
  "reqwest 0.12.28",
+ "rstest",
  "serde_json",
  "tokio",
  "wiremock",
@@ -2767,8 +2774,8 @@ name = "litellm-cache-memory"
 version = "0.1.0"
 dependencies = [
  "litellm-cache",
+ "litellm-cache-testing",
  "rstest",
- "serde_json",
  "tokio",
 ]
 
@@ -2776,9 +2783,10 @@ dependencies = [
 name = "litellm-cache-qdrant-semantic"
 version = "0.1.0"
 dependencies = [
+ "futures-executor",
  "futures-util",
  "litellm-cache",
- "litellm-cache-response",
+ "litellm-cache-testing",
  "qdrant-client",
  "reqwest 0.12.28",
  "rstest",
@@ -2797,9 +2805,11 @@ name = "litellm-cache-redis"
 version = "0.1.0"
 dependencies = [
  "litellm-cache",
+ "litellm-cache-testing",
  "r2d2",
  "redis",
  "redis-test",
+ "rstest",
  "serde_json",
  "tokio",
 ]
@@ -2810,10 +2820,10 @@ version = "0.1.0"
 dependencies = [
  "litellm-cache",
  "litellm-cache-redis",
- "litellm-cache-response",
- "r2d2",
+ "litellm-cache-testing",
  "redis",
  "redis-test",
+ "rstest",
  "serde_json",
  "sha2 0.10.9",
  "tokio",
@@ -2829,6 +2839,7 @@ dependencies = [
  "py_literal",
  "redis",
  "redis-test",
+ "rstest",
  "serde",
  "serde_json",
  "sha2 0.10.9",
@@ -2841,22 +2852,35 @@ version = "0.1.0"
 dependencies = [
  "aws-credential-types",
  "aws-sdk-s3",
+ "aws-smithy-runtime-api",
  "aws-smithy-types",
  "aws-types",
+ "futures-util",
+ "http 1.4.2",
  "litellm-auth-aws",
  "litellm-cache",
+ "litellm-cache-testing",
+ "reqwest 0.12.28",
+ "rstest",
  "serde_json",
  "tokio",
  "wiremock",
 ]
 
+[[package]]
+name = "litellm-cache-testing"
+version = "0.1.0"
+dependencies = [
+ "litellm-cache",
+]
+
 [[package]]
 name = "litellm-cache-valkey-semantic"
 version = "0.1.0"
 dependencies = [
  "litellm-cache",
  "litellm-cache-redis",
- "litellm-cache-response",
+ "litellm-cache-testing",
  "redis",
  "redis-test",
  "rstest",
@@ -2922,6 +2946,7 @@ name = "litellm-core-utils"
 version = "0.1.0"
 dependencies = [
  "fancy-regex 0.19.2",
+ "litellm-tracing",
  "litellm-types",
  "rstest",
  "serde",
@@ -2932,6 +2957,14 @@ dependencies = [
  "url",
 ]
 
+[[package]]
+name = "litellm-cost"
+version = "0.1.0"
+dependencies = [
+ "criterion",
+ "proptest",
+]
+
 [[package]]
 name = "litellm-framing"
 version = "0.1.0"
@@ -3022,6 +3055,20 @@ dependencies = [
  "url",
 ]
 
+[[package]]
+name = "litellm-model-catalog"
+version = "0.1.0"
+dependencies = [
+ "criterion",
+ "indexmap 2.14.0",
+ "litellm-model-catalog",
+ "rstest",
+ "schemars 1.2.2",
+ "serde",
+ "serde_json",
+ "thiserror 2.0.19",
+]
+
 [[package]]
 name = "litellm-python-bridge"
 version = "0.1.0"
@@ -3029,6 +3076,7 @@ dependencies = [
  "aws-sdk-secretsmanager",
  "bytes",
  "criterion",
+ "fancy-regex 0.19.2",
  "futures-util",
  "litellm-auth",
  "litellm-auth-aws",
@@ -3047,6 +3095,7 @@ dependencies = [
  "litellm-callbacks-legacy-python",
  "litellm-core",
  "litellm-core-utils",
+ "litellm-host",
  "litellm-host-python",
  "litellm-http",
  "litellm-llms",
@@ -3054,6 +3103,7 @@ dependencies = [
  "litellm-secrets-aws",
  "litellm-secrets-types",
  "litellm-token-counter",
+ "litellm-tracing",
  "litellm-types",
  "pyo3",
  "pyo3-async-runtimes",
@@ -3072,6 +3122,21 @@ dependencies = [
  "wiremock",
 ]
 
+[[package]]
+name = "litellm-python-compat"
+version = "0.1.0"
+dependencies = [
+ "criterion",
+ "hex",
+ "num-bigint 0.4.8",
+ "num-traits",
+ "rstest",
+ "serde",
+ "serde-pickle",
+ "serde_json",
+ "thiserror 2.0.19",
+]
+
 [[package]]
 name = "litellm-secrets"
 version = "0.1.0"
@@ -3111,11 +3176,11 @@ dependencies = [
  "litellm-auth-aws",
  "litellm-core-utils",
  "litellm-secrets-types",
+ "litellm-tracing",
  "rstest",
  "serde_json",
  "thiserror 2.0.19",
  "tokio",
- "tracing",
  "veil",
  "wiremock",
 ]
@@ -3147,6 +3212,7 @@ dependencies = [
  "base64 0.22.1",
  "litellm-core-utils",
  "litellm-secrets-types",
+ "litellm-tracing",
  "moka",
  "percent-encoding",
  "reqwest 0.12.28",
@@ -3155,7 +3221,6 @@ dependencies = [
  "serde_json",
  "thiserror 2.0.19",
  "tokio",
- "tracing",
  "veil",
  "wiremock",
 ]
@@ -3270,6 +3335,19 @@ dependencies = [
  "tiktoken-rs",
 ]
 
+[[package]]
+name = "litellm-tracing"
+version = "0.1.0"
+dependencies = [
+ "fancy-regex 0.19.2",
+ "percent-encoding",
+ "rstest",
+ "serde_json",
+ "tokio",
+ "tracing",
+ "tracing-subscriber",
+]
+
 [[package]]
 name = "litellm-types"
 version = "0.1.0"
@@ -4709,10 +4787,23 @@ checksum = "687274d293b6cdc6e73e0fee520bf2049650090d7164f87672d212a3c530cf4a"
 dependencies = [
  "dyn-clone",
  "ref-cast",
+ "schemars_derive",
  "serde",
  "serde_json",
 ]
 
+[[package]]
+name = "schemars_derive"
+version = "1.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d98c67716b46af2f0b8cf752abc930f6f9aecfbf671ecfb531db8a31dbe4e2ba"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "serde_derive_internals",
+ "syn 3.0.0",
+]
+
 [[package]]
 name = "scopeguard"
 version = "1.2.0"
@@ -4801,6 +4892,17 @@ dependencies = [
  "syn 3.0.0",
 ]
 
+[[package]]
+name = "serde_derive_internals"
+version = "0.30.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f852137cce035d6a4df67ccce505ff6b3e9fd3a10e3e52b24dc71e650bb1a9bd"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 3.0.0",
+]
+
 [[package]]
 name = "serde_json"
 version = "1.0.150"

litellm-rust/Cargo.toml

@@ -9,6 +9,8 @@ license = "MIT"
 repository = "https://github.com/BerriAI/litellm"
 
 [workspace.dependencies]
+litellm-tracing = { path = "crates/tracing" }
+tracing = "0.1"
 litellm-core = { path = "crates/core" }
 litellm-host = { path = "crates/host" }
 litellm-callbacks-legacy-python = { path = "crates/callbacks-legacy-python" }
@@ -39,6 +41,7 @@ litellm-cache-disk = { path = "crates/cache-disk" }
 litellm-cache-redis-semantic = { path = "crates/cache-redis-semantic" }
 litellm-cache-response = { path = "crates/cache-response" }
 litellm-cache-qdrant-semantic = { path = "crates/cache-qdrant-semantic" }
+litellm-cache-testing = { path = "crates/cache-testing" }
 litellm-token-counter = { path = "crates/token-counter" }
 litellm-token-counter-fast = { path = "crates/token-counter-fast" }
 litellm-token-counter-huggingface = { path = "crates/token-counter-huggingface" }

litellm-rust/crates/cache-azure-blob/Cargo.toml

@@ -14,9 +14,14 @@ async-trait = "0.1"
 azure_core = "1.1.0"
 azure_storage_blob = "1.1.0"
 futures-util.workspace = true
+reqwest.workspace = true
 tokio.workspace = true
 url.workspace = true
 
 [dev-dependencies]
 litellm-cache-response.workspace = true
+litellm-cache-testing.workspace = true
+rstest.workspace = true
 serde_json.workspace = true
+tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }
+wiremock = "0.6.5"

litellm-rust/crates/cache-azure-blob/src/cache.rs

@@ -3,7 +3,7 @@ use std::{sync::Arc, time::Duration};
 use azure_core::{
     credentials::TokenCredential,
     error::ErrorKind,
-    http::{ClientOptions, RequestContent},
+    http::{ClientOptions, RequestContent, Transport},
 };
 use azure_storage_blob::{
     BlobContainerClient, BlobContainerClientOptions,
@@ -11,13 +11,12 @@ use azure_storage_blob::{
 };
 use futures_util::{TryStreamExt, future::try_join_all};
 use litellm_cache::{
-    BaseCache, BatchCache, CacheCodec, CacheConnectionResult, CacheConnectionStatus, Error,
-    ExactCacheContext, FlushCache,
+    BaseCache, BatchCache, CacheCodec, DisconnectCache, Error, ExactCacheContext, FlushCache,
 };
 use tokio::runtime::Handle;
 use url::Url;
 
-use crate::credential::AzureBlobCredential;
+use crate::{credential::AzureBlobCredential, transport::ReqwestTransport};
 
 pub struct AzureBlobCache<C> {
     container: BlobContainerClient,
@@ -28,9 +27,11 @@ pub struct AzureBlobCache<C> {
 }
 
 impl<C: CacheCodec> AzureBlobCache<C> {
+    /// `http` is the host's pooled client; the SDK sends every request through it.
     pub async fn connect(
         account_url: &str,
         container: &str,
+        http: reqwest::Client,
         codec: C,
         runtime: Handle,
     ) -> Result<Self, Error> {
@@ -38,7 +39,10 @@ impl<C: CacheCodec> AzureBlobCache<C> {
             account_url,
             container,
             Some(Arc::new(AzureBlobCredential::default())),
-            ClientOptions::default(),
+            ClientOptions {
+                transport: Some(Transport::new(Arc::new(ReqwestTransport(http)))),
+                ..ClientOptions::default()
+            },
             codec,
             runtime,
         )
@@ -152,7 +156,11 @@ impl<C: CacheCodec> AzureBlobCache<C> {
     }
 
     fn block_on<T>(&self, future: impl Future<Output = T>) -> T {
-        self.runtime.block_on(future)
+        if Handle::try_current().is_ok() {
+            tokio::task::block_in_place(|| self.runtime.block_on(future))
+        } else {
+            self.runtime.block_on(future)
+        }
     }
 }
 
@@ -217,25 +225,6 @@ impl<C: CacheCodec> BaseCache for AzureBlobCache<C> {
         .await
         .map(drop)
     }
-
-    async fn disconnect(&self) -> Result<(), Error> {
-        Ok(())
-    }
-
-    async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
-        Ok(match self.container.get_properties(None).await {
-            Ok(_) => CacheConnectionResult {
-                status: CacheConnectionStatus::Success,
-                message: "Azure Blob cache connection test successful".into(),
-                error: None,
-            },
-            Err(error) => CacheConnectionResult {
-                status: CacheConnectionStatus::Failed,
-                message: format!("Azure Blob connection failed: {error}"),
-                error: Some(error.to_string()),
-            },
-        })
-    }
 }
 
 impl<C: CacheCodec> BatchCache for AzureBlobCache<C> {}
@@ -250,5 +239,10 @@ impl<C: CacheCodec> FlushCache for AzureBlobCache<C> {
     }
 }
 
-#[cfg(test)]
-mod tests;
+impl<C: CacheCodec> DisconnectCache for AzureBlobCache<C> {
+    /// Python closes its two SDK clients; the Rust clients hold no connection of their own
+    /// (the pooled transport belongs to the host), so there is nothing to release.
+    async fn disconnect(&self) -> Result<(), Error> {
+        Ok(())
+    }
+}

litellm-rust/crates/cache-azure-blob/src/cache/tests.rs

@@ -1,746 +0,0 @@
-use std::{
-    collections::BTreeMap,
-    sync::{Arc, Mutex},
-    time::Duration,
-};
-
-use azure_core::http::{
-    AsyncRawResponse, Body, ClientOptions, HttpClient, Method, Request, StatusCode, Transport,
-    headers::{HeaderName, Headers},
-};
-use litellm_cache::{
-    BaseCache, BatchCache, BatchEntry, CacheConnectionStatus, Error, ExactCacheContext, FlushCache,
-};
-use litellm_cache_response::{
-    CacheEntry, CacheKeyField, CacheKeyInput, ResponseCache, ResponseCacheCodec,
-    ResponseCacheRequest, cache_key,
-};
-use serde_json::json;
-use tokio::runtime::Runtime;
-
-use super::AzureBlobCache;
-
-const ACCOUNT_URL: &str = "https://example.blob.core.windows.net";
-const CONTAINER: &str = "litellm-cache";
-const IF_NONE_MATCH: HeaderName = HeaderName::from_static("if-none-match");
-const ERROR_CODE: HeaderName = HeaderName::from_static("x-ms-error-code");
-
-#[derive(Clone, Debug, PartialEq, Eq)]
-struct RecordedRequest {
-    method: Method,
-    path: String,
-    query: String,
-    if_none_match: Option<String>,
-}
-
-#[derive(Default)]
-struct FakeState {
-    container_exists: bool,
-    blobs: BTreeMap<String, Vec<u8>>,
-    requests: Vec<RecordedRequest>,
-    failing: bool,
-    precondition_conflicts: bool,
-}
-
-#[derive(Clone, Default)]
-struct FakeBlobService {
-    state: Arc<Mutex<FakeState>>,
-}
-
-impl std::fmt::Debug for FakeBlobService {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        f.write_str("FakeBlobService")
-    }
-}
-
-impl FakeBlobService {
-    fn with_existing_container() -> Self {
-        let service = Self::default();
-        service.state.lock().unwrap().container_exists = true;
-        service
-    }
-
-    fn blob(&self, name: &str) -> Option<Vec<u8>> {
-        self.state.lock().unwrap().blobs.get(name).cloned()
-    }
-
-    fn blob_names(&self) -> Vec<String> {
-        self.state.lock().unwrap().blobs.keys().cloned().collect()
-    }
-
-    fn seed_blob(&self, name: &str, bytes: &[u8]) {
-        self.state
-            .lock()
-            .unwrap()
-            .blobs
-            .insert(name.to_string(), bytes.to_vec());
-    }
-
-    fn set_failing(&self, failing: bool) {
-        self.state.lock().unwrap().failing = failing;
-    }
-
-    fn set_precondition_conflicts(&self, enabled: bool) {
-        self.state.lock().unwrap().precondition_conflicts = enabled;
-    }
-
-    fn requests(&self) -> Vec<RecordedRequest> {
-        self.state.lock().unwrap().requests.clone()
-    }
-
-    fn container_exists(&self) -> bool {
-        self.state.lock().unwrap().container_exists
-    }
-
-    fn respond(status: StatusCode, error_code: Option<&str>, body: Vec<u8>) -> AsyncRawResponse {
-        let mut headers = Headers::new();
-        if let Some(code) = error_code {
-            headers.insert(ERROR_CODE, code.to_string());
-        }
-        AsyncRawResponse::from_bytes(status, headers, body)
-    }
-
-    fn list_body(state: &FakeState) -> Vec<u8> {
-        let mut xml = String::from(
-            r#"<?xml version="1.0" encoding="utf-8"?><EnumerationResults ServiceEndpoint="https://example.blob.core.windows.net/" ContainerName="litellm-cache"><Blobs>"#,
-        );
-        for name in state.blobs.keys() {
-            xml.push_str(&format!(
-                "<Blob><Name>{name}</Name><Properties><BlobType>BlockBlob</BlobType></Properties></Blob>"
-            ));
-        }
-        xml.push_str("</Blobs><NextMarker /></EnumerationResults>");
-        xml.into_bytes()
-    }
-}
-
-#[async_trait::async_trait]
-impl HttpClient for FakeBlobService {
-    async fn execute_request(&self, request: &Request) -> azure_core::Result<AsyncRawResponse> {
-        let mut state = self.state.lock().unwrap();
-        let path = request.url().path().to_string();
-        let query = request.url().query().unwrap_or_default().to_string();
-        let if_none_match = request
-            .headers()
-            .get_optional_str(&IF_NONE_MATCH)
-            .map(str::to_owned);
-        state.requests.push(RecordedRequest {
-            method: request.method(),
-            path: path.clone(),
-            query: query.clone(),
-            if_none_match: if_none_match.clone(),
-        });
-        if state.failing {
-            return Ok(Self::respond(
-                StatusCode::Forbidden,
-                Some("AuthorizationFailure"),
-                Vec::new(),
-            ));
-        }
-        let container_path = format!("/{CONTAINER}");
-        let blob_name = path
-            .strip_prefix(&format!("{container_path}/"))
-            .map(str::to_owned);
-        let is_container = path == container_path && query.contains("restype=container");
-        let response = match (request.method(), is_container, blob_name) {
-            (Method::Put, true, None) if state.container_exists => Self::respond(
-                StatusCode::Conflict,
-                Some("ContainerAlreadyExists"),
-                Vec::new(),
-            ),
-            (Method::Put, true, None) => {
-                state.container_exists = true;
-                Self::respond(StatusCode::Created, None, Vec::new())
-            }
-            (Method::Get, true, None) if query.contains("comp=list") => {
-                Self::respond(StatusCode::Ok, None, Self::list_body(&state))
-            }
-            (Method::Get, true, None) if state.container_exists => {
-                Self::respond(StatusCode::Ok, None, Vec::new())
-            }
-            (Method::Get, true, None) => {
-                Self::respond(StatusCode::NotFound, Some("ContainerNotFound"), Vec::new())
-            }
-            (Method::Put, false, Some(name)) => {
-                if if_none_match.as_deref() == Some("*") && state.blobs.contains_key(&name) {
-                    if state.precondition_conflicts {
-                        Self::respond(
-                            StatusCode::PreconditionFailed,
-                            Some("ConditionNotMet"),
-                            Vec::new(),
-                        )
-                    } else {
-                        Self::respond(StatusCode::Conflict, Some("BlobAlreadyExists"), Vec::new())
-                    }
-                } else {
-                    let bytes = match request.body() {
-                        Body::Bytes(bytes) => bytes.to_vec(),
-                        Body::SeekableStream(_) => panic!("unexpected streaming upload"),
-                    };
-                    state.blobs.insert(name, bytes);
-                    Self::respond(StatusCode::Created, None, Vec::new())
-                }
-            }
-            (Method::Get, false, Some(name)) => match state.blobs.get(&name) {
-                Some(bytes) => Self::respond(StatusCode::Ok, None, bytes.clone()),
-                None => Self::respond(StatusCode::NotFound, Some("BlobNotFound"), Vec::new()),
-            },
-            (Method::Delete, false, Some(name)) => match state.blobs.remove(&name) {
-                Some(_) => Self::respond(StatusCode::Accepted, None, Vec::new()),
-                None => Self::respond(StatusCode::NotFound, Some("BlobNotFound"), Vec::new()),
-            },
-            (method, _, _) => panic!("unexpected request {method:?} {path}?{query}"),
-        };
-        Ok(response)
-    }
-}
-
-struct Fixture {
-    runtime: Runtime,
-    service: FakeBlobService,
-    cache: Arc<AzureBlobCache<ResponseCacheCodec>>,
-}
-
-impl Fixture {
-    fn new(service: FakeBlobService) -> Self {
-        let runtime = Runtime::new().unwrap();
-        let cache = runtime
-            .block_on(Self::connect(&service, runtime.handle().clone()))
-            .unwrap();
-        Self {
-            runtime,
-            service,
-            cache: Arc::new(cache),
-        }
-    }
-
-    async fn connect(
-        service: &FakeBlobService,
-        handle: tokio::runtime::Handle,
-    ) -> Result<AzureBlobCache<ResponseCacheCodec>, Error> {
-        AzureBlobCache::connect_with_options(
-            ACCOUNT_URL,
-            CONTAINER,
-            None,
-            ClientOptions {
-                transport: Some(Transport::new(Arc::new(service.clone()))),
-                ..ClientOptions::default()
-            },
-            ResponseCacheCodec,
-            handle,
-        )
-        .await
-    }
-
-    fn response_cache(&self) -> ResponseCache<AzureBlobCache<ResponseCacheCodec>> {
-        ResponseCache::new(self.cache.clone())
-    }
-
-    fn stored_json(&self, key: &str) -> serde_json::Value {
-        serde_json::from_slice(&self.service.blob(key).expect("blob should exist")).unwrap()
-    }
-}
-
-fn request(model: &str) -> ResponseCacheRequest {
-    ResponseCacheRequest::new(CacheKeyInput {
-        fields: vec![CacheKeyField {
-            name: "model".into(),
-            value: Some(model.into()),
-            api_parameter: true,
-            internal_parameter: false,
-        }],
-        preset: None,
-        namespace: None,
-        include_provider_parameters: false,
-    })
-}
-
-fn now() -> Duration {
-    Duration::from_secs(1_700_000_000)
-}
-
-fn entry(value: serde_json::Value) -> CacheEntry {
-    CacheEntry {
-        timestamp: Some(1_700_000_000.5),
-        response: value,
-    }
-}
-
-fn no_ttl() -> ExactCacheContext {
-    ExactCacheContext::default()
-}
-
-fn with_ttl(seconds: u64) -> ExactCacheContext {
-    ExactCacheContext {
-        ttl: Some(Duration::from_secs(seconds)),
-    }
-}
-
-#[test]
-fn connect_creates_the_container_once() {
-    let fixture = Fixture::new(FakeBlobService::default());
-    assert!(fixture.service.container_exists());
-    assert_eq!(
-        fixture.service.requests(),
-        vec![RecordedRequest {
-            method: Method::Put,
-            path: format!("/{CONTAINER}"),
-            query: "restype=container".into(),
-            if_none_match: None,
-        }]
-    );
-    assert_eq!(fixture.cache.account_url(), ACCOUNT_URL);
-    assert_eq!(fixture.cache.container_name(), CONTAINER);
-}
-
-#[test]
-fn connect_accepts_an_existing_container() {
-    let fixture = Fixture::new(FakeBlobService::with_existing_container());
-    assert!(fixture.service.container_exists());
-    assert_eq!(fixture.service.requests().len(), 1);
-}
-
-#[test]
-fn connect_accepts_account_urls_with_trailing_slash() {
-    let runtime = Runtime::new().unwrap();
-    let service = FakeBlobService::default();
-    let cache = runtime
-        .block_on(AzureBlobCache::connect_with_options(
-            "https://example.blob.core.windows.net/",
-            CONTAINER,
-            None,
-            ClientOptions {
-                transport: Some(Transport::new(Arc::new(service.clone()))),
-                ..ClientOptions::default()
-            },
-            ResponseCacheCodec,
-            runtime.handle().clone(),
-        ))
-        .unwrap();
-    assert_eq!(service.requests()[0].path, format!("/{CONTAINER}"));
-    assert_eq!(cache.account_url(), "https://example.blob.core.windows.net");
-}
-
-#[test]
-fn connect_keeps_account_url_query_parameters_on_the_container_path() {
-    let runtime = Runtime::new().unwrap();
-    let service = FakeBlobService::default();
-    runtime
-        .block_on(AzureBlobCache::connect_with_options(
-            "https://example.blob.core.windows.net/?sv=2024-01-01&sig=abc",
-            CONTAINER,
-            None,
-            ClientOptions {
-                transport: Some(Transport::new(Arc::new(service.clone()))),
-                ..ClientOptions::default()
-            },
-            ResponseCacheCodec,
-            runtime.handle().clone(),
-        ))
-        .unwrap();
-    let create = &service.requests()[0];
-    assert_eq!(create.path, format!("/{CONTAINER}"));
-    assert!(create.query.contains("sig=abc"));
-}
-
-#[test]
-fn connect_surfaces_service_failures() {
-    let runtime = Runtime::new().unwrap();
-    let service = FakeBlobService::default();
-    service.set_failing(true);
-    let result = runtime.block_on(Fixture::connect(&service, runtime.handle().clone()));
-    assert!(matches!(result, Err(Error::Unavailable)));
-}
-
-#[test]
-fn sync_set_and_get_round_trip_python_json_shape() {
-    let fixture = Fixture::new(FakeBlobService::default());
-    let value = entry(json!({"choices": [{"message": {"content": "héllo 🌍"}}]}));
-    fixture
-        .cache
-        .set_cache("key-1", value.clone(), &no_ttl())
-        .unwrap();
-
-    assert_eq!(
-        fixture.stored_json("key-1"),
-        json!({
-            "timestamp": 1_700_000_000.5,
-            "response": {"choices": [{"message": {"content": "héllo 🌍"}}]}
-        })
-    );
-    assert_eq!(
-        fixture.cache.get_cache("key-1", &no_ttl()).unwrap(),
-        Some(value)
-    );
-}
-
-#[test]
-fn sync_set_does_not_overwrite_an_existing_blob() {
-    let fixture = Fixture::new(FakeBlobService::default());
-    fixture
-        .cache
-        .set_cache("key", entry(json!({"v": "first"})), &no_ttl())
-        .unwrap();
-    fixture
-        .cache
-        .set_cache("key", entry(json!({"v": "second"})), &no_ttl())
-        .unwrap();
-
-    assert_eq!(
-        fixture.stored_json("key")["response"],
-        json!({"v": "first"})
-    );
-    let uploads: Vec<_> = fixture
-        .service
-        .requests()
-        .into_iter()
-        .filter(|request| request.method == Method::Put && request.path.ends_with("/key"))
-        .collect();
-    assert_eq!(uploads.len(), 2);
-    assert!(
-        uploads
-            .iter()
-            .all(|request| request.if_none_match.as_deref() == Some("*"))
-    );
-}
-
-#[test]
-fn sync_set_treats_a_precondition_conflict_as_an_existing_blob() {
-    let fixture = Fixture::new(FakeBlobService::default());
-    fixture.service.set_precondition_conflicts(true);
-    fixture
-        .cache
-        .set_cache("key", entry(json!({"v": "first"})), &no_ttl())
-        .unwrap();
-    fixture
-        .cache
-        .set_cache("key", entry(json!({"v": "second"})), &no_ttl())
-        .unwrap();
-
-    assert_eq!(
-        fixture.stored_json("key")["response"],
-        json!({"v": "first"})
-    );
-}
-
-#[test]
-fn async_set_overwrites_an_existing_blob() {
-    let fixture = Fixture::new(FakeBlobService::default());
-    fixture.runtime.block_on(async {
-        fixture
-            .cache
-            .async_set_cache("key", entry(json!({"v": "first"})), no_ttl())
-            .await
-            .unwrap();
-        fixture
-            .cache
-            .async_set_cache("key", entry(json!({"v": "second"})), no_ttl())
-            .await
-            .unwrap();
-        assert_eq!(
-            fixture
-                .cache
-                .async_get_cache("key", &no_ttl())
-                .await
-                .unwrap(),
-            Some(entry(json!({"v": "second"})))
-        );
-    });
-    assert_eq!(
-        fixture.stored_json("key")["response"],
-        json!({"v": "second"})
-    );
-    assert!(
-        fixture
-            .service
-            .requests()
-            .iter()
-            .filter(|request| request.method == Method::Put && request.path.ends_with("/key"))
-            .all(|request| request.if_none_match.is_none())
-    );
-}
-
-#[test]
-fn missing_blobs_are_misses() {
-    let fixture = Fixture::new(FakeBlobService::default());
-    assert_eq!(fixture.cache.get_cache("absent", &no_ttl()).unwrap(), None);
-    assert_eq!(
-        fixture
-            .runtime
-            .block_on(fixture.cache.async_get_cache("absent", &no_ttl()))
-            .unwrap(),
-        None
-    );
-}
-
-#[test]
-fn ttl_is_ignored_and_entries_never_expire() {
-    let fixture = Fixture::new(FakeBlobService::default());
-    assert_eq!(fixture.cache.get_ttl(&with_ttl(1)), None);
-    assert_eq!(fixture.cache.get_ttl(&no_ttl()), None);
-
-    fixture
-        .cache
-        .set_cache("key", entry(json!("value")), &with_ttl(1))
-        .unwrap();
-    std::thread::sleep(Duration::from_millis(1100));
-    assert_eq!(
-        fixture.cache.get_cache("key", &with_ttl(1)).unwrap(),
-        Some(entry(json!("value")))
-    );
-    assert!(
-        fixture
-            .service
-            .requests()
-            .iter()
-            .all(|request| !request.query.contains("expiry"))
-    );
-}
-
-#[test]
-fn malformed_blobs_are_invalid_entries_and_response_cache_misses() {
-    let fixture = Fixture::new(FakeBlobService::default());
-    fixture.service.seed_blob("broken-json", b"{not json");
-    fixture
-        .service
-        .seed_blob("broken-utf8", &[0xff, 0xfe, 0x22]);
-    fixture
-        .service
-        .seed_blob("wrong-shape", br#"{"timestamp": "yesterday"}"#);
-
-    for key in ["broken-json", "broken-utf8", "wrong-shape"] {
-        assert!(matches!(
-            fixture.cache.get_cache(key, &no_ttl()),
-            Err(Error::InvalidEntry)
-        ));
-    }
-
-    let response_cache = fixture.response_cache();
-    let broken = request("broken");
-    fixture
-        .service
-        .seed_blob(&cache_key(&broken.key), b"{not json");
-    assert_eq!(response_cache.lookup(&broken, now()).unwrap(), None);
-    assert_eq!(
-        fixture
-            .runtime
-            .block_on(response_cache.async_lookup(&broken, now()))
-            .unwrap(),
-        None
-    );
-}
-
-#[test]
-fn batch_get_preserves_order_and_marks_misses_and_invalid_entries() {
-    let fixture = Fixture::new(FakeBlobService::default());
-    fixture
-        .cache
-        .set_cache("a", entry(json!("A")), &no_ttl())
-        .unwrap();
-    fixture
-        .cache
-        .set_cache("c", entry(json!("C")), &no_ttl())
-        .unwrap();
-    fixture.service.seed_blob("bad", b"nope");
-    let keys = ["c", "missing", "a", "bad"].map(String::from);
-
-    let sync = fixture.cache.batch_get_cache(&keys, &no_ttl()).unwrap();
-    assert_eq!(
-        sync,
-        vec![
-            BatchEntry::Hit(entry(json!("C"))),
-            BatchEntry::Miss,
-            BatchEntry::Hit(entry(json!("A"))),
-            BatchEntry::Invalid,
-        ]
-    );
-
-    let asynchronous = fixture
-        .runtime
-        .block_on(fixture.cache.async_batch_get_cache(keys.to_vec(), no_ttl()))
-        .unwrap();
-    assert_eq!(asynchronous, sync);
-
-    let response_cache = fixture.response_cache();
-    let requests = [request("hit"), request("missing"), request("bad")];
-    response_cache
-        .store(&requests[0], json!("HIT"), now())
-        .unwrap();
-    fixture
-        .service
-        .seed_blob(&cache_key(&requests[2].key), b"nope");
-    let hits = response_cache.lookup_batch(&requests, now()).unwrap();
-    assert_eq!(hits.values, vec![Some(json!("HIT")), None, None]);
-    assert_eq!(hits.missing_indices, vec![1, 2]);
-    let async_hits = fixture
-        .runtime
-        .block_on(response_cache.async_lookup_batch(&requests, now()))
-        .unwrap();
-    assert_eq!(async_hits.values, hits.values);
-}
-
-#[test]
-fn async_pipeline_writes_every_entry_with_overwrite() {
-    let fixture = Fixture::new(FakeBlobService::default());
-    fixture.service.seed_blob("k2", b"stale");
-    fixture
-        .runtime
-        .block_on(fixture.cache.async_set_cache_pipeline(
-            vec![
-                ("k1".into(), entry(json!({"n": 1}))),
-                ("k2".into(), entry(json!({"n": 2}))),
-                ("k3".into(), entry(json!({"n": 3}))),
-            ],
-            with_ttl(30),
-        ))
-        .unwrap();
-    assert_eq!(fixture.service.blob_names(), ["k1", "k2", "k3"]);
-    assert_eq!(fixture.stored_json("k2")["response"], json!({"n": 2}));
-}
-
-#[test]
-fn flush_deletes_every_blob_in_the_container() {
-    let fixture = Fixture::new(FakeBlobService::default());
-    for key in ["x", "y", "z"] {
-        fixture
-            .cache
-            .set_cache(key, entry(json!(key)), &no_ttl())
-            .unwrap();
-    }
-    fixture.cache.flush_cache().unwrap();
-    assert!(fixture.service.blob_names().is_empty());
-    assert!(fixture.service.container_exists());
-
-    fixture
-        .cache
-        .set_cache("again", entry(json!(1)), &no_ttl())
-        .unwrap();
-    fixture
-        .runtime
-        .block_on(fixture.cache.async_flush_cache())
-        .unwrap();
-    assert!(fixture.service.blob_names().is_empty());
-}
-
-#[test]
-fn service_failures_map_to_unavailable() {
-    let fixture = Fixture::new(FakeBlobService::default());
-    fixture.service.set_failing(true);
-    assert!(matches!(
-        fixture.cache.get_cache("key", &no_ttl()),
-        Err(Error::Unavailable)
-    ));
-    assert!(matches!(
-        fixture.cache.set_cache("key", entry(json!(1)), &no_ttl()),
-        Err(Error::Unavailable)
-    ));
-    assert!(matches!(
-        fixture.cache.flush_cache(),
-        Err(Error::Unavailable)
-    ));
-    assert!(matches!(
-        fixture.runtime.block_on(
-            fixture
-                .cache
-                .async_set_cache_pipeline(vec![("k".into(), entry(json!(1)))], no_ttl())
-        ),
-        Err(Error::Unavailable)
-    ));
-}
-
-#[test]
-fn test_connection_reports_container_reachability() {
-    let fixture = Fixture::new(FakeBlobService::default());
-    let ok = fixture
-        .runtime
-        .block_on(fixture.cache.test_connection())
-        .unwrap();
-    assert_eq!(ok.status, CacheConnectionStatus::Success);
-    assert!(ok.error.is_none());
-
-    fixture.service.set_failing(true);
-    let failed = fixture
-        .runtime
-        .block_on(fixture.cache.test_connection())
-        .unwrap();
-    assert_eq!(failed.status, CacheConnectionStatus::Failed);
-    assert!(failed.error.is_some());
-}
-
-#[test]
-fn disconnect_is_idempotent_and_keeps_data() {
-    let fixture = Fixture::new(FakeBlobService::default());
-    fixture
-        .cache
-        .set_cache("key", entry(json!(1)), &no_ttl())
-        .unwrap();
-    fixture.runtime.block_on(async {
-        fixture.cache.disconnect().await.unwrap();
-        fixture.cache.disconnect().await.unwrap();
-    });
-    assert_eq!(
-        fixture.cache.get_cache("key", &no_ttl()).unwrap(),
-        Some(entry(json!(1)))
-    );
-}
-
-#[test]
-fn response_cache_stores_and_reads_through_the_backend() {
-    let fixture = Fixture::new(FakeBlobService::default());
-    let response_cache = fixture.response_cache();
-    let mut request = request("gpt");
-    request.context = with_ttl(60);
-    let response = json!({"id": "chatcmpl-1"});
-    response_cache
-        .store(&request, response.clone(), now())
-        .unwrap();
-    assert_eq!(
-        fixture.stored_json(&cache_key(&request.key)),
-        json!({"timestamp": 1_700_000_000.0, "response": {"id": "chatcmpl-1"}})
-    );
-    assert_eq!(
-        response_cache
-            .lookup(&request, now() + Duration::from_secs(3600))
-            .unwrap(),
-        Some(response.clone())
-    );
-    assert_eq!(
-        fixture
-            .runtime
-            .block_on(response_cache.async_lookup(&request, now() + Duration::from_secs(3600)))
-            .unwrap(),
-        Some(response.clone())
-    );
-    fixture.runtime.block_on(async {
-        response_cache
-            .async_store(&request, json!("replaced"), now())
-            .await
-            .unwrap();
-        assert_eq!(
-            response_cache.async_lookup(&request, now()).await.unwrap(),
-            Some(json!("replaced"))
-        );
-        response_cache.async_flush().await.unwrap();
-        assert_eq!(
-            response_cache.async_lookup(&request, now()).await.unwrap(),
-            None
-        );
-    });
-}
-
-#[test]
-fn non_object_responses_are_written_serialized_like_python() {
-    let fixture = Fixture::new(FakeBlobService::default());
-    fixture
-        .cache
-        .set_cache("s", entry(json!("plain")), &no_ttl())
-        .unwrap();
-    assert_eq!(
-        fixture.stored_json("s"),
-        json!({"timestamp": 1_700_000_000.5, "response": "\"plain\""})
-    );
-    assert_eq!(
-        fixture.cache.get_cache("s", &no_ttl()).unwrap(),
-        Some(entry(json!("plain")))
-    );
-}

litellm-rust/crates/cache-azure-blob/src/lib.rs

@@ -1,5 +1,7 @@
 mod cache;
 mod credential;
+mod transport;
 
 pub use cache::AzureBlobCache;
 pub use credential::AzureBlobCredential;
+pub use transport::ReqwestTransport;

litellm-rust/crates/cache-azure-blob/src/transport.rs

@@ -0,0 +1,49 @@
+use azure_core::{
+    error::ErrorKind,
+    http::{
+        AsyncRawResponse, Body, HttpClient, Request,
+        headers::{HeaderName, HeaderValue, Headers},
+    },
+};
+use futures_util::TryStreamExt;
+
+#[derive(Debug)]
+pub struct ReqwestTransport(pub reqwest::Client);
+
+#[async_trait::async_trait]
+impl HttpClient for ReqwestTransport {
+    async fn execute_request(&self, request: &Request) -> azure_core::Result<AsyncRawResponse> {
+        let method = reqwest::Method::from_bytes(request.method().as_ref().as_bytes())
+            .map_err(|error| azure_core::Error::new(ErrorKind::Other, error))?;
+        let mut outgoing = self.0.request(method, request.url().as_str());
+        for (name, value) in request.headers().iter() {
+            outgoing = outgoing.header(name.as_str(), value.as_str());
+        }
+        let outgoing = match request.body().clone() {
+            Body::Bytes(bytes) => outgoing.body(bytes),
+            Body::SeekableStream(stream) => outgoing.body(reqwest::Body::wrap_stream(stream)),
+        };
+        let response = outgoing.send().await.map_err(|error| {
+            let kind = if error.is_connect() {
+                ErrorKind::Connection
+            } else {
+                ErrorKind::Io
+            };
+            azure_core::Error::new(kind, error)
+        })?;
+        let status = response.status().as_u16().into();
+        let mut headers = Headers::new();
+        for (name, value) in response.headers() {
+            if let Ok(value) = value.to_str() {
+                headers.insert(
+                    HeaderName::from(name.as_str().to_owned()),
+                    HeaderValue::from(value.to_owned()),
+                );
+            }
+        }
+        let body = response
+            .bytes_stream()
+            .map_err(|error| azure_core::Error::new(ErrorKind::Io, error));
+        Ok(AsyncRawResponse::new(status, headers, Box::pin(body)))
+    }
+}

litellm-rust/crates/cache-azure-blob/tests/cache.rs

@@ -0,0 +1,494 @@
+mod support;
+
+use std::{sync::Arc, time::Duration};
+
+use azure_core::http::Method;
+use litellm_cache::{
+    BaseCache, BatchCache, BatchEntry, DisconnectCache, Error, ExactCacheContext, FlushCache,
+};
+use litellm_cache_azure_blob::AzureBlobCache;
+use litellm_cache_response::{
+    CacheEntry, CacheKeyField, CacheKeyInput, ResponseCache, ResponseCacheCodec,
+    ResponseCacheRequest, cache_key,
+};
+use rstest::{fixture, rstest};
+use serde_json::json;
+use support::{ACCOUNT_URL, CONTAINER, FakeBlobService, RecordedRequest};
+use tokio::runtime::Runtime;
+
+type Fixture = support::Fixture<ResponseCacheCodec>;
+
+#[fixture]
+fn fixture() -> Fixture {
+    Fixture::new(FakeBlobService::default(), ResponseCacheCodec)
+}
+
+fn response_cache(fixture: &Fixture) -> ResponseCache<AzureBlobCache<ResponseCacheCodec>> {
+    ResponseCache::new(fixture.cache.clone())
+}
+
+fn request(model: &str) -> ResponseCacheRequest {
+    ResponseCacheRequest::new(CacheKeyInput {
+        fields: vec![CacheKeyField {
+            name: "model".into(),
+            value: Some(model.into()),
+            api_parameter: true,
+            internal_parameter: false,
+        }],
+        preset: None,
+        namespace: None,
+        include_provider_parameters: false,
+    })
+}
+
+fn now() -> Duration {
+    Duration::from_secs(1_700_000_000)
+}
+
+fn entry(value: serde_json::Value) -> CacheEntry {
+    CacheEntry {
+        timestamp: Some(1_700_000_000.5),
+        response: value,
+    }
+}
+
+fn no_ttl() -> ExactCacheContext {
+    ExactCacheContext::default()
+}
+
+fn with_ttl(seconds: u64) -> ExactCacheContext {
+    ExactCacheContext {
+        ttl: Some(Duration::from_secs(seconds)),
+    }
+}
+
+fn connect_to(account_url: &str) -> (FakeBlobService, AzureBlobCache<ResponseCacheCodec>) {
+    let runtime = Runtime::new().unwrap();
+    let service = FakeBlobService::default();
+    let cache = runtime
+        .block_on(support::connect(
+            &service,
+            account_url,
+            ResponseCacheCodec,
+            runtime.handle().clone(),
+        ))
+        .unwrap();
+    (service, cache)
+}
+
+#[rstest]
+fn connect_creates_the_container_once(fixture: Fixture) {
+    assert!(fixture.service.container_exists());
+    assert_eq!(
+        fixture.service.requests(),
+        vec![RecordedRequest {
+            method: Method::Put,
+            path: format!("/{CONTAINER}"),
+            query: "restype=container".into(),
+            if_none_match: None,
+        }]
+    );
+    assert_eq!(fixture.cache.account_url(), ACCOUNT_URL);
+    assert_eq!(fixture.cache.container_name(), CONTAINER);
+}
+
+#[rstest]
+fn connect_accepts_an_existing_container() {
+    let fixture = Fixture::new(
+        FakeBlobService::with_existing_container(),
+        ResponseCacheCodec,
+    );
+    assert!(fixture.service.container_exists());
+    assert_eq!(fixture.service.requests().len(), 1);
+}
+
+#[rstest]
+fn connect_accepts_account_urls_with_trailing_slash() {
+    let (service, cache) = connect_to("https://example.blob.core.windows.net/");
+    assert_eq!(service.requests()[0].path, format!("/{CONTAINER}"));
+    assert_eq!(cache.account_url(), "https://example.blob.core.windows.net");
+}
+
+#[rstest]
+fn connect_keeps_account_url_query_parameters_on_the_container_path() {
+    let (service, _) = connect_to("https://example.blob.core.windows.net/?sv=2024-01-01&sig=abc");
+    let create = &service.requests()[0];
+    assert_eq!(create.path, format!("/{CONTAINER}"));
+    assert!(create.query.contains("sig=abc"));
+}
+
+#[rstest]
+fn connect_surfaces_service_failures() {
+    let runtime = Runtime::new().unwrap();
+    let service = FakeBlobService::default();
+    service.set_failing(true);
+    let result = runtime.block_on(support::connect(
+        &service,
+        ACCOUNT_URL,
+        ResponseCacheCodec,
+        runtime.handle().clone(),
+    ));
+    assert!(matches!(result, Err(Error::Unavailable)));
+}
+
+#[rstest]
+fn sync_set_and_get_round_trip_python_json_shape(fixture: Fixture) {
+    let value = entry(json!({"choices": [{"message": {"content": "héllo 🌍"}}]}));
+    fixture
+        .cache
+        .set_cache("key-1", value.clone(), &no_ttl())
+        .unwrap();
+
+    assert_eq!(
+        fixture.stored_json("key-1"),
+        json!({
+            "timestamp": 1_700_000_000.5,
+            "response": {"choices": [{"message": {"content": "héllo 🌍"}}]}
+        })
+    );
+    assert_eq!(
+        fixture.cache.get_cache("key-1", &no_ttl()).unwrap(),
+        Some(value)
+    );
+}
+
+#[rstest]
+#[case::blob_already_exists(false)]
+#[case::precondition_conflict(true)]
+fn sync_set_does_not_overwrite_an_existing_blob(fixture: Fixture, #[case] precondition: bool) {
+    fixture.service.set_precondition_conflicts(precondition);
+    fixture
+        .cache
+        .set_cache("key", entry(json!({"v": "first"})), &no_ttl())
+        .unwrap();
+    fixture
+        .cache
+        .set_cache("key", entry(json!({"v": "second"})), &no_ttl())
+        .unwrap();
+
+    assert_eq!(
+        fixture.stored_json("key")["response"],
+        json!({"v": "first"})
+    );
+    let uploads: Vec<_> = fixture
+        .service
+        .requests()
+        .into_iter()
+        .filter(|request| request.method == Method::Put && request.path.ends_with("/key"))
+        .collect();
+    assert_eq!(uploads.len(), 2);
+    assert!(
+        uploads
+            .iter()
+            .all(|request| request.if_none_match.as_deref() == Some("*"))
+    );
+}
+
+#[rstest]
+fn async_set_overwrites_an_existing_blob(fixture: Fixture) {
+    fixture.runtime.block_on(async {
+        fixture
+            .cache
+            .async_set_cache("key", entry(json!({"v": "first"})), no_ttl())
+            .await
+            .unwrap();
+        fixture
+            .cache
+            .async_set_cache("key", entry(json!({"v": "second"})), no_ttl())
+            .await
+            .unwrap();
+        assert_eq!(
+            fixture
+                .cache
+                .async_get_cache("key", &no_ttl())
+                .await
+                .unwrap(),
+            Some(entry(json!({"v": "second"})))
+        );
+    });
+    assert_eq!(
+        fixture.stored_json("key")["response"],
+        json!({"v": "second"})
+    );
+    assert!(
+        fixture
+            .service
+            .requests()
+            .iter()
+            .filter(|request| request.method == Method::Put && request.path.ends_with("/key"))
+            .all(|request| request.if_none_match.is_none())
+    );
+}
+
+#[rstest]
+fn missing_blobs_are_misses(fixture: Fixture) {
+    assert_eq!(fixture.cache.get_cache("absent", &no_ttl()).unwrap(), None);
+    assert_eq!(
+        fixture
+            .runtime
+            .block_on(fixture.cache.async_get_cache("absent", &no_ttl()))
+            .unwrap(),
+        None
+    );
+}
+
+#[rstest]
+fn ttl_is_ignored_and_entries_never_expire(fixture: Fixture) {
+    assert_eq!(fixture.cache.get_ttl(&with_ttl(1)), None);
+    assert_eq!(fixture.cache.get_ttl(&no_ttl()), None);
+
+    fixture
+        .cache
+        .set_cache("key", entry(json!("value")), &with_ttl(1))
+        .unwrap();
+    std::thread::sleep(Duration::from_millis(1100));
+    assert_eq!(
+        fixture.cache.get_cache("key", &with_ttl(1)).unwrap(),
+        Some(entry(json!("value")))
+    );
+    assert!(
+        fixture
+            .service
+            .requests()
+            .iter()
+            .all(|request| !request.query.contains("expiry"))
+    );
+}
+
+#[rstest]
+#[case::broken_json("broken-json", b"{not json".as_slice())]
+#[case::broken_utf8("broken-utf8", &[0xff, 0xfe, 0x22])]
+#[case::wrong_shape("wrong-shape", br#"{"timestamp": "yesterday"}"#.as_slice())]
+fn malformed_blobs_are_invalid_entries(fixture: Fixture, #[case] key: &str, #[case] bytes: &[u8]) {
+    fixture.service.seed_blob(key, bytes);
+    assert!(matches!(
+        fixture.cache.get_cache(key, &no_ttl()),
+        Err(Error::InvalidEntry)
+    ));
+}
+
+#[rstest]
+fn malformed_blobs_are_response_cache_misses(fixture: Fixture) {
+    let response_cache = response_cache(&fixture);
+    let broken = request("broken");
+    fixture
+        .service
+        .seed_blob(&cache_key(&broken.key), b"{not json");
+    assert_eq!(response_cache.lookup(&broken, now()).unwrap(), None);
+    assert_eq!(
+        fixture
+            .runtime
+            .block_on(response_cache.async_lookup(&broken, now()))
+            .unwrap(),
+        None
+    );
+}
+
+#[rstest]
+fn batch_get_preserves_order_and_marks_misses_and_invalid_entries(fixture: Fixture) {
+    fixture
+        .cache
+        .set_cache("a", entry(json!("A")), &no_ttl())
+        .unwrap();
+    fixture
+        .cache
+        .set_cache("c", entry(json!("C")), &no_ttl())
+        .unwrap();
+    fixture.service.seed_blob("bad", b"nope");
+    let keys = ["c", "missing", "a", "bad"].map(String::from);
+
+    let sync = fixture.cache.batch_get_cache(&keys, &no_ttl()).unwrap();
+    assert_eq!(
+        sync,
+        vec![
+            BatchEntry::Hit(entry(json!("C"))),
+            BatchEntry::Miss,
+            BatchEntry::Hit(entry(json!("A"))),
+            BatchEntry::Invalid,
+        ]
+    );
+
+    let asynchronous = fixture
+        .runtime
+        .block_on(fixture.cache.async_batch_get_cache(keys.to_vec(), no_ttl()))
+        .unwrap();
+    assert_eq!(asynchronous, sync);
+
+    let response_cache = response_cache(&fixture);
+    let requests = [request("hit"), request("missing"), request("bad")];
+    response_cache
+        .store(&requests[0], json!("HIT"), now())
+        .unwrap();
+    fixture
+        .service
+        .seed_blob(&cache_key(&requests[2].key), b"nope");
+    let hits = response_cache.lookup_batch(&requests, now()).unwrap();
+    assert_eq!(hits.values, vec![Some(json!("HIT")), None, None]);
+    assert_eq!(hits.missing_indices, vec![1, 2]);
+    let async_hits = fixture
+        .runtime
+        .block_on(response_cache.async_lookup_batch(&requests, now()))
+        .unwrap();
+    assert_eq!(async_hits.values, hits.values);
+}
+
+#[rstest]
+fn async_pipeline_writes_every_entry_with_overwrite(fixture: Fixture) {
+    fixture.service.seed_blob("k2", b"stale");
+    fixture
+        .runtime
+        .block_on(fixture.cache.async_set_cache_pipeline(
+            vec![
+                ("k1".into(), entry(json!({"n": 1}))),
+                ("k2".into(), entry(json!({"n": 2}))),
+                ("k3".into(), entry(json!({"n": 3}))),
+            ],
+            with_ttl(30),
+        ))
+        .unwrap();
+    assert_eq!(fixture.service.blob_names(), ["k1", "k2", "k3"]);
+    assert_eq!(fixture.stored_json("k2")["response"], json!({"n": 2}));
+}
+
+#[rstest]
+fn flush_deletes_every_blob_in_the_container(fixture: Fixture) {
+    for key in ["x", "y", "z"] {
+        fixture
+            .cache
+            .set_cache(key, entry(json!(key)), &no_ttl())
+            .unwrap();
+    }
+    fixture.cache.flush_cache().unwrap();
+    assert!(fixture.service.blob_names().is_empty());
+    assert!(fixture.service.container_exists());
+
+    fixture
+        .cache
+        .set_cache("again", entry(json!(1)), &no_ttl())
+        .unwrap();
+    fixture
+        .runtime
+        .block_on(fixture.cache.async_flush_cache())
+        .unwrap();
+    assert!(fixture.service.blob_names().is_empty());
+}
+
+#[rstest]
+fn service_failures_map_to_unavailable(fixture: Fixture) {
+    fixture.service.set_failing(true);
+    assert!(matches!(
+        fixture.cache.get_cache("key", &no_ttl()),
+        Err(Error::Unavailable)
+    ));
+    assert!(matches!(
+        fixture.cache.set_cache("key", entry(json!(1)), &no_ttl()),
+        Err(Error::Unavailable)
+    ));
+    assert!(matches!(
+        fixture.cache.flush_cache(),
+        Err(Error::Unavailable)
+    ));
+    assert!(matches!(
+        fixture.runtime.block_on(
+            fixture
+                .cache
+                .async_set_cache_pipeline(vec![("k".into(), entry(json!(1)))], no_ttl())
+        ),
+        Err(Error::Unavailable)
+    ));
+}
+
+#[rstest]
+fn disconnect_is_idempotent_and_keeps_data(fixture: Fixture) {
+    fixture
+        .cache
+        .set_cache("key", entry(json!(1)), &no_ttl())
+        .unwrap();
+    fixture.runtime.block_on(async {
+        fixture.cache.disconnect().await.unwrap();
+        fixture.cache.disconnect().await.unwrap();
+    });
+    assert_eq!(
+        fixture.cache.get_cache("key", &no_ttl()).unwrap(),
+        Some(entry(json!(1)))
+    );
+}
+
+#[rstest]
+fn response_cache_stores_and_reads_through_the_backend(fixture: Fixture) {
+    let response_cache = response_cache(&fixture);
+    let mut request = request("gpt");
+    request.context = with_ttl(60);
+    let response = json!({"id": "chatcmpl-1"});
+    response_cache
+        .store(&request, response.clone(), now())
+        .unwrap();
+    assert_eq!(
+        fixture.stored_json(&cache_key(&request.key)),
+        json!({"timestamp": 1_700_000_000.0, "response": {"id": "chatcmpl-1"}})
+    );
+    assert_eq!(
+        response_cache
+            .lookup(&request, now() + Duration::from_secs(3600))
+            .unwrap(),
+        Some(response.clone())
+    );
+    assert_eq!(
+        fixture
+            .runtime
+            .block_on(response_cache.async_lookup(&request, now() + Duration::from_secs(3600)))
+            .unwrap(),
+        Some(response.clone())
+    );
+    fixture.runtime.block_on(async {
+        response_cache
+            .async_store(&request, json!("replaced"), now())
+            .await
+            .unwrap();
+        assert_eq!(
+            response_cache.async_lookup(&request, now()).await.unwrap(),
+            Some(json!("replaced"))
+        );
+        response_cache.async_flush().await.unwrap();
+        assert_eq!(
+            response_cache.async_lookup(&request, now()).await.unwrap(),
+            None
+        );
+    });
+}
+
+#[rstest]
+fn non_object_responses_are_written_serialized_like_python(fixture: Fixture) {
+    fixture
+        .cache
+        .set_cache("s", entry(json!("plain")), &no_ttl())
+        .unwrap();
+    assert_eq!(
+        fixture.stored_json("s"),
+        json!({"timestamp": 1_700_000_000.5, "response": "\"plain\""})
+    );
+    assert_eq!(
+        fixture.cache.get_cache("s", &no_ttl()).unwrap(),
+        Some(entry(json!("plain")))
+    );
+}
+
+#[rstest]
+#[tokio::test(flavor = "multi_thread")]
+async fn sync_methods_block_inside_a_multi_thread_runtime() {
+    let service = FakeBlobService::default();
+    let cache = support::connect(
+        &service,
+        ACCOUNT_URL,
+        ResponseCacheCodec,
+        tokio::runtime::Handle::current(),
+    )
+    .await
+    .map(Arc::new)
+    .unwrap();
+    cache.set_cache("key", entry(json!(1)), &no_ttl()).unwrap();
+    assert_eq!(
+        cache.get_cache("key", &no_ttl()).unwrap(),
+        Some(entry(json!(1)))
+    );
+}

litellm-rust/crates/cache-azure-blob/tests/contract.rs

@@ -0,0 +1,81 @@
+mod support;
+
+use litellm_cache::{ExactCacheContext, JsonCodec};
+use litellm_cache_azure_blob::AzureBlobCache;
+use litellm_cache_testing as contract;
+use rstest::{fixture, rstest};
+use serde_json::{Value, json};
+use support::{ACCOUNT_URL, FakeBlobService};
+use tokio::runtime::Handle;
+
+#[fixture]
+async fn azure() -> AzureBlobCache<JsonCodec<Value>> {
+    support::connect(
+        &FakeBlobService::default(),
+        ACCOUNT_URL,
+        JsonCodec::new(),
+        Handle::current(),
+    )
+    .await
+    .unwrap()
+}
+
+#[fixture]
+fn context() -> ExactCacheContext {
+    ExactCacheContext::default()
+}
+
+const PREFIX: &str = "contract:";
+
+// `overwrite_replaces` does not apply: sync `set_cache` never overwrites a blob, as in Python.
+
+#[rstest]
+#[tokio::test(flavor = "multi_thread")]
+async fn hit_and_miss(
+    #[future(awt)] azure: AzureBlobCache<JsonCodec<Value>>,
+    context: ExactCacheContext,
+) {
+    contract::hit_and_miss(&azure, context, PREFIX, json!({"answer": 42})).await;
+}
+
+#[rstest]
+#[tokio::test(flavor = "multi_thread")]
+async fn sync_async_equivalence(
+    #[future(awt)] azure: AzureBlobCache<JsonCodec<Value>>,
+    context: ExactCacheContext,
+) {
+    contract::sync_async_equivalence(&azure, context, PREFIX, json!("first"), json!([2])).await;
+}
+
+#[rstest]
+#[tokio::test(flavor = "multi_thread")]
+async fn pipeline_writes_every_entry(
+    #[future(awt)] azure: AzureBlobCache<JsonCodec<Value>>,
+    context: ExactCacheContext,
+) {
+    contract::pipeline_writes_every_entry(
+        &azure,
+        context,
+        PREFIX,
+        vec![json!("a"), json!(2), json!({"c": true})],
+    )
+    .await;
+}
+
+#[rstest]
+#[tokio::test(flavor = "multi_thread")]
+async fn batch_preserves_order(
+    #[future(awt)] azure: AzureBlobCache<JsonCodec<Value>>,
+    context: ExactCacheContext,
+) {
+    contract::batch_preserves_order(&azure, context, PREFIX, json!("first"), json!(2)).await;
+}
+
+#[rstest]
+#[tokio::test(flavor = "multi_thread")]
+async fn flush_clears(
+    #[future(awt)] azure: AzureBlobCache<JsonCodec<Value>>,
+    context: ExactCacheContext,
+) {
+    contract::flush_clears(&azure, context, PREFIX, json!("value")).await;
+}

litellm-rust/crates/cache-azure-blob/tests/support/mod.rs

@@ -0,0 +1,239 @@
+#![allow(dead_code)]
+
+use std::{
+    collections::BTreeMap,
+    sync::{Arc, Mutex},
+};
+
+use azure_core::http::{
+    AsyncRawResponse, Body, ClientOptions, HttpClient, Method, Request, StatusCode, Transport,
+    headers::{HeaderName, Headers},
+};
+use litellm_cache::{CacheCodec, Error};
+use litellm_cache_azure_blob::AzureBlobCache;
+use tokio::runtime::{Handle, Runtime};
+
+pub const ACCOUNT_URL: &str = "https://example.blob.core.windows.net";
+pub const CONTAINER: &str = "litellm-cache";
+const IF_NONE_MATCH: HeaderName = HeaderName::from_static("if-none-match");
+const ERROR_CODE: HeaderName = HeaderName::from_static("x-ms-error-code");
+
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct RecordedRequest {
+    pub method: Method,
+    pub path: String,
+    pub query: String,
+    pub if_none_match: Option<String>,
+}
+
+#[derive(Default)]
+struct FakeState {
+    container_exists: bool,
+    blobs: BTreeMap<String, Vec<u8>>,
+    requests: Vec<RecordedRequest>,
+    failing: bool,
+    precondition_conflicts: bool,
+}
+
+#[derive(Clone, Default)]
+pub struct FakeBlobService {
+    state: Arc<Mutex<FakeState>>,
+}
+
+impl std::fmt::Debug for FakeBlobService {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_str("FakeBlobService")
+    }
+}
+
+impl FakeBlobService {
+    pub fn with_existing_container() -> Self {
+        let service = Self::default();
+        service.state.lock().unwrap().container_exists = true;
+        service
+    }
+
+    pub fn blob(&self, name: &str) -> Option<Vec<u8>> {
+        self.state.lock().unwrap().blobs.get(name).cloned()
+    }
+
+    pub fn blob_names(&self) -> Vec<String> {
+        self.state.lock().unwrap().blobs.keys().cloned().collect()
+    }
+
+    pub fn seed_blob(&self, name: &str, bytes: &[u8]) {
+        self.state
+            .lock()
+            .unwrap()
+            .blobs
+            .insert(name.to_string(), bytes.to_vec());
+    }
+
+    pub fn set_failing(&self, failing: bool) {
+        self.state.lock().unwrap().failing = failing;
+    }
+
+    pub fn set_precondition_conflicts(&self, enabled: bool) {
+        self.state.lock().unwrap().precondition_conflicts = enabled;
+    }
+
+    pub fn requests(&self) -> Vec<RecordedRequest> {
+        self.state.lock().unwrap().requests.clone()
+    }
+
+    pub fn container_exists(&self) -> bool {
+        self.state.lock().unwrap().container_exists
+    }
+
+    fn respond(status: StatusCode, error_code: Option<&str>, body: Vec<u8>) -> AsyncRawResponse {
+        let mut headers = Headers::new();
+        if let Some(code) = error_code {
+            headers.insert(ERROR_CODE, code.to_string());
+        }
+        AsyncRawResponse::from_bytes(status, headers, body)
+    }
+
+    fn list_body(state: &FakeState) -> Vec<u8> {
+        let mut xml = String::from(
+            r#"<?xml version="1.0" encoding="utf-8"?><EnumerationResults ServiceEndpoint="https://example.blob.core.windows.net/" ContainerName="litellm-cache"><Blobs>"#,
+        );
+        for name in state.blobs.keys() {
+            xml.push_str(&format!(
+                "<Blob><Name>{name}</Name><Properties><BlobType>BlockBlob</BlobType></Properties></Blob>"
+            ));
+        }
+        xml.push_str("</Blobs><NextMarker /></EnumerationResults>");
+        xml.into_bytes()
+    }
+}
+
+#[async_trait::async_trait]
+impl HttpClient for FakeBlobService {
+    async fn execute_request(&self, request: &Request) -> azure_core::Result<AsyncRawResponse> {
+        let mut state = self.state.lock().unwrap();
+        let path = request.url().path().to_string();
+        let query = request.url().query().unwrap_or_default().to_string();
+        let if_none_match = request
+            .headers()
+            .get_optional_str(&IF_NONE_MATCH)
+            .map(str::to_owned);
+        state.requests.push(RecordedRequest {
+            method: request.method(),
+            path: path.clone(),
+            query: query.clone(),
+            if_none_match: if_none_match.clone(),
+        });
+        if state.failing {
+            return Ok(Self::respond(
+                StatusCode::Forbidden,
+                Some("AuthorizationFailure"),
+                Vec::new(),
+            ));
+        }
+        let container_path = format!("/{CONTAINER}");
+        let blob_name = path
+            .strip_prefix(&format!("{container_path}/"))
+            .map(str::to_owned);
+        let is_container = path == container_path && query.contains("restype=container");
+        let response = match (request.method(), is_container, blob_name) {
+            (Method::Put, true, None) if state.container_exists => Self::respond(
+                StatusCode::Conflict,
+                Some("ContainerAlreadyExists"),
+                Vec::new(),
+            ),
+            (Method::Put, true, None) => {
+                state.container_exists = true;
+                Self::respond(StatusCode::Created, None, Vec::new())
+            }
+            (Method::Get, true, None) if query.contains("comp=list") => {
+                Self::respond(StatusCode::Ok, None, Self::list_body(&state))
+            }
+            (Method::Get, true, None) if state.container_exists => {
+                Self::respond(StatusCode::Ok, None, Vec::new())
+            }
+            (Method::Get, true, None) => {
+                Self::respond(StatusCode::NotFound, Some("ContainerNotFound"), Vec::new())
+            }
+            (Method::Put, false, Some(name)) => {
+                if if_none_match.as_deref() == Some("*") && state.blobs.contains_key(&name) {
+                    if state.precondition_conflicts {
+                        Self::respond(
+                            StatusCode::PreconditionFailed,
+                            Some("ConditionNotMet"),
+                            Vec::new(),
+                        )
+                    } else {
+                        Self::respond(StatusCode::Conflict, Some("BlobAlreadyExists"), Vec::new())
+                    }
+                } else {
+                    let bytes = match request.body() {
+                        Body::Bytes(bytes) => bytes.to_vec(),
+                        Body::SeekableStream(_) => panic!("unexpected streaming upload"),
+                    };
+                    state.blobs.insert(name, bytes);
+                    Self::respond(StatusCode::Created, None, Vec::new())
+                }
+            }
+            (Method::Get, false, Some(name)) => match state.blobs.get(&name) {
+                Some(bytes) => Self::respond(StatusCode::Ok, None, bytes.clone()),
+                None => Self::respond(StatusCode::NotFound, Some("BlobNotFound"), Vec::new()),
+            },
+            (Method::Delete, false, Some(name)) => match state.blobs.remove(&name) {
+                Some(_) => Self::respond(StatusCode::Accepted, None, Vec::new()),
+                None => Self::respond(StatusCode::NotFound, Some("BlobNotFound"), Vec::new()),
+            },
+            (method, _, _) => panic!("unexpected request {method:?} {path}?{query}"),
+        };
+        Ok(response)
+    }
+}
+
+pub async fn connect<C: CacheCodec>(
+    service: &FakeBlobService,
+    account_url: &str,
+    codec: C,
+    handle: Handle,
+) -> Result<AzureBlobCache<C>, Error> {
+    AzureBlobCache::connect_with_options(
+        account_url,
+        CONTAINER,
+        None,
+        ClientOptions {
+            transport: Some(Transport::new(Arc::new(service.clone()))),
+            ..ClientOptions::default()
+        },
+        codec,
+        handle,
+    )
+    .await
+}
+
+/// A cache on its own fake service and runtime, so sync methods run outside any runtime.
+pub struct Fixture<C> {
+    pub runtime: Runtime,
+    pub service: FakeBlobService,
+    pub cache: Arc<AzureBlobCache<C>>,
+}
+
+impl<C: CacheCodec> Fixture<C> {
+    pub fn new(service: FakeBlobService, codec: C) -> Self {
+        let runtime = Runtime::new().unwrap();
+        let cache = runtime
+            .block_on(connect(
+                &service,
+                ACCOUNT_URL,
+                codec,
+                runtime.handle().clone(),
+            ))
+            .unwrap();
+        Self {
+            runtime,
+            service,
+            cache: Arc::new(cache),
+        }
+    }
+
+    pub fn stored_json(&self, key: &str) -> serde_json::Value {
+        serde_json::from_slice(&self.service.blob(key).expect("blob should exist")).unwrap()
+    }
+}

litellm-rust/crates/cache-azure-blob/tests/transport.rs

@@ -0,0 +1,90 @@
+use std::sync::Arc;
+
+use azure_core::http::{ClientOptions, Transport};
+use litellm_cache::{BaseCache, ExactCacheContext, JsonCodec};
+use litellm_cache_azure_blob::{AzureBlobCache, ReqwestTransport};
+use rstest::{fixture, rstest};
+use serde_json::{Value, json};
+use tokio::runtime::Handle;
+use wiremock::{
+    Mock, MockServer, ResponseTemplate,
+    matchers::{body_json, header, method, path, query_param},
+};
+
+#[fixture]
+async fn server() -> MockServer {
+    let server = MockServer::start().await;
+    Mock::given(method("PUT"))
+        .and(path("/litellm-cache"))
+        .and(query_param("restype", "container"))
+        .respond_with(ResponseTemplate::new(201))
+        .expect(1)
+        .mount(&server)
+        .await;
+    server
+}
+
+async fn connect(server: &MockServer) -> AzureBlobCache<JsonCodec<Value>> {
+    AzureBlobCache::connect_with_options(
+        &server.uri(),
+        "litellm-cache",
+        None,
+        ClientOptions {
+            transport: Some(Transport::new(Arc::new(ReqwestTransport(
+                reqwest::Client::new(),
+            )))),
+            ..ClientOptions::default()
+        },
+        JsonCodec::new(),
+        Handle::current(),
+    )
+    .await
+    .unwrap()
+}
+
+#[rstest]
+#[tokio::test(flavor = "multi_thread")]
+async fn uploads_go_through_the_host_client(#[future(awt)] server: MockServer) {
+    Mock::given(method("PUT"))
+        .and(path("/litellm-cache/key"))
+        .and(header("if-none-match", "*"))
+        .and(body_json(json!({"answer": 1})))
+        .respond_with(ResponseTemplate::new(201))
+        .expect(1)
+        .mount(&server)
+        .await;
+    connect(&server)
+        .await
+        .set_cache("key", json!({"answer": 1}), &ExactCacheContext::default())
+        .unwrap();
+}
+
+#[rstest]
+#[case::hit(
+    ResponseTemplate::new(200).set_body_json(json!({"answer": 2})),
+    Some(json!({"answer": 2}))
+)]
+#[case::blob_not_found(
+    ResponseTemplate::new(404).insert_header("x-ms-error-code", "BlobNotFound"),
+    None
+)]
+#[tokio::test(flavor = "multi_thread")]
+async fn downloads_map_the_host_client_response(
+    #[future(awt)] server: MockServer,
+    #[case] response: ResponseTemplate,
+    #[case] expected: Option<Value>,
+) {
+    Mock::given(method("GET"))
+        .and(path("/litellm-cache/key"))
+        .respond_with(response)
+        .mount(&server)
+        .await;
+    assert_eq!(
+        connect(&server)
+            .await
+            .async_get_cache("key", &ExactCacheContext::default())
+            .await
+            .unwrap(),
+        expected
+    );
+}

litellm-rust/crates/cache-disk/Cargo.toml

@@ -15,5 +15,6 @@ serde_json.workspace = true
 tokio.workspace = true
 
 [dev-dependencies]
+litellm-cache-testing.workspace = true
 rstest.workspace = true
 tempfile = "3.27.0"

litellm-rust/crates/cache-disk/src/cache.rs

@@ -5,8 +5,8 @@ use std::{
 };
 
 use litellm_cache::{
-    BaseCache, BatchCache, BatchEntry, CacheCodec, CacheConnectionResult, CacheConnectionStatus,
-    CounterCache, DeleteCache, Error, ExactCacheContext, FlushCache,
+    BaseCache, BatchCache, BatchEntry, CacheCodec, CounterCache, DeleteCache, DisconnectCache,
+    Error, ExactCacheContext, FlushCache,
 };
 
 use crate::{DiskStore, DiskcacheSqliteStore, PythonDiskCacheAdapter, StoredValue, ValueAdapter};
@@ -150,29 +150,6 @@ impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> BaseCache for DiskCache<S, D,
         })
         .await
     }
-
-    async fn disconnect(&self) -> Result<(), Error> {
-        Ok(())
-    }
-
-    async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
-        let result = Self::run_blocking(Arc::clone(&self.store), |store| {
-            store.probe().map(|_| CacheConnectionResult {
-                status: CacheConnectionStatus::Success,
-                message: "Disk cache connection test successful".into(),
-                error: None,
-            })
-        })
-        .await;
-        Ok(match result {
-            Ok(result) => result,
-            Err(error) => CacheConnectionResult {
-                status: CacheConnectionStatus::Failed,
-                message: format!("Disk cache connection failed: {error}"),
-                error: Some(error.to_string()),
-            },
-        })
-    }
 }
 
 impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> BatchCache for DiskCache<S, D, A> {
@@ -241,9 +218,13 @@ impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> FlushCache for DiskCache<S, D
     }
 }
 
-impl<S: CacheCodec<Value = f64>, D: DiskStore, A: ValueAdapter> CounterCache
-    for DiskCache<S, D, A>
-{
+impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> DisconnectCache for DiskCache<S, D, A> {
+    async fn disconnect(&self) -> Result<(), Error> {
+        Ok(())
+    }
+}
+
+impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> CounterCache for DiskCache<S, D, A> {
     fn increment_cache(
         &self,
         key: &str,
@@ -264,6 +245,7 @@ impl<S: CacheCodec<Value = f64>, D: DiskStore, A: ValueAdapter> CounterCache
         key: &str,
         amount: f64,
         context: ExactCacheContext,
+        _refresh_ttl: bool,
     ) -> Result<f64, Error> {
         let key = key.to_string();
         let adapter = Arc::clone(&self.adapter);

litellm-rust/crates/cache-disk/src/sqlite.rs

@@ -544,18 +544,6 @@ impl DiskStore for DiskcacheSqliteStore {
             }
         }
     }
-
-    fn probe(&self) -> Result<(), Error> {
-        let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
-        connection
-            .query_row(
-                "SELECT value FROM Settings WHERE key = 'count'",
-                [],
-                |row| row.get::<_, i64>(0),
-            )
-            .map(|_| ())
-            .map_err(|_| Error::Unavailable)
-    }
 }
 
 fn default_settings() -> HashMap<String, Value> {

litellm-rust/crates/cache-disk/src/store.rs

@@ -29,5 +29,4 @@ pub trait DiskStore: Send + Sync + 'static {
         now: f64,
         apply: &mut dyn FnMut(Option<StoredValue>) -> Result<(StoredValue, Option<f64>), Error>,
     ) -> Result<(), Error>;
-    fn probe(&self) -> Result<(), Error>;
 }

litellm-rust/crates/cache-disk/tests/cache.rs

@@ -7,8 +7,8 @@ use std::{
 };
 
 use litellm_cache::{
-    BaseCache, BatchCache, BatchEntry, CacheCodec, CounterCache, DeleteCache, ExactCacheContext,
-    FlushCache, JsonCodec,
+    BaseCache, BatchCache, BatchEntry, CacheCodec, CounterCache, DeleteCache, DisconnectCache,
+    ExactCacheContext, FlushCache, JsonCodec,
 };
 use litellm_cache_disk::{DiskCache, DiskStore, DiskcacheSqliteStore, StoredValue, ValueAdapter};
 use rstest::{fixture, rstest};
@@ -395,7 +395,7 @@ fn delete_flush_and_spilled_file_replacement_clean_up_storage(sandbox: Sandbox)
 
 #[rstest]
 #[tokio::test]
-async fn async_operations_connection_and_delete_match_sync_operations(sandbox: Sandbox) {
+async fn async_operations_disconnect_and_delete_match_sync_operations(sandbox: Sandbox) {
     let cache = sandbox.cache::<Value>();
     let context = ExactCacheContext {
         ttl: Some(Duration::from_secs(60)),
@@ -424,8 +424,101 @@ async fn async_operations_connection_and_delete_match_sync_operations(sandbox: S
     );
     cache.async_delete_cache("a").await.unwrap();
     cache.async_flush_cache().await.unwrap();
+    cache.disconnect().await.unwrap();
+}
+
+#[derive(Clone, Copy, Debug)]
+enum Increment {
+    Sync,
+    Async { refresh_ttl: bool },
+}
+
+impl Increment {
+    async fn apply(
+        self,
+        cache: &DiskCache<JsonCodec<Value>>,
+        key: &str,
+        amount: f64,
+        context: ExactCacheContext,
+    ) -> f64 {
+        match self {
+            Self::Sync => cache.increment_cache(key, amount, context).unwrap(),
+            Self::Async { refresh_ttl } => cache
+                .async_increment(key, amount, context, refresh_ttl)
+                .await
+                .unwrap(),
+        }
+    }
+}
+
+#[rstest]
+#[case::sync_missing(Increment::Sync, None, 3.0, 3.0)]
+#[case::sync_existing_int(Increment::Sync, Some(json!(7)), 5.0, 12.0)]
+#[case::sync_non_int(Increment::Sync, Some(json!("not-a-number")), 4.0, 4.0)]
+#[case::async_missing(Increment::Async { refresh_ttl: false }, None, 2.0, 2.0)]
+#[case::async_existing_int(Increment::Async { refresh_ttl: false }, Some(json!(10)), 5.0, 15.0)]
+#[case::async_non_int(Increment::Async { refresh_ttl: false }, Some(json!("corrupt")), 9.0, 9.0)]
+#[case::async_refresh_ttl_is_ignored(Increment::Async { refresh_ttl: true }, Some(json!(1)), 1.0, 2.0)]
+#[tokio::test]
+async fn increments_read_back_through_get_cache(
+    sandbox: Sandbox,
+    #[case] increment: Increment,
+    #[case] initial: Option<Value>,
+    #[case] amount: f64,
+    #[case] expected: f64,
+) {
+    let cache = sandbox.cache::<Value>();
+    let context = ExactCacheContext::default();
+    if let Some(initial) = initial {
+        cache
+            .async_set_cache("counter", initial, context.clone())
+            .await
+            .unwrap();
+    }
     assert_eq!(
-        cache.test_connection().await.unwrap().status,
-        litellm_cache::CacheConnectionStatus::Success
+        increment
+            .apply(&cache, "counter", amount, context.clone())
+            .await,
+        expected
+    );
+    assert_eq!(
+        cache.get_cache("counter", &context).unwrap(),
+        Some(json!(expected as i64))
     );
 }
+
+#[rstest]
+#[case::without_refresh(false)]
+#[case::with_refresh(true)]
+#[tokio::test]
+async fn async_increment_rewrites_ttl_on_every_write(sandbox: Sandbox, #[case] refresh_ttl: bool) {
+    let cache = sandbox.cache::<Value>();
+    let expiry = || {
+        sandbox
+            .db()
+            .query_row(
+                "SELECT expire_time IS NOT NULL FROM Cache WHERE key = 'counter'",
+                [],
+                |row| row.get::<_, bool>(0),
+            )
+            .unwrap()
+    };
+    let ttl = ExactCacheContext {
+        ttl: Some(Duration::from_secs(60)),
+    };
+    cache
+        .async_increment("counter", 1.0, ttl.clone(), refresh_ttl)
+        .await
+        .unwrap();
+    assert!(expiry());
+    cache
+        .async_increment("counter", 1.0, ExactCacheContext::default(), refresh_ttl)
+        .await
+        .unwrap();
+    assert!(!expiry());
+    cache
+        .async_increment("counter", 1.0, ttl, refresh_ttl)
+        .await
+        .unwrap();
+    assert!(expiry());
+}

litellm-rust/crates/cache-disk/tests/contract.rs

@@ -0,0 +1,82 @@
+use litellm_cache::{ExactCacheContext, JsonCodec};
+use litellm_cache_disk::DiskCache;
+use litellm_cache_testing as contract;
+use rstest::{fixture, rstest};
+use serde_json::{Value, json};
+use tempfile::TempDir;
+
+struct Disk {
+    cache: DiskCache<JsonCodec<Value>>,
+    _directory: TempDir,
+}
+
+#[fixture]
+fn disk() -> Disk {
+    let directory = tempfile::tempdir().unwrap();
+    Disk {
+        cache: DiskCache::open(directory.path(), JsonCodec::new()).unwrap(),
+        _directory: directory,
+    }
+}
+
+#[fixture]
+fn context() -> ExactCacheContext {
+    ExactCacheContext::default()
+}
+
+const PREFIX: &str = "contract:";
+
+#[rstest]
+#[tokio::test]
+async fn hit_and_miss(disk: Disk, context: ExactCacheContext) {
+    contract::hit_and_miss(&disk.cache, context, PREFIX, json!({"answer": 42})).await;
+}
+
+#[rstest]
+#[tokio::test]
+async fn sync_async_equivalence(disk: Disk, context: ExactCacheContext) {
+    contract::sync_async_equivalence(&disk.cache, context, PREFIX, json!("first"), json!([2]))
+        .await;
+}
+
+#[rstest]
+#[tokio::test]
+async fn overwrite_replaces(disk: Disk, context: ExactCacheContext) {
+    contract::overwrite_replaces(&disk.cache, context, PREFIX, json!(1), json!({"b": 2})).await;
+}
+
+#[rstest]
+#[tokio::test]
+async fn pipeline_writes_every_entry(disk: Disk, context: ExactCacheContext) {
+    contract::pipeline_writes_every_entry(
+        &disk.cache,
+        context,
+        PREFIX,
+        vec![json!("a"), json!(2), json!({"c": true})],
+    )
+    .await;
+}
+
+#[rstest]
+#[tokio::test]
+async fn batch_preserves_order(disk: Disk, context: ExactCacheContext) {
+    contract::batch_preserves_order(&disk.cache, context, PREFIX, json!("first"), json!(2)).await;
+}
+
+#[rstest]
+#[tokio::test]
+async fn delete_removes_key(disk: Disk, context: ExactCacheContext) {
+    contract::delete_removes_key(&disk.cache, context, PREFIX, json!("value")).await;
+}
+
+#[rstest]
+#[tokio::test]
+async fn flush_clears(disk: Disk, context: ExactCacheContext) {
+    contract::flush_clears(&disk.cache, context, PREFIX, json!("value")).await;
+}
+
+#[rstest]
+#[tokio::test]
+async fn counter_accumulates(disk: Disk, context: ExactCacheContext) {
+    contract::counter_accumulates(&disk.cache, context, PREFIX).await;
+}

litellm-rust/crates/cache-gcs/Cargo.toml

@@ -15,6 +15,8 @@ reqwest.workspace = true
 tokio.workspace = true
 
 [dev-dependencies]
+litellm-cache-testing.workspace = true
+rstest.workspace = true
 serde_json.workspace = true
 tokio.workspace = true
 wiremock = "0.6.5"

litellm-rust/crates/cache-gcs/src/cache.rs

@@ -2,7 +2,7 @@ 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,
+    BaseCache, BatchCache, BatchEntry, CacheCodec, DisconnectCache, Error, ExactCacheContext,
     FlushCache,
 };
 use percent_encoding::{AsciiSet, NON_ALPHANUMERIC, percent_encode};
@@ -53,25 +53,25 @@ pub struct GcsCache<S: CacheCodec> {
 }
 
 impl<S: CacheCodec> GcsCache<S> {
-    pub fn new(config: GcsConfig, codec: S) -> Result<Self, Error> {
+    pub fn new(config: GcsConfig, client: Client, codec: S) -> Self {
         let token = Arc::new(GcpTokenSource::new(config.path_service_account.clone()));
-        Self::with_token_source(config, codec, token)
+        Self::with_token_source(config, client, codec, token)
     }
 
     pub fn with_token_source(
         config: GcsConfig,
+        client: Client,
         codec: S,
         token: Arc<dyn TokenSource>,
-    ) -> Result<Self, Error> {
-        let client = Client::builder().build().map_err(|_| Error::Unavailable)?;
+    ) -> Self {
         let key_prefix = key_prefix(config.gcs_path.as_deref());
-        Ok(Self {
+        Self {
             config,
             key_prefix,
             client,
             token,
             codec,
-        })
+        }
     }
 
     pub fn bucket_name(&self) -> &str {
@@ -154,26 +154,26 @@ impl<S: CacheCodec> GcsCache<S> {
         F: Future<Output = Result<T, Error>> + Send,
         T: Send,
     {
-        let run = || {
+        let run = |future: F| {
             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);
+        match tokio::runtime::Handle::try_current() {
+            Ok(handle) if handle.runtime_flavor() == tokio::runtime::RuntimeFlavor::MultiThread => {
+                tokio::task::block_in_place(|| handle.block_on(future))
             }
-            return std::thread::scope(|scope| {
+            Ok(_) => std::thread::scope(|scope| {
                 scope
-                    .spawn(run)
+                    .spawn(|| run(future))
                     .join()
                     .map_err(|_| Error::Unavailable)
                     .and_then(|result| result)
-            });
+            }),
+            Err(_) => run(future),
         }
-        run()
     }
 }
 
@@ -222,14 +222,12 @@ impl<S: CacheCodec> BaseCache for GcsCache<S> {
         .await
         .map(|_| ())
     }
+}
 
+impl<S: CacheCodec> DisconnectCache for GcsCache<S> {
     async fn disconnect(&self) -> Result<(), Error> {
         Ok(())
     }
-
-    async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
-        Err(Error::UnsupportedOperation)
-    }
 }
 
 impl<S: CacheCodec> BatchCache for GcsCache<S> {

litellm-rust/crates/cache-gcs/tests/cache.rs

@@ -1,37 +1,32 @@
-use std::{sync::Arc, time::Duration};
+mod support;
+
+use std::{future::Future, pin::Pin, sync::Arc, time::Duration};
 
 use litellm_cache::{
-    BaseCache, BatchCache, BatchEntry, CacheContext, Error, ExactCacheContext, FlushCache,
-    JsonCodec,
+    BaseCache, BatchCache, BatchEntry, CacheContext, DisconnectCache, Error, ExactCacheContext,
+    FlushCache,
 };
-use litellm_cache_gcs::{GcsCache, GcsConfig, StaticTokenSource, TokenSource, key_prefix};
-use serde_json::json;
+use litellm_cache_gcs::{GcsCache, GcsConfig, TokenSource, key_prefix};
+use rstest::{fixture, rstest};
+use serde_json::{Value, json};
+use support::FakeBucket;
 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(),
-    }
+#[fixture]
+async fn server() -> MockServer {
+    MockServer::start().await
 }
 
-fn cache(server: &MockServer, gcs_path: Option<&str>) -> GcsCache<JsonCodec<serde_json::Value>> {
-    GcsCache::with_token_source(
-        config(server, gcs_path),
-        JsonCodec::new(),
-        Arc::new(StaticTokenSource("tok".into())),
-    )
-    .unwrap()
+fn context() -> ExactCacheContext {
+    ExactCacheContext::default()
 }
 
+#[rstest]
 #[tokio::test]
-async fn set_writes_encoded_object_and_headers() {
-    let server = MockServer::start().await;
+async fn set_writes_encoded_object_and_headers(#[future(awt)] server: MockServer) {
     Mock::given(method("POST"))
         .and(path("/upload/storage/v1/b/bucket/o"))
         .and(query_param("uploadType", "media"))
@@ -42,12 +37,8 @@ async fn set_writes_encoded_object_and_headers() {
         .expect(1)
         .mount(&server)
         .await;
-    cache(&server, Some("cache/"))
-        .set_cache(
-            "team:a b/c",
-            json!({"value": "entry"}),
-            &ExactCacheContext::default(),
-        )
+    support::cache(&server, Some("cache/"))
+        .set_cache("team:a b/c", json!({"value": "entry"}), &context())
         .unwrap();
     let requests = server.received_requests().await.unwrap();
     assert_eq!(requests.len(), 1);
@@ -57,103 +48,108 @@ async fn set_writes_encoded_object_and_headers() {
     );
 }
 
+#[rstest]
+#[case::hit(
+    "hit",
+    ResponseTemplate::new(200).set_body_json(json!({"value": "entry"})),
+    Ok(Some(json!({"value": "entry"})))
+)]
+#[case::missing("missing", ResponseTemplate::new(404), Ok(None))]
+#[case::server_error("server-error", ResponseTemplate::new(500), Err(Error::Unavailable))]
+#[case::invalid(
+    "invalid",
+    ResponseTemplate::new(200).set_body_string("not json"),
+    Err(Error::InvalidEntry)
+)]
 #[tokio::test]
-async fn get_maps_statuses_and_decode_failures() {
-    let server = MockServer::start().await;
+async fn get_maps_statuses_and_decode_failures(
+    #[future(awt)] server: MockServer,
+    #[case] key: &str,
+    #[case] response: ResponseTemplate,
+    #[case] expected: Result<Option<Value>, Error>,
+) {
     Mock::given(method("GET"))
-        .and(path("/storage/v1/b/bucket/o/hit"))
+        .and(path(format!("/storage/v1/b/bucket/o/{key}")))
         .and(query_param("alt", "media"))
-        .respond_with(ResponseTemplate::new(200).set_body_json(json!({"value": "entry"})))
+        .respond_with(response)
         .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
-    );
+    let cache = support::cache(&server, None);
+    assert_eq!(cache.get_cache(key, &context()), expected);
+    assert_eq!(cache.async_get_cache(key, &context()).await, expected);
 }
 
-#[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("")), "");
+#[rstest]
+#[case::none(None, "")]
+#[case::trailing_slash(Some("a/b/"), "a/b/")]
+#[case::no_trailing_slash(Some("a/b"), "a/b/")]
+#[case::empty(Some(""), "")]
+fn key_prefix_normalizes_paths(#[case] gcs_path: Option<&str>, #[case] expected: &str) {
+    assert_eq!(key_prefix(gcs_path), expected);
 }
 
+#[rstest]
 #[tokio::test]
-async fn object_names_use_python_quote_encoding() {
-    let server = MockServer::start().await;
+async fn cache_exposes_its_configuration(#[future(awt)] server: MockServer) {
+    let cache = GcsCache::new(
+        GcsConfig {
+            path_service_account: Some("/secrets/sa.json".into()),
+            ..support::config(&server, Some("folder"))
+        },
+        reqwest::Client::new(),
+        litellm_cache::JsonCodec::<Value>::new(),
+    );
+    assert_eq!(cache.bucket_name(), "bucket");
+    assert_eq!(cache.key_prefix(), "folder/");
+    assert_eq!(cache.path_service_account(), Some("/secrets/sa.json"));
+    assert_eq!(cache.object_name("k"), "folder/k");
+}
+
+#[rstest]
+#[case::punctuation("a~b-c_d.e/f g%h", "uploadType=media&name=p%2Fa~b-c_d.e%2Ff%20g%25h")]
+#[case::utf8("ключ", "uploadType=media&name=p%2F%D0%BA%D0%BB%D1%8E%D1%87")]
+#[tokio::test]
+async fn object_names_use_python_quote_encoding(
+    #[future(awt)] server: MockServer,
+    #[case] key: &str,
+    #[case] query: &str,
+) {
     Mock::given(method("POST"))
         .and(path("/upload/storage/v1/b/bucket/o"))
         .and(query_param("uploadType", "media"))
         .respond_with(ResponseTemplate::new(200))
+        .expect(1)
+        .mount(&server)
+        .await;
+    support::cache(&server, Some("p/"))
+        .async_set_cache(key, json!({"value": key}), context())
+        .await
+        .unwrap();
+    let requests = server.received_requests().await.unwrap();
+    assert_eq!(requests[0].url.query(), Some(query));
+}
+
+#[rstest]
+#[tokio::test]
+async fn object_names_are_encoded_in_the_download_path(#[future(awt)] server: MockServer) {
+    Mock::given(method("GET"))
+        .and(path("/storage/v1/b/bucket/o/p%2Fa%3Ab%20c"))
+        .and(query_param("alt", "media"))
+        .respond_with(ResponseTemplate::new(200).set_body_json(json!(1)))
         .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"));
+    let cache = support::cache(&server, Some("p"));
+    assert_eq!(cache.get_cache("a:b c", &context()), Ok(Some(json!(1))));
+    assert_eq!(
+        cache.async_get_cache("a:b c", &context()).await,
+        Ok(Some(json!(1)))
+    );
 }
 
+#[rstest]
 #[tokio::test]
-async fn ignores_ttl_and_writes_pipeline_concurrently() {
-    let server = MockServer::start().await;
+async fn ignores_ttl_and_writes_pipeline_concurrently(#[future(awt)] server: MockServer) {
     for key in ["one", "two", "three"] {
         Mock::given(method("POST"))
             .and(path("/upload/storage/v1/b/bucket/o"))
@@ -164,12 +160,10 @@ async fn ignores_ttl_and_writes_pipeline_concurrently() {
             .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
-    );
+    let cache = support::cache(&server, None);
+    let with_ttl = context().with_ttl(Some(Duration::from_secs(5)));
+    assert_eq!(cache.get_ttl(&context()), None);
+    assert_eq!(cache.get_ttl(&with_ttl), None);
     cache
         .async_set_cache_pipeline(
             vec![
@@ -177,15 +171,15 @@ async fn ignores_ttl_and_writes_pipeline_concurrently() {
                 ("two".into(), json!({"key": "two"})),
                 ("three".into(), json!({"key": "three"})),
             ],
-            ExactCacheContext::default().with_ttl(Some(Duration::from_secs(5))),
+            with_ttl,
         )
         .await
         .unwrap();
 }
 
+#[rstest]
 #[tokio::test]
-async fn async_batch_get_preserves_hits_misses_and_invalid_entries() {
-    let server = MockServer::start().await;
+async fn batch_get_preserves_hits_misses_and_invalid_entries(#[future(awt)] server: MockServer) {
     Mock::given(method("GET"))
         .and(path("/storage/v1/b/bucket/o/hit"))
         .respond_with(ResponseTemplate::new(200).set_body_json(json!({"value": "entry"})))
@@ -201,122 +195,83 @@ async fn async_batch_get_preserves_hits_misses_and_invalid_entries() {
         .respond_with(ResponseTemplate::new(200).set_body_string("not json"))
         .mount(&server)
         .await;
+    let cache = support::cache(&server, None);
+    let keys = vec!["hit".to_string(), "missing".into(), "invalid".into()];
+    let expected = vec![
+        BatchEntry::Hit(json!({"value": "entry"})),
+        BatchEntry::Miss,
+        BatchEntry::Invalid,
+    ];
 
+    assert_eq!(cache.batch_get_cache(&keys, &context()).unwrap(), expected);
     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,
-        ]
+        cache.async_batch_get_cache(keys, context()).await.unwrap(),
+        expected
     );
 }
 
+#[rstest]
 #[tokio::test]
-async fn lifecycle_operations_are_noops_and_connection_test_is_unsupported() {
-    let server = MockServer::start().await;
-    let cache = cache(&server, None);
+async fn flush_and_disconnect_are_noops_like_python(#[future(awt)] server: MockServer) {
+    let cache = support::cache(&server, None);
     assert_eq!(cache.flush_cache(), Ok(()));
+    assert_eq!(cache.async_flush_cache().await, Ok(()));
     assert_eq!(cache.disconnect().await, Ok(()));
-    assert_eq!(
-        cache.test_connection().await,
-        Err(Error::UnsupportedOperation)
-    );
+    assert!(server.received_requests().await.unwrap().is_empty());
 }
 
-#[test]
+fn round_trip(cache: &support::JsonGcsCache) -> Result<Option<Value>, Error> {
+    cache.set_cache("key", json!({"value": "entry"}), &context())?;
+    cache.get_cache("key", &context())
+}
+
+#[rstest]
 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"}))
-    );
+    let server = runtime.block_on(FakeBucket::serve());
+    let cache = support::cache(&server, None);
+    assert_eq!(round_trip(&cache), Ok(Some(json!({"value": "entry"}))));
 }
 
+#[rstest]
 #[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"}))
-    );
+    let server = FakeBucket::serve().await;
+    let cache = support::cache(&server, None);
+    assert_eq!(round_trip(&cache), Ok(Some(json!({"value": "entry"}))));
+}
+
+#[rstest]
+#[tokio::test]
+async fn sync_operations_work_inside_a_current_thread_runtime() {
+    let server = FakeBucket::serve().await;
+    let cache = support::cache(&server, None);
+    assert_eq!(round_trip(&cache), Ok(Some(json!({"value": "entry"}))));
 }
 
 struct FailingTokenSource;
 
 impl TokenSource for FailingTokenSource {
-    fn bearer_token(
-        &self,
-    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<String, Error>> + Send + '_>>
-    {
+    fn bearer_token(&self) -> Pin<Box<dyn Future<Output = Result<String, Error>> + Send + '_>> {
         Box::pin(async { Err(Error::Unavailable) })
     }
 }
 
+#[rstest]
 #[tokio::test]
-async fn token_source_failure_skips_http() {
-    let server = MockServer::start().await;
-    let cache = GcsCache::with_token_source(
-        config(&server, None),
-        JsonCodec::<serde_json::Value>::new(),
-        Arc::new(FailingTokenSource),
-    )
-    .unwrap();
+async fn token_source_failure_skips_http(#[future(awt)] server: MockServer) {
+    let cache = support::cache_with_token(&server, None, Arc::new(FailingTokenSource));
+    assert_eq!(
+        cache.get_cache("key", &context()).unwrap_err(),
+        Error::Unavailable
+    );
     assert_eq!(
         cache
-            .get_cache("key", &ExactCacheContext::default())
+            .async_set_cache("key", json!(1), context())
+            .await
             .unwrap_err(),
         Error::Unavailable
     );

litellm-rust/crates/cache-gcs/tests/contract.rs

@@ -0,0 +1,65 @@
+mod support;
+
+use litellm_cache::ExactCacheContext;
+use litellm_cache_testing as contract;
+use rstest::{fixture, rstest};
+use serde_json::json;
+use support::{FakeBucket, JsonGcsCache};
+use wiremock::MockServer;
+
+struct Gcs {
+    cache: JsonGcsCache,
+    _server: MockServer,
+}
+
+#[fixture]
+async fn gcs() -> Gcs {
+    let server = FakeBucket::serve().await;
+    Gcs {
+        cache: support::cache(&server, Some("contract")),
+        _server: server,
+    }
+}
+
+#[fixture]
+fn context() -> ExactCacheContext {
+    ExactCacheContext::default()
+}
+
+const PREFIX: &str = "contract:";
+
+#[rstest]
+#[tokio::test(flavor = "multi_thread")]
+async fn hit_and_miss(#[future(awt)] gcs: Gcs, context: ExactCacheContext) {
+    contract::hit_and_miss(&gcs.cache, context, PREFIX, json!({"answer": 42})).await;
+}
+
+#[rstest]
+#[tokio::test(flavor = "multi_thread")]
+async fn sync_async_equivalence(#[future(awt)] gcs: Gcs, context: ExactCacheContext) {
+    contract::sync_async_equivalence(&gcs.cache, context, PREFIX, json!("first"), json!([2])).await;
+}
+
+#[rstest]
+#[tokio::test(flavor = "multi_thread")]
+async fn overwrite_replaces(#[future(awt)] gcs: Gcs, context: ExactCacheContext) {
+    contract::overwrite_replaces(&gcs.cache, context, PREFIX, json!(1), json!({"b": 2})).await;
+}
+
+#[rstest]
+#[tokio::test(flavor = "multi_thread")]
+async fn pipeline_writes_every_entry(#[future(awt)] gcs: Gcs, context: ExactCacheContext) {
+    contract::pipeline_writes_every_entry(
+        &gcs.cache,
+        context,
+        PREFIX,
+        vec![json!("a"), json!(2), json!({"c": true})],
+    )
+    .await;
+}
+
+#[rstest]
+#[tokio::test(flavor = "multi_thread")]
+async fn batch_preserves_order(#[future(awt)] gcs: Gcs, context: ExactCacheContext) {
+    contract::batch_preserves_order(&gcs.cache, context, PREFIX, json!("first"), json!(2)).await;
+}

litellm-rust/crates/cache-gcs/tests/support/mod.rs

@@ -0,0 +1,87 @@
+#![allow(dead_code)]
+
+use std::{
+    collections::HashMap,
+    sync::{Arc, Mutex},
+};
+
+use litellm_cache::JsonCodec;
+use litellm_cache_gcs::{GcsCache, GcsConfig, StaticTokenSource, TokenSource};
+use percent_encoding::percent_decode_str;
+use serde_json::Value;
+use wiremock::{Mock, MockServer, Request, Respond, ResponseTemplate, http::Method, matchers::any};
+
+pub type JsonGcsCache = GcsCache<JsonCodec<Value>>;
+
+pub 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(),
+    }
+}
+
+pub fn cache_with_token(
+    server: &MockServer,
+    gcs_path: Option<&str>,
+    token: Arc<dyn TokenSource>,
+) -> JsonGcsCache {
+    GcsCache::with_token_source(
+        config(server, gcs_path),
+        reqwest::Client::new(),
+        JsonCodec::new(),
+        token,
+    )
+}
+
+pub fn cache(server: &MockServer, gcs_path: Option<&str>) -> JsonGcsCache {
+    cache_with_token(server, gcs_path, Arc::new(StaticTokenSource("tok".into())))
+}
+
+/// An in-memory bucket speaking the JSON API's media upload and `alt=media` download.
+#[derive(Clone, Default)]
+pub struct FakeBucket {
+    objects: Arc<Mutex<HashMap<String, Vec<u8>>>>,
+}
+
+impl FakeBucket {
+    pub async fn serve() -> MockServer {
+        let server = MockServer::start().await;
+        Mock::given(any())
+            .respond_with(Self::default())
+            .mount(&server)
+            .await;
+        server
+    }
+}
+
+impl Respond for FakeBucket {
+    fn respond(&self, request: &Request) -> ResponseTemplate {
+        let mut objects = self.objects.lock().unwrap();
+        match request.method {
+            Method::POST => {
+                let name = request
+                    .url
+                    .query_pairs()
+                    .find_map(|(key, value)| (key == "name").then(|| value.into_owned()))
+                    .expect("uploads carry the object name");
+                objects.insert(name, request.body.clone());
+                ResponseTemplate::new(200)
+            }
+            Method::GET => {
+                let encoded = request
+                    .url
+                    .path()
+                    .strip_prefix("/storage/v1/b/bucket/o/")
+                    .expect("downloads address an object");
+                let name = percent_decode_str(encoded).decode_utf8().unwrap();
+                match objects.get(name.as_ref()) {
+                    Some(body) => ResponseTemplate::new(200).set_body_bytes(body.clone()),
+                    None => ResponseTemplate::new(404),
+                }
+            }
+            _ => ResponseTemplate::new(405),
+        }
+    }
+}

litellm-rust/crates/cache-memory/Cargo.toml

@@ -9,6 +9,6 @@ repository.workspace = true
 litellm-cache.workspace = true
 
 [dev-dependencies]
-serde_json.workspace = true
+litellm-cache-testing.workspace = true
 rstest.workspace = true
 tokio.workspace = true

litellm-rust/crates/cache-memory/src/cache.rs

@@ -7,8 +7,8 @@ use std::{
 };
 
 use litellm_cache::{
-    BaseCache, BatchCache, CacheConnectionResult, CacheConnectionStatus, ClaimCache, CounterCache,
-    DeleteCache, Error, ExactCacheContext, FlushCache, IncrementOperation, SetCache, TtlCache,
+    BaseCache, BatchCache, ClaimCache, CounterCache, DeleteCache, DisconnectCache, Error,
+    ExactCacheContext, FlushCache, SetCache, TtlCache,
 };
 
 const DEFAULT_MAX_SIZE_IN_MEMORY: usize = 200;
@@ -75,7 +75,9 @@ impl<V: Clone> InMemoryCache<V> {
                 expiration_heap: BinaryHeap::new(),
             }),
             max_size_in_memory: max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY),
-            default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
+            default_ttl: default_ttl
+                .filter(|ttl| !ttl.is_zero())
+                .unwrap_or(DEFAULT_TTL),
             max_entry_bytes,
             measure_value,
             now: Arc::new(now),
@@ -91,21 +93,9 @@ impl<V: Clone> InMemoryCache<V> {
         if self.max_size_in_memory == 0 {
             return Ok(CacheWrite::Disabled);
         }
-        if let (Some(limit), Some(measure)) = (self.max_entry_bytes, &self.measure_value)
-            && measure(&value)? > limit
-        {
-            return Ok(CacheWrite::TooLarge);
-        }
         let now = (self.now)();
         let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
-        let key = key.into();
-        Self::evict(&mut state, self.max_size_in_memory, now, &key);
-        let expiration = state.expirations.get(&key).copied();
-        if expiration.is_none_or(|expiration| expiration < now) {
-            Self::set_expiration(&mut state, &key, now + ttl.unwrap_or(self.default_ttl));
-        }
-        state.values.insert(key, value);
-        Ok(CacheWrite::Stored)
+        self.store(&mut state, key.into(), value, ttl, now)
     }
 
     pub fn get_cache(&self, key: &str) -> Result<Option<V>, Error> {
@@ -121,6 +111,70 @@ impl<V: Clone> InMemoryCache<V> {
         Ok(state.values.get(key).cloned())
     }
 
+    /// `check_value_size`: whether `value` fits `max_entry_bytes`. Always `true` without a
+    /// limit and a measure, since typed values have no generic size.
+    pub fn check_value_size(&self, value: &V) -> Result<bool, Error> {
+        match (self.max_entry_bytes, &self.measure_value) {
+            (Some(limit), Some(measure)) => Ok(measure(value)? <= limit),
+            _ => Ok(true),
+        }
+    }
+
+    /// `evict_cache`: drops expired entries, then the earliest-expiring ones until a new key
+    /// fits.
+    pub fn evict_cache(&self) -> Result<(), Error> {
+        let now = (self.now)();
+        let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
+        Self::evict(&mut state, self.max_size_in_memory, now, None);
+        Ok(())
+    }
+
+    /// `evict_element_if_expired`: `true` when `key` had expired and was removed.
+    pub fn evict_element_if_expired(&self, key: &str) -> Result<bool, Error> {
+        let now = (self.now)();
+        let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
+        let expired = state
+            .expirations
+            .get(key)
+            .is_some_and(|expiration| *expiration < now);
+        if expired {
+            Self::remove(&mut state, key);
+        }
+        Ok(expired)
+    }
+
+    /// `allow_ttl_override`: a write may set the TTL when the key has none or it has passed.
+    pub fn allow_ttl_override(&self, key: &str) -> Result<bool, Error> {
+        let now = (self.now)();
+        Ok(self
+            .expires_at(key)?
+            .is_none_or(|expiration| expiration < now))
+    }
+
+    /// The number of stored entries, expired ones included until they are evicted.
+    pub fn len(&self) -> Result<usize, Error> {
+        Ok(self
+            .state
+            .lock()
+            .map_err(|_| Error::Unavailable)?
+            .values
+            .len())
+    }
+
+    pub fn is_empty(&self) -> Result<bool, Error> {
+        Ok(self.len()? == 0)
+    }
+
+    /// Entries in the expiration heap, stale ones included; bounded by eviction.
+    pub fn expiration_heap_len(&self) -> Result<usize, Error> {
+        Ok(self
+            .state
+            .lock()
+            .map_err(|_| Error::Unavailable)?
+            .expiration_heap
+            .len())
+    }
+
     pub fn max_size_in_memory(&self) -> usize {
         self.max_size_in_memory
     }
@@ -172,7 +226,9 @@ impl<V: Clone> InMemoryCache<V> {
         Ok(())
     }
 
-    fn evict(state: &mut CacheState<V>, capacity: usize, now: Duration, key: &str) {
+    /// Writing an existing `key` never evicts another entry, unlike Python, which pops the
+    /// earliest-expiring entry whenever the cache is full.
+    fn evict(state: &mut CacheState<V>, capacity: usize, now: Duration, key: Option<&str>) {
         while let Some(Reverse((expiration, key))) = state.expiration_heap.peek().cloned() {
             if state.expirations.get(&key).copied() != Some(expiration) {
                 state.expiration_heap.pop();
@@ -183,7 +239,7 @@ impl<V: Clone> InMemoryCache<V> {
                 break;
             }
         }
-        if state.values.contains_key(key) {
+        if key.is_some_and(|key| state.values.contains_key(key)) {
             return;
         }
         while state.values.len() >= capacity {
@@ -209,6 +265,40 @@ impl<V: Clone> InMemoryCache<V> {
         state.values.remove(key);
         state.expirations.remove(key);
     }
+
+    /// `get_cache` under the held lock: an expired entry is removed and reads as missing.
+    fn live(state: &mut CacheState<V>, key: &str, now: Duration) -> Option<V> {
+        if state
+            .expirations
+            .get(key)
+            .is_some_and(|expiration| *expiration < now)
+        {
+            Self::remove(state, key);
+        }
+        state.values.get(key).cloned()
+    }
+
+    /// Python `set_cache` under the held lock: evict first (even when `key` already exists),
+    /// then skip oversized values, then write, keeping a live key's expiry.
+    fn store(
+        &self,
+        state: &mut CacheState<V>,
+        key: String,
+        value: V,
+        ttl: Option<Duration>,
+        now: Duration,
+    ) -> Result<CacheWrite, Error> {
+        Self::evict(state, self.max_size_in_memory, now, None);
+        if !self.check_value_size(&value)? {
+            return Ok(CacheWrite::TooLarge);
+        }
+        let expiration = state.expirations.get(&key).copied();
+        if expiration.is_none_or(|expiration| expiration < now) {
+            Self::set_expiration(state, &key, now + ttl.unwrap_or(self.default_ttl));
+        }
+        state.values.insert(key, value);
+        Ok(CacheWrite::Stored)
+    }
 }
 
 impl<V> ClaimCache for InMemoryCache<V>
@@ -227,7 +317,7 @@ where
         }
         let now = (self.now)();
         let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
-        Self::evict(&mut state, self.max_size_in_memory, now, key);
+        Self::evict(&mut state, self.max_size_in_memory, now, Some(key));
         let existing = state
             .values
             .get(key)
@@ -262,38 +352,12 @@ impl CounterCache for InMemoryCache<f64> {
         }
         let now = (self.now)();
         let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
-        Self::evict(&mut state, self.max_size_in_memory, now, key);
-        let value = state.values.get(key).copied().unwrap_or_default() + amount;
-        if !state.expirations.contains_key(key) {
-            Self::set_expiration(
-                &mut state,
-                key,
-                now + self.get_ttl(&context).unwrap_or(self.default_ttl),
-            );
-        }
-        state.values.insert(key.into(), value);
+        let value = Self::live(&mut state, key, now).unwrap_or_default() + amount;
+        self.store(&mut state, key.into(), value, self.get_ttl(&context), now)?;
         Ok(value)
     }
 }
 
-impl InMemoryCache<f64> {
-    pub async fn async_increment_pipeline(
-        &self,
-        operations: Vec<IncrementOperation>,
-    ) -> Result<Vec<f64>, Error> {
-        operations
-            .into_iter()
-            .map(|operation| {
-                self.increment_cache(
-                    &operation.key,
-                    operation.amount,
-                    ExactCacheContext { ttl: operation.ttl },
-                )
-            })
-            .collect()
-    }
-}
-
 impl<V: Clone + Send + Sync + 'static> BaseCache for InMemoryCache<V> {
     type Value = V;
     type Context = ExactCacheContext;
@@ -315,18 +379,12 @@ impl<V: Clone + Send + Sync + 'static> BaseCache for InMemoryCache<V> {
     fn get_cache(&self, key: &str, _: &ExactCacheContext) -> Result<Option<Self::Value>, Error> {
         self.get_cache(key)
     }
+}
 
+impl<V: Clone + Send + Sync + 'static> DisconnectCache for InMemoryCache<V> {
     async fn disconnect(&self) -> Result<(), Error> {
         Ok(())
     }
-
-    async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
-        Ok(CacheConnectionResult {
-            status: CacheConnectionStatus::Success,
-            message: "In-memory cache connection test successful".into(),
-            error: None,
-        })
-    }
 }
 
 impl<V: Clone + Send + Sync + 'static> BatchCache for InMemoryCache<V> {}
@@ -367,34 +425,9 @@ where
         }
         let now = (self.now)();
         let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
-        Self::evict(&mut state, self.max_size_in_memory, now, key);
-        let mut stored = state.values.get(key).cloned().unwrap_or_default();
+        let mut stored = Self::live(&mut state, key, now).unwrap_or_default();
         stored.extend(values.iter().cloned());
-        if let (Some(limit), Some(measure)) = (self.max_entry_bytes, &self.measure_value)
-            && measure(&stored)? > limit
-        {
-            return Ok(values);
-        }
-        if !state.expirations.contains_key(key) {
-            Self::set_expiration(&mut state, key, now + ttl.unwrap_or(self.default_ttl));
-        }
-        state.values.insert(key.into(), stored);
+        self.store(&mut state, key.into(), stored, ttl, now)?;
         Ok(values)
     }
 }
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[test]
-    fn repeated_increments_keep_one_heap_entry_per_expiration() {
-        let cache = InMemoryCache::<f64>::new(Some(4), None);
-        for _ in 0..100 {
-            cache
-                .increment_cache("counter", 1.0, ExactCacheContext::default())
-                .unwrap();
-        }
-        assert_eq!(cache.state.lock().unwrap().expiration_heap.len(), 1);
-    }
-}

litellm-rust/crates/cache-memory/tests/cache.rs

@@ -8,111 +8,296 @@ use std::{
 };
 
 use litellm_cache::{
-    BaseCache, CacheBackend, CacheConnectionStatus, ClaimCache, CounterCache, DeleteCache, Error,
-    ExactCacheContext, IncrementOperation, SetCache, get_cache, set_cache,
+    BaseCache, BatchCache, BatchEntry, CacheBackend, ClaimCache, CounterCache, DeleteCache,
+    DisconnectCache, Error, ExactCacheContext, FlushCache, IncrementOperation, SetCache, TtlCache,
+    get_cache, set_cache,
 };
 use litellm_cache_memory::{CacheWrite, InMemoryCache};
 use rstest::{fixture, rstest};
 
+type Clock = Arc<AtomicU64>;
+
 #[fixture]
-fn clock() -> Arc<AtomicU64> {
+fn clock() -> Clock {
     Arc::new(AtomicU64::new(100))
 }
 
-fn cache(clock: Arc<AtomicU64>, capacity: usize) -> InMemoryCache<String> {
+fn cache_with<V: Clone>(clock: &Clock, capacity: usize) -> InMemoryCache<V> {
+    let clock = clock.clone();
     InMemoryCache::with_clock(Some(capacity), Some(Duration::from_secs(60)), move || {
-        Duration::from_secs(clock.load(Ordering::SeqCst))
+        Duration::from_millis(clock.load(Ordering::SeqCst) * 1000)
     })
 }
 
+fn cache(clock: &Clock, capacity: usize) -> InMemoryCache<String> {
+    cache_with(clock, capacity)
+}
+
+fn at(clock: &Clock, seconds: u64) {
+    clock.store(seconds, Ordering::SeqCst);
+}
+
+fn secs(seconds: u64) -> Option<Duration> {
+    Some(Duration::from_secs(seconds))
+}
+
+fn ttl(seconds: u64) -> ExactCacheContext {
+    ExactCacheContext { ttl: secs(seconds) }
+}
+
+fn measured(capacity: usize) -> InMemoryCache<String> {
+    InMemoryCache::with_clock_and_size_measurement(
+        Some(capacity),
+        secs(60),
+        Some(4),
+        Some(Arc::new(|value: &String| {
+            if value.is_empty() {
+                return Err(Error::InvalidEntry);
+            }
+            Ok(value.len())
+        })),
+        || Duration::from_secs(100),
+    )
+}
+
 #[rstest]
-fn default_explicit_and_override_ttls_follow_python_rules(clock: Arc<AtomicU64>) {
-    let cache = cache(clock.clone(), 4);
+fn default_explicit_and_override_ttls_follow_python_rules(clock: Clock) {
+    let cache = cache(&clock, 4);
     cache.set_cache("key", "first".into(), None).unwrap();
-    assert_eq!(
-        cache.expires_at("key").unwrap(),
-        Some(Duration::from_secs(160))
-    );
-    cache
-        .set_cache("key", "second".into(), Some(Duration::from_secs(10)))
-        .unwrap();
-    assert_eq!(
-        cache.expires_at("key").unwrap(),
-        Some(Duration::from_secs(160))
-    );
-    clock.store(160, Ordering::SeqCst);
+    assert_eq!(cache.expires_at("key").unwrap(), secs(160));
+    cache.set_cache("key", "second".into(), secs(10)).unwrap();
+    assert_eq!(cache.expires_at("key").unwrap(), secs(160));
+    at(&clock, 160);
     assert_eq!(cache.get_cache("key").unwrap(), Some("second".into()));
-    clock.store(161, Ordering::SeqCst);
+    at(&clock, 161);
     assert_eq!(cache.get_cache("key").unwrap(), None);
-    cache
-        .set_cache("key", "third".into(), Some(Duration::from_secs(10)))
-        .unwrap();
-    assert_eq!(
-        cache.expires_at("key").unwrap(),
-        Some(Duration::from_secs(171))
-    );
+    assert_eq!(cache.expires_at("key").unwrap(), None);
+    cache.set_cache("key", "third".into(), secs(10)).unwrap();
+    assert_eq!(cache.expires_at("key").unwrap(), secs(171));
 }
 
 #[rstest]
-fn write_at_expiry_boundary_refreshes_ttl(clock: Arc<AtomicU64>) {
-    let cache = cache(clock.clone(), 4);
-    cache
-        .set_cache("key", "first".into(), Some(Duration::from_secs(10)))
-        .unwrap();
-    clock.store(110, Ordering::SeqCst);
-    cache
-        .set_cache("key", "second".into(), Some(Duration::from_secs(10)))
-        .unwrap();
-    assert_eq!(
-        cache.expires_at("key").unwrap(),
-        Some(Duration::from_secs(120))
-    );
-    clock.store(115, Ordering::SeqCst);
+#[case::unset(None, secs(600))]
+#[case::zero_falls_back_like_python_or(Some(Duration::ZERO), secs(600))]
+#[case::explicit(secs(5), secs(5))]
+fn default_ttl_falls_back_to_ten_minutes(
+    #[case] default_ttl: Option<Duration>,
+    #[case] expected: Option<Duration>,
+) {
+    let cache = InMemoryCache::<String>::with_clock(None, default_ttl, || Duration::ZERO);
+    assert_eq!(cache.get_ttl(&ExactCacheContext::default()), expected);
+    cache.set_cache("key", "value".into(), None).unwrap();
+    assert_eq!(cache.expires_at("key").unwrap(), expected);
+    assert_eq!(cache.max_size_in_memory(), 200);
+}
+
+#[rstest]
+fn write_at_expiry_boundary_refreshes_ttl(clock: Clock) {
+    let cache = cache(&clock, 4);
+    cache.set_cache("key", "first".into(), secs(10)).unwrap();
+    at(&clock, 110);
+    cache.set_cache("key", "second".into(), secs(10)).unwrap();
+    assert_eq!(cache.expires_at("key").unwrap(), secs(120));
+    at(&clock, 115);
     assert_eq!(cache.get_cache("key").unwrap(), Some("second".into()));
 }
 
 #[rstest]
-fn capacity_evicts_earliest_and_ignores_stale_heap_entries(clock: Arc<AtomicU64>) {
-    let cache = cache(clock, 2);
-    cache
-        .set_cache("early", "a".into(), Some(Duration::from_secs(10)))
-        .unwrap();
-    cache
-        .set_cache("late", "b".into(), Some(Duration::from_secs(20)))
-        .unwrap();
+fn expired_key_without_a_read_allows_a_ttl_override(clock: Clock) {
+    let cache = cache(&clock, 4);
+    cache.set_cache("key", "first".into(), secs(1)).unwrap();
+    assert_eq!(cache.allow_ttl_override("key"), Ok(false));
+    at(&clock, 102);
+    assert_eq!(cache.allow_ttl_override("key"), Ok(true));
+    cache.set_cache("key", "second".into(), secs(1)).unwrap();
+    assert_eq!(cache.expires_at("key").unwrap(), secs(103));
+    assert_eq!(cache.allow_ttl_override("missing"), Ok(true));
+}
+
+#[rstest]
+fn capacity_evicts_earliest_and_ignores_stale_heap_entries(clock: Clock) {
+    let cache = cache(&clock, 2);
+    cache.set_cache("early", "a".into(), secs(10)).unwrap();
+    cache.set_cache("late", "b".into(), secs(20)).unwrap();
     cache.delete_cache("early").unwrap();
-    cache
-        .set_cache("new", "c".into(), Some(Duration::from_secs(30)))
-        .unwrap();
+    cache.set_cache("new", "c".into(), secs(30)).unwrap();
     assert_eq!(cache.get_cache("late").unwrap(), Some("b".into()));
-    cache
-        .set_cache("last", "d".into(), Some(Duration::from_secs(40)))
-        .unwrap();
+    cache.set_cache("last", "d".into(), secs(40)).unwrap();
     assert_eq!(cache.get_cache("late").unwrap(), None);
 }
 
-#[test]
-fn disabled_size_limited_and_validated_writes_are_observable() {
-    let cache = |capacity| {
-        InMemoryCache::with_clock_and_size_measurement(
-            Some(capacity),
-            Some(Duration::from_secs(60)),
-            Some(4),
-            Some(Arc::new(|value: &String| {
-                if value.is_empty() {
-                    return Err(Error::InvalidEntry);
-                }
-                Ok(value.len())
-            })),
-            || Duration::from_secs(100),
-        )
-    };
-    let disabled = cache(0);
+#[rstest]
+fn max_size_is_respected_when_every_item_has_a_long_ttl(clock: Clock) {
+    let cache = cache(&clock, 3);
+    for index in 0..3 {
+        at(&clock, 100 + index);
+        cache
+            .set_cache(
+                format!("key_{index}"),
+                format!("value_{index}"),
+                secs(86_400),
+            )
+            .unwrap();
+    }
+    assert_eq!(cache.len(), Ok(3));
+    cache
+        .set_cache("key_3", "value_3".into(), secs(86_400))
+        .unwrap();
+    assert_eq!(cache.len(), Ok(3));
+    assert_eq!(cache.get_cache("key_0").unwrap(), None);
+    assert_eq!(cache.expires_at("key_0").unwrap(), None);
+    for key in ["key_1", "key_2", "key_3"] {
+        assert!(cache.get_cache(key).unwrap().is_some(), "{key}");
+    }
+}
+
+#[rstest]
+fn expired_items_are_evicted_before_live_ones(clock: Clock) {
+    let cache = cache(&clock, 3);
+    cache.set_cache("expired_1", "1".into(), secs(1)).unwrap();
+    cache.set_cache("expired_2", "2".into(), secs(1)).unwrap();
+    cache
+        .set_cache("long_lived", "3".into(), secs(86_400))
+        .unwrap();
+    assert_eq!(cache.len(), Ok(3));
+    at(&clock, 102);
+    cache
+        .set_cache("new_item", "4".into(), secs(86_400))
+        .unwrap();
+    assert_eq!(cache.len(), Ok(2));
+    assert_eq!(cache.get_cache("long_lived").unwrap(), Some("3".into()));
+    assert_eq!(cache.get_cache("new_item").unwrap(), Some("4".into()));
+    for key in ["expired_1", "expired_2"] {
+        assert_eq!(cache.expires_at(key).unwrap(), None, "{key}");
+    }
+}
+
+#[rstest]
+fn injected_clock_controls_expiry_and_eviction(clock: Clock) {
+    let cache = cache(&clock, 2);
+    at(&clock, 0);
+    cache
+        .set_cache("first", "original".into(), secs(10))
+        .unwrap();
+    at(&clock, 9);
+    cache.set_cache("second", "survivor".into(), None).unwrap();
+    assert_eq!(cache.get_cache("first").unwrap(), Some("original".into()));
+    at(&clock, 11);
+    assert_eq!(cache.get_cache("first").unwrap(), None);
+    cache
+        .set_cache("third", "replacement".into(), None)
+        .unwrap();
+    assert_eq!(cache.get_cache("second").unwrap(), Some("survivor".into()));
+    at(&clock, 70);
+    cache.set_cache("fourth", "new".into(), None).unwrap();
+    assert_eq!(cache.get_cache("second").unwrap(), None);
     assert_eq!(
-        disabled.set_cache("a", "x".into(), None).unwrap(),
+        cache.get_cache("third").unwrap(),
+        Some("replacement".into())
+    );
+    assert_eq!(cache.get_cache("fourth").unwrap(), Some("new".into()));
+}
+
+#[rstest]
+fn rewriting_one_key_keeps_one_heap_entry(clock: Clock) {
+    let cache = cache(&clock, 10);
+    for index in 0..1_000 {
+        cache
+            .set_cache("hot_key", format!("value_{index}"), secs(60))
+            .unwrap();
+    }
+    assert_eq!(cache.expiration_heap_len(), Ok(1));
+}
+
+#[rstest]
+fn repeated_increments_keep_one_heap_entry_per_expiration() {
+    let cache = InMemoryCache::<f64>::new(Some(4), None);
+    for _ in 0..100 {
+        cache
+            .increment_cache("counter", 1.0, ExactCacheContext::default())
+            .unwrap();
+    }
+    assert_eq!(cache.expiration_heap_len(), Ok(1));
+}
+
+#[rstest]
+fn reinserting_expired_keys_below_capacity_prunes_the_heap(clock: Clock) {
+    let cache = cache(&clock, 200);
+    for cycle in 0..3 {
+        for index in 0..5 {
+            cache
+                .set_cache(format!("key_{index}"), format!("value_{cycle}"), secs(1))
+                .unwrap();
+        }
+        at(&clock, 100 + 2 * (cycle + 1));
+    }
+    for index in 0..5 {
+        cache
+            .set_cache(format!("key_{index}"), "final".into(), secs(1))
+            .unwrap();
+    }
+    assert_eq!(cache.len(), Ok(5));
+    assert_eq!(cache.expiration_heap_len(), Ok(5));
+}
+
+#[rstest]
+fn evict_cache_drops_expired_entries_then_makes_room(clock: Clock) {
+    let cache = cache(&clock, 2);
+    assert_eq!(cache.is_empty(), Ok(true));
+    cache.set_cache("short", "a".into(), secs(1)).unwrap();
+    cache.set_cache("long", "b".into(), secs(50)).unwrap();
+    at(&clock, 102);
+    cache.evict_cache().unwrap();
+    assert_eq!(cache.len(), Ok(1));
+    assert_eq!(cache.expires_at("short").unwrap(), None);
+    cache.set_cache("longer", "c".into(), secs(90)).unwrap();
+    cache.evict_cache().unwrap();
+    assert_eq!(cache.len(), Ok(1));
+    assert_eq!(cache.get_cache("long").unwrap(), None);
+    assert_eq!(cache.get_cache("longer").unwrap(), Some("c".into()));
+}
+
+#[rstest]
+fn evict_element_if_expired_reports_removal(clock: Clock) {
+    let cache = cache(&clock, 4);
+    cache.set_cache("key", "value".into(), secs(10)).unwrap();
+    assert_eq!(cache.evict_element_if_expired("key"), Ok(false));
+    assert_eq!(cache.evict_element_if_expired("missing"), Ok(false));
+    at(&clock, 110);
+    assert_eq!(cache.evict_element_if_expired("key"), Ok(false));
+    at(&clock, 111);
+    assert_eq!(cache.evict_element_if_expired("key"), Ok(true));
+    assert_eq!(cache.len(), Ok(0));
+    assert_eq!(cache.expires_at("key").unwrap(), None);
+}
+
+#[rstest]
+#[case::fits("ok", Ok(true))]
+#[case::at_limit("four", Ok(true))]
+#[case::too_large("oversized", Ok(false))]
+#[case::measure_error("", Err(Error::InvalidEntry))]
+fn check_value_size_applies_the_entry_limit(
+    #[case] value: &str,
+    #[case] expected: Result<bool, Error>,
+) {
+    assert_eq!(measured(2).check_value_size(&value.to_string()), expected);
+}
+
+#[rstest]
+fn values_are_unbounded_without_a_measure() {
+    let cache = InMemoryCache::<String>::default();
+    assert_eq!(cache.max_entry_bytes(), None);
+    assert_eq!(cache.check_value_size(&"x".repeat(1 << 20)), Ok(true));
+}
+
+#[rstest]
+fn disabled_size_limited_and_validated_writes_are_observable() {
+    assert_eq!(
+        measured(0).set_cache("a", "x".into(), None).unwrap(),
         CacheWrite::Disabled
     );
-    let cache = cache(2);
+    let cache = measured(2);
+    assert_eq!(cache.max_entry_bytes(), Some(4));
     assert_eq!(
         cache.set_cache("large", "oversized".into(), None).unwrap(),
         CacheWrite::TooLarge
@@ -132,30 +317,21 @@ fn disabled_size_limited_and_validated_writes_are_observable() {
     assert_eq!(cache.get_cache("small").unwrap(), None);
 }
 
+#[rstest]
 #[tokio::test]
-async fn connection_test_matches_python_result_contract() {
+async fn disconnect_is_a_no_op_that_keeps_entries() {
     let cache = InMemoryCache::<String>::default();
-    let result = BaseCache::test_connection(&cache).await.unwrap();
-    assert_eq!(result.status, CacheConnectionStatus::Success);
-    assert_eq!(result.message, "In-memory cache connection test successful");
-    assert_eq!(result.error, None);
-    assert_eq!(
-        serde_json::to_value(result).unwrap(),
-        serde_json::json!({
-            "status": "success",
-            "message": "In-memory cache connection test successful"
-        })
-    );
+    cache.set_cache("key", "value".into(), None).unwrap();
+    cache.disconnect().await.unwrap();
+    assert_eq!(cache.get_cache("key").unwrap(), Some("value".into()));
 }
 
+#[rstest]
 #[tokio::test]
-async fn generic_consumers_share_typed_values_and_honor_expiration() {
-    let clock = clock();
-    let cache: CacheBackend<InMemoryCache<String>> = Arc::new(cache(clock.clone(), 4));
+async fn generic_consumers_share_typed_values_and_honor_expiration(clock: Clock) {
+    let cache: CacheBackend<InMemoryCache<String>> = Arc::new(self::cache(&clock, 4));
     let reader = Arc::clone(&cache);
-    let context = ExactCacheContext {
-        ttl: Some(Duration::from_secs(5)),
-    };
+    let context = ttl(5);
     set_cache(cache.as_ref(), "sync", "first".into(), &context).unwrap();
     assert_eq!(
         get_cache(reader.as_ref(), "sync", &context).unwrap(),
@@ -181,7 +357,7 @@ async fn generic_consumers_share_typed_values_and_honor_expiration() {
         reader.async_get_cache("async", &context).await.unwrap(),
         None
     );
-    clock.store(106, Ordering::SeqCst);
+    at(&clock, 106);
     assert_eq!(get_cache(reader.as_ref(), "sync", &context).unwrap(), None);
     assert_eq!(
         reader.async_get_cache("batch", &context).await.unwrap(),
@@ -189,34 +365,95 @@ async fn generic_consumers_share_typed_values_and_honor_expiration() {
     );
 }
 
-#[test]
-fn claims_are_atomic_and_refresh_eligible_winners() {
-    let clock = clock();
-    let cache = InMemoryCache::with_clock(Some(4), Some(Duration::from_secs(60)), {
-        let clock = clock.clone();
-        move || Duration::from_secs(clock.load(Ordering::SeqCst))
-    });
-    let context = ExactCacheContext {
-        ttl: Some(Duration::from_secs(10)),
-    };
+#[rstest]
+#[case::context_ttl(ttl(5), secs(105))]
+#[case::default_ttl(ExactCacheContext::default(), secs(160))]
+#[tokio::test]
+async fn pipeline_writes_use_the_context_ttl_or_the_default(
+    clock: Clock,
+    #[case] context: ExactCacheContext,
+    #[case] expected: Option<Duration>,
+) {
+    let cache = cache(&clock, 4);
+    cache
+        .async_set_cache_pipeline(
+            vec![("a".into(), "1".into()), ("b".into(), "2".into())],
+            context,
+        )
+        .await
+        .unwrap();
+    assert_eq!(cache.expires_at("a").unwrap(), expected);
+    assert_eq!(cache.expires_at("b").unwrap(), expected);
+}
+
+#[rstest]
+#[tokio::test]
+async fn batch_reads_return_one_entry_per_key_and_drop_expired_ones(clock: Clock) {
+    let cache = cache(&clock, 4);
+    cache.set_cache("short", "a".into(), secs(1)).unwrap();
+    cache.set_cache("long", "b".into(), secs(50)).unwrap();
+    let keys = vec!["short".to_string(), "missing".into(), "long".into()];
+    assert_eq!(
+        cache
+            .batch_get_cache(&keys, &ExactCacheContext::default())
+            .unwrap(),
+        [
+            BatchEntry::Hit("a".to_string()),
+            BatchEntry::Miss,
+            BatchEntry::Hit("b".into()),
+        ]
+    );
+    at(&clock, 102);
+    assert_eq!(
+        cache
+            .async_batch_get_cache(keys, ExactCacheContext::default())
+            .await
+            .unwrap(),
+        [
+            BatchEntry::Miss,
+            BatchEntry::Miss,
+            BatchEntry::Hit("b".into())
+        ]
+    );
+}
+
+#[rstest]
+#[tokio::test]
+async fn flush_clears_values_and_expirations(clock: Clock) {
+    let cache = cache(&clock, 4);
+    cache.set_cache("a", "1".into(), None).unwrap();
+    cache.set_cache("b", "2".into(), None).unwrap();
+    cache.flush_cache().unwrap();
+    assert_eq!(cache.len(), Ok(0));
+    assert_eq!(cache.expiration_heap_len(), Ok(0));
+    cache.set_cache("c", "3".into(), None).unwrap();
+    FlushCache::async_flush_cache(&cache).await.unwrap();
+    assert_eq!(cache.is_empty(), Ok(true));
+    assert_eq!(
+        cache.async_get_oldest_n_keys(5).await.unwrap(),
+        Vec::<String>::new()
+    );
+}
+
+#[rstest]
+fn claims_are_atomic_and_refresh_eligible_winners(clock: Clock) {
+    let cache = cache(&clock, 4);
+    let context = ttl(10);
     assert_eq!(
         cache
             .claim_cache("affinity", "first".to_string(), &[], context.clone())
             .unwrap(),
         "first"
     );
-    clock.store(103, Ordering::SeqCst);
+    at(&clock, 103);
     assert_eq!(
         cache
             .claim_cache("affinity", "second".to_string(), &[], context.clone())
             .unwrap(),
         "first"
     );
-    assert_eq!(
-        cache.expires_at("affinity").unwrap(),
-        Some(Duration::from_secs(110))
-    );
-    clock.store(105, Ordering::SeqCst);
+    assert_eq!(cache.expires_at("affinity").unwrap(), secs(110));
+    at(&clock, 105);
     assert_eq!(
         cache
             .claim_cache(
@@ -228,13 +465,10 @@ fn claims_are_atomic_and_refresh_eligible_winners() {
             .unwrap(),
         "first"
     );
-    assert_eq!(
-        cache.expires_at("affinity").unwrap(),
-        Some(Duration::from_secs(115))
-    );
+    assert_eq!(cache.expires_at("affinity").unwrap(), secs(115));
 }
 
-#[test]
+#[rstest]
 fn counters_increment_under_one_lock() {
     let cache = InMemoryCache::<f64>::default();
     assert_eq!(
@@ -250,44 +484,107 @@ fn counters_increment_under_one_lock() {
 }
 
 #[rstest]
-fn rewriting_an_existing_key_at_capacity_keeps_other_entries(clock: Arc<AtomicU64>) {
-    let cache = cache(clock, 2);
-    cache
-        .set_cache("hot", "1".into(), Some(Duration::from_secs(10)))
-        .unwrap();
-    cache
-        .set_cache("cold", "2".into(), Some(Duration::from_secs(20)))
-        .unwrap();
+fn concurrent_increments_are_atomic() {
+    let cache = Arc::new(InMemoryCache::<f64>::default());
+    cache.set_cache("counter", 1000.0, None).unwrap();
+    let threads = (0..8)
+        .map(|_| {
+            let cache = cache.clone();
+            std::thread::spawn(move || {
+                cache
+                    .increment_cache("counter", 1.0, ExactCacheContext::default())
+                    .unwrap()
+            })
+        })
+        .collect::<Vec<_>>();
+    for thread in threads {
+        thread.join().unwrap();
+    }
+    assert_eq!(cache.get_cache("counter").unwrap(), Some(1008.0));
+}
+
+#[rstest]
+#[case::window_semantics(false)]
+#[case::refresh_ttl_is_ignored(true)]
+#[tokio::test]
+async fn async_increment_delegates_to_the_locked_sync_path(
+    clock: Clock,
+    #[case] refresh_ttl: bool,
+) {
+    let cache = cache_with::<f64>(&clock, 4);
+    assert_eq!(
+        cache
+            .async_increment("counter", 2.0, ttl(10), refresh_ttl)
+            .await,
+        Ok(2.0)
+    );
+    at(&clock, 105);
+    assert_eq!(
+        cache
+            .async_increment("counter", 3.0, ttl(10), refresh_ttl)
+            .await,
+        Ok(5.0)
+    );
+    assert_eq!(cache.get_cache("counter").unwrap(), Some(5.0));
+    assert_eq!(cache.expires_at("counter").unwrap(), secs(110));
+}
+
+#[rstest]
+fn expired_counters_restart_from_zero_with_a_new_ttl(clock: Clock) {
+    let cache = cache_with::<f64>(&clock, 4);
+    cache.increment_cache("counter", 2.0, ttl(10)).unwrap();
+    at(&clock, 111);
+    assert_eq!(cache.increment_cache("counter", 1.0, ttl(10)), Ok(1.0));
+    assert_eq!(cache.expires_at("counter").unwrap(), secs(121));
+}
+
+/// Python `InMemoryCache.set_cache` runs `evict_cache()` before every insert, and step 2 evicts
+/// the earliest expiry while `len(cache_dict) >= max_size_in_memory`, even when the key being
+/// written already exists.
+#[rstest]
+fn overwriting_an_existing_key_at_capacity_evicts_the_earliest_expiry_like_python(clock: Clock) {
+    let cache = cache(&clock, 2);
+    cache.set_cache("hot", "1".into(), secs(10)).unwrap();
+    cache.set_cache("cold", "2".into(), secs(20)).unwrap();
 
     cache.set_cache("cold", "3".into(), None).unwrap();
-    assert_eq!(cache.get_cache("hot").unwrap(), Some("1".into()));
+
+    assert_eq!(cache.get_cache("hot").unwrap(), None);
     assert_eq!(cache.get_cache("cold").unwrap(), Some("3".into()));
+}
+
+/// `claim_cache` has no Python counterpart; it never evicts another entry for a key it holds.
+#[rstest]
+fn claiming_an_existing_key_at_capacity_keeps_other_entries(clock: Clock) {
+    let cache = cache(&clock, 2);
+    cache.set_cache("hot", "1".into(), secs(10)).unwrap();
+    cache.set_cache("cold", "2".into(), secs(20)).unwrap();
 
     cache
         .claim_cache("cold", "4".into(), &[], ExactCacheContext::default())
         .unwrap();
     assert_eq!(cache.get_cache("hot").unwrap(), Some("1".into()));
-
-    cache.set_cache("new", "5".into(), None).unwrap();
-    assert_eq!(cache.get_cache("hot").unwrap(), None);
-    assert_eq!(cache.get_cache("cold").unwrap(), Some("3".into()));
-    assert_eq!(cache.get_cache("new").unwrap(), Some("5".into()));
+    assert_eq!(cache.get_cache("cold").unwrap(), Some("2".into()));
 }
 
-#[test]
-fn incrementing_an_existing_counter_at_capacity_keeps_every_counter() {
-    let cache = InMemoryCache::<f64>::new(Some(2), None);
+/// Python `increment_cache` is `get_cache` then `set_cache`, so at capacity the write evicts
+/// the earliest expiry first: equal expiries tie-break on the key, and the value read before
+/// eviction is the one written back.
+#[rstest]
+fn incrementing_at_capacity_evicts_the_earliest_expiry_like_python(clock: Clock) {
+    let cache = cache_with::<f64>(&clock, 2);
     for key in ["a", "b", "a", "b"] {
         cache
             .increment_cache(key, 1.0, ExactCacheContext::default())
             .unwrap();
     }
-    assert_eq!(cache.get_cache("a").unwrap(), Some(2.0));
+    assert_eq!(cache.get_cache("a").unwrap(), None);
     assert_eq!(cache.get_cache("b").unwrap(), Some(2.0));
 }
 
-#[test]
-fn disabled_cache_does_not_retain_claims_or_counters() {
+#[rstest]
+#[tokio::test]
+async fn disabled_cache_does_not_retain_claims_counters_or_sets() {
     let claims = InMemoryCache::<String>::new(Some(0), None);
     assert_eq!(
         claims
@@ -305,62 +602,128 @@ fn disabled_cache_does_not_retain_claims_or_counters() {
         2.0
     );
     assert_eq!(counters.get_cache("key").unwrap(), None);
+
+    let sets = InMemoryCache::<HashSet<String>>::new(Some(0), None);
+    assert_eq!(
+        sets.async_set_cache_sadd("key", vec!["a".into()], None)
+            .await
+            .unwrap(),
+        ["a"]
+    );
+    assert_eq!(sets.get_cache("key").unwrap(), None);
 }
 
+#[rstest]
 #[tokio::test]
-async fn ttl_and_oldest_key_operations_use_the_stored_expirations() {
-    let clock = Arc::new(AtomicU64::new(100));
-    let cache = cache(clock, 3);
-    cache
-        .set_cache("later", "2".into(), Some(Duration::from_secs(20)))
-        .unwrap();
-    cache
-        .set_cache("first", "1".into(), Some(Duration::from_secs(10)))
-        .unwrap();
+async fn ttl_and_oldest_key_operations_use_the_stored_expirations(clock: Clock) {
+    let cache = cache(&clock, 3);
+    cache.set_cache("later", "2".into(), secs(20)).unwrap();
+    cache.set_cache("first", "1".into(), secs(10)).unwrap();
+    cache.set_cache("latest", "3".into(), secs(30)).unwrap();
 
+    assert_eq!(cache.async_get_ttl("first").await.unwrap(), secs(110));
     assert_eq!(
-        cache.async_get_ttl("first").await.unwrap(),
-        Some(Duration::from_secs(110))
+        TtlCache::async_get_ttl(&cache, "later").await.unwrap(),
+        secs(120)
     );
     assert_eq!(cache.async_get_oldest_n_keys(1).await.unwrap(), ["first"]);
+    assert_eq!(
+        cache.async_get_oldest_n_keys(10).await.unwrap(),
+        ["first", "later", "latest"]
+    );
+    assert_eq!(
+        cache.async_get_oldest_n_keys(0).await.unwrap(),
+        Vec::<String>::new()
+    );
     assert_eq!(cache.async_get_ttl("missing").await.unwrap(), None);
 }
 
+#[rstest]
 #[tokio::test]
-async fn increment_pipeline_preserves_operation_order() {
-    let cache = InMemoryCache::<f64>::new(Some(3), None);
+async fn increment_pipeline_preserves_operation_order(clock: Clock) {
+    let cache = cache_with::<f64>(&clock, 3);
+    let operation = |key: &str, amount, ttl| IncrementOperation {
+        key: key.into(),
+        amount,
+        ttl: secs(ttl),
+    };
     assert_eq!(
         cache
             .async_increment_pipeline(vec![
-                IncrementOperation {
-                    key: "a".into(),
-                    amount: 1.0,
-                    ttl: Some(Duration::from_secs(10)),
-                },
-                IncrementOperation {
-                    key: "a".into(),
-                    amount: 2.0,
-                    ttl: Some(Duration::from_secs(20)),
-                },
+                operation("a", 1.0, 10),
+                operation("b", 5.0, 30),
+                operation("a", 2.0, 20),
             ])
             .await
             .unwrap(),
-        [1.0, 3.0]
+        [1.0, 5.0, 3.0]
     );
     assert_eq!(cache.get_cache("a").unwrap(), Some(3.0));
+    assert_eq!(cache.expires_at("a").unwrap(), secs(110));
+    assert_eq!(cache.expires_at("b").unwrap(), secs(130));
+    assert_eq!(
+        cache.async_increment_pipeline(Vec::new()).await.unwrap(),
+        Vec::<f64>::new()
+    );
 }
 
+#[rstest]
 #[tokio::test]
-async fn set_capability_preserves_python_result_and_deduplicates_storage() {
-    let cache = InMemoryCache::<HashSet<String>>::new(None, None);
-    let inserted = vec!["a".into(), "a".into(), "b".into()];
+async fn set_capability_preserves_python_result_and_deduplicates_storage(clock: Clock) {
+    let cache = cache_with::<HashSet<String>>(&clock, 4);
+    let inserted = vec!["a".to_string(), "a".into(), "b".into()];
     assert_eq!(
         cache
-            .async_set_cache_sadd("members", inserted.clone(), None)
+            .async_set_cache_sadd("members", inserted.clone(), secs(10))
             .await
             .unwrap(),
         inserted
     );
+    assert_eq!(
+        cache
+            .async_set_cache_sadd("members", vec!["c".into()], secs(99))
+            .await
+            .unwrap(),
+        ["c"]
+    );
+    assert_eq!(
+        cache.get_cache("members").unwrap(),
+        Some(HashSet::from(["a".into(), "b".into(), "c".into()]))
+    );
+    assert_eq!(cache.expires_at("members").unwrap(), secs(110));
+    at(&clock, 111);
+    cache
+        .async_set_cache_sadd("members", vec!["d".into()], None)
+        .await
+        .unwrap();
+    assert_eq!(
+        cache.get_cache("members").unwrap(),
+        Some(HashSet::from(["d".into()]))
+    );
+    assert_eq!(cache.expires_at("members").unwrap(), secs(171));
+}
+
+#[rstest]
+#[tokio::test]
+async fn oversized_set_additions_are_not_stored() {
+    let cache = InMemoryCache::<HashSet<String>>::with_clock_and_size_measurement(
+        Some(4),
+        None,
+        Some(2),
+        Some(Arc::new(|value: &HashSet<String>| Ok(value.len()))),
+        || Duration::ZERO,
+    );
+    cache
+        .async_set_cache_sadd("members", vec!["a".into(), "b".into()], None)
+        .await
+        .unwrap();
+    assert_eq!(
+        cache
+            .async_set_cache_sadd("members", vec!["c".into()], None)
+            .await
+            .unwrap(),
+        ["c"]
+    );
     assert_eq!(
         cache.get_cache("members").unwrap(),
         Some(HashSet::from(["a".into(), "b".into()]))

litellm-rust/crates/cache-memory/tests/contract.rs

@@ -0,0 +1,98 @@
+use std::time::Duration;
+
+use litellm_cache::ExactCacheContext;
+use litellm_cache_memory::InMemoryCache;
+use litellm_cache_testing as contract;
+use rstest::{fixture, rstest};
+
+#[fixture]
+fn strings() -> InMemoryCache<String> {
+    InMemoryCache::new(Some(16), None)
+}
+
+#[fixture]
+fn counters() -> InMemoryCache<f64> {
+    InMemoryCache::new(Some(16), None)
+}
+
+#[fixture]
+fn context() -> ExactCacheContext {
+    ExactCacheContext {
+        ttl: Some(Duration::from_secs(60)),
+    }
+}
+
+#[rstest]
+#[tokio::test]
+async fn hit_and_miss(strings: InMemoryCache<String>, context: ExactCacheContext) {
+    contract::hit_and_miss(&strings, context, "memory:", "value".into()).await;
+}
+
+#[rstest]
+#[tokio::test]
+async fn sync_async_equivalence(strings: InMemoryCache<String>, context: ExactCacheContext) {
+    contract::sync_async_equivalence(
+        &strings,
+        context,
+        "memory:",
+        "first".into(),
+        "second".into(),
+    )
+    .await;
+}
+
+#[rstest]
+#[tokio::test]
+async fn overwrite_replaces(strings: InMemoryCache<String>, context: ExactCacheContext) {
+    contract::overwrite_replaces(
+        &strings,
+        context,
+        "memory:",
+        "first".into(),
+        "second".into(),
+    )
+    .await;
+}
+
+#[rstest]
+#[tokio::test]
+async fn pipeline_writes_every_entry(strings: InMemoryCache<String>, context: ExactCacheContext) {
+    contract::pipeline_writes_every_entry(
+        &strings,
+        context,
+        "memory:",
+        vec!["a".into(), "b".into(), "c".into()],
+    )
+    .await;
+}
+
+#[rstest]
+#[tokio::test]
+async fn batch_preserves_order(strings: InMemoryCache<String>, context: ExactCacheContext) {
+    contract::batch_preserves_order(
+        &strings,
+        context,
+        "memory:",
+        "first".into(),
+        "second".into(),
+    )
+    .await;
+}
+
+#[rstest]
+#[tokio::test]
+async fn delete_removes_key(strings: InMemoryCache<String>, context: ExactCacheContext) {
+    contract::delete_removes_key(&strings, context, "memory:", "value".into()).await;
+}
+
+#[rstest]
+#[tokio::test]
+async fn flush_clears(strings: InMemoryCache<String>, context: ExactCacheContext) {
+    contract::flush_clears(&strings, context, "memory:", "value".into()).await;
+}
+
+#[rstest]
+#[tokio::test]
+async fn counter_accumulates(counters: InMemoryCache<f64>, context: ExactCacheContext) {
+    contract::counter_accumulates(&counters, context, "memory:").await;
+}

litellm-rust/crates/cache-qdrant-semantic/Cargo.toml

@@ -17,7 +17,8 @@ tokio.workspace = true
 uuid.workspace = true
 
 [dev-dependencies]
-litellm-cache-response.workspace = true
+futures-executor = "0.3"
+litellm-cache-testing.workspace = true
 rstest.workspace = true
 tonic = "0.14"
 tonic-prost = "0.14"

litellm-rust/crates/cache-qdrant-semantic/src/cache.rs

@@ -1,7 +1,8 @@
-use std::future::Future;
-
 use futures_util::future::try_join_all;
-use litellm_cache::{BaseCache, CacheCodec, CacheConnectionResult, Error, SemanticCacheContext};
+use litellm_cache::{
+    BaseCache, CacheCodec, Error, SemanticCacheContext,
+    semantic::{Embedder, SemanticCache, SemanticLookup, prompt_from_messages},
+};
 use qdrant_client::{
     Payload, Qdrant,
     qdrant::{
@@ -14,26 +15,7 @@ use qdrant_client::{
 use serde_json::{Map, Value, json};
 use uuid::Uuid;
 
-use crate::prompt_from_messages;
-
-pub trait Embedder: Send + Sync + 'static {
-    fn model(&self) -> &str;
-    fn embed(&self, input: &str) -> impl Future<Output = Result<Vec<f32>, Error>> + Send;
-}
-
-#[derive(Clone, Debug, PartialEq)]
-pub enum Quantization {
-    Binary,
-    Scalar,
-    Product,
-}
-
-pub struct QdrantSemanticConfig {
-    pub collection_name: String,
-    pub similarity_threshold: f64,
-    pub vector_size: u64,
-    pub quantization: Quantization,
-}
+use crate::{QdrantSemanticConfig, Quantization};
 
 pub struct QdrantSemanticCache<E: Embedder, C: CacheCodec> {
     client: Qdrant,
@@ -100,14 +82,9 @@ impl<E: Embedder, C: CacheCodec> QdrantSemanticCache<E, C> {
         &self.embedder
     }
 
+    /// Python reads `kwargs["messages"]` unguarded, so a request without messages fails.
     fn prompt(context: &SemanticCacheContext) -> Result<String, Error> {
-        let Some(messages) = context.messages.as_ref().and_then(Value::as_array) else {
-            return Err(Error::MissingPrompt);
-        };
-        if messages.is_empty() {
-            return Err(Error::MissingPrompt);
-        }
-        Ok(prompt_from_messages(messages))
+        prompt_from_messages(context).ok_or(Error::MissingPrompt)
     }
 
     async fn set(
@@ -117,7 +94,10 @@ impl<E: Embedder, C: CacheCodec> QdrantSemanticCache<E, C> {
         context: &SemanticCacheContext,
     ) -> Result<(), Error> {
         let prompt = Self::prompt(context)?;
-        let vector = self.embedder.embed(&prompt).await?;
+        let vector = self
+            .embedder
+            .async_embed(&prompt, context.metadata.as_ref())
+            .await?;
         let response =
             String::from_utf8(self.codec.encode(&value)?).map_err(|_| Error::InvalidEntry)?;
         let payload = Payload::try_from(json!({
@@ -147,9 +127,12 @@ impl<E: Embedder, C: CacheCodec> QdrantSemanticCache<E, C> {
         &self,
         key: &str,
         context: &SemanticCacheContext,
-    ) -> Result<Option<C::Value>, Error> {
+    ) -> Result<SemanticLookup<C::Value>, Error> {
         let prompt = Self::prompt(context)?;
-        let vector = self.embedder.embed(&prompt).await?;
+        let vector = self
+            .embedder
+            .async_embed(&prompt, context.metadata.as_ref())
+            .await?;
         let result = self
             .client
             .search_points(
@@ -171,20 +154,27 @@ impl<E: Embedder, C: CacheCodec> QdrantSemanticCache<E, C> {
             .await
             .map_err(|_| Error::Unavailable)?;
         let Some(point) = result.result.into_iter().next() else {
-            return Ok(None);
+            return Ok(SemanticLookup::miss(Some(0.0)));
         };
         let payload: Map<String, Value> = Payload::from(point.payload).into();
-        if payload.get("litellm_cache_key").and_then(Value::as_str) != Some(key) {
-            return Ok(None);
+        if !payload
+            .get("litellm_cache_key")
+            .is_some_and(|cached| python_str(cached).as_deref() == Some(key))
+        {
+            return Ok(SemanticLookup::miss(Some(0.0)));
         }
-        if f64::from(point.score) < self.config.similarity_threshold {
-            return Ok(None);
+        let similarity = f64::from(point.score);
+        if similarity < self.config.similarity_threshold {
+            return Ok(SemanticLookup::miss(Some(similarity)));
         }
         let response = payload
             .get("response")
             .and_then(Value::as_str)
             .ok_or(Error::InvalidEntry)?;
-        self.codec.decode(response.as_bytes()).map(Some)
+        Ok(SemanticLookup {
+            value: Some(self.codec.decode(response.as_bytes())?),
+            similarity: Some(similarity),
+        })
     }
 }
 
@@ -219,7 +209,8 @@ impl<E: Embedder, C: CacheCodec> BaseCache for QdrantSemanticCache<E, C> {
     }
 
     fn get_cache(&self, key: &str, context: &Self::Context) -> Result<Option<Self::Value>, Error> {
-        self.runtime.block_on(self.get(key, context))
+        self.get_cache_with_similarity(key, context)
+            .map(|lookup| lookup.value)
     }
 
     async fn async_set_cache(
@@ -236,7 +227,7 @@ impl<E: Embedder, C: CacheCodec> BaseCache for QdrantSemanticCache<E, C> {
         key: &str,
         context: &Self::Context,
     ) -> Result<Option<Self::Value>, Error> {
-        self.get(key, context).await
+        self.get(key, context).await.map(|lookup| lookup.value)
     }
 
     async fn async_set_cache_pipeline(
@@ -251,12 +242,36 @@ impl<E: Embedder, C: CacheCodec> BaseCache for QdrantSemanticCache<E, C> {
         .await
         .map(|_| ())
     }
+}
 
-    async fn disconnect(&self) -> Result<(), Error> {
-        Ok(())
+/// Python stamps the top point's score, even below the threshold, and `0.0` when there is no
+/// point or it belongs to another key. A request without messages fails before any search.
+impl<E: Embedder, C: CacheCodec> SemanticCache for QdrantSemanticCache<E, C> {
+    fn get_cache_with_similarity(
+        &self,
+        key: &str,
+        context: &Self::Context,
+    ) -> Result<SemanticLookup<Self::Value>, Error> {
+        self.runtime.block_on(self.get(key, context))
     }
 
-    async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
-        Err(Error::UnsupportedOperation)
+    async fn async_get_cache_with_similarity(
+        &self,
+        key: &str,
+        context: &Self::Context,
+    ) -> Result<SemanticLookup<Self::Value>, Error> {
+        self.get(key, context).await
+    }
+}
+
+/// `str(value)` for the scalar payload values `_payload_matches_cache_key` compares; `None` for
+/// null (a pre-isolation point without a key) and for containers, which never equal a key.
+fn python_str(value: &Value) -> Option<String> {
+    match value {
+        Value::String(text) => Some(text.clone()),
+        Value::Number(number) => Some(number.to_string()),
+        Value::Bool(true) => Some("True".into()),
+        Value::Bool(false) => Some("False".into()),
+        Value::Null | Value::Array(_) | Value::Object(_) => None,
     }
 }

litellm-rust/crates/cache-qdrant-semantic/src/config.rs

@@ -0,0 +1,13 @@
+#[derive(Clone, Debug, PartialEq)]
+pub enum Quantization {
+    Binary,
+    Scalar,
+    Product,
+}
+
+pub struct QdrantSemanticConfig {
+    pub collection_name: String,
+    pub similarity_threshold: f64,
+    pub vector_size: u64,
+    pub quantization: Quantization,
+}

litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs

@@ -1,11 +1,9 @@
 use std::time::Duration;
 
-use litellm_cache::Error;
+use litellm_cache::{Error, semantic::Embedder};
 use reqwest::Client;
 use serde_json::Value;
 
-use crate::Embedder;
-
 pub struct OpenAiEmbedder {
     client: Client,
     api_base: String,
@@ -31,14 +29,16 @@ impl OpenAiEmbedder {
             timeout: config.timeout,
         }
     }
-}
 
-impl Embedder for OpenAiEmbedder {
-    fn model(&self) -> &str {
+    pub fn model(&self) -> &str {
         &self.model
     }
+}
 
-    async fn embed(&self, input: &str) -> Result<Vec<f32>, Error> {
+/// An OpenAI-compatible `/embeddings` call. It has no router to route on, so `metadata` is
+/// unused, and it only embeds asynchronously: sync cache calls block on the cache's runtime.
+impl Embedder for OpenAiEmbedder {
+    async fn async_embed(&self, input: &str, _metadata: Option<&Value>) -> Result<Vec<f32>, Error> {
         let request = self
             .client
             .post(format!("{}/embeddings", self.api_base))

litellm-rust/crates/cache-qdrant-semantic/src/lib.rs

@@ -1,7 +1,7 @@
+mod cache;
+mod config;
 mod embedder;
-mod prompt;
-mod semantic;
 
+pub use cache::QdrantSemanticCache;
+pub use config::{QdrantSemanticConfig, Quantization};
 pub use embedder::{OpenAiEmbedder, OpenAiEmbedderConfig};
-pub use prompt::prompt_from_messages;
-pub use semantic::{Embedder, QdrantSemanticCache, QdrantSemanticConfig, Quantization};

litellm-rust/crates/cache-qdrant-semantic/src/prompt.rs

@@ -1,59 +0,0 @@
-use serde_json::Value;
-
-fn search_results_text(search_results: Option<&Value>) -> String {
-    let Some(Value::Array(results)) = search_results else {
-        return String::new();
-    };
-    results
-        .iter()
-        .filter_map(Value::as_object)
-        .flat_map(|result| {
-            let source = result
-                .get("source")
-                .and_then(Value::as_str)
-                .map(str::to_owned);
-            let title = result
-                .get("title")
-                .and_then(Value::as_str)
-                .map(str::to_owned);
-            let content = result
-                .get("content")
-                .and_then(Value::as_array)
-                .into_iter()
-                .flatten()
-                .filter_map(Value::as_object)
-                .filter_map(|block| block.get("text").and_then(Value::as_str).map(str::to_owned));
-            let citations = result
-                .get("citations")
-                .filter(|value| !value.is_null())
-                .map(|value| serde_json::to_string(value).unwrap_or_default());
-            source
-                .into_iter()
-                .chain(title)
-                .chain(content)
-                .chain(citations)
-        })
-        .collect()
-}
-
-pub fn prompt_from_messages(messages: &[Value]) -> String {
-    messages
-        .iter()
-        .filter_map(Value::as_object)
-        .map(|message| {
-            let content = match message.get("content") {
-                Some(Value::String(content)) => content.clone(),
-                Some(Value::Array(parts)) => parts
-                    .iter()
-                    .filter_map(Value::as_object)
-                    .filter_map(|part| part.get("text").and_then(Value::as_str))
-                    .collect(),
-                _ => String::new(),
-            };
-            format!(
-                "{content}{}",
-                search_results_text(message.get("search_results"))
-            )
-        })
-        .collect()
-}

litellm-rust/crates/cache-qdrant-semantic/tests/contract.rs

@@ -0,0 +1,91 @@
+//! `overwrite_replaces` does not apply: like Python, every write upserts a new `uuid4` point,
+//! so a second write with the same prompt adds a tie instead of replacing the first.
+
+mod support;
+
+use std::future::Future;
+
+use litellm_cache::{JsonCodec, SemanticCacheContext, semantic::PreparedEmbedding};
+use litellm_cache_qdrant_semantic::{QdrantSemanticCache, QdrantSemanticConfig, Quantization};
+use litellm_cache_testing as contract;
+use qdrant_client::Qdrant;
+use rstest::{fixture, rstest};
+use serde_json::{Value, json};
+use support::{FakeQdrant, FakeState};
+
+type Cache = QdrantSemanticCache<PreparedEmbedding, JsonCodec<Value>>;
+
+const PREFIX: &str = "contract:";
+
+#[fixture]
+fn context() -> SemanticCacheContext {
+    SemanticCacheContext {
+        messages: Some(json!([{"role": "user", "content": "contract prompt"}])),
+        ..Default::default()
+    }
+}
+
+/// Runs a contract against a fresh fake Qdrant. The sync cache methods block on the runtime, so
+/// the contract is polled on a blocking thread outside the runtime's own executor.
+async fn run<F, Fut>(check: F)
+where
+    F: FnOnce(Cache) -> Fut + Send + 'static,
+    Fut: Future<Output = ()>,
+{
+    let server = FakeQdrant::start(FakeState::default()).await;
+    let runtime = tokio::runtime::Handle::current();
+    let cache = QdrantSemanticCache::connect(
+        Qdrant::from_url(&server.url()).build().unwrap(),
+        PreparedEmbedding(vec![0.6, 0.8]),
+        JsonCodec::new(),
+        QdrantSemanticConfig {
+            collection_name: "contract".to_owned(),
+            similarity_threshold: 0.9,
+            vector_size: 2,
+            quantization: Quantization::Binary,
+        },
+        runtime.clone(),
+    )
+    .await
+    .unwrap();
+    tokio::task::spawn_blocking(move || {
+        let _guard = runtime.enter();
+        futures_executor::block_on(check(cache));
+    })
+    .await
+    .unwrap();
+    server.stop();
+}
+
+#[rstest]
+#[tokio::test(flavor = "multi_thread")]
+async fn hit_and_miss(context: SemanticCacheContext) {
+    run(|cache| async move {
+        contract::hit_and_miss(&cache, context, PREFIX, json!({"answer": 42})).await;
+    })
+    .await;
+}
+
+#[rstest]
+#[tokio::test(flavor = "multi_thread")]
+async fn sync_async_equivalence(context: SemanticCacheContext) {
+    run(|cache| async move {
+        contract::sync_async_equivalence(&cache, context, PREFIX, json!("first"), json!([2])).await;
+    })
+    .await;
+}
+
+#[rstest]
+#[tokio::test(flavor = "multi_thread")]
+async fn pipeline_writes_every_entry(context: SemanticCacheContext) {
+    run(|cache| async move {
+        contract::pipeline_writes_every_entry(
+            &cache,
+            context,
+            PREFIX,
+            vec![json!("a"), json!(2), json!({"c": true})],
+        )
+        .await;
+    })
+    .await;
+}

litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs

@@ -3,9 +3,10 @@ use std::{
     time::Duration,
 };
 
-use litellm_cache::Error;
-use litellm_cache_qdrant_semantic::{Embedder, OpenAiEmbedder, OpenAiEmbedderConfig};
-use serde_json::Value;
+use litellm_cache::{Error, semantic::Embedder};
+use litellm_cache_qdrant_semantic::{OpenAiEmbedder, OpenAiEmbedderConfig};
+use rstest::rstest;
+use serde_json::{Value, json};
 use tokio::{
     io::{AsyncReadExt, AsyncWriteExt},
     net::TcpListener,
@@ -98,6 +99,7 @@ fn config(base: String, timeout: Option<Duration>) -> OpenAiEmbedderConfig {
     }
 }
 
+#[rstest]
 #[tokio::test]
 async fn posts_embeddings_request_and_parses_vector() {
     let server = TestHttpServer::response("200 OK", r#"{"data":[{"embedding":[0.1,0.2]}]}"#).await;
@@ -108,7 +110,14 @@ async fn posts_embeddings_request_and_parses_vector() {
             Some(Duration::from_secs(1)),
         ),
     );
-    assert_eq!(embedder.embed("hello").await.unwrap(), vec![0.1, 0.2]);
+    assert_eq!(embedder.model(), "test-model");
+    assert_eq!(
+        embedder
+            .async_embed("hello", Some(&json!({"ignored": true})))
+            .await
+            .unwrap(),
+        vec![0.1, 0.2]
+    );
     let request = server.request.lock().unwrap().clone().unwrap();
     let request_text = String::from_utf8(request).unwrap();
     assert!(request_text.starts_with("POST /embeddings HTTP/1.1\r\n"));
@@ -120,37 +129,50 @@ async fn posts_embeddings_request_and_parses_vector() {
     assert_eq!(body["encoding_format"], "float");
 }
 
+#[rstest]
+#[case::error_status("500 Internal Server Error", "{}", 0, None, Err(Error::Unavailable))]
+#[case::timed_out(
+    "200 OK",
+    r#"{"data":[{"embedding":[0.1,0.2]}]}"#,
+    500,
+    Some(Duration::from_millis(200)),
+    Err(Error::Unavailable)
+)]
+#[case::within_timeout(
+    "200 OK",
+    r#"{"data":[{"embedding":[0.1,0.2]}]}"#,
+    100,
+    Some(Duration::from_secs(1)),
+    Ok(vec![0.1, 0.2])
+)]
+#[case::missing_embedding("200 OK", r#"{"data":[]}"#, 0, None, Err(Error::Unavailable))]
 #[tokio::test]
-async fn status_and_timeout_errors_are_unavailable() {
-    let server = TestHttpServer::response("500 Internal Server Error", "{}").await;
-    let embedder = OpenAiEmbedder::new(reqwest::Client::new(), config(server.base_url(), None));
-    assert_eq!(embedder.embed("hello").await, Err(Error::Unavailable));
-
-    let server = TestHttpServer::response_after(
-        "200 OK",
-        r#"{"data":[{"embedding":[0.1,0.2]}]}"#,
-        Duration::from_millis(500),
-    )
-    .await;
-    let embedder = OpenAiEmbedder::new(
-        reqwest::Client::new(),
-        config(server.base_url(), Some(Duration::from_millis(200))),
-    );
-    assert_eq!(embedder.embed("hello").await, Err(Error::Unavailable));
-
-    let server = TestHttpServer::response_after(
-        "200 OK",
-        r#"{"data":[{"embedding":[0.1,0.2]}]}"#,
-        Duration::from_millis(100),
-    )
-    .await;
-    let embedder = OpenAiEmbedder::new(
-        reqwest::Client::new(),
-        config(server.base_url(), Some(Duration::from_secs(1))),
-    );
-    assert_eq!(embedder.embed("hello").await.unwrap(), vec![0.1, 0.2]);
+async fn status_timeout_and_body_errors_are_unavailable(
+    #[case] status: &str,
+    #[case] body: &str,
+    #[case] delay_ms: u64,
+    #[case] timeout: Option<Duration>,
+    #[case] expected: Result<Vec<f32>, Error>,
+) {
+    let server =
+        TestHttpServer::response_after(status, body, Duration::from_millis(delay_ms)).await;
+    let embedder = OpenAiEmbedder::new(reqwest::Client::new(), config(server.base_url(), timeout));
+    assert_eq!(embedder.async_embed("hello", None).await, expected);
 }
 
+#[rstest]
+fn sync_embedding_is_unsupported() {
+    let embedder = OpenAiEmbedder::new(
+        reqwest::Client::new(),
+        config("http://127.0.0.1:9".to_owned(), None),
+    );
+    assert_eq!(
+        embedder.embed("hello", None),
+        Err(Error::UnsupportedOperation)
+    );
+}
+
+#[rstest]
 #[tokio::test]
 async fn uses_the_injected_client() {
     let server = TestHttpServer::response("200 OK", r#"{"data":[{"embedding":[0.1,0.2]}]}"#).await;
@@ -159,7 +181,10 @@ async fn uses_the_injected_client() {
         .build()
         .unwrap();
     let embedder = OpenAiEmbedder::new(client, config(server.base_url(), None));
-    assert_eq!(embedder.embed("hello").await.unwrap(), vec![0.1, 0.2]);
+    assert_eq!(
+        embedder.async_embed("hello", None).await.unwrap(),
+        vec![0.1, 0.2]
+    );
     let request = server.request.lock().unwrap().clone().unwrap();
     let request_text = String::from_utf8(request).unwrap();
     assert!(request_text.contains("\r\nuser-agent: litellm-embedder-test\r\n"));

litellm-rust/crates/cache-qdrant-semantic/tests/prompt.rs

@@ -1,38 +0,0 @@
-use litellm_cache_qdrant_semantic::prompt_from_messages;
-use serde_json::json;
-
-#[test]
-fn prompt_matches_python_message_content_rules() {
-    let messages = vec![
-        json!({"role": "user", "content": "hello"}),
-        json!({
-            "role": "user",
-            "content": [
-                {"type": "text", "text": "world"},
-                {"type": "image_url", "image_url": {"url": "ignored"}},
-                {"type": "text", "text": "!"},
-            ],
-        }),
-    ];
-
-    assert_eq!(prompt_from_messages(&messages), "helloworld!");
-}
-
-#[test]
-fn prompt_includes_search_result_text_and_compact_citations() {
-    let messages = vec![json!({
-        "role": "tool",
-        "content": null,
-        "search_results": [{
-            "source": "source",
-            "title": "title",
-            "content": [{"text": "body"}],
-            "citations": {"page": 1, "section": "intro"},
-        }],
-    })];
-
-    assert_eq!(
-        prompt_from_messages(&messages),
-        r#"sourcetitlebody{"page":1,"section":"intro"}"#
-    );
-}

litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs

@@ -1,27 +1,32 @@
-#[path = "support/mod.rs"]
 mod support;
 
-use std::{collections::HashMap, sync::Arc, time::Duration};
+use std::{
+    collections::HashMap,
+    sync::{Arc, Mutex},
+    time::Duration,
+};
 
-use litellm_cache::{BaseCache, CacheCodec, CacheContext, Error, SemanticCacheContext};
-use litellm_cache_qdrant_semantic::{
-    Embedder, QdrantSemanticCache, QdrantSemanticConfig, Quantization,
+use litellm_cache::{
+    BaseCache, CacheContext, Error, JsonCodec, SemanticCacheContext,
+    semantic::{Embedder, SemanticCache, SemanticLookup},
 };
-use litellm_cache_response::{
-    CacheEntry, CacheKeyInput, ResponseCache, ResponseCacheCodec, ResponseCacheRequest,
-};
-use qdrant_client::Payload;
+use litellm_cache_qdrant_semantic::{QdrantSemanticCache, QdrantSemanticConfig, Quantization};
 use qdrant_client::{
-    Qdrant,
+    Payload, Qdrant,
     qdrant::{self, CompressionRatio, Distance, PointId, QuantizationType, Value, VectorParams},
 };
+use rstest::{fixture, rstest};
 use serde_json::{Value as JsonValue, json};
-
 use support::{FakeQdrant, FakeState, StoredPoint};
 
+type Calls = Arc<Mutex<Vec<(String, Option<JsonValue>)>>>;
+type Cache = QdrantSemanticCache<FixedEmbedder, JsonCodec<JsonValue>>;
+
+/// Embeds known prompts, fails on anything else, and records every call.
 #[derive(Clone)]
 struct FixedEmbedder {
     vectors: Arc<HashMap<String, Vec<f32>>>,
+    calls: Calls,
 }
 
 impl FixedEmbedder {
@@ -33,16 +38,21 @@ impl FixedEmbedder {
                     .map(|(prompt, vector)| (prompt.to_owned(), vector))
                     .collect(),
             ),
+            calls: Calls::default(),
         }
     }
 }
 
 impl Embedder for FixedEmbedder {
-    fn model(&self) -> &str {
-        "fixed"
-    }
-
-    async fn embed(&self, input: &str) -> Result<Vec<f32>, Error> {
+    async fn async_embed(
+        &self,
+        input: &str,
+        metadata: Option<&JsonValue>,
+    ) -> Result<Vec<f32>, Error> {
+        self.calls
+            .lock()
+            .unwrap()
+            .push((input.to_owned(), metadata.cloned()));
         self.vectors.get(input).cloned().ok_or(Error::Unavailable)
     }
 }
@@ -63,22 +73,20 @@ fn context(prompt: &str) -> SemanticCacheContext {
     }
 }
 
-fn value(response: JsonValue) -> CacheEntry {
-    CacheEntry {
-        timestamp: Some(1.0),
-        response,
-    }
+#[fixture]
+fn entry() -> JsonValue {
+    json!({"timestamp": 1.0, "response": {"answer": 42}})
 }
 
 async fn connect(
     server: &FakeQdrant,
     vectors: impl IntoIterator<Item = (&'static str, Vec<f32>)>,
-) -> QdrantSemanticCache<FixedEmbedder, ResponseCacheCodec> {
+) -> Cache {
     let client = Qdrant::from_url(&server.url()).build().unwrap();
     QdrantSemanticCache::connect(
         client,
         FixedEmbedder::new(vectors),
-        ResponseCacheCodec,
+        JsonCodec::new(),
         config(Quantization::Binary),
         tokio::runtime::Handle::current(),
     )
@@ -86,72 +94,70 @@ async fn connect(
     .unwrap()
 }
 
+#[rstest]
+#[case::binary(Quantization::Binary)]
+#[case::scalar(Quantization::Scalar)]
+#[case::product(Quantization::Product)]
 #[tokio::test(flavor = "multi_thread")]
-#[expect(
-    deprecated,
-    reason = "the test verifies Qdrant's legacy always_ram quantization contract"
-)]
-async fn connect_sets_collection_quantization_and_index() {
-    for (quantization, expected) in [
-        (Quantization::Binary, 0),
-        (Quantization::Scalar, 1),
-        (Quantization::Product, 2),
-    ] {
-        let server = FakeQdrant::start(FakeState::default()).await;
-        let client = Qdrant::from_url(&server.url()).build().unwrap();
-        QdrantSemanticCache::connect(
-            client,
-            FixedEmbedder::new([]),
-            ResponseCacheCodec,
-            config(quantization),
-            tokio::runtime::Handle::current(),
-        )
-        .await
+async fn connect_sets_collection_quantization_and_index(#[case] quantization: Quantization) {
+    let server = FakeQdrant::start(FakeState::default()).await;
+    let client = Qdrant::from_url(&server.url()).build().unwrap();
+    QdrantSemanticCache::connect(
+        client,
+        FixedEmbedder::new([]),
+        JsonCodec::<JsonValue>::new(),
+        config(quantization.clone()),
+        tokio::runtime::Handle::current(),
+    )
+    .await
+    .unwrap();
+    let state = server.state.lock().unwrap();
+    let request = &state.created_collections[0];
+    let Some(qdrant::vectors_config::Config::Params(VectorParams { size, distance, .. })) = request
+        .vectors_config
+        .as_ref()
+        .and_then(|config| config.config.clone())
+    else {
+        panic!("missing vector params");
+    };
+    assert_eq!(size, 2);
+    assert_eq!(distance, Distance::Cosine as i32);
+    let quantization_config = request
+        .quantization_config
+        .as_ref()
+        .unwrap()
+        .quantization
         .unwrap();
-        let state = server.state.lock().unwrap();
-        let request = &state.created_collections[0];
-        let Some(qdrant::vectors_config::Config::Params(VectorParams { size, distance, .. })) =
-            request
-                .vectors_config
-                .as_ref()
-                .and_then(|config| config.config.clone())
-        else {
-            panic!("missing vector params");
-        };
-        assert_eq!(size, 2);
-        assert_eq!(distance, Distance::Cosine as i32);
-        let quantization_config = request
-            .quantization_config
-            .as_ref()
-            .unwrap()
-            .quantization
-            .unwrap();
-        match (expected, quantization_config) {
-            (0, qdrant::quantization_config::Quantization::Binary(binary)) => {
-                assert_eq!(binary.always_ram, Some(false));
-            }
-            (1, qdrant::quantization_config::Quantization::Scalar(scalar)) => {
-                assert_eq!(scalar.r#type, QuantizationType::Int8 as i32);
-                assert_eq!(scalar.quantile, Some(0.99));
-                assert_eq!(scalar.always_ram, Some(false));
-            }
-            (2, qdrant::quantization_config::Quantization::Product(product)) => {
-                assert_eq!(product.compression, CompressionRatio::X16 as i32);
-                assert_eq!(product.always_ram, Some(false));
-            }
-            _ => panic!("unexpected quantization"),
+    #[expect(
+        deprecated,
+        reason = "the test verifies Qdrant's legacy always_ram quantization contract"
+    )]
+    match (quantization, quantization_config) {
+        (Quantization::Binary, qdrant::quantization_config::Quantization::Binary(binary)) => {
+            assert_eq!(binary.always_ram, Some(false));
         }
-        assert!(state.index_creations >= 1);
-        assert_eq!(state.field_indexes[0].collection_name, "semantic");
-        assert_eq!(state.field_indexes[0].field_name, "litellm_cache_key");
-        assert_eq!(
-            state.field_indexes[0].field_type,
-            Some(qdrant::FieldType::Keyword as i32)
-        );
-        server.stop();
+        (Quantization::Scalar, qdrant::quantization_config::Quantization::Scalar(scalar)) => {
+            assert_eq!(scalar.r#type, QuantizationType::Int8 as i32);
+            assert_eq!(scalar.quantile, Some(0.99));
+            assert_eq!(scalar.always_ram, Some(false));
+        }
+        (Quantization::Product, qdrant::quantization_config::Quantization::Product(product)) => {
+            assert_eq!(product.compression, CompressionRatio::X16 as i32);
+            assert_eq!(product.always_ram, Some(false));
+        }
+        _ => panic!("unexpected quantization"),
     }
+    assert!(state.index_creations >= 1);
+    assert_eq!(state.field_indexes[0].collection_name, "semantic");
+    assert_eq!(state.field_indexes[0].field_name, "litellm_cache_key");
+    assert_eq!(
+        state.field_indexes[0].field_type,
+        Some(qdrant::FieldType::Keyword as i32)
+    );
+    server.stop();
 }
 
+#[rstest]
 #[tokio::test(flavor = "multi_thread")]
 async fn existing_collection_skips_create_and_index_failure_is_non_fatal() {
     let server = FakeQdrant::start(FakeState {
@@ -160,19 +166,25 @@ async fn existing_collection_skips_create_and_index_failure_is_non_fatal() {
         ..Default::default()
     })
     .await;
-    let _cache = connect(&server, [("hello", vec![1.0, 0.0])]).await;
+    let cache = connect(&server, [("hello", vec![1.0, 0.0])]).await;
+    assert_eq!(cache.collection_name(), "semantic");
+    assert_eq!(cache.similarity_threshold(), 0.9);
+    assert_eq!(cache.vector_size(), 2);
     let state = server.state.lock().unwrap();
     assert!(state.created_collections.is_empty());
     assert!(state.index_creations >= 1);
     server.stop();
 }
 
+#[rstest]
 #[tokio::test(flavor = "multi_thread")]
-async fn async_and_sync_set_get_store_exact_payload() {
+async fn async_and_sync_set_get_store_exact_payload(entry: JsonValue) {
     let server = FakeQdrant::start(FakeState::default()).await;
     let cache = Arc::new(connect(&server, [("hello", vec![1.0, 0.0])]).await);
-    let ctx = context("hello");
-    let entry = value(json!({"answer": 42}));
+    let ctx = SemanticCacheContext {
+        metadata: Some(json!({"tenant": "team"})),
+        ..context("hello")
+    };
     cache
         .async_set_cache("key", entry.clone(), ctx.clone())
         .await
@@ -188,10 +200,8 @@ async fn async_and_sync_set_get_store_exact_payload() {
         payload_keys.sort();
         assert_eq!(payload_keys, ["litellm_cache_key", "response", "text"]);
         assert_eq!(payload["litellm_cache_key"], Value::from("key"));
-        assert_eq!(
-            payload["response"],
-            Value::from(String::from_utf8(ResponseCacheCodec.encode(&entry).unwrap()).unwrap())
-        );
+        assert_eq!(payload["text"], Value::from("hello"));
+        assert_eq!(payload["response"], Value::from(entry.to_string()));
     }
     let sync_entry = entry.clone();
     let sync_cache = cache.clone();
@@ -207,208 +217,276 @@ async fn async_and_sync_set_get_store_exact_payload() {
     })
     .await
     .unwrap();
+    assert_eq!(
+        *cache.embedder().calls.lock().unwrap(),
+        vec![("hello".to_owned(), ctx.metadata.clone()); 4]
+    );
     server.stop();
 }
 
+#[rstest]
+#[case::content_parts_skip_images(
+    json!([
+        {"role": "user", "content": "hello"},
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "world"},
+                {"type": "image_url", "image_url": {"url": "ignored"}},
+                {"type": "text", "text": "!"},
+            ],
+        },
+    ]),
+    "helloworld!"
+)]
+#[case::search_results_and_compact_citations(
+    json!([{
+        "role": "tool",
+        "content": null,
+        "search_results": [{
+            "source": "source",
+            "title": "title",
+            "content": [{"text": "body"}],
+            "citations": {"page": 1, "section": "intro"},
+        }],
+    }]),
+    r#"sourcetitlebody{"page":1,"section":"intro"}"#
+)]
 #[tokio::test(flavor = "multi_thread")]
-async fn misses_and_payload_validation_are_safe() {
+async fn prompt_matches_python_message_rules(
+    #[case] messages: JsonValue,
+    #[case] prompt: &'static str,
+    entry: JsonValue,
+) {
+    let server = FakeQdrant::start(FakeState::default()).await;
+    let cache = connect(&server, [(prompt, vec![1.0, 0.0])]).await;
+    let context = SemanticCacheContext {
+        messages: Some(messages),
+        ..Default::default()
+    };
+
+    cache.async_set_cache("key", entry, context).await.unwrap();
+
+    assert_eq!(cache.embedder().calls.lock().unwrap()[0].0, prompt);
+    assert_eq!(
+        server.state.lock().unwrap().points[0].payload["text"],
+        Value::from(prompt)
+    );
+    server.stop();
+}
+
+#[rstest]
+#[case::no_messages(SemanticCacheContext::default())]
+#[case::empty_messages(SemanticCacheContext { messages: Some(json!([])), ..Default::default() })]
+#[case::responses_input_is_not_read(SemanticCacheContext { input: Some(json!("hello")), ..Default::default() })]
+#[tokio::test(flavor = "multi_thread")]
+async fn requests_without_messages_are_missing_a_prompt(
+    #[case] context: SemanticCacheContext,
+    entry: JsonValue,
+) {
+    let server = FakeQdrant::start(FakeState::default()).await;
+    let cache = connect(&server, [("hello", vec![1.0, 0.0])]).await;
+
+    assert_eq!(
+        cache.async_set_cache("key", entry, context.clone()).await,
+        Err(Error::MissingPrompt)
+    );
+    assert_eq!(
+        cache.async_get_cache("key", &context).await,
+        Err(Error::MissingPrompt)
+    );
+    assert!(cache.embedder().calls.lock().unwrap().is_empty());
+    server.stop();
+}
+
+#[rstest]
+#[case::other_key("other", "hello", None)]
+#[case::below_similarity_threshold("key", "near", None)]
+#[tokio::test(flavor = "multi_thread")]
+async fn misses_and_payload_validation_are_safe(
+    #[case] key: &str,
+    #[case] prompt: &str,
+    #[case] numeric_key_point: Option<u64>,
+    entry: JsonValue,
+) {
     let server = FakeQdrant::start(FakeState::default()).await;
     let cache = connect(
         &server,
         [("hello", vec![1.0, 0.0]), ("near", vec![0.7, 0.71414286])],
     )
     .await;
-    let entry = value(json!({"answer": 1}));
     cache
         .async_set_cache("key", entry, context("hello"))
         .await
         .unwrap();
+    if let Some(id) = numeric_key_point {
+        server.insert_point(StoredPoint {
+            id: Some(PointId::from(id)),
+            vector: vec![1.0, 0.0],
+            payload: Payload::try_from(json!({
+                "litellm_cache_key": id,
+                "response": "{}",
+            }))
+            .unwrap()
+            .into(),
+        });
+    }
+
     assert_eq!(
-        cache
-            .async_get_cache("other", &context("hello"))
-            .await
-            .unwrap(),
-        None
-    );
-    assert_eq!(
-        cache
-            .async_get_cache("key", &context("near"))
-            .await
-            .unwrap(),
-        None
-    );
-    server.insert_point(StoredPoint {
-        id: Some(PointId::from(99_u64)),
-        vector: vec![1.0, 0.0],
-        payload: Payload::try_from(json!({
-            "litellm_cache_key": 99,
-            "response": "{}",
-        }))
-        .unwrap()
-        .into(),
-    });
-    assert_eq!(
-        cache
-            .async_get_cache("99", &context("hello"))
-            .await
-            .unwrap(),
+        cache.async_get_cache(key, &context(prompt)).await.unwrap(),
         None
     );
     server.stop();
 }
 
+#[rstest]
+#[case::hit("key", context("hello"), Ok((true, Some(1.0))))]
+#[case::below_similarity_threshold("key", context("near"), Ok((false, Some(0.7))))]
+#[case::no_results("other", context("hello"), Ok((false, Some(0.0))))]
+#[case::no_prompt("key", SemanticCacheContext::default(), Err(Error::MissingPrompt))]
 #[tokio::test(flavor = "multi_thread")]
-async fn decoding_errors_missing_prompt_pipeline_and_ttl_behave_as_required() {
+async fn lookup_reports_python_semantic_similarity(
+    #[case] key: &'static str,
+    #[case] context: SemanticCacheContext,
+    #[case] expected: Result<(bool, Option<f64>), Error>,
+    #[values(false, true)] use_async: bool,
+    entry: JsonValue,
+) {
     let server = FakeQdrant::start(FakeState::default()).await;
-    let cache = connect(&server, [("one", vec![1.0, 0.0]), ("two", vec![0.0, 1.0])]).await;
-    let empty = SemanticCacheContext::default();
-    assert_eq!(
-        cache
-            .async_set_cache("key", value(json!({})), empty.clone())
-            .await,
-        Err(Error::MissingPrompt)
+    let cache = Arc::new(
+        connect(
+            &server,
+            [("hello", vec![1.0, 0.0]), ("near", vec![0.7, 0.71414286])],
+        )
+        .await,
     );
+    cache
+        .async_set_cache("key", entry.clone(), self::context("hello"))
+        .await
+        .unwrap();
+    server.insert_point(StoredPoint {
+        id: Some(PointId::from(99_u64)),
+        vector: vec![1.0, 0.0],
+        payload: Payload::try_from(json!({"litellm_cache_key": 99, "response": "{}"}))
+            .unwrap()
+            .into(),
+    });
+
+    let lookup = if use_async {
+        cache.async_get_cache_with_similarity(key, &context).await
+    } else {
+        let cache = Arc::clone(&cache);
+        tokio::task::spawn_blocking(move || cache.get_cache_with_similarity(key, &context))
+            .await
+            .unwrap()
+    };
+
+    match (lookup, expected) {
+        (Ok(SemanticLookup { value, similarity }), Ok((hit, expected))) => {
+            assert_eq!(value, hit.then_some(entry));
+            assert_eq!(similarity.is_some(), expected.is_some());
+            if let (Some(similarity), Some(expected)) = (similarity, expected) {
+                assert!((similarity - expected).abs() < 1e-6, "{similarity}");
+            }
+        }
+        (lookup, expected) => assert_eq!(lookup.map(|_| ()), expected.map(|_| ())),
+    }
+    server.stop();
+}
+
+#[rstest]
+#[case::codec_decodes_the_payload(Some(json!("{\"a\":1}")), Ok(Some(json!({"a": 1}))))]
+#[case::undecodable_response(Some(json!("not json")), Err(Error::InvalidEntry))]
+#[case::non_string_response(Some(json!(1)), Err(Error::InvalidEntry))]
+#[case::missing_response(None, Err(Error::InvalidEntry))]
+#[tokio::test(flavor = "multi_thread")]
+async fn stored_responses_go_through_the_codec(
+    #[case] response: Option<JsonValue>,
+    #[case] expected: Result<Option<JsonValue>, Error>,
+) {
+    let server = FakeQdrant::start(FakeState::default()).await;
+    let cache = connect(&server, [("hello", vec![1.0, 0.0])]).await;
+    let mut payload = serde_json::Map::new();
+    payload.insert("litellm_cache_key".to_owned(), json!("key"));
+    if let Some(response) = response {
+        payload.insert("response".to_owned(), response);
+    }
+    server.insert_point(StoredPoint {
+        id: Some(PointId::from(1_u64)),
+        vector: vec![1.0, 0.0],
+        payload: Payload::try_from(JsonValue::Object(payload))
+            .unwrap()
+            .into(),
+    });
+
     assert_eq!(
-        cache.async_get_cache("key", &empty).await,
-        Err(Error::MissingPrompt)
+        cache.async_get_cache("key", &context("hello")).await,
+        expected
     );
+    server.stop();
+}
+
+#[rstest]
+#[tokio::test(flavor = "multi_thread")]
+async fn embedding_failures_propagate() {
+    let server = FakeQdrant::start(FakeState::default()).await;
+    let cache = connect(&server, []).await;
+
     assert_eq!(
         cache.async_get_cache("key", &context("unknown")).await,
         Err(Error::Unavailable)
     );
+    server.stop();
+}
+
+#[rstest]
+#[tokio::test(flavor = "multi_thread")]
+async fn ttl_is_ignored_and_entries_do_not_expire(entry: JsonValue) {
+    let server = FakeQdrant::start(FakeState::default()).await;
+    let cache = connect(&server, [("one", vec![1.0, 0.0])]).await;
+    let ctx = context("one").with_ttl(Some(Duration::from_secs(1)));
+
+    assert_eq!(cache.get_ttl(&ctx), None);
     cache
-        .async_set_cache(
-            "ttl",
-            value(json!({"ttl": true})),
-            context("one").with_ttl(Some(Duration::from_secs(1))),
-        )
+        .async_set_cache("ttl", entry, ctx.clone())
         .await
         .unwrap();
     tokio::time::sleep(Duration::from_millis(1_100)).await;
-    assert!(
-        cache
-            .async_get_cache(
-                "ttl",
-                &context("one").with_ttl(Some(Duration::from_secs(1))),
-            )
-            .await
-            .unwrap()
-            .is_some()
-    );
+    assert!(cache.async_get_cache("ttl", &ctx).await.unwrap().is_some());
+    server.stop();
+}
+
+#[rstest]
+#[tokio::test(flavor = "multi_thread")]
+async fn pipeline_upserts_each_entry_and_waits_for_indexing() {
+    let server = FakeQdrant::start(FakeState::default()).await;
+    let cache = connect(&server, [("one", vec![1.0, 0.0])]).await;
+
     cache
         .async_set_cache_pipeline(
             vec![
-                ("one".to_owned(), value(json!({"n": 1}))),
-                ("two".to_owned(), value(json!({"n": 2}))),
+                ("one".to_owned(), json!({"n": 1})),
+                ("two".to_owned(), json!({"n": 2})),
             ],
             context("one"),
         )
         .await
         .unwrap();
-    assert!(
-        cache
-            .async_get_cache("one", &context("one"))
-            .await
-            .unwrap()
-            .is_some()
-    );
-    assert!(
-        cache
-            .async_get_cache("two", &context("one"))
-            .await
-            .unwrap()
-            .is_some()
-    );
+
+    for (key, value) in [("one", json!({"n": 1})), ("two", json!({"n": 2}))] {
+        assert_eq!(
+            cache.async_get_cache(key, &context("one")).await.unwrap(),
+            Some(value)
+        );
+    }
     assert_eq!(
         server.state.lock().unwrap().upsert_waits,
-        vec![Some(true), Some(true), Some(true)]
-    );
-    assert_eq!(cache.get_ttl(&context("one")), None);
-    assert_eq!(
-        cache.test_connection().await,
-        Err(Error::UnsupportedOperation)
-    );
-    server.stop();
-}
-
-#[tokio::test(flavor = "multi_thread")]
-async fn response_payloads_decode_and_invalid_entries_fail() {
-    let server = FakeQdrant::start(FakeState::default()).await;
-    let cache = connect(&server, [("hello", vec![1.0, 0.0])]).await;
-    for (key, response) in [
-        ("python", json!("{'timestamp': 1.0, 'response': {'a': 1}}")),
-        ("garbage", json!("not json")),
-        ("missing", json!("unused")),
-    ] {
-        let mut payload = serde_json::Map::new();
-        payload.insert("litellm_cache_key".to_owned(), json!(key));
-        if key != "missing" {
-            payload.insert("response".to_owned(), response);
-        }
-        server.insert_point(StoredPoint {
-            id: Some(PointId::from(key.len() as u64)),
-            vector: vec![1.0, 0.0],
-            payload: Payload::try_from(JsonValue::Object(payload))
-                .unwrap()
-                .into(),
-        });
-    }
-    assert_eq!(
-        cache
-            .async_get_cache("python", &context("hello"))
-            .await
-            .unwrap(),
-        Some(value(json!({"a": 1})))
-    );
-    assert_eq!(
-        cache.async_get_cache("garbage", &context("hello")).await,
-        Err(Error::InvalidEntry)
-    );
-    assert_eq!(
-        cache.async_get_cache("missing", &context("hello")).await,
-        Err(Error::InvalidEntry)
-    );
-    server.stop();
-}
-
-#[tokio::test(flavor = "multi_thread")]
-async fn response_cache_facade_turns_invalid_entry_into_miss() {
-    let server = FakeQdrant::start(FakeState::default()).await;
-    let cache = Arc::new(connect(&server, [("hello", vec![1.0, 0.0])]).await);
-    let request = ResponseCacheRequest::<SemanticCacheContext>::new(CacheKeyInput {
-        preset: Some("key".to_owned()),
-        ..Default::default()
-    })
-    .with_context(context("hello"));
-    let response = json!({"answer": 42});
-    let facade = ResponseCache::new(cache.clone());
-    facade
-        .async_store(&request, response.clone(), Duration::from_secs(1))
-        .await
-        .unwrap();
-    assert_eq!(
-        facade
-            .async_lookup(&request, Duration::from_secs(1))
-            .await
-            .unwrap(),
-        Some(response)
-    );
-    {
-        let mut state = server.state.lock().unwrap();
-        state.points[0]
-            .payload
-            .insert("response".to_owned(), Value::from("not json"));
-    }
-    assert_eq!(
-        facade
-            .async_lookup(&request, Duration::from_secs(1))
-            .await
-            .unwrap(),
-        None
+        vec![Some(true), Some(true)]
     );
     server.stop();
 }
 
+#[rstest]
 #[tokio::test(flavor = "multi_thread")]
 async fn stopped_qdrant_server_maps_to_unavailable() {
     let server = FakeQdrant::start(FakeState::default()).await;
@@ -420,3 +498,28 @@ async fn stopped_qdrant_server_maps_to_unavailable() {
         Err(Error::Unavailable)
     );
 }
+
+/// `_payload_matches_cache_key` compares `str(cached_key) == str(key)`, so a point whose stored
+/// key is the number 99 answers a lookup for `"99"`.
+#[rstest]
+#[tokio::test(flavor = "multi_thread")]
+async fn numeric_stored_cache_keys_match_like_python_str() {
+    let server = FakeQdrant::start(FakeState::default()).await;
+    let cache = connect(&server, [("hello", vec![1.0, 0.0])]).await;
+    server.insert_point(StoredPoint {
+        id: Some(PointId::from(99_u64)),
+        vector: vec![1.0, 0.0],
+        payload: Payload::try_from(json!({"litellm_cache_key": 99, "response": "{}"}))
+            .unwrap()
+            .into(),
+    });
+
+    let lookup = cache
+        .async_get_cache_with_similarity("99", &context("hello"))
+        .await
+        .unwrap();
+
+    assert_eq!(lookup.value, Some(json!({})));
+    assert!((lookup.similarity.unwrap() - 1.0).abs() < 1e-6);
+    server.stop();
+}

litellm-rust/crates/cache-qdrant-semantic/tests/support/mod.rs

@@ -1,15 +1,16 @@
+#![allow(dead_code)]
+
 use std::{
     collections::{HashMap, HashSet},
     net::SocketAddr,
     sync::{Arc, Mutex},
 };
 
-use qdrant_client::qdrant::collections_server::CollectionsServer;
 use qdrant_client::qdrant::{
     self, CollectionExists, CollectionExistsRequest, CollectionExistsResponse,
     CollectionOperationResponse, CreateCollection, CreateFieldIndexCollection, Filter, PointId,
     PointsOperationResponse, ScoredPoint, SearchPoints, SearchResponse, Value, Vector, Vectors,
-    collections_server::Collections,
+    collections_server::{Collections, CollectionsServer},
     points_server::{Points, PointsServer},
 };
 use tokio::sync::oneshot;

litellm-rust/crates/cache-redis-semantic/Cargo.toml

@@ -8,14 +8,12 @@ repository.workspace = true
 [dependencies]
 litellm-cache.workspace = true
 litellm-cache-redis.workspace = true
-litellm-cache-response.workspace = true
 redis = { version = "1.7.0", features = ["tls-rustls"] }
-r2d2 = "0.8.10"
-serde_json.workspace = true
 sha2.workspace = true
-tokio.workspace = true
 
 [dev-dependencies]
+litellm-cache-testing.workspace = true
 redis-test = "1.0.4"
+rstest.workspace = true
 serde_json.workspace = true
 tokio.workspace = true

litellm-rust/crates/cache-redis-semantic/src/cache.rs

@@ -1,105 +1,36 @@
 use std::{
-    future::Future,
-    sync::{Arc, OnceLock},
+    sync::Arc,
     time::{Duration, SystemTime, UNIX_EPOCH},
 };
 
 use litellm_cache::{
-    BaseCache, CacheCodec, CacheConnectionResult, CacheConnectionStatus, Error,
-    SemanticCacheContext,
+    BaseCache, CacheCodec, Error, SemanticCacheContext,
+    semantic::{Embedder, SemanticCache, SemanticLookup, prompt_from_context},
 };
 use litellm_cache_redis::{
     RedisTopology,
     connection::{ConnectionRef, Connections},
 };
-use litellm_cache_response::{CacheEntry, ResponseCacheCodec};
-use serde_json::Value;
 use sha2::{Digest, Sha256};
 
-use crate::prompt::prompt_from_context;
-
-const CACHE_KEY_FIELD: &str = "litellm_cache_key";
-const VECTOR_FIELD: &str = "prompt_vector";
-
-pub trait Embedder: Send + Sync + 'static {
-    fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result<Vec<f32>, Error>;
-
-    fn async_embed(
-        &self,
-        prompt: &str,
-        metadata: Option<&Value>,
-    ) -> impl Future<Output = Result<Vec<f32>, Error>> + Send;
-}
-
-#[derive(Clone, Debug)]
-pub struct RedisSemanticConfig {
-    pub index_name: String,
-    pub similarity_threshold: f32,
-}
+use crate::{
+    RedisSemanticConfig,
+    index::{CACHE_KEY_FIELD, Index, VECTOR_FIELD},
+    reply::{bytes_field, first_document, number_field, string_field},
+};
 
 struct Inner {
-    index_name: String,
+    index: Index,
     distance_threshold: f64,
-    resolved_index: OnceLock<String>,
-    codec: ResponseCacheCodec,
     clock: fn() -> f64,
 }
 
 impl Inner {
-    fn new(config: RedisSemanticConfig) -> Self {
+    fn new(config: RedisSemanticConfig, clock: fn() -> f64) -> Self {
         Self {
-            index_name: config.index_name,
+            index: Index::new(config.index_name),
             distance_threshold: 1.0 - f64::from(config.similarity_threshold),
-            resolved_index: OnceLock::new(),
-            codec: ResponseCacheCodec,
-            clock: timestamp,
-        }
-    }
-
-    fn ensure_index(
-        &self,
-        connection: &mut ConnectionRef<'_>,
-        dims: usize,
-    ) -> Result<String, Error> {
-        if let Some(name) = self.resolved_index.get() {
-            return Ok(name.clone());
-        }
-        let name = match index_compatible(connection, &self.index_name, dims)? {
-            Some(true) => self.index_name.clone(),
-            Some(false) => self.isolated_index(connection, dims)?,
-            None => match create_index(connection, &self.index_name, dims) {
-                Ok(()) => self.index_name.clone(),
-                Err(_) => match index_compatible(connection, &self.index_name, dims)? {
-                    Some(true) => self.index_name.clone(),
-                    Some(false) => self.isolated_index(connection, dims)?,
-                    None => return Err(Error::Unavailable),
-                },
-            },
-        };
-        let _ = self.resolved_index.set(name.clone());
-        Ok(name)
-    }
-
-    fn isolated_index(
-        &self,
-        connection: &mut ConnectionRef<'_>,
-        dims: usize,
-    ) -> Result<String, Error> {
-        let name = format!("{}_isolated", self.index_name);
-        match index_compatible(connection, &name, dims)? {
-            Some(true) => Ok(name),
-            Some(false) => {
-                redis::cmd("FT.DROPINDEX")
-                    .arg(&name)
-                    .query::<()>(connection)
-                    .map_err(|_| Error::Unavailable)?;
-                create_index(connection, &name, dims)?;
-                Ok(name)
-            }
-            None => {
-                create_index(connection, &name, dims)?;
-                Ok(name)
-            }
+            clock,
         }
     }
 
@@ -107,15 +38,14 @@ impl Inner {
         &self,
         connection: &mut ConnectionRef<'_>,
         tag: &str,
-        value: &CacheEntry,
+        response: Vec<u8>,
         prompt: &str,
         vector: &[f32],
         ttl: Option<Duration>,
     ) -> Result<(), Error> {
-        let index = self.ensure_index(connection, vector.len())?;
+        let index = self.index.ensure(connection, vector.len())?;
         let entry_id = entry_id(prompt, tag);
         let hash_key = format!("{index}:{entry_id}");
-        let response = self.codec.encode(value)?;
         redis::cmd("HSET")
             .arg(&hash_key)
             .arg("entry_id")
@@ -149,8 +79,8 @@ impl Inner {
         connection: &mut ConnectionRef<'_>,
         tag: &str,
         vector: &[f32],
-    ) -> Result<Option<CacheEntry>, Error> {
-        let index = self.ensure_index(connection, vector.len())?;
+    ) -> Result<SemanticLookup<Vec<u8>>, Error> {
+        let index = self.index.ensure(connection, vector.len())?;
         let query = format!(
             "(@{CACHE_KEY_FIELD}:{{{}}})=>[KNN 1 @{VECTOR_FIELD} $vector AS vector_distance]",
             escape_tag(tag)
@@ -183,57 +113,80 @@ impl Inner {
             .query::<redis::Value>(connection)
             .map_err(|_| Error::Unavailable)?;
         let Some(fields) = first_document(&result) else {
-            return Ok(None);
+            return Ok(SemanticLookup::miss(Some(0.0)));
         };
         if string_field(fields, CACHE_KEY_FIELD).as_deref() != Some(tag) {
-            return Ok(None);
+            return Ok(SemanticLookup::miss(Some(0.0)));
         }
-        if number_field(fields, "vector_distance")
-            .is_none_or(|distance| distance > self.distance_threshold)
-        {
-            return Ok(None);
-        }
-        let Some(response) = bytes_field(fields, "response") else {
-            return Ok(None);
+        // redisvl's range query only returns entries within the distance threshold, so a
+        // farther hit reads as no result.
+        let Some(distance) = number_field(fields, "vector_distance")
+            .filter(|distance| *distance <= self.distance_threshold)
+        else {
+            return Ok(SemanticLookup::miss(Some(0.0)));
         };
-        self.codec.decode(&response).map(Some)
-    }
-}
-
-pub struct RedisSemanticCache<E: Embedder, C = redis::Connection> {
-    connections: Arc<Connections<C>>,
-    embedder: E,
-    inner: Arc<Inner>,
-}
-
-impl<E: Embedder> RedisSemanticCache<E> {
-    pub fn new(url: &str, embedder: E, config: RedisSemanticConfig) -> Result<Self, Error> {
-        Ok(Self {
-            connections: Arc::new(Connections::open(url, &RedisTopology::Standalone)?),
-            embedder,
-            inner: Arc::new(Inner::new(config)),
+        let Some(response) = bytes_field(fields, "response") else {
+            return Ok(SemanticLookup::miss(Some(0.0)));
+        };
+        Ok(SemanticLookup {
+            value: Some(response),
+            similarity: Some(1.0 - distance),
         })
     }
 }
 
-impl<E: Embedder, C: redis::ConnectionLike + Send + 'static> RedisSemanticCache<E, C> {
-    pub fn with_connection(connection: C, embedder: E, config: RedisSemanticConfig) -> Self {
+/// `RedisSemanticCache`: a redisvl-compatible semantic index on Redis Stack. Values go through
+/// the injected codec, so the response layer decides what a cached entry is.
+pub struct RedisSemanticCache<E, S, C = redis::Connection> {
+    connections: Arc<Connections<C>>,
+    embedder: E,
+    codec: S,
+    inner: Arc<Inner>,
+}
+
+impl<E: Embedder, S: CacheCodec> RedisSemanticCache<E, S> {
+    pub fn new(
+        url: &str,
+        embedder: E,
+        codec: S,
+        config: RedisSemanticConfig,
+    ) -> Result<Self, Error> {
+        Ok(Self {
+            connections: Arc::new(Connections::open(url, &RedisTopology::Standalone)?),
+            embedder,
+            codec,
+            inner: Arc::new(Inner::new(config, timestamp)),
+        })
+    }
+}
+
+impl<E, S, C> RedisSemanticCache<E, S, C>
+where
+    E: Embedder,
+    S: CacheCodec,
+    C: redis::ConnectionLike + Send + 'static,
+{
+    pub fn with_connection(
+        connection: C,
+        embedder: E,
+        codec: S,
+        config: RedisSemanticConfig,
+    ) -> Self {
         Self {
             connections: Arc::new(Connections::fixed(connection)),
             embedder,
-            inner: Arc::new(Inner::new(config)),
+            codec,
+            inner: Arc::new(Inner::new(config, timestamp)),
         }
     }
 
     pub fn with_clock(self, clock: fn() -> f64) -> Self {
+        let config = RedisSemanticConfig {
+            index_name: self.index_name().to_owned(),
+            similarity_threshold: self.similarity_threshold(),
+        };
         Self {
-            inner: Arc::new(Inner {
-                index_name: self.inner.index_name.clone(),
-                distance_threshold: self.inner.distance_threshold,
-                resolved_index: OnceLock::new(),
-                codec: self.inner.codec,
-                clock,
-            }),
+            inner: Arc::new(Inner::new(config, clock)),
             ..self
         }
     }
@@ -243,7 +196,7 @@ impl<E: Embedder, C: redis::ConnectionLike + Send + 'static> RedisSemanticCache<
     }
 
     pub fn index_name(&self) -> &str {
-        &self.inner.index_name
+        self.inner.index.name()
     }
 
     pub fn similarity_threshold(&self) -> f32 {
@@ -253,12 +206,25 @@ impl<E: Embedder, C: redis::ConnectionLike + Send + 'static> RedisSemanticCache<
     fn tag<'a>(key: &'a str, context: &'a SemanticCacheContext) -> &'a str {
         context.scope.as_deref().unwrap_or(key)
     }
+
+    fn decode(&self, lookup: SemanticLookup<Vec<u8>>) -> Result<SemanticLookup<S::Value>, Error> {
+        Ok(SemanticLookup {
+            value: lookup
+                .value
+                .map(|bytes| self.codec.decode(&bytes))
+                .transpose()?,
+            similarity: lookup.similarity,
+        })
+    }
 }
 
-impl<E: Embedder, C: redis::ConnectionLike + Send + 'static> BaseCache
-    for RedisSemanticCache<E, C>
+impl<E, S, C> BaseCache for RedisSemanticCache<E, S, C>
+where
+    E: Embedder,
+    S: CacheCodec,
+    C: redis::ConnectionLike + Send + 'static,
 {
-    type Value = CacheEntry;
+    type Value = S::Value;
     type Context = SemanticCacheContext;
 
     fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
@@ -274,22 +240,18 @@ impl<E: Embedder, C: redis::ConnectionLike + Send + 'static> BaseCache
         let Some(prompt) = prompt_from_context(context) else {
             return Ok(());
         };
+        let response = self.codec.encode(&value)?;
         let vector = self.embedder.embed(&prompt, context.metadata.as_ref())?;
-        let tag = Self::tag(key, context).to_string();
+        let tag = Self::tag(key, context);
         self.connections.execute(|connection| {
             self.inner
-                .store(connection, &tag, &value, &prompt, &vector, context.ttl)
+                .store(connection, tag, response, &prompt, &vector, context.ttl)
         })
     }
 
     fn get_cache(&self, key: &str, context: &Self::Context) -> Result<Option<Self::Value>, Error> {
-        let Some(prompt) = prompt_from_context(context) else {
-            return Ok(None);
-        };
-        let vector = self.embedder.embed(&prompt, context.metadata.as_ref())?;
-        let tag = Self::tag(key, context).to_string();
-        self.connections
-            .execute(|connection| self.inner.lookup(connection, &tag, &vector))
+        self.get_cache_with_similarity(key, context)
+            .map(|lookup| lookup.value)
     }
 
     async fn async_set_cache(
@@ -301,14 +263,15 @@ impl<E: Embedder, C: redis::ConnectionLike + Send + 'static> BaseCache
         let Some(prompt) = prompt_from_context(&context) else {
             return Ok(());
         };
+        let response = self.codec.encode(&value)?;
         let vector = self
             .embedder
             .async_embed(&prompt, context.metadata.as_ref())
             .await?;
-        let tag = Self::tag(key, &context).to_string();
+        let tag = Self::tag(key, &context).to_owned();
         let inner = Arc::clone(&self.inner);
         Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
-            inner.store(connection, &tag, &value, &prompt, &vector, context.ttl)
+            inner.store(connection, &tag, response, &prompt, &vector, context.ttl)
         })
         .await
     }
@@ -318,49 +281,54 @@ impl<E: Embedder, C: redis::ConnectionLike + Send + 'static> BaseCache
         key: &str,
         context: &Self::Context,
     ) -> Result<Option<Self::Value>, Error> {
+        self.async_get_cache_with_similarity(key, context)
+            .await
+            .map(|lookup| lookup.value)
+    }
+}
+
+/// Python stamps a similarity of `0.0` when there is no prompt or no hit in the key's scope.
+impl<E, S, C> SemanticCache for RedisSemanticCache<E, S, C>
+where
+    E: Embedder,
+    S: CacheCodec,
+    C: redis::ConnectionLike + Send + 'static,
+{
+    fn get_cache_with_similarity(
+        &self,
+        key: &str,
+        context: &Self::Context,
+    ) -> Result<SemanticLookup<Self::Value>, Error> {
         let Some(prompt) = prompt_from_context(context) else {
-            return Ok(None);
+            return Ok(SemanticLookup::miss(Some(0.0)));
+        };
+        let vector = self.embedder.embed(&prompt, context.metadata.as_ref())?;
+        let tag = Self::tag(key, context);
+        let lookup = self
+            .connections
+            .execute(|connection| self.inner.lookup(connection, tag, &vector))?;
+        self.decode(lookup)
+    }
+
+    async fn async_get_cache_with_similarity(
+        &self,
+        key: &str,
+        context: &Self::Context,
+    ) -> Result<SemanticLookup<Self::Value>, Error> {
+        let Some(prompt) = prompt_from_context(context) else {
+            return Ok(SemanticLookup::miss(Some(0.0)));
         };
         let vector = self
             .embedder
             .async_embed(&prompt, context.metadata.as_ref())
             .await?;
-        let tag = Self::tag(key, context).to_string();
+        let tag = Self::tag(key, context).to_owned();
         let inner = Arc::clone(&self.inner);
-        Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
+        let lookup = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
             inner.lookup(connection, &tag, &vector)
         })
-        .await
-    }
-
-    async fn disconnect(&self) -> Result<(), Error> {
-        Ok(())
-    }
-
-    async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
-        match Connections::run_blocking(Arc::clone(&self.connections), |connection| {
-            Ok(match redis::cmd("PING").query::<String>(connection) {
-                Ok(_) => CacheConnectionResult {
-                    status: CacheConnectionStatus::Success,
-                    message: "Redis cache connection test successful".into(),
-                    error: None,
-                },
-                Err(error) => CacheConnectionResult {
-                    status: CacheConnectionStatus::Failed,
-                    message: format!("Redis connection failed: {error}"),
-                    error: Some(error.to_string()),
-                },
-            })
-        })
-        .await
-        {
-            Ok(result) => Ok(result),
-            Err(error) => Ok(CacheConnectionResult {
-                status: CacheConnectionStatus::Failed,
-                message: format!("Redis connection failed: {error}"),
-                error: Some(error.to_string()),
-            }),
-        }
+        .await?;
+        self.decode(lookup)
     }
 }
 
@@ -387,228 +355,46 @@ fn vector_buffer(vector: &[f32]) -> Vec<u8> {
 }
 
 fn escape_tag(value: &str) -> String {
-    value
-        .chars()
-        .flat_map(|ch| {
-            if matches!(
-                ch,
-                ',' | '.'
-                    | '<'
-                    | '>'
-                    | '{'
-                    | '}'
-                    | '['
-                    | ']'
-                    | '\\'
-                    | '"'
-                    | '\''
-                    | ':'
-                    | ';'
-                    | '!'
-                    | '@'
-                    | '#'
-                    | '$'
-                    | '%'
-                    | '^'
-                    | '&'
-                    | '*'
-                    | '('
-                    | ')'
-                    | '-'
-                    | '+'
-                    | '='
-                    | '~'
-                    | '|'
-                    | '/'
-                    | ' '
-                    | '?'
-            ) {
-                vec!['\\', ch]
-            } else {
-                vec![ch]
-            }
-        })
-        .collect()
-}
-
-fn create_index(connection: &mut ConnectionRef<'_>, name: &str, dims: usize) -> Result<(), Error> {
-    redis::cmd("FT.CREATE")
-        .arg(name)
-        .arg("ON")
-        .arg("HASH")
-        .arg("PREFIX")
-        .arg(1)
-        .arg(name)
-        .arg("SCORE")
-        .arg(1.0)
-        .arg("SCHEMA")
-        .arg("prompt")
-        .arg("TEXT")
-        .arg("WEIGHT")
-        .arg(1)
-        .arg("response")
-        .arg("TEXT")
-        .arg("WEIGHT")
-        .arg(1)
-        .arg("inserted_at")
-        .arg("NUMERIC")
-        .arg("updated_at")
-        .arg("NUMERIC")
-        .arg(VECTOR_FIELD)
-        .arg("VECTOR")
-        .arg("FLAT")
-        .arg(6)
-        .arg("TYPE")
-        .arg("FLOAT32")
-        .arg("DIM")
-        .arg(dims)
-        .arg("DISTANCE_METRIC")
-        .arg("COSINE")
-        .arg(CACHE_KEY_FIELD)
-        .arg("TAG")
-        .arg("SEPARATOR")
-        .arg(",")
-        .query::<()>(connection)
-        .map_err(|_| Error::Unavailable)
-}
-
-fn index_compatible(
-    connection: &mut ConnectionRef<'_>,
-    name: &str,
-    dims: usize,
-) -> Result<Option<bool>, Error> {
-    let info = match redis::cmd("FT.INFO")
-        .arg(name)
-        .query::<redis::Value>(connection)
-    {
-        Ok(info) => info,
-        Err(error) if unknown_index(&error) => return Ok(None),
-        Err(_) => return Err(Error::Unavailable),
-    };
-    Ok(Some(schema_compatible(&info, dims)))
-}
-
-fn unknown_index(error: &redis::RedisError) -> bool {
-    let message = error.to_string().to_lowercase();
-    message.contains("unknown") && message.contains("index")
-}
-
-fn schema_compatible(info: &redis::Value, dims: usize) -> bool {
-    let redis::Value::Array(entries) = info else {
-        return false;
-    };
-    let attributes = entries
-        .as_chunks::<2>()
-        .0
-        .iter()
-        .find(|pair| string_value(&pair[0]).as_deref() == Some("attributes"))
-        .map(|pair| &pair[1]);
-    let Some(redis::Value::Array(attributes)) = attributes else {
-        return false;
-    };
-    let fields = attributes
-        .iter()
-        .map(|attribute| {
-            let redis::Value::Array(attribute) = attribute else {
-                return (None, None, None, None, None);
-            };
-            let mut name = None;
-            let mut field_type = None;
-            let mut dim = None;
-            let mut data_type = None;
-            let mut distance_metric = None;
-            for pair in attribute.as_chunks::<2>().0 {
-                match string_value(&pair[0]).as_deref() {
-                    Some("identifier") => name = string_value(&pair[1]),
-                    Some("type") => field_type = string_value(&pair[1]),
-                    Some("dim") => dim = number_value(&pair[1]),
-                    Some("data_type") => data_type = string_value(&pair[1]),
-                    Some("distance_metric") => distance_metric = string_value(&pair[1]),
-                    _ => {}
-                }
-            }
-            (name, field_type, dim, data_type, distance_metric)
-        })
-        .collect::<Vec<_>>();
-    let has_field = |name: &str, field_type: &str| {
-        fields
-            .iter()
-            .any(|(n, t, ..)| n.as_deref() == Some(name) && t.as_deref() == Some(field_type))
-    };
-    has_field("prompt", "TEXT")
-        && has_field("response", "TEXT")
-        && has_field("inserted_at", "NUMERIC")
-        && has_field("updated_at", "NUMERIC")
-        && has_field(CACHE_KEY_FIELD, "TAG")
-        && fields.iter().any(|(n, t, d, data, metric)| {
-            n.as_deref() == Some(VECTOR_FIELD)
-                && t.as_deref() == Some("VECTOR")
-                && *d == Some(dims as f64)
-                && data
-                    .as_deref()
-                    .is_some_and(|data| data.eq_ignore_ascii_case("float32"))
-                && metric
-                    .as_deref()
-                    .is_some_and(|metric| metric.eq_ignore_ascii_case("cosine"))
-        })
-}
-
-fn string_value(value: &redis::Value) -> Option<String> {
-    match value {
-        redis::Value::BulkString(bytes) => String::from_utf8(bytes.clone()).ok(),
-        redis::Value::SimpleString(text) => Some(text.clone()),
-        redis::Value::VerbatimString { text, .. } => Some(text.clone()),
-        _ => None,
-    }
-}
-
-fn number_value(value: &redis::Value) -> Option<f64> {
-    match value {
-        redis::Value::Int(number) => Some(*number as f64),
-        redis::Value::Double(number) => Some(*number),
-        _ => string_value(value).and_then(|text| text.parse().ok()),
-    }
-}
-
-fn first_document(result: &redis::Value) -> Option<&[redis::Value]> {
-    let redis::Value::Array(items) = result else {
-        return None;
-    };
-    let [count, _document_id, fields, ..] = items.as_slice() else {
-        return None;
-    };
-    if !matches!(count, redis::Value::Int(count) if *count > 0) {
-        return None;
-    }
-    match fields {
-        redis::Value::Array(fields) => Some(fields.as_slice()),
-        _ => None,
-    }
-}
-
-fn field_value<'a>(fields: &'a [redis::Value], name: &str) -> Option<&'a redis::Value> {
-    fields
-        .as_chunks::<2>()
-        .0
-        .iter()
-        .find(|pair| string_value(&pair[0]).as_deref() == Some(name))
-        .map(|pair| &pair[1])
-}
-
-fn string_field(fields: &[redis::Value], name: &str) -> Option<String> {
-    field_value(fields, name).and_then(string_value)
-}
-
-fn number_field(fields: &[redis::Value], name: &str) -> Option<f64> {
-    field_value(fields, name).and_then(number_value)
-}
-
-fn bytes_field(fields: &[redis::Value], name: &str) -> Option<Vec<u8>> {
-    match field_value(fields, name)? {
-        redis::Value::BulkString(bytes) => Some(bytes.clone()),
-        redis::Value::SimpleString(text) => Some(text.clone().into_bytes()),
-        _ => None,
+    let mut escaped = String::with_capacity(value.len());
+    for ch in value.chars() {
+        if matches!(
+            ch,
+            ',' | '.'
+                | '<'
+                | '>'
+                | '{'
+                | '}'
+                | '['
+                | ']'
+                | '\\'
+                | '"'
+                | '\''
+                | ':'
+                | ';'
+                | '!'
+                | '@'
+                | '#'
+                | '$'
+                | '%'
+                | '^'
+                | '&'
+                | '*'
+                | '('
+                | ')'
+                | '-'
+                | '+'
+                | '='
+                | '~'
+                | '|'
+                | '/'
+                | ' '
+                | '?'
+        ) {
+            escaped.push('\\');
+        }
+        escaped.push(ch);
     }
+    escaped
 }
 
 fn ttl_seconds(ttl: Duration) -> u64 {

litellm-rust/crates/cache-redis-semantic/src/config.rs

@@ -0,0 +1,8 @@
+/// `RedisSemanticCache.DEFAULT_REDIS_INDEX_NAME`.
+pub const DEFAULT_INDEX_NAME: &str = "litellm_semantic_cache_index";
+
+#[derive(Clone, Debug)]
+pub struct RedisSemanticConfig {
+    pub index_name: String,
+    pub similarity_threshold: f32,
+}

litellm-rust/crates/cache-redis-semantic/src/index.rs

@@ -0,0 +1,205 @@
+use std::sync::OnceLock;
+
+use litellm_cache::Error;
+use litellm_cache_redis::connection::ConnectionRef;
+
+use crate::reply::{number_value, string_value};
+
+pub(crate) const CACHE_KEY_FIELD: &str = "litellm_cache_key";
+pub(crate) const VECTOR_FIELD: &str = "prompt_vector";
+
+/// The redisvl `SemanticCache` index, resolved once per cache: the configured name when its
+/// schema fits, else `<name>_isolated`, recreated when that one is stale too.
+pub(crate) struct Index {
+    name: String,
+    resolved: OnceLock<String>,
+}
+
+impl Index {
+    pub(crate) fn new(name: String) -> Self {
+        Self {
+            name,
+            resolved: OnceLock::new(),
+        }
+    }
+
+    pub(crate) fn name(&self) -> &str {
+        &self.name
+    }
+
+    pub(crate) fn ensure(
+        &self,
+        connection: &mut ConnectionRef<'_>,
+        dims: usize,
+    ) -> Result<String, Error> {
+        if let Some(name) = self.resolved.get() {
+            return Ok(name.clone());
+        }
+        let name = match index_compatible(connection, &self.name, dims)? {
+            Some(true) => self.name.clone(),
+            Some(false) => self.isolated(connection, dims)?,
+            None => match create_index(connection, &self.name, dims) {
+                Ok(()) => self.name.clone(),
+                Err(_) => match index_compatible(connection, &self.name, dims)? {
+                    Some(true) => self.name.clone(),
+                    Some(false) => self.isolated(connection, dims)?,
+                    None => return Err(Error::Unavailable),
+                },
+            },
+        };
+        let _ = self.resolved.set(name.clone());
+        Ok(name)
+    }
+
+    fn isolated(&self, connection: &mut ConnectionRef<'_>, dims: usize) -> Result<String, Error> {
+        let name = format!("{}_isolated", self.name);
+        match index_compatible(connection, &name, dims)? {
+            Some(true) => Ok(name),
+            Some(false) => {
+                redis::cmd("FT.DROPINDEX")
+                    .arg(&name)
+                    .query::<()>(connection)
+                    .map_err(|_| Error::Unavailable)?;
+                create_index(connection, &name, dims)?;
+                Ok(name)
+            }
+            None => {
+                create_index(connection, &name, dims)?;
+                Ok(name)
+            }
+        }
+    }
+}
+
+fn create_index(connection: &mut ConnectionRef<'_>, name: &str, dims: usize) -> Result<(), Error> {
+    redis::cmd("FT.CREATE")
+        .arg(name)
+        .arg("ON")
+        .arg("HASH")
+        .arg("PREFIX")
+        .arg(1)
+        .arg(name)
+        .arg("SCORE")
+        .arg(1.0)
+        .arg("SCHEMA")
+        .arg("prompt")
+        .arg("TEXT")
+        .arg("WEIGHT")
+        .arg(1)
+        .arg("response")
+        .arg("TEXT")
+        .arg("WEIGHT")
+        .arg(1)
+        .arg("inserted_at")
+        .arg("NUMERIC")
+        .arg("updated_at")
+        .arg("NUMERIC")
+        .arg(VECTOR_FIELD)
+        .arg("VECTOR")
+        .arg("FLAT")
+        .arg(6)
+        .arg("TYPE")
+        .arg("FLOAT32")
+        .arg("DIM")
+        .arg(dims)
+        .arg("DISTANCE_METRIC")
+        .arg("COSINE")
+        .arg(CACHE_KEY_FIELD)
+        .arg("TAG")
+        .arg("SEPARATOR")
+        .arg(",")
+        .query::<()>(connection)
+        .map_err(|_| Error::Unavailable)
+}
+
+fn index_compatible(
+    connection: &mut ConnectionRef<'_>,
+    name: &str,
+    dims: usize,
+) -> Result<Option<bool>, Error> {
+    let info = match redis::cmd("FT.INFO")
+        .arg(name)
+        .query::<redis::Value>(connection)
+    {
+        Ok(info) => info,
+        Err(error) if unknown_index(&error) => return Ok(None),
+        Err(_) => return Err(Error::Unavailable),
+    };
+    Ok(Some(schema_compatible(&info, dims)))
+}
+
+fn unknown_index(error: &redis::RedisError) -> bool {
+    let message = error.to_string().to_lowercase();
+    message.contains("unknown") && message.contains("index")
+}
+
+struct Attribute {
+    name: Option<String>,
+    field_type: Option<String>,
+    dim: Option<f64>,
+    data_type: Option<String>,
+    distance_metric: Option<String>,
+}
+
+fn attribute(value: &redis::Value) -> Option<Attribute> {
+    let redis::Value::Array(pairs) = value else {
+        return None;
+    };
+    let mut attribute = Attribute {
+        name: None,
+        field_type: None,
+        dim: None,
+        data_type: None,
+        distance_metric: None,
+    };
+    for pair in pairs.as_chunks::<2>().0 {
+        match string_value(&pair[0]).as_deref() {
+            Some("identifier") => attribute.name = string_value(&pair[1]),
+            Some("type") => attribute.field_type = string_value(&pair[1]),
+            Some("dim") => attribute.dim = number_value(&pair[1]),
+            Some("data_type") => attribute.data_type = string_value(&pair[1]),
+            Some("distance_metric") => attribute.distance_metric = string_value(&pair[1]),
+            _ => {}
+        }
+    }
+    Some(attribute)
+}
+
+fn schema_compatible(info: &redis::Value, dims: usize) -> bool {
+    let redis::Value::Array(entries) = info else {
+        return false;
+    };
+    let attributes = entries
+        .as_chunks::<2>()
+        .0
+        .iter()
+        .find(|pair| string_value(&pair[0]).as_deref() == Some("attributes"))
+        .map(|pair| &pair[1]);
+    let Some(redis::Value::Array(attributes)) = attributes else {
+        return false;
+    };
+    let fields = attributes.iter().filter_map(attribute).collect::<Vec<_>>();
+    let has_field = |name: &str, field_type: &str| {
+        fields.iter().any(|field| {
+            field.name.as_deref() == Some(name) && field.field_type.as_deref() == Some(field_type)
+        })
+    };
+    has_field("prompt", "TEXT")
+        && has_field("response", "TEXT")
+        && has_field("inserted_at", "NUMERIC")
+        && has_field("updated_at", "NUMERIC")
+        && has_field(CACHE_KEY_FIELD, "TAG")
+        && fields.iter().any(|field| {
+            field.name.as_deref() == Some(VECTOR_FIELD)
+                && field.field_type.as_deref() == Some("VECTOR")
+                && field.dim == Some(dims as f64)
+                && field
+                    .data_type
+                    .as_deref()
+                    .is_some_and(|data| data.eq_ignore_ascii_case("float32"))
+                && field
+                    .distance_metric
+                    .as_deref()
+                    .is_some_and(|metric| metric.eq_ignore_ascii_case("cosine"))
+        })
+}

litellm-rust/crates/cache-redis-semantic/src/lib.rs

@@ -1,5 +1,7 @@
 mod cache;
-mod prompt;
+mod config;
+mod index;
+mod reply;
 
-pub use cache::{Embedder, RedisSemanticCache, RedisSemanticConfig};
-pub use prompt::prompt_from_context;
+pub use cache::RedisSemanticCache;
+pub use config::{DEFAULT_INDEX_NAME, RedisSemanticConfig};

litellm-rust/crates/cache-redis-semantic/src/prompt.rs

@@ -1,97 +0,0 @@
-use litellm_cache::SemanticCacheContext;
-use serde_json::Value;
-
-pub fn prompt_from_context(context: &SemanticCacheContext) -> Option<String> {
-    if let Some(messages) = context.messages.as_ref().and_then(Value::as_array)
-        && !messages.is_empty()
-    {
-        return Some(messages_text(messages));
-    }
-    let input = context.input.as_ref()?;
-    let mut parts = Vec::new();
-    collect_input_text(input, &mut parts);
-    let prompt = parts.join("\n").trim().to_string();
-    (!prompt.is_empty()).then_some(prompt)
-}
-
-fn messages_text(messages: &[Value]) -> String {
-    let mut text = String::new();
-    for message in messages {
-        let Some(message) = message.as_object() else {
-            continue;
-        };
-        match message.get("content") {
-            Some(Value::String(content)) => text.push_str(content),
-            Some(Value::Array(parts)) => {
-                for part in parts {
-                    if let Some(text_content) = part.get("text").and_then(Value::as_str) {
-                        text.push_str(text_content);
-                    }
-                }
-            }
-            _ => {}
-        }
-        text.push_str(&search_results_text(message.get("search_results")));
-    }
-    text
-}
-
-fn search_results_text(search_results: Option<&Value>) -> String {
-    let Some(Value::Array(results)) = search_results else {
-        return String::new();
-    };
-    let mut text = String::new();
-    for result in results {
-        let Some(result) = result.as_object() else {
-            continue;
-        };
-        for key in ["source", "title"] {
-            if let Some(value) = result.get(key).and_then(Value::as_str) {
-                text.push_str(value);
-            }
-        }
-        if let Some(Value::Array(content)) = result.get("content") {
-            for block in content {
-                if let Some(value) = block.get("text").and_then(Value::as_str) {
-                    text.push_str(value);
-                }
-            }
-        }
-        if let Some(citations) = result.get("citations") {
-            text.push_str(&citations.to_string());
-        }
-    }
-    text
-}
-
-fn collect_input_text(value: &Value, parts: &mut Vec<String>) {
-    match value {
-        Value::String(text) => {
-            let trimmed = text.trim();
-            if !trimmed.is_empty() {
-                parts.push(trimmed.to_string());
-            }
-        }
-        Value::Array(items) => {
-            for item in items {
-                collect_input_text(item, parts);
-            }
-        }
-        Value::Object(map) => {
-            if let Some(content) = map.get("content").filter(|content| !content.is_null()) {
-                collect_input_text(content, parts);
-                return;
-            }
-            for key in ["text", "output", "input_text", "output_text"] {
-                if let Some(Value::String(text)) = map.get(key) {
-                    let trimmed = text.trim();
-                    if !trimmed.is_empty() {
-                        parts.push(trimmed.to_string());
-                        return;
-                    }
-                }
-            }
-        }
-        _ => {}
-    }
-}

litellm-rust/crates/cache-redis-semantic/src/reply.rs

@@ -0,0 +1,57 @@
+pub(crate) fn string_value(value: &redis::Value) -> Option<String> {
+    match value {
+        redis::Value::BulkString(bytes) => String::from_utf8(bytes.clone()).ok(),
+        redis::Value::SimpleString(text) => Some(text.clone()),
+        redis::Value::VerbatimString { text, .. } => Some(text.clone()),
+        _ => None,
+    }
+}
+
+pub(crate) fn number_value(value: &redis::Value) -> Option<f64> {
+    match value {
+        redis::Value::Int(number) => Some(*number as f64),
+        redis::Value::Double(number) => Some(*number),
+        _ => string_value(value).and_then(|text| text.parse().ok()),
+    }
+}
+
+pub(crate) fn first_document(result: &redis::Value) -> Option<&[redis::Value]> {
+    let redis::Value::Array(items) = result else {
+        return None;
+    };
+    let [count, _document_id, fields, ..] = items.as_slice() else {
+        return None;
+    };
+    if !matches!(count, redis::Value::Int(count) if *count > 0) {
+        return None;
+    }
+    match fields {
+        redis::Value::Array(fields) => Some(fields.as_slice()),
+        _ => None,
+    }
+}
+
+fn field_value<'a>(fields: &'a [redis::Value], name: &str) -> Option<&'a redis::Value> {
+    fields
+        .as_chunks::<2>()
+        .0
+        .iter()
+        .find(|pair| string_value(&pair[0]).as_deref() == Some(name))
+        .map(|pair| &pair[1])
+}
+
+pub(crate) fn string_field(fields: &[redis::Value], name: &str) -> Option<String> {
+    field_value(fields, name).and_then(string_value)
+}
+
+pub(crate) fn number_field(fields: &[redis::Value], name: &str) -> Option<f64> {
+    field_value(fields, name).and_then(number_value)
+}
+
+pub(crate) fn bytes_field(fields: &[redis::Value], name: &str) -> Option<Vec<u8>> {
+    match field_value(fields, name)? {
+        redis::Value::BulkString(bytes) => Some(bytes.clone()),
+        redis::Value::SimpleString(text) => Some(text.clone().into_bytes()),
+        _ => None,
+    }
+}

litellm-rust/crates/cache-redis-semantic/tests/contract.rs

@@ -0,0 +1,63 @@
+mod support;
+
+use litellm_cache::{JsonCodec, SemanticCacheContext, semantic::PreparedEmbedding};
+use litellm_cache_redis_semantic::{DEFAULT_INDEX_NAME, RedisSemanticCache, RedisSemanticConfig};
+use litellm_cache_testing as contract;
+use rstest::{fixture, rstest};
+use serde_json::{Value, json};
+use support::FakeSearch;
+
+type Cache = RedisSemanticCache<PreparedEmbedding, JsonCodec<Value>, FakeSearch>;
+
+const PREFIX: &str = "contract:";
+
+#[fixture]
+fn cache() -> Cache {
+    RedisSemanticCache::with_connection(
+        FakeSearch::default(),
+        PreparedEmbedding(vec![0.6, 0.8]),
+        JsonCodec::new(),
+        RedisSemanticConfig {
+            index_name: DEFAULT_INDEX_NAME.into(),
+            similarity_threshold: 0.9,
+        },
+    )
+}
+
+#[fixture]
+fn context() -> SemanticCacheContext {
+    SemanticCacheContext {
+        messages: Some(json!([{"role": "user", "content": "contract prompt"}])),
+        ..Default::default()
+    }
+}
+
+#[rstest]
+#[tokio::test]
+async fn hit_and_miss(cache: Cache, context: SemanticCacheContext) {
+    contract::hit_and_miss(&cache, context, PREFIX, json!({"answer": 42})).await;
+}
+
+#[rstest]
+#[tokio::test]
+async fn sync_async_equivalence(cache: Cache, context: SemanticCacheContext) {
+    contract::sync_async_equivalence(&cache, context, PREFIX, json!("first"), json!([2])).await;
+}
+
+#[rstest]
+#[tokio::test]
+async fn overwrite_replaces(cache: Cache, context: SemanticCacheContext) {
+    contract::overwrite_replaces(&cache, context, PREFIX, json!(1), json!({"b": 2})).await;
+}
+
+#[rstest]
+#[tokio::test]
+async fn pipeline_writes_every_entry(cache: Cache, context: SemanticCacheContext) {
+    contract::pipeline_writes_every_entry(
+        &cache,
+        context,
+        PREFIX,
+        vec![json!("a"), json!(2), json!({"c": true})],
+    )
+    .await;
+}

litellm-rust/crates/cache-redis-semantic/tests/support/mod.rs

@@ -0,0 +1,299 @@
+#![allow(dead_code)]
+
+use std::{
+    collections::{BTreeMap, HashMap},
+    sync::{Arc, Mutex},
+};
+
+use litellm_cache::{Error, semantic::Embedder};
+use serde_json::Value;
+
+pub type EmbedCalls = Arc<Mutex<Vec<(String, Option<Value>)>>>;
+
+/// Embeds known prompts to fixed vectors, anything else to `[0.1, 0.2, 0.3]`, and records every
+/// prompt with its metadata.
+pub struct FakeEmbedder {
+    vectors: HashMap<String, Vec<f32>>,
+    pub calls: EmbedCalls,
+}
+
+impl FakeEmbedder {
+    pub fn new(vectors: &[(&str, &[f32])]) -> Self {
+        Self {
+            vectors: vectors
+                .iter()
+                .map(|(prompt, vector)| ((*prompt).to_owned(), vector.to_vec()))
+                .collect(),
+            calls: EmbedCalls::default(),
+        }
+    }
+}
+
+impl Embedder for FakeEmbedder {
+    fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result<Vec<f32>, Error> {
+        self.calls
+            .lock()
+            .unwrap()
+            .push((prompt.to_owned(), metadata.cloned()));
+        Ok(self
+            .vectors
+            .get(prompt)
+            .cloned()
+            .unwrap_or_else(|| vec![0.1, 0.2, 0.3]))
+    }
+
+    async fn async_embed(&self, prompt: &str, metadata: Option<&Value>) -> Result<Vec<f32>, Error> {
+        self.embed(prompt, metadata)
+    }
+}
+
+struct FakeIndex {
+    prefix: Vec<u8>,
+    dims: usize,
+    vector_field: String,
+}
+
+#[derive(Default)]
+struct SearchState {
+    indexes: HashMap<String, FakeIndex>,
+    hashes: BTreeMap<Vec<u8>, BTreeMap<String, Vec<u8>>>,
+}
+
+/// An in-memory Redis Stack speaking the `FT.*`, `HSET` and `EXPIRE` subset the semantic cache
+/// sends, with exact cosine KNN over the hashes under an index prefix.
+#[derive(Clone, Default)]
+pub struct FakeSearch {
+    state: Arc<Mutex<SearchState>>,
+}
+
+impl FakeSearch {
+    fn run(&self, args: Vec<Vec<u8>>) -> redis::RedisResult<redis::Value> {
+        let mut state = self.state.lock().unwrap();
+        let text = |index: usize| String::from_utf8_lossy(&args[index]).into_owned();
+        match text(0).to_uppercase().as_str() {
+            "FT.CREATE" => {
+                let name = text(1);
+                if state.indexes.contains_key(&name) {
+                    return Err(error("Index already exists"));
+                }
+                let position = |token: &str| args.iter().position(|arg| arg == token.as_bytes());
+                let prefix = args[position("PREFIX").unwrap() + 2].clone();
+                let dims = text(position("DIM").unwrap() + 1).parse().unwrap();
+                let vector_field = text(position("VECTOR").unwrap() - 1);
+                state.indexes.insert(
+                    name,
+                    FakeIndex {
+                        prefix,
+                        dims,
+                        vector_field,
+                    },
+                );
+                Ok(redis::Value::Okay)
+            }
+            "FT.INFO" => {
+                let index = state
+                    .indexes
+                    .get(&text(1))
+                    .ok_or_else(|| error("Unknown index name"))?;
+                Ok(index_info(index))
+            }
+            "FT.DROPINDEX" => {
+                state.indexes.remove(&text(1));
+                Ok(redis::Value::Okay)
+            }
+            "HSET" => {
+                let hash = state.hashes.entry(args[1].clone()).or_default();
+                for pair in args[2..].chunks(2) {
+                    hash.insert(
+                        String::from_utf8_lossy(&pair[0]).into_owned(),
+                        pair[1].clone(),
+                    );
+                }
+                Ok(redis::Value::Int(((args.len() - 2) / 2) as i64))
+            }
+            "EXPIRE" => Ok(redis::Value::Int(i64::from(
+                state.hashes.contains_key(&args[1]),
+            ))),
+            "FT.SEARCH" => {
+                let index = state
+                    .indexes
+                    .get(&text(1))
+                    .ok_or_else(|| error("no such index"))?;
+                let query = text(2);
+                let tag = query_tag(&query);
+                let params = args.iter().position(|arg| arg == b"PARAMS").unwrap();
+                let vector = floats(&args[params + 3]);
+                let best = state
+                    .hashes
+                    .iter()
+                    .filter(|(key, _)| key.starts_with(&index.prefix))
+                    .filter(|(_, fields)| {
+                        fields.get("litellm_cache_key").map(Vec::as_slice) == Some(tag.as_bytes())
+                    })
+                    .filter_map(|(key, fields)| {
+                        let stored = floats(fields.get(&index.vector_field)?);
+                        (stored.len() == index.dims)
+                            .then(|| (key, fields, 1.0 - cosine(&vector, &stored)))
+                    })
+                    .min_by(|left, right| left.2.total_cmp(&right.2));
+                let Some((key, fields, distance)) = best else {
+                    return Ok(redis::Value::Array(vec![redis::Value::Int(0)]));
+                };
+                let mut reply = fields
+                    .iter()
+                    .filter(|(name, _)| **name != index.vector_field)
+                    .flat_map(|(name, value)| [bulk(name.as_bytes()), bulk(value)])
+                    .collect::<Vec<_>>();
+                reply.extend([
+                    bulk(b"vector_distance"),
+                    bulk(distance.to_string().as_bytes()),
+                ]);
+                Ok(redis::Value::Array(vec![
+                    redis::Value::Int(1),
+                    bulk(key),
+                    redis::Value::Array(reply),
+                ]))
+            }
+            "PING" => Ok(redis::Value::SimpleString("PONG".into())),
+            _ => Err(error("unsupported command")),
+        }
+    }
+}
+
+impl redis::ConnectionLike for FakeSearch {
+    fn req_packed_command(&mut self, command: &[u8]) -> redis::RedisResult<redis::Value> {
+        let mut commands = parse_commands(command);
+        self.run(commands.remove(0))
+    }
+
+    fn req_packed_commands(
+        &mut self,
+        commands: &[u8],
+        offset: usize,
+        count: usize,
+    ) -> redis::RedisResult<Vec<redis::Value>> {
+        let replies = parse_commands(commands)
+            .into_iter()
+            .map(|args| self.run(args))
+            .collect::<redis::RedisResult<Vec<_>>>()?;
+        Ok(replies.into_iter().skip(offset).take(count).collect())
+    }
+
+    fn get_db(&self) -> i64 {
+        0
+    }
+
+    fn check_connection(&mut self) -> bool {
+        true
+    }
+
+    fn is_open(&self) -> bool {
+        true
+    }
+}
+
+fn error(message: &'static str) -> redis::RedisError {
+    redis::RedisError::from((redis::ErrorKind::Extension, message))
+}
+
+fn bulk(bytes: &[u8]) -> redis::Value {
+    redis::Value::BulkString(bytes.to_vec())
+}
+
+fn index_info(index: &FakeIndex) -> redis::Value {
+    let attribute = |name: &str, field_type: &str| {
+        redis::Value::Array(vec![
+            bulk(b"identifier"),
+            bulk(name.as_bytes()),
+            bulk(b"type"),
+            bulk(field_type.as_bytes()),
+        ])
+    };
+    redis::Value::Array(vec![
+        bulk(b"attributes"),
+        redis::Value::Array(vec![
+            attribute("prompt", "TEXT"),
+            attribute("response", "TEXT"),
+            attribute("inserted_at", "NUMERIC"),
+            attribute("updated_at", "NUMERIC"),
+            attribute("litellm_cache_key", "TAG"),
+            redis::Value::Array(vec![
+                bulk(b"identifier"),
+                bulk(index.vector_field.as_bytes()),
+                bulk(b"type"),
+                bulk(b"VECTOR"),
+                bulk(b"dim"),
+                redis::Value::Int(index.dims as i64),
+                bulk(b"data_type"),
+                bulk(b"FLOAT32"),
+                bulk(b"distance_metric"),
+                bulk(b"COSINE"),
+            ]),
+        ]),
+    ])
+}
+
+/// The tag inside `@litellm_cache_key:{...}`, with query escapes removed.
+fn query_tag(query: &str) -> String {
+    let start = query.find("@litellm_cache_key:{").unwrap() + "@litellm_cache_key:{".len();
+    let mut tag = String::new();
+    let mut characters = query[start..].chars();
+    while let Some(character) = characters.next() {
+        match character {
+            '\\' => tag.extend(characters.next()),
+            '}' => break,
+            character => tag.push(character),
+        }
+    }
+    tag
+}
+
+fn floats(bytes: &[u8]) -> Vec<f32> {
+    bytes
+        .as_chunks::<4>()
+        .0
+        .iter()
+        .map(|chunk| f32::from_le_bytes(*chunk))
+        .collect()
+}
+
+fn cosine(left: &[f32], right: &[f32]) -> f64 {
+    let dot = left
+        .iter()
+        .zip(right)
+        .map(|(left, right)| f64::from(*left) * f64::from(*right))
+        .sum::<f64>();
+    let norm = |vector: &[f32]| {
+        vector
+            .iter()
+            .map(|value| f64::from(*value).powi(2))
+            .sum::<f64>()
+            .sqrt()
+    };
+    dot / (norm(left) * norm(right))
+}
+
+/// Splits a packed RESP request into each command's arguments.
+fn parse_commands(mut bytes: &[u8]) -> Vec<Vec<Vec<u8>>> {
+    let line = |bytes: &mut &[u8]| {
+        let end = bytes
+            .windows(2)
+            .position(|window| window == b"\r\n")
+            .unwrap();
+        let text = String::from_utf8(bytes[1..end].to_vec()).unwrap();
+        *bytes = &bytes[end + 2..];
+        text.parse::<usize>().unwrap()
+    };
+    let mut commands = Vec::new();
+    while !bytes.is_empty() {
+        let count = line(&mut bytes);
+        let mut args = Vec::with_capacity(count);
+        for _ in 0..count {
+            let length = line(&mut bytes);
+            args.push(bytes[..length].to_vec());
+            bytes = &bytes[length + 2..];
+        }
+        commands.push(args);
+    }
+    commands
+}

litellm-rust/crates/cache-redis/Cargo.toml

@@ -12,5 +12,7 @@ r2d2 = "0.8.10"
 tokio.workspace = true
 
 [dev-dependencies]
+litellm-cache-testing.workspace = true
 redis-test = "1.0.4"
+rstest.workspace = true
 serde_json.workspace = true

litellm-rust/crates/cache-redis/src/cache.rs

@@ -1,110 +1,25 @@
 use std::{
-    sync::{Arc, Mutex},
+    sync::{Arc, OnceLock},
     time::Duration,
 };
 
-use litellm_cache::{
-    BaseCache, BatchCache, BatchEntry, CacheCodec, CacheConnectionResult, CacheConnectionStatus,
-    ClaimCache, CounterCache, DeleteCache, Error, ExactCacheContext, FlushCache,
-};
-use redis::Commands;
+use litellm_cache::{BatchEntry, CacheCodec, Error};
 
-use crate::topology::RedisTopology;
-
-mod connection;
-mod operations;
-
-pub use connection::ConnectionRef;
-use connection::{ClusterConnectionManager, ConnectionManager};
-
-pub use operations::{
-    RedisArg, RedisLpopOperation, RedisLpopResult, RedisRpushOperation, RedisScript,
+use crate::{
+    connection::{ConnectionRef, Connections},
+    topology::RedisTopology,
 };
 
 const DEFAULT_TTL: Duration = Duration::from_secs(600);
-const REDIS_TIMEOUT: Duration = Duration::from_secs(5);
-const REDIS_POOL_SIZE: u32 = 16;
-
-const INCREMENT_SCRIPT: &str = concat!(
-    "local value = redis.call('INCRBYFLOAT', KEYS[1], ARGV[1]); ",
-    "if redis.call('TTL', KEYS[1]) == -1 then ",
-    "redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return value"
-);
-
-const CLAIM_SCRIPT: &str = concat!(
-    "local current = redis.call('GET', KEYS[1]); ",
-    "if ARGV[1] == '' then if current ~= false and current ~= '' then return 0; end; ",
-    "elseif current ~= ARGV[1] then return 0; end; ",
-    "if ARGV[3] ~= '' then redis.call('SET', KEYS[1], ARGV[3], 'EX', ARGV[2]); ",
-    "elseif ARGV[4] == '1' then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return 1"
-);
-const CLAIM_ATTEMPTS: usize = 8;
-
-#[allow(private_interfaces)]
-pub enum Connections<C> {
-    Pool(r2d2::Pool<ConnectionManager>),
-    Cluster(r2d2::Pool<ClusterConnectionManager>),
-    Fixed(Mutex<C>),
-}
-
-impl<C> Connections<C>
-where
-    C: redis::ConnectionLike + Send + 'static,
-{
-    pub fn execute<T>(
-        &self,
-        operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error>,
-    ) -> Result<T, Error> {
-        match self {
-            Self::Pool(pool) => {
-                let mut pooled = pool.get().map_err(|_| Error::Unavailable)?;
-                let result = operation(&mut ConnectionRef::Node(&mut pooled.connection));
-                pooled.failed = matches!(result, Err(Error::Unavailable));
-                result
-            }
-            Self::Cluster(pool) => {
-                let mut pooled = pool.get().map_err(|_| Error::Unavailable)?;
-                let result = operation(&mut ConnectionRef::Cluster(&mut pooled.connection));
-                pooled.failed = matches!(result, Err(Error::Unavailable));
-                result
-            }
-            Self::Fixed(connection) => {
-                let mut connection = connection.lock().map_err(|_| Error::Unavailable)?;
-                operation(&mut ConnectionRef::Node(&mut *connection))
-            }
-        }
-    }
-
-    pub async fn run_blocking<T, F>(connections: Arc<Self>, operation: F) -> Result<T, Error>
-    where
-        T: Send + 'static,
-        F: FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error> + Send + 'static,
-    {
-        tokio::task::spawn_blocking(move || connections.execute(operation))
-            .await
-            .map_err(|_| Error::Unavailable)?
-    }
-
-    pub fn fixed(connection: C) -> Self {
-        Self::Fixed(Mutex::new(connection))
-    }
-
-    pub fn open(url: &str, topology: &RedisTopology) -> Result<Self, Error> {
-        match topology {
-            RedisTopology::Standalone => Ok(Self::Pool(pool(ConnectionManager::open(url)?)?)),
-            RedisTopology::Cluster { startup_nodes } => Ok(Self::Cluster(pool(
-                ClusterConnectionManager::open(url, startup_nodes)?,
-            )?)),
-        }
-    }
-}
 
 pub struct RedisCache<S, C = redis::Connection> {
-    connections: Arc<Connections<C>>,
-    default_ttl: Duration,
-    codec: S,
-    namespace: Option<String>,
-    topology: RedisTopology,
+    pub(crate) connections: Arc<Connections<C>>,
+    pub(crate) default_ttl: Duration,
+    pub(crate) codec: S,
+    pub(crate) namespace: Option<String>,
+    pub(crate) topology: RedisTopology,
+    /// The server's major version, read from `INFO` once, like Python's `redis_version`.
+    pub(crate) major_version: Arc<OnceLock<u32>>,
 }
 
 impl<S: CacheCodec> RedisCache<S> {
@@ -125,20 +40,11 @@ impl<S: CacheCodec> RedisCache<S> {
             codec,
             namespace: None,
             topology: topology.clone(),
+            major_version: Arc::default(),
         })
     }
 }
 
-fn pool<M: r2d2::ManageConnection>(manager: M) -> Result<r2d2::Pool<M>, Error> {
-    r2d2::Pool::builder()
-        .max_size(REDIS_POOL_SIZE)
-        .min_idle(Some(0))
-        .connection_timeout(REDIS_TIMEOUT)
-        .test_on_check_out(false)
-        .build(manager)
-        .map_err(|_| Error::Unavailable)
-}
-
 impl<S, C> RedisCache<S, C>
 where
     S: CacheCodec,
@@ -151,6 +57,7 @@ where
             codec,
             namespace: None,
             topology: RedisTopology::Standalone,
+            major_version: Arc::default(),
         }
     }
 
@@ -169,11 +76,52 @@ where
         &self.topology
     }
 
-    fn namespaced_key(&self, key: &str) -> String {
+    pub(crate) fn namespaced_key(&self, key: &str) -> String {
         namespaced_key(self.namespace.as_deref(), key)
     }
 
-    fn namespaced_pattern(&self) -> Result<String, Error> {
+    pub(crate) fn namespaced_keys(&self, keys: &[String]) -> Vec<String> {
+        keys.iter().map(|key| self.namespaced_key(key)).collect()
+    }
+
+    /// Whole seconds for `ttl`, falling back to the default TTL like Python's `get_ttl`.
+    pub(crate) fn ttl_or_default(&self, ttl: Option<Duration>) -> u64 {
+        ttl_seconds(ttl.unwrap_or(self.default_ttl))
+    }
+
+    pub(crate) fn execute<T>(
+        &self,
+        operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error>,
+    ) -> Result<T, Error> {
+        self.connections.execute(operation)
+    }
+
+    /// `_parse_redis_major_version`: the major version from `INFO`, or
+    /// `DEFAULT_REDIS_MAJOR_VERSION` when `INFO` fails or its version does not parse. The first
+    /// answer is kept, as Python reads `redis_version` once at construction.
+    pub(crate) async fn major_version(&self) -> u32 {
+        if let Some(version) = self.major_version.get() {
+            return *version;
+        }
+        let info = self
+            .run(|connection| connection.node_text(&redis::cmd("INFO")))
+            .await;
+        let version = info
+            .ok()
+            .and_then(|info| parse_major_version(&info))
+            .unwrap_or_else(default_major_version);
+        *self.major_version.get_or_init(|| version)
+    }
+
+    pub(crate) async fn run<T, F>(&self, operation: F) -> Result<T, Error>
+    where
+        T: Send + 'static,
+        F: FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error> + Send + 'static,
+    {
+        Connections::run_blocking(Arc::clone(&self.connections), operation).await
+    }
+
+    pub(crate) fn namespaced_pattern(&self) -> Result<String, Error> {
         let namespace = self.namespace.as_ref().ok_or(Error::UnscopedFlush)?;
         let escaped: String = namespace
             .chars()
@@ -188,18 +136,7 @@ where
         Ok(format!("{escaped}:*"))
     }
 
-    fn flush_matching(connection: &mut ConnectionRef<'_>, pattern: &str) -> Result<(), Error> {
-        connection.scan(pattern, 1000, |connection, keys| {
-            if !keys.is_empty() {
-                connection
-                    .del::<_, usize>(keys)
-                    .map_err(|_| Error::Unavailable)?;
-            }
-            Ok(true)
-        })
-    }
-
-    fn decode_response(&self, value: redis::Value) -> Result<Option<S::Value>, Error> {
+    pub(crate) fn decode_response(&self, value: redis::Value) -> Result<Option<S::Value>, Error> {
         match value {
             redis::Value::Nil => Ok(None),
             redis::Value::BulkString(bytes) => self.codec.decode(&bytes).map(Some),
@@ -208,7 +145,10 @@ where
         }
     }
 
-    fn decode_batch_response(&self, value: redis::Value) -> Result<BatchEntry<S::Value>, Error> {
+    pub(crate) fn decode_batch_response(
+        &self,
+        value: redis::Value,
+    ) -> Result<BatchEntry<S::Value>, Error> {
         match self.decode_response(value) {
             Ok(Some(value)) => Ok(BatchEntry::Hit(value)),
             Ok(None) => Ok(BatchEntry::Miss),
@@ -216,15 +156,9 @@ where
             Err(error) => Err(error),
         }
     }
-
-    fn ttl_seconds(ttl: Duration) -> u64 {
-        ttl.as_secs()
-            .saturating_add(u64::from(ttl.subsec_nanos() > 0))
-            .max(1)
-    }
 }
 
-fn namespaced_key(namespace: Option<&str>, key: &str) -> String {
+pub(crate) fn namespaced_key(namespace: Option<&str>, key: &str) -> String {
     match namespace {
         Some(namespace) if !key.starts_with(&format!("{namespace}:")) => {
             format!("{namespace}:{key}")
@@ -233,469 +167,26 @@ fn namespaced_key(namespace: Option<&str>, key: &str) -> String {
     }
 }
 
-impl<S, C> BaseCache for RedisCache<S, C>
-where
-    S: CacheCodec,
-    C: redis::ConnectionLike + Send + 'static,
-{
-    type Value = S::Value;
-    type Context = ExactCacheContext;
+pub(crate) fn ttl_seconds(ttl: Duration) -> u64 {
+    ttl.as_secs()
+        .saturating_add(u64::from(ttl.subsec_nanos() > 0))
+        .max(1)
+}
 
-    fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
-        context.ttl.or(Some(self.default_ttl))
-    }
-
-    fn set_cache(
-        &self,
-        key: &str,
-        value: Self::Value,
-        context: &ExactCacheContext,
-    ) -> Result<(), Error> {
-        let payload = self.codec.encode(&value)?;
-        let ttl = Self::ttl_seconds(self.get_ttl(context).unwrap_or(self.default_ttl));
-        let key = self.namespaced_key(key);
-        self.connections.execute(|connection| {
-            connection
-                .set_ex::<_, _, ()>(key, payload, ttl)
-                .map_err(|_| Error::Unavailable)
-        })
-    }
-
-    fn get_cache(&self, key: &str, _: &ExactCacheContext) -> Result<Option<Self::Value>, Error> {
-        let key = self.namespaced_key(key);
-        let value = self.connections.execute(|connection| {
-            connection
-                .get::<_, redis::Value>(key)
-                .map_err(|_| Error::Unavailable)
-        })?;
-        self.decode_response(value)
-    }
-
-    async fn async_set_cache(
-        &self,
-        key: &str,
-        value: Self::Value,
-        context: ExactCacheContext,
-    ) -> Result<(), Error> {
-        let payload = self.codec.encode(&value)?;
-        let key = self.namespaced_key(key);
-        let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
-        Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
-            connection
-                .set_ex::<_, _, ()>(key, payload, ttl)
-                .map_err(|_| Error::Unavailable)
-        })
-        .await
-    }
-
-    async fn async_get_cache(
-        &self,
-        key: &str,
-        _: &ExactCacheContext,
-    ) -> Result<Option<Self::Value>, Error> {
-        let key = self.namespaced_key(key);
-        let value = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
-            connection
-                .get::<_, redis::Value>(key)
-                .map_err(|_| Error::Unavailable)
-        })
-        .await?;
-        self.decode_response(value)
-    }
-
-    async fn async_set_cache_pipeline(
-        &self,
-        cache_list: Vec<(String, Self::Value)>,
-        context: ExactCacheContext,
-    ) -> Result<(), Error> {
-        let entries = cache_list
-            .into_iter()
-            .map(|(key, value)| {
-                self.codec
-                    .encode(&value)
-                    .map(|payload| (self.namespaced_key(&key), payload))
-            })
-            .collect::<Result<Vec<_>, _>>()?;
-        let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
-        if entries.is_empty() {
-            return Ok(());
-        }
-        Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
-            let commands = entries
-                .into_iter()
-                .map(|(key, payload)| {
-                    let mut command = redis::cmd("SETEX");
-                    command.arg(key).arg(ttl).arg(payload);
-                    command
-                })
-                .collect();
-            connection.pipeline(commands).map(drop)
-        })
-        .await
-    }
-
-    async fn disconnect(&self) -> Result<(), Error> {
-        Ok(())
-    }
-
-    async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
-        match Connections::run_blocking(Arc::clone(&self.connections), |connection| {
-            Ok(match connection.ping() {
-                Ok(_) => CacheConnectionResult {
-                    status: CacheConnectionStatus::Success,
-                    message: "Redis cache connection test successful".into(),
-                    error: None,
-                },
-                Err(error) => CacheConnectionResult {
-                    status: CacheConnectionStatus::Failed,
-                    message: format!("Redis connection failed: {error}"),
-                    error: Some(error.to_string()),
-                },
-            })
-        })
-        .await
-        {
-            Ok(result) => Ok(result),
-            Err(error) => Ok(CacheConnectionResult {
-                status: CacheConnectionStatus::Failed,
-                message: format!("Redis connection failed: {error}"),
-                error: Some(error.to_string()),
-            }),
-        }
+fn parse_major_version(info: &str) -> Option<u32> {
+    let version = info
+        .lines()
+        .find_map(|line| line.trim().strip_prefix("redis_version:"))?
+        .trim();
+    match version.split_once('.') {
+        Some((major, _)) => major.parse().ok(),
+        None => version.parse::<f64>().ok().map(|major| major as u32),
     }
 }
 
-impl<S, C> BatchCache for RedisCache<S, C>
-where
-    S: CacheCodec,
-    C: redis::ConnectionLike + Send + 'static,
-{
-    fn batch_get_cache(
-        &self,
-        keys: &[String],
-        _: &ExactCacheContext,
-    ) -> Result<Vec<BatchEntry<Self::Value>>, Error> {
-        let keys = keys
-            .iter()
-            .map(|key| self.namespaced_key(key))
-            .collect::<Vec<_>>();
-        let values = self.connections.execute(|connection| {
-            redis::cmd("MGET")
-                .arg(keys)
-                .query::<Vec<redis::Value>>(connection)
-                .map_err(|_| Error::Unavailable)
-        })?;
-        values
-            .into_iter()
-            .map(|value| self.decode_batch_response(value))
-            .collect()
-    }
-
-    async fn async_batch_get_cache(
-        &self,
-        keys: Vec<String>,
-        _: ExactCacheContext,
-    ) -> Result<Vec<BatchEntry<Self::Value>>, Error> {
-        let keys = keys
-            .iter()
-            .map(|key| self.namespaced_key(key))
-            .collect::<Vec<_>>();
-        let values = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
-            redis::cmd("MGET")
-                .arg(keys)
-                .query::<Vec<redis::Value>>(connection)
-                .map_err(|_| Error::Unavailable)
-        })
-        .await?;
-        values
-            .into_iter()
-            .map(|value| self.decode_batch_response(value))
-            .collect()
-    }
-}
-
-impl<S, C> DeleteCache for RedisCache<S, C>
-where
-    S: CacheCodec,
-    C: redis::ConnectionLike + Send + 'static,
-{
-    fn delete_cache(&self, key: &str) -> Result<(), Error> {
-        let key = self.namespaced_key(key);
-        self.connections
-            .execute(|connection| connection.del::<_, ()>(key).map_err(|_| Error::Unavailable))
-    }
-
-    async fn async_delete_cache(&self, key: &str) -> Result<(), Error> {
-        let key = self.namespaced_key(key);
-        Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
-            connection.del::<_, ()>(key).map_err(|_| Error::Unavailable)
-        })
-        .await
-    }
-}
-
-impl<S, C> FlushCache for RedisCache<S, C>
-where
-    S: CacheCodec,
-    C: redis::ConnectionLike + Send + 'static,
-{
-    fn flush_cache(&self) -> Result<(), Error> {
-        let pattern = self.namespaced_pattern()?;
-        self.connections
-            .execute(|connection| Self::flush_matching(connection, &pattern))
-    }
-
-    async fn async_flush_cache(&self) -> Result<(), Error> {
-        let pattern = self.namespaced_pattern()?;
-        Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
-            Self::flush_matching(connection, &pattern)
-        })
-        .await
-    }
-}
-
-impl<S, C> CounterCache for RedisCache<S, C>
-where
-    S: CacheCodec<Value = f64>,
-    C: redis::ConnectionLike + Send + 'static,
-{
-    fn increment_cache(
-        &self,
-        key: &str,
-        amount: f64,
-        context: ExactCacheContext,
-    ) -> Result<f64, Error> {
-        let key = self.namespaced_key(key);
-        let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
-        self.connections
-            .execute(|connection| increment(connection, key, amount, ttl))
-    }
-
-    async fn async_increment(
-        &self,
-        key: &str,
-        amount: f64,
-        context: ExactCacheContext,
-    ) -> Result<f64, Error> {
-        let key = self.namespaced_key(key);
-        let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
-        Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
-            increment(connection, key, amount, ttl)
-        })
-        .await
-    }
-}
-
-fn increment(
-    connection: &mut ConnectionRef<'_>,
-    key: String,
-    amount: f64,
-    ttl: u64,
-) -> Result<f64, Error> {
-    redis::cmd("EVAL")
-        .arg(INCREMENT_SCRIPT)
-        .arg(1)
-        .arg(key)
-        .arg(amount)
-        .arg(ttl)
-        .query(connection)
-        .map_err(|_| Error::Unavailable)
-}
-
-fn stored_bytes(value: redis::Value) -> Result<Option<Vec<u8>>, Error> {
-    match value {
-        redis::Value::Nil => Ok(None),
-        redis::Value::BulkString(bytes) => Ok(Some(bytes)),
-        redis::Value::SimpleString(text) => Ok(Some(text.into_bytes())),
-        _ => Err(Error::InvalidEntry),
-    }
-}
-
-/// Eligibility is decided on decoded values, so a pin written by another encoder (Python's
-/// `json.dumps` spacing or key order) still matches. The write is a compare-and-set on the
-/// bytes that decision was made on, retried when another claimant wins the race.
-fn claim<S: CacheCodec>(
-    connection: &mut ConnectionRef<'_>,
-    codec: &S,
-    key: &str,
-    candidate: S::Value,
-    eligible: &[S::Value],
-    ttl: u64,
-) -> Result<S::Value, Error>
-where
-    S::Value: PartialEq,
-{
-    let payload = codec.encode(&candidate)?;
-    if payload.is_empty() {
-        return Err(Error::InvalidEntry);
-    }
-    for _ in 0..CLAIM_ATTEMPTS {
-        let current = stored_bytes(
-            connection
-                .get::<_, redis::Value>(key)
-                .map_err(|_| Error::Unavailable)?,
-        )?
-        .filter(|bytes| !bytes.is_empty());
-        let existing = current
-            .as_deref()
-            .and_then(|bytes| codec.decode(bytes).ok())
-            .filter(|existing| eligible.is_empty() || eligible.contains(existing));
-        let refresh = existing
-            .as_ref()
-            .is_some_and(|existing| !eligible.is_empty() || *existing == candidate);
-        let write: &[u8] = if existing.is_some() { b"" } else { &payload };
-        let applied = redis::cmd("EVAL")
-            .arg(CLAIM_SCRIPT)
-            .arg(1)
-            .arg(key)
-            .arg(current.as_deref().unwrap_or_default())
-            .arg(ttl)
-            .arg(write)
-            .arg(u8::from(refresh))
-            .query::<bool>(connection)
-            .map_err(|_| Error::Unavailable)?;
-        if applied {
-            return Ok(existing.unwrap_or(candidate));
-        }
-    }
-    Err(Error::Unavailable)
-}
-
-impl<S, C> ClaimCache for RedisCache<S, C>
-where
-    S: CacheCodec + Clone + 'static,
-    S::Value: PartialEq,
-    C: redis::ConnectionLike + Send + 'static,
-{
-    fn claim_cache(
-        &self,
-        key: &str,
-        candidate: S::Value,
-        eligible: &[S::Value],
-        context: ExactCacheContext,
-    ) -> Result<S::Value, Error> {
-        let key = self.namespaced_key(key);
-        let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
-        self.connections
-            .execute(|connection| claim(connection, &self.codec, &key, candidate, eligible, ttl))
-    }
-
-    async fn async_claim_cache(
-        &self,
-        key: &str,
-        candidate: S::Value,
-        eligible: Vec<S::Value>,
-        context: ExactCacheContext,
-    ) -> Result<S::Value, Error> {
-        let key = self.namespaced_key(key);
-        let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
-        let codec = self.codec.clone();
-        Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
-            claim(connection, &codec, &key, candidate, &eligible, ttl)
-        })
-        .await
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use std::time::Duration;
-
-    use litellm_cache::{
-        BaseCache, CacheCodec, DeleteCache, ExactCacheContext, FlushCache, JsonCodec,
-    };
-    use redis_test::{MockCmd, MockRedisConnection};
-    use serde_json::json;
-
-    use super::RedisCache;
-
-    fn entry() -> serde_json::Value {
-        json!({"deployment": "model-a", "cooldown_seconds": 30})
-    }
-
-    #[test]
-    fn ttl_seconds_rounds_up_and_keeps_expiration_positive() {
-        assert_eq!(
-            RedisCache::<JsonCodec<serde_json::Value>>::ttl_seconds(Duration::ZERO),
-            1
-        );
-        assert_eq!(
-            RedisCache::<JsonCodec<serde_json::Value>>::ttl_seconds(Duration::from_millis(1500)),
-            2
-        );
-        assert_eq!(
-            RedisCache::<JsonCodec<serde_json::Value>>::ttl_seconds(Duration::from_secs(15)),
-            15
-        );
-    }
-
-    #[test]
-    fn redis_commands_round_trip_entries_and_delete_only_namespaced_keys() {
-        let value = entry();
-        let payload = JsonCodec::<serde_json::Value>::new()
-            .encode(&value)
-            .unwrap();
-        let connection = MockRedisConnection::new([
-            MockCmd::new(
-                redis::cmd("SETEX")
-                    .arg("litellm-cache:key")
-                    .arg(600)
-                    .arg(payload.clone()),
-                Ok("OK"),
-            ),
-            MockCmd::new(redis::cmd("GET").arg("litellm-cache:key"), Ok(payload)),
-            MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)),
-        ])
-        .assert_all_commands_consumed();
-        let cache =
-            RedisCache::with_connection(connection, None, JsonCodec::<serde_json::Value>::new())
-                .with_namespace(Some("litellm-cache".into()));
-
-        cache
-            .set_cache("key", value.clone(), &ExactCacheContext::default())
-            .unwrap();
-        assert_eq!(
-            cache
-                .get_cache("key", &ExactCacheContext::default())
-                .unwrap(),
-            Some(value)
-        );
-        cache.delete_cache("key").unwrap();
-    }
-
-    #[test]
-    fn flush_scans_and_deletes_only_cache_keys() {
-        let connection = MockRedisConnection::new([
-            MockCmd::new(
-                redis::cmd("SCAN")
-                    .cursor_arg(0)
-                    .arg("MATCH")
-                    .arg("litellm-cache:*")
-                    .arg("COUNT")
-                    .arg(1000),
-                Ok(redis_test::redis_value!(["0", ["litellm-cache:key"]])),
-            ),
-            MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)),
-        ])
-        .assert_all_commands_consumed();
-        let cache =
-            RedisCache::with_connection(connection, None, JsonCodec::<serde_json::Value>::new())
-                .with_namespace(Some("litellm-cache".into()));
-
-        cache.flush_cache().unwrap();
-    }
-
-    #[tokio::test]
-    async fn test_connection_runs_ping_off_executor() {
-        let connection = MockRedisConnection::new([MockCmd::new(redis::cmd("PING"), Ok("PONG"))])
-            .assert_all_commands_consumed();
-        let cache =
-            RedisCache::with_connection(connection, None, JsonCodec::<serde_json::Value>::new())
-                .with_namespace(Some("litellm-cache".into()));
-
-        assert_eq!(
-            cache.test_connection().await.unwrap().status,
-            litellm_cache::CacheConnectionStatus::Success
-        );
-    }
+fn default_major_version() -> u32 {
+    std::env::var("DEFAULT_REDIS_MAJOR_VERSION")
+        .ok()
+        .and_then(|value| value.parse().ok())
+        .unwrap_or(7)
 }

litellm-rust/crates/cache-redis/src/cache/operations.rs

@@ -1,632 +0,0 @@
-use std::{sync::Arc, time::Duration};
-
-use litellm_cache::{
-    CacheCodec, CacheScript, ClientInfoCache, Error, IncrementOperation, QueueCache, ScanCache,
-    ScriptCache, SetCache, TtlCache,
-};
-use redis::Commands;
-
-use super::{ConnectionRef, Connections, RedisCache, namespaced_key};
-
-const INCREMENT_WITH_FLOOR_SCRIPT: &str = concat!(
-    "local count = redis.call('INCRBY', KEYS[1], ARGV[1]); ",
-    "if count < 0 then count = redis.call('INCRBY', KEYS[1], -count); end; ",
-    "if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ",
-    "return count"
-);
-const SET_MAX_SCRIPT: &str = concat!(
-    "local current = redis.call('GET', KEYS[1]); ",
-    "if current == false or tonumber(current) < tonumber(ARGV[1]) then ",
-    "redis.call('SET', KEYS[1], ARGV[1]); ",
-    "if tonumber(ARGV[2]) > 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ",
-    "return ARGV[1]; end; return current"
-);
-
-#[derive(Clone, Debug, PartialEq)]
-pub enum RedisArg {
-    Bytes(Vec<u8>),
-    Integer(i64),
-    Float(f64),
-}
-
-impl From<&str> for RedisArg {
-    fn from(value: &str) -> Self {
-        Self::Bytes(value.as_bytes().to_vec())
-    }
-}
-
-impl From<String> for RedisArg {
-    fn from(value: String) -> Self {
-        Self::Bytes(value.into_bytes())
-    }
-}
-
-impl From<Vec<u8>> for RedisArg {
-    fn from(value: Vec<u8>) -> Self {
-        Self::Bytes(value)
-    }
-}
-
-impl From<i64> for RedisArg {
-    fn from(value: i64) -> Self {
-        Self::Integer(value)
-    }
-}
-
-impl From<f64> for RedisArg {
-    fn from(value: f64) -> Self {
-        Self::Float(value)
-    }
-}
-
-impl redis::ToRedisArgs for RedisArg {
-    fn write_redis_args<W>(&self, out: &mut W)
-    where
-        W: ?Sized + redis::RedisWrite,
-    {
-        match self {
-            Self::Bytes(value) => value.write_redis_args(out),
-            Self::Integer(value) => value.write_redis_args(out),
-            Self::Float(value) => value.write_redis_args(out),
-        }
-    }
-}
-
-#[derive(Clone, Debug, PartialEq)]
-pub struct RedisRpushOperation {
-    pub key: String,
-    pub values: Vec<RedisArg>,
-}
-
-#[derive(Clone, Debug, PartialEq, Eq)]
-pub struct RedisLpopOperation {
-    pub key: String,
-    pub count: Option<usize>,
-}
-
-#[derive(Clone, Debug, PartialEq, Eq)]
-pub enum RedisLpopResult {
-    Missing,
-    Value(Vec<u8>),
-    Values(Vec<Vec<u8>>),
-}
-
-pub struct RedisScript<C> {
-    connections: Arc<Connections<C>>,
-    namespace: Option<String>,
-    source: String,
-}
-
-impl<C> CacheScript for RedisScript<C>
-where
-    C: redis::ConnectionLike + Send + 'static,
-{
-    type Argument = RedisArg;
-    type Output = redis::Value;
-
-    async fn invoke(
-        &self,
-        keys: Vec<String>,
-        arguments: Vec<Self::Argument>,
-    ) -> Result<Self::Output, Error> {
-        let keys = keys
-            .into_iter()
-            .map(|key| namespaced_key(self.namespace.as_deref(), &key))
-            .collect::<Vec<_>>();
-        let connections = Arc::clone(&self.connections);
-        let source = self.source.clone();
-        tokio::task::spawn_blocking(move || {
-            connections.execute(|connection| {
-                redis::cmd("EVAL")
-                    .arg(source)
-                    .arg(keys.len())
-                    .arg(keys)
-                    .arg(arguments)
-                    .query(connection)
-                    .map_err(|_| Error::Unavailable)
-            })
-        })
-        .await
-        .map_err(|_| Error::Unavailable)?
-    }
-}
-
-impl<S, C> RedisCache<S, C>
-where
-    S: CacheCodec,
-    C: redis::ConnectionLike + Send + 'static,
-{
-    pub async fn delete_cache_keys(&self, keys: Vec<String>) -> Result<usize, Error> {
-        if keys.is_empty() {
-            return Ok(0);
-        }
-        let keys = keys
-            .into_iter()
-            .map(|key| self.namespaced_key(&key))
-            .collect::<Vec<_>>();
-        Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
-            connection.del(keys).map_err(|_| Error::Unavailable)
-        })
-        .await
-    }
-
-    pub fn batch_get_counts(&self, keys: &[String]) -> Result<Vec<Option<i64>>, Error> {
-        let keys = keys
-            .iter()
-            .map(|key| self.namespaced_key(key))
-            .collect::<Vec<_>>();
-        let values = self.connections.execute(|connection| {
-            redis::cmd("MGET")
-                .arg(keys)
-                .query::<Vec<redis::Value>>(connection)
-                .map_err(|_| Error::Unavailable)
-        })?;
-        values.into_iter().map(count).collect()
-    }
-
-    pub async fn async_batch_get_counts(
-        &self,
-        keys: Vec<String>,
-    ) -> Result<Vec<Option<i64>>, Error> {
-        let keys = keys
-            .iter()
-            .map(|key| self.namespaced_key(key))
-            .collect::<Vec<_>>();
-        let values = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
-            redis::cmd("MGET")
-                .arg(keys)
-                .query::<Vec<redis::Value>>(connection)
-                .map_err(|_| Error::Unavailable)
-        })
-        .await?;
-        values.into_iter().map(count).collect()
-    }
-
-    pub fn sync_ping(&self) -> Result<bool, Error> {
-        self.connections
-            .execute(|connection| connection.ping().map_err(|_| Error::Unavailable))
-    }
-
-    pub async fn ping(&self) -> Result<bool, Error> {
-        Connections::run_blocking(Arc::clone(&self.connections), |connection| {
-            connection.ping().map_err(|_| Error::Unavailable)
-        })
-        .await
-    }
-
-    pub async fn async_get_ttl(&self, key: &str) -> Result<Option<i64>, Error> {
-        let key = self.namespaced_key(key);
-        let ttl = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
-            redis::cmd("TTL")
-                .arg(key)
-                .query::<i64>(connection)
-                .map_err(|_| Error::Unavailable)
-        })
-        .await?;
-        Ok((ttl >= 0).then_some(ttl))
-    }
-
-    pub async fn async_scan_iter(&self, pattern: &str, count: usize) -> Result<Vec<String>, Error> {
-        let pattern = format!("{}*", self.namespaced_key(pattern));
-        Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
-            let mut matches = Vec::new();
-            connection.scan(&pattern, count, |_, keys| {
-                matches.extend(keys);
-                Ok(matches.len() < count)
-            })?;
-            matches.truncate(count);
-            Ok(matches)
-        })
-        .await
-    }
-
-    pub async fn async_set_cache_sadd(
-        &self,
-        key: &str,
-        values: Vec<RedisArg>,
-        ttl: Option<Duration>,
-    ) -> Result<usize, Error> {
-        if values.is_empty() {
-            return Err(Error::InvalidEntry);
-        }
-        let key = self.namespaced_key(key);
-        let ttl = Self::ttl_seconds(ttl.unwrap_or(self.default_ttl));
-        Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
-            let mut sadd = redis::cmd("SADD");
-            sadd.arg(&key).arg(values);
-            let mut expire = redis::cmd("EXPIRE");
-            expire.arg(&key).arg(ttl);
-            let replies = connection.pipeline(vec![sadd, expire])?;
-            replies
-                .into_iter()
-                .next()
-                .map(redis::from_redis_value::<usize>)
-                .transpose()
-                .map_err(|_| Error::Unavailable)?
-                .ok_or(Error::Unavailable)
-        })
-        .await
-    }
-
-    pub async fn async_rpush(&self, key: &str, values: Vec<RedisArg>) -> Result<usize, Error> {
-        if values.is_empty() {
-            return Err(Error::InvalidEntry);
-        }
-        let key = self.namespaced_key(key);
-        Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
-            redis::cmd("RPUSH")
-                .arg(key)
-                .arg(values)
-                .query(connection)
-                .map_err(|_| Error::Unavailable)
-        })
-        .await
-    }
-
-    pub async fn async_rpush_pipeline(
-        &self,
-        operations: Vec<RedisRpushOperation>,
-    ) -> Result<Vec<usize>, Error> {
-        let operations = operations
-            .into_iter()
-            .map(|operation| {
-                if operation.values.is_empty() {
-                    return Err(Error::InvalidEntry);
-                }
-                Ok((self.namespaced_key(&operation.key), operation.values))
-            })
-            .collect::<Result<Vec<_>, _>>()?;
-        if operations.is_empty() {
-            return Ok(Vec::new());
-        }
-        Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
-            let commands = operations
-                .into_iter()
-                .map(|(key, values)| {
-                    let mut command = redis::cmd("RPUSH");
-                    command.arg(key).arg(values);
-                    command
-                })
-                .collect();
-            connection
-                .pipeline(commands)?
-                .into_iter()
-                .map(|value| redis::from_redis_value(value).map_err(|_| Error::Unavailable))
-                .collect()
-        })
-        .await
-    }
-
-    pub async fn async_lpop(
-        &self,
-        key: &str,
-        count: Option<usize>,
-    ) -> Result<RedisLpopResult, Error> {
-        let key = self.namespaced_key(key);
-        let multiple = count.is_some();
-        let value = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
-            let mut command = redis::cmd("LPOP");
-            command.arg(key);
-            if let Some(count) = count {
-                command.arg(count);
-            }
-            command
-                .query::<redis::Value>(connection)
-                .map_err(|_| Error::Unavailable)
-        })
-        .await?;
-        lpop_result(value, multiple)
-    }
-
-    pub async fn async_lpop_pipeline(
-        &self,
-        operations: Vec<RedisLpopOperation>,
-    ) -> Result<Vec<RedisLpopResult>, Error> {
-        let operations = operations
-            .into_iter()
-            .map(|operation| (self.namespaced_key(&operation.key), operation.count))
-            .collect::<Vec<_>>();
-        if operations.is_empty() {
-            return Ok(Vec::new());
-        }
-        let multiple = operations
-            .iter()
-            .map(|(_, count)| count.is_some())
-            .collect::<Vec<_>>();
-        let values = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
-            let commands = operations
-                .into_iter()
-                .map(|(key, count)| {
-                    let mut command = redis::cmd("LPOP");
-                    command.arg(key);
-                    if let Some(count) = count {
-                        command.arg(count);
-                    }
-                    command
-                })
-                .collect();
-            connection.pipeline(commands)
-        })
-        .await?;
-        values
-            .into_iter()
-            .zip(multiple)
-            .map(|(value, multiple)| lpop_result(value, multiple))
-            .collect()
-    }
-
-    pub async fn async_eval(
-        &self,
-        script: String,
-        keys: Vec<String>,
-        arguments: Vec<RedisArg>,
-    ) -> Result<redis::Value, Error> {
-        let keys = keys
-            .into_iter()
-            .map(|key| self.namespaced_key(&key))
-            .collect::<Vec<_>>();
-        Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
-            redis::cmd("EVAL")
-                .arg(script)
-                .arg(keys.len())
-                .arg(keys)
-                .arg(arguments)
-                .query(connection)
-                .map_err(|_| Error::Unavailable)
-        })
-        .await
-    }
-
-    pub fn client_list(&self) -> Result<String, Error> {
-        self.connections
-            .execute(|connection| connection.node_text(redis::cmd("CLIENT").arg("LIST")))
-    }
-
-    pub fn info(&self) -> Result<String, Error> {
-        self.connections
-            .execute(|connection| connection.node_text(&redis::cmd("INFO")))
-    }
-
-    pub fn flushall(&self) -> Result<(), Error> {
-        self.connections.execute(|connection| connection.flushall())
-    }
-}
-
-impl<S, C> RedisCache<S, C>
-where
-    S: CacheCodec<Value = f64>,
-    C: redis::ConnectionLike + Send + 'static,
-{
-    pub fn increment_with_floor(
-        &self,
-        key: &str,
-        amount: i64,
-        ttl: Duration,
-    ) -> Result<i64, Error> {
-        let key = self.namespaced_key(key);
-        let ttl = Self::ttl_seconds(ttl);
-        self.connections
-            .execute(|connection| increment_with_floor(connection, key, amount, ttl))
-    }
-
-    pub async fn async_increment_pipeline(
-        &self,
-        operations: Vec<IncrementOperation>,
-    ) -> Result<Vec<f64>, Error> {
-        let operations = operations
-            .into_iter()
-            .map(|operation| {
-                (
-                    self.namespaced_key(&operation.key),
-                    operation.amount,
-                    operation.ttl.map(Self::ttl_seconds),
-                )
-            })
-            .collect::<Vec<_>>();
-        if operations.is_empty() {
-            return Ok(Vec::new());
-        }
-        Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
-            let mut commands = Vec::with_capacity(operations.len() * 2);
-            let mut increments = Vec::with_capacity(operations.len());
-            for (key, amount, ttl) in operations {
-                let mut increment = redis::cmd("INCRBYFLOAT");
-                increment.arg(&key).arg(amount);
-                increments.push(commands.len());
-                commands.push(increment);
-                if let Some(ttl) = ttl {
-                    let mut expire = redis::cmd("EXPIRE");
-                    expire.arg(key).arg(ttl);
-                    commands.push(expire);
-                }
-            }
-            let mut replies = connection.pipeline(commands)?;
-            increments
-                .into_iter()
-                .map(|index| {
-                    redis::from_redis_value(std::mem::take(&mut replies[index]))
-                        .map_err(|_| Error::Unavailable)
-                })
-                .collect()
-        })
-        .await
-    }
-
-    pub async fn async_increment_with_floor(
-        &self,
-        key: &str,
-        amount: i64,
-        ttl: Duration,
-    ) -> Result<i64, Error> {
-        let key = self.namespaced_key(key);
-        let ttl = Self::ttl_seconds(ttl);
-        Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
-            increment_with_floor(connection, key, amount, ttl)
-        })
-        .await
-    }
-
-    pub async fn async_set_max(
-        &self,
-        key: &str,
-        value: f64,
-        ttl: Option<Duration>,
-    ) -> Result<f64, Error> {
-        let key = self.namespaced_key(key);
-        let ttl = Self::ttl_seconds(ttl.unwrap_or(self.default_ttl));
-        Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
-            redis::cmd("EVAL")
-                .arg(SET_MAX_SCRIPT)
-                .arg(1)
-                .arg(key)
-                .arg(value)
-                .arg(ttl)
-                .query(connection)
-                .map_err(|_| Error::Unavailable)
-        })
-        .await
-    }
-}
-
-fn redis_bytes(value: redis::Value) -> Result<Vec<u8>, Error> {
-    match value {
-        redis::Value::BulkString(bytes) => Ok(bytes),
-        redis::Value::SimpleString(text) => Ok(text.into_bytes()),
-        _ => Err(Error::InvalidEntry),
-    }
-}
-
-fn lpop_result(value: redis::Value, multiple: bool) -> Result<RedisLpopResult, Error> {
-    match value {
-        redis::Value::Nil => Ok(RedisLpopResult::Missing),
-        redis::Value::Array(values) if multiple => values
-            .into_iter()
-            .map(redis_bytes)
-            .collect::<Result<Vec<_>, _>>()
-            .map(RedisLpopResult::Values),
-        value if !multiple => redis_bytes(value).map(RedisLpopResult::Value),
-        _ => Err(Error::InvalidEntry),
-    }
-}
-
-fn count(value: redis::Value) -> Result<Option<i64>, Error> {
-    match value {
-        redis::Value::Nil => Ok(None),
-        redis::Value::Int(value) => Ok(Some(value)),
-        redis::Value::BulkString(value) => std::str::from_utf8(&value)
-            .ok()
-            .and_then(|value| value.parse().ok())
-            .map(Some)
-            .ok_or(Error::InvalidEntry),
-        redis::Value::SimpleString(value) => {
-            value.parse().map(Some).map_err(|_| Error::InvalidEntry)
-        }
-        _ => Err(Error::InvalidEntry),
-    }
-}
-
-fn increment_with_floor(
-    connection: &mut ConnectionRef<'_>,
-    key: String,
-    amount: i64,
-    ttl: u64,
-) -> Result<i64, Error> {
-    redis::cmd("EVAL")
-        .arg(INCREMENT_WITH_FLOOR_SCRIPT)
-        .arg(1)
-        .arg(key)
-        .arg(amount)
-        .arg(ttl)
-        .query(connection)
-        .map_err(|_| Error::Unavailable)
-}
-
-impl<S, C> TtlCache for RedisCache<S, C>
-where
-    S: CacheCodec,
-    C: redis::ConnectionLike + Send + 'static,
-{
-    async fn async_get_ttl(&self, key: &str) -> Result<Option<Duration>, Error> {
-        RedisCache::async_get_ttl(self, key)
-            .await
-            .map(|ttl| ttl.map(|seconds| Duration::from_secs(seconds as u64)))
-    }
-}
-
-impl<S, C> ScanCache for RedisCache<S, C>
-where
-    S: CacheCodec,
-    C: redis::ConnectionLike + Send + 'static,
-{
-    async fn async_scan_iter(&self, pattern: &str, count: usize) -> Result<Vec<String>, Error> {
-        RedisCache::async_scan_iter(self, pattern, count).await
-    }
-}
-
-impl<S, C> ClientInfoCache for RedisCache<S, C>
-where
-    S: CacheCodec,
-    C: redis::ConnectionLike + Send + 'static,
-{
-    type ClientList = String;
-    type Info = String;
-
-    fn client_list(&self) -> Result<Self::ClientList, Error> {
-        RedisCache::client_list(self)
-    }
-
-    fn info(&self) -> Result<Self::Info, Error> {
-        RedisCache::info(self)
-    }
-}
-
-impl<S, C> SetCache for RedisCache<S, C>
-where
-    S: CacheCodec,
-    C: redis::ConnectionLike + Send + 'static,
-{
-    type SetValue = RedisArg;
-    type SetResult = usize;
-
-    async fn async_set_cache_sadd(
-        &self,
-        key: &str,
-        values: Vec<Self::SetValue>,
-        ttl: Option<Duration>,
-    ) -> Result<Self::SetResult, Error> {
-        RedisCache::async_set_cache_sadd(self, key, values, ttl).await
-    }
-}
-
-impl<S, C> QueueCache for RedisCache<S, C>
-where
-    S: CacheCodec,
-    C: redis::ConnectionLike + Send + 'static,
-{
-    type QueueValue = RedisArg;
-    type PopResult = RedisLpopResult;
-
-    async fn async_rpush(&self, key: &str, values: Vec<Self::QueueValue>) -> Result<usize, Error> {
-        RedisCache::async_rpush(self, key, values).await
-    }
-
-    async fn async_lpop(&self, key: &str, count: Option<usize>) -> Result<Self::PopResult, Error> {
-        RedisCache::async_lpop(self, key, count).await
-    }
-}
-
-impl<S, C> ScriptCache for RedisCache<S, C>
-where
-    S: CacheCodec,
-    C: redis::ConnectionLike + Send + 'static,
-{
-    type Script = RedisScript<C>;
-
-    fn async_register_script(&self, source: String) -> Self::Script {
-        RedisScript {
-            connections: Arc::clone(&self.connections),
-            namespace: self.namespace.clone(),
-            source,
-        }
-    }
-}

litellm-rust/crates/cache-redis/src/claim.rs

@@ -0,0 +1,105 @@
+use litellm_cache::{CacheCodec, ClaimCache, Error, ExactCacheContext};
+use redis::Commands;
+
+use crate::{cache::RedisCache, connection::ConnectionRef};
+
+const CLAIM_SCRIPT: &str = concat!(
+    "local current = redis.call('GET', KEYS[1]); ",
+    "if ARGV[1] == '' then if current ~= false and current ~= '' then return 0; end; ",
+    "elseif current ~= ARGV[1] then return 0; end; ",
+    "if ARGV[3] ~= '' then redis.call('SET', KEYS[1], ARGV[3], 'EX', ARGV[2]); ",
+    "elseif ARGV[4] == '1' then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return 1"
+);
+const CLAIM_ATTEMPTS: usize = 8;
+
+fn stored_bytes(value: redis::Value) -> Result<Option<Vec<u8>>, Error> {
+    match value {
+        redis::Value::Nil => Ok(None),
+        redis::Value::BulkString(bytes) => Ok(Some(bytes)),
+        redis::Value::SimpleString(text) => Ok(Some(text.into_bytes())),
+        _ => Err(Error::InvalidEntry),
+    }
+}
+
+/// Eligibility is decided on decoded values, so a pin written by another encoder (Python's
+/// `json.dumps` spacing or key order) still matches. The write is a compare-and-set on the
+/// bytes that decision was made on, retried when another claimant wins the race.
+fn claim<S: CacheCodec>(
+    connection: &mut ConnectionRef<'_>,
+    codec: &S,
+    key: &str,
+    candidate: S::Value,
+    eligible: &[S::Value],
+    ttl: u64,
+) -> Result<S::Value, Error>
+where
+    S::Value: PartialEq,
+{
+    let payload = codec.encode(&candidate)?;
+    if payload.is_empty() {
+        return Err(Error::InvalidEntry);
+    }
+    for _ in 0..CLAIM_ATTEMPTS {
+        let current = stored_bytes(
+            connection
+                .get::<_, redis::Value>(key)
+                .map_err(|_| Error::Unavailable)?,
+        )?
+        .filter(|bytes| !bytes.is_empty());
+        let existing = current
+            .as_deref()
+            .and_then(|bytes| codec.decode(bytes).ok())
+            .filter(|existing| eligible.is_empty() || eligible.contains(existing));
+        let refresh = existing
+            .as_ref()
+            .is_some_and(|existing| !eligible.is_empty() || *existing == candidate);
+        let write: &[u8] = if existing.is_some() { b"" } else { &payload };
+        let applied = redis::cmd("EVAL")
+            .arg(CLAIM_SCRIPT)
+            .arg(1)
+            .arg(key)
+            .arg(current.as_deref().unwrap_or_default())
+            .arg(ttl)
+            .arg(write)
+            .arg(u8::from(refresh))
+            .query::<bool>(connection)
+            .map_err(|_| Error::Unavailable)?;
+        if applied {
+            return Ok(existing.unwrap_or(candidate));
+        }
+    }
+    Err(Error::Unavailable)
+}
+
+impl<S, C> ClaimCache for RedisCache<S, C>
+where
+    S: CacheCodec + Clone + 'static,
+    S::Value: PartialEq,
+    C: redis::ConnectionLike + Send + 'static,
+{
+    fn claim_cache(
+        &self,
+        key: &str,
+        candidate: S::Value,
+        eligible: &[S::Value],
+        context: ExactCacheContext,
+    ) -> Result<S::Value, Error> {
+        let key = self.namespaced_key(key);
+        let ttl = self.ttl_or_default(context.ttl);
+        self.execute(|connection| claim(connection, &self.codec, &key, candidate, eligible, ttl))
+    }
+
+    async fn async_claim_cache(
+        &self,
+        key: &str,
+        candidate: S::Value,
+        eligible: Vec<S::Value>,
+        context: ExactCacheContext,
+    ) -> Result<S::Value, Error> {
+        let key = self.namespaced_key(key);
+        let ttl = self.ttl_or_default(context.ttl);
+        let codec = self.codec.clone();
+        self.run(move |connection| claim(connection, &codec, &key, candidate, &eligible, ttl))
+            .await
+    }
+}

litellm-rust/crates/cache-redis/src/connection.rs

@@ -1,29 +1,126 @@
-use std::collections::HashMap;
+use std::{
+    sync::{Arc, Mutex},
+    time::Duration,
+};
 
 use litellm_cache::Error;
 use redis::{
-    ConnectionAddr, ConnectionInfo, ConnectionLike, IntoConnectionInfo,
-    cluster::{ClusterClient, ClusterClientBuilder, ClusterConnection, NodeAddress},
+    ConnectionAddr, ConnectionInfo, IntoConnectionInfo,
+    cluster::{
+        ClusterClient, ClusterClientBuilder, ClusterConnection, ClusterPipeline, NodeAddress,
+    },
     cluster_routing::{
-        MultipleNodeRoutingInfo, ResponsePolicy, RoutingInfo, SingleNodeRoutingInfo, Slot,
+        MultipleNodeRoutingInfo, ResponsePolicy, RoutingInfo, SingleNodeRoutingInfo,
     },
 };
 
-use super::REDIS_TIMEOUT;
-use crate::topology::RedisNode;
+use crate::topology::{RedisNode, RedisTopology};
 
-pub struct PooledConnection<C> {
-    pub(super) connection: C,
-    pub(super) failed: bool,
+pub(crate) const REDIS_TIMEOUT: Duration = Duration::from_secs(5);
+const REDIS_POOL_SIZE: u32 = 16;
+
+#[allow(private_interfaces)]
+pub enum Connections<C> {
+    Pool(r2d2::Pool<ConnectionManager>),
+    Cluster(r2d2::Pool<ClusterConnectionManager>),
+    Fixed(Mutex<C>),
+}
+
+impl<C> Connections<C>
+where
+    C: redis::ConnectionLike + Send + 'static,
+{
+    pub fn execute<T>(
+        &self,
+        operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error>,
+    ) -> Result<T, Error> {
+        match self {
+            Self::Pool(pool) => {
+                let mut pooled = pool.get().map_err(|_| Error::Unavailable)?;
+                let result = operation(&mut ConnectionRef::Node(&mut pooled.connection));
+                pooled.failed = matches!(result, Err(Error::Unavailable));
+                result
+            }
+            Self::Cluster(pool) => {
+                let mut pooled = pool.get().map_err(|_| Error::Unavailable)?;
+                let result = operation(&mut ConnectionRef::Cluster(&mut pooled.connection));
+                pooled.failed = matches!(result, Err(Error::Unavailable));
+                result
+            }
+            Self::Fixed(connection) => {
+                let mut connection = connection.lock().map_err(|_| Error::Unavailable)?;
+                operation(&mut ConnectionRef::Node(&mut *connection))
+            }
+        }
+    }
+
+    pub async fn run_blocking<T, F>(connections: Arc<Self>, operation: F) -> Result<T, Error>
+    where
+        T: Send + 'static,
+        F: FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error> + Send + 'static,
+    {
+        tokio::task::spawn_blocking(move || connections.execute(operation))
+            .await
+            .map_err(|_| Error::Unavailable)?
+    }
+
+    pub fn fixed(connection: C) -> Self {
+        Self::Fixed(Mutex::new(connection))
+    }
+
+    pub fn open(url: &str, topology: &RedisTopology) -> Result<Self, Error> {
+        match topology {
+            RedisTopology::Standalone => Ok(Self::Pool(pool(ConnectionManager::open(url)?)?)),
+            RedisTopology::Cluster { startup_nodes } => Ok(Self::Cluster(pool(
+                ClusterConnectionManager::open(url, startup_nodes)?,
+            )?)),
+        }
+    }
+
+    /// Closes every idle pooled connection; the next operation opens a fresh one. Connections
+    /// checked out right now return to the pool, and a caller-owned connection stays open.
+    pub fn disconnect(&self) {
+        match self {
+            Self::Pool(pool) => close_idle(pool),
+            Self::Cluster(pool) => close_idle(pool),
+            Self::Fixed(_) => {}
+        }
+    }
+}
+
+fn pool<M: r2d2::ManageConnection>(manager: M) -> Result<r2d2::Pool<M>, Error> {
+    r2d2::Pool::builder()
+        .max_size(REDIS_POOL_SIZE)
+        .min_idle(Some(0))
+        .connection_timeout(REDIS_TIMEOUT)
+        .test_on_check_out(false)
+        .build(manager)
+        .map_err(|_| Error::Unavailable)
+}
+
+fn close_idle<M, T>(pool: &r2d2::Pool<M>)
+where
+    M: r2d2::ManageConnection<Connection = PooledConnection<T>>,
+{
+    let mut idle = Vec::new();
+    while let Some(mut connection) = pool.try_get() {
+        connection.failed = true;
+        idle.push(connection);
+    }
+}
+
+pub(crate) struct PooledConnection<C> {
+    connection: C,
+    failed: bool,
 }
 
 /// Pools connections without a checkout PING, which would double every operation's round trips.
 /// A timed-out command leaves its reply on the socket while redis still reports the connection
 /// open, so any connection whose operation failed is discarded instead of being reused.
-pub struct ConnectionManager(redis::Client);
+pub(crate) struct ConnectionManager(redis::Client);
 
 impl ConnectionManager {
-    pub(super) fn open(url: &str) -> Result<Self, Error> {
+    fn open(url: &str) -> Result<Self, Error> {
         redis::Client::open(url)
             .map(Self)
             .map_err(|_| Error::Unavailable)
@@ -54,10 +151,10 @@ impl r2d2::ManageConnection for ConnectionManager {
     }
 }
 
-pub struct ClusterConnectionManager(ClusterClient);
+pub(crate) struct ClusterConnectionManager(ClusterClient);
 
 impl ClusterConnectionManager {
-    pub(super) fn open(url: &str, startup_nodes: &[RedisNode]) -> Result<Self, Error> {
+    fn open(url: &str, startup_nodes: &[RedisNode]) -> Result<Self, Error> {
         if startup_nodes.is_empty() {
             return Err(Error::Unavailable);
         }
@@ -172,43 +269,36 @@ impl redis::ConnectionLike for ConnectionRef<'_> {
 }
 
 impl ConnectionRef<'_> {
-    pub(crate) fn pipeline(
+    /// Runs `pipeline` and decodes its non-ignored replies as `T`. A cluster connection refuses
+    /// `Pipeline::query`, so there a transaction goes to its keys' slot as one MULTI/EXEC and
+    /// anything else is split per node by `ClusterPipeline`; either way the raw replies are
+    /// handed back to `pipeline` to decode.
+    pub(crate) fn query_pipeline<T: redis::FromRedisValue>(
         &mut self,
-        commands: Vec<redis::Cmd>,
-    ) -> Result<Vec<redis::Value>, Error> {
+        pipeline: &redis::Pipeline,
+    ) -> Result<T, Error> {
         match self {
-            Self::Node(connection) => {
-                let mut pipeline = redis::pipe();
-                for command in &commands {
-                    pipeline.add_command(command.clone());
-                }
-                pipeline
-                    .query::<Vec<redis::Value>>(*connection)
-                    .map_err(|_| Error::Unavailable)
+            Self::Node(connection) => pipeline.query(*connection),
+            Self::Cluster(connection) if pipeline.is_transaction() => {
+                redis::ConnectionLike::req_packed_commands(
+                    *connection,
+                    &pipeline.get_packed_pipeline(),
+                    pipeline.len() + 1,
+                    1,
+                )
+                .and_then(|replies| pipeline.query(&mut Replies(Some(replies))))
             }
             Self::Cluster(connection) => {
-                let mut replies: Vec<Option<redis::Value>> = vec![None; commands.len()];
-                for indices in slot_groups(&commands).into_values() {
-                    let mut pipeline = redis::pipe();
-                    for index in &indices {
-                        pipeline.add_command(commands[*index].clone());
-                    }
-                    let values = connection
-                        .req_packed_commands(&pipeline.get_packed_pipeline(), 0, indices.len())
-                        .map_err(|_| Error::Unavailable)?;
-                    if values.len() != indices.len() {
-                        return Err(Error::Unavailable);
-                    }
-                    for (index, value) in indices.into_iter().zip(values) {
-                        replies[index] = Some(value);
-                    }
+                let mut cluster = ClusterPipeline::with_capacity(pipeline.len());
+                for command in pipeline.cmd_iter() {
+                    cluster.add_command(command.clone());
                 }
-                replies
-                    .into_iter()
-                    .collect::<Option<Vec<_>>>()
-                    .ok_or(Error::Unavailable)
+                cluster
+                    .query(connection)
+                    .and_then(|replies| pipeline.query(&mut Replies(Some(replies))))
             }
         }
+        .map_err(|_| Error::Unavailable)
     }
 
     pub(crate) fn scan(
@@ -379,14 +469,35 @@ fn scan_command(cursor: u64, pattern: &str, count: usize) -> redis::Cmd {
     command
 }
 
-fn slot_groups(commands: &[redis::Cmd]) -> HashMap<Slot, Vec<usize>> {
-    let mut groups: HashMap<Slot, Vec<usize>> = HashMap::new();
-    for (index, command) in commands.iter().enumerate() {
-        let key = match command.args_iter().nth(1) {
-            Some(redis::Arg::Simple(key)) => key,
-            _ => b"",
-        };
-        groups.entry(Slot::for_key(key)).or_default().push(index);
+/// Hands already received pipeline replies to `Pipeline::query`, so it applies its own
+/// ignore and error handling to replies a cluster pipeline gathered from several nodes.
+struct Replies(Option<Vec<redis::Value>>);
+
+impl redis::ConnectionLike for Replies {
+    fn req_packed_command(&mut self, _: &[u8]) -> redis::RedisResult<redis::Value> {
+        Err((redis::ErrorKind::Client, "replies hold a pipeline only").into())
+    }
+
+    fn req_packed_commands(
+        &mut self,
+        _: &[u8],
+        _: usize,
+        _: usize,
+    ) -> redis::RedisResult<Vec<redis::Value>> {
+        self.0
+            .take()
+            .ok_or_else(|| (redis::ErrorKind::Client, "replies were already read").into())
+    }
+
+    fn get_db(&self) -> i64 {
+        0
+    }
+
+    fn check_connection(&mut self) -> bool {
+        true
+    }
+
+    fn is_open(&self) -> bool {
+        true
     }
-    groups
 }

litellm-rust/crates/cache-redis/src/counter.rs

@@ -0,0 +1,205 @@
+use std::time::Duration;
+
+use litellm_cache::{
+    BoundedCounterCache, CacheCodec, CountReadCache, CounterCache, Error, ExactCacheContext,
+    IncrementOperation,
+};
+
+use crate::{
+    cache::{RedisCache, ttl_seconds},
+    connection::ConnectionRef,
+    store::mget,
+};
+
+const INCREMENT_SCRIPT: &str = concat!(
+    "local value = redis.call('INCRBYFLOAT', KEYS[1], ARGV[1]); ",
+    "if redis.call('TTL', KEYS[1]) == -1 then ",
+    "redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return value"
+);
+const INCREMENT_WITH_FLOOR_SCRIPT: &str = concat!(
+    "local count = redis.call('INCRBY', KEYS[1], ARGV[1]); ",
+    "if count < 0 then count = redis.call('INCRBY', KEYS[1], -count); end; ",
+    "if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ",
+    "return count"
+);
+const SET_MAX_SCRIPT: &str = concat!(
+    "local current = redis.call('GET', KEYS[1]); ",
+    "if current == false or tonumber(current) < tonumber(ARGV[1]) then ",
+    "redis.call('SET', KEYS[1], ARGV[1]); ",
+    "if tonumber(ARGV[2]) > 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ",
+    "return ARGV[1]; end; return current"
+);
+
+impl<S, C> CounterCache for RedisCache<S, C>
+where
+    S: CacheCodec,
+    C: redis::ConnectionLike + Send + 'static,
+{
+    fn increment_cache(
+        &self,
+        key: &str,
+        amount: f64,
+        context: ExactCacheContext,
+    ) -> Result<f64, Error> {
+        let key = self.namespaced_key(key);
+        let ttl = self.ttl_or_default(context.ttl);
+        self.execute(|connection| increment(connection, key, amount, ttl, false))
+    }
+
+    /// Python `_incrbyfloat_with_ttl`: without `refresh_ttl` the TTL is set only on a key that
+    /// has none, in one atomic script; with it, every increment re-arms the TTL.
+    async fn async_increment(
+        &self,
+        key: &str,
+        amount: f64,
+        context: ExactCacheContext,
+        refresh_ttl: bool,
+    ) -> Result<f64, Error> {
+        let key = self.namespaced_key(key);
+        let ttl = self.ttl_or_default(context.ttl);
+        self.run(move |connection| increment(connection, key, amount, ttl, refresh_ttl))
+            .await
+    }
+
+    async fn async_increment_pipeline(
+        &self,
+        operations: Vec<IncrementOperation>,
+    ) -> Result<Vec<f64>, Error> {
+        if operations.is_empty() {
+            return Ok(Vec::new());
+        }
+        let mut pipeline = redis::pipe();
+        for operation in operations {
+            let key = self.namespaced_key(&operation.key);
+            pipeline.cmd("INCRBYFLOAT").arg(&key).arg(operation.amount);
+            if let Some(ttl) = operation.ttl {
+                pipeline
+                    .cmd("EXPIRE")
+                    .arg(key)
+                    .arg(ttl_seconds(ttl))
+                    .ignore();
+            }
+        }
+        self.run(move |connection| connection.query_pipeline(&pipeline))
+            .await
+    }
+}
+
+fn increment(
+    connection: &mut ConnectionRef<'_>,
+    key: String,
+    amount: f64,
+    ttl: u64,
+    refresh_ttl: bool,
+) -> Result<f64, Error> {
+    if !refresh_ttl {
+        return redis::cmd("EVAL")
+            .arg(INCREMENT_SCRIPT)
+            .arg(1)
+            .arg(key)
+            .arg(amount)
+            .arg(ttl)
+            .query(connection)
+            .map_err(|_| Error::Unavailable);
+    }
+    connection
+        .query_pipeline(
+            redis::pipe()
+                .cmd("INCRBYFLOAT")
+                .arg(&key)
+                .arg(amount)
+                .cmd("EXPIRE")
+                .arg(&key)
+                .arg(ttl)
+                .ignore(),
+        )
+        .map(|(value,)| value)
+}
+
+impl<S, C> CountReadCache for RedisCache<S, C>
+where
+    S: CacheCodec,
+    C: redis::ConnectionLike + Send + 'static,
+{
+    fn batch_get_counts(&self, keys: &[String]) -> Result<Vec<Option<i64>>, Error> {
+        let keys = self.namespaced_keys(keys);
+        self.execute(|connection| mget(connection, keys))?
+            .into_iter()
+            .map(count)
+            .collect()
+    }
+
+    async fn async_batch_get_counts(&self, keys: Vec<String>) -> Result<Vec<Option<i64>>, Error> {
+        let keys = self.namespaced_keys(&keys);
+        self.run(move |connection| mget(connection, keys))
+            .await?
+            .into_iter()
+            .map(count)
+            .collect()
+    }
+}
+
+fn count(value: redis::Value) -> Result<Option<i64>, Error> {
+    redis::from_redis_value(value).map_err(|_| Error::InvalidEntry)
+}
+
+impl<S, C> BoundedCounterCache for RedisCache<S, C>
+where
+    S: CacheCodec,
+    C: redis::ConnectionLike + Send + 'static,
+{
+    fn increment_with_floor(&self, key: &str, amount: i64, ttl: Duration) -> Result<i64, Error> {
+        let key = self.namespaced_key(key);
+        let ttl = ttl_seconds(ttl);
+        self.execute(|connection| increment_with_floor(connection, key, amount, ttl))
+    }
+
+    async fn async_increment_with_floor(
+        &self,
+        key: &str,
+        amount: i64,
+        ttl: Duration,
+    ) -> Result<i64, Error> {
+        let key = self.namespaced_key(key);
+        let ttl = ttl_seconds(ttl);
+        self.run(move |connection| increment_with_floor(connection, key, amount, ttl))
+            .await
+    }
+
+    async fn async_set_max(
+        &self,
+        key: &str,
+        value: f64,
+        ttl: Option<Duration>,
+    ) -> Result<f64, Error> {
+        let key = self.namespaced_key(key);
+        let ttl = self.ttl_or_default(ttl);
+        self.run(move |connection| {
+            redis::cmd("EVAL")
+                .arg(SET_MAX_SCRIPT)
+                .arg(1)
+                .arg(key)
+                .arg(value)
+                .arg(ttl)
+                .query(connection)
+                .map_err(|_| Error::Unavailable)
+        })
+        .await
+    }
+}
+
+fn increment_with_floor(
+    connection: &mut ConnectionRef<'_>,
+    key: String,
+    amount: i64,
+    ttl: u64,
+) -> Result<i64, Error> {
+    redis::cmd("EVAL")
+        .arg(INCREMENT_WITH_FLOOR_SCRIPT)
+        .arg(1)
+        .arg(key)
+        .arg(amount)
+        .arg(ttl)
+        .query(connection)
+        .map_err(|_| Error::Unavailable)
+}

litellm-rust/crates/cache-redis/src/keys.rs

@@ -0,0 +1,63 @@
+use std::time::Duration;
+
+use litellm_cache::{CacheCodec, Error, RefreshTtlCache, ScanCache, TtlCache};
+
+use crate::cache::RedisCache;
+
+impl<S, C> TtlCache for RedisCache<S, C>
+where
+    S: CacheCodec,
+    C: redis::ConnectionLike + Send + 'static,
+{
+    async fn async_get_ttl(&self, key: &str) -> Result<Option<Duration>, Error> {
+        let key = self.namespaced_key(key);
+        let ttl = self
+            .run(move |connection| {
+                redis::cmd("TTL")
+                    .arg(key)
+                    .query::<i64>(connection)
+                    .map_err(|_| Error::Unavailable)
+            })
+            .await?;
+        Ok(u64::try_from(ttl).ok().map(Duration::from_secs))
+    }
+}
+
+impl<S, C> RefreshTtlCache for RedisCache<S, C>
+where
+    S: CacheCodec,
+    C: redis::ConnectionLike + Send + 'static,
+{
+    async fn async_refresh_ttl(&self, key: &str, ttl: Option<Duration>) -> Result<bool, Error> {
+        let key = self.namespaced_key(key);
+        let ttl = self.ttl_or_default(ttl);
+        self.run(move |connection| {
+            redis::cmd("EXPIRE")
+                .arg(key)
+                .arg(ttl)
+                .query(connection)
+                .map_err(|_| Error::Unavailable)
+        })
+        .await
+    }
+}
+
+impl<S, C> ScanCache for RedisCache<S, C>
+where
+    S: CacheCodec,
+    C: redis::ConnectionLike + Send + 'static,
+{
+    async fn async_scan_iter(&self, pattern: &str, count: usize) -> Result<Vec<String>, Error> {
+        let pattern = format!("{}*", self.namespaced_key(pattern));
+        self.run(move |connection| {
+            let mut matches = Vec::new();
+            connection.scan(&pattern, count, |_, keys| {
+                matches.extend(keys);
+                Ok(matches.len() < count)
+            })?;
+            matches.truncate(count);
+            Ok(matches)
+        })
+        .await
+    }
+}

litellm-rust/crates/cache-redis/src/lib.rs

@@ -1,11 +1,15 @@
 mod cache;
+mod claim;
+pub mod connection;
+mod counter;
+mod keys;
+mod lifecycle;
+mod queue;
+mod script;
+mod store;
 mod topology;
 
-pub mod connection {
-    pub use crate::cache::{ConnectionRef, Connections};
-}
-
-pub use cache::{
-    RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisRpushOperation, RedisScript,
-};
+pub use cache::RedisCache;
+pub use queue::{RedisLpopOperation, RedisLpopResult, RedisRpushOperation};
+pub use script::{RedisArg, RedisScript};
 pub use topology::{RedisNode, RedisTopology};

litellm-rust/crates/cache-redis/src/lifecycle.rs

@@ -0,0 +1,85 @@
+use litellm_cache::{
+    CacheCodec, CacheConnectionResult, CacheConnectionStatus, ClientInfoCache, ConnectionCache,
+    DisconnectCache, Error, PingCache,
+};
+
+use crate::{cache::RedisCache, topology::RedisTopology};
+
+impl<S, C> PingCache for RedisCache<S, C>
+where
+    S: CacheCodec,
+    C: redis::ConnectionLike + Send + 'static,
+{
+    fn sync_ping(&self) -> Result<bool, Error> {
+        self.execute(|connection| connection.ping().map_err(|_| Error::Unavailable))
+    }
+
+    async fn ping(&self) -> Result<bool, Error> {
+        self.run(|connection| connection.ping().map_err(|_| Error::Unavailable))
+            .await
+    }
+}
+
+impl<S, C> ConnectionCache for RedisCache<S, C>
+where
+    S: CacheCodec,
+    C: redis::ConnectionLike + Send + 'static,
+{
+    /// Python `RedisCache.test_connection`, or `RedisClusterCache.test_connection` for a
+    /// cluster topology, which differs only in its messages.
+    async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
+        let label = match self.topology {
+            RedisTopology::Standalone => "Redis",
+            RedisTopology::Cluster { .. } => "Redis Cluster",
+        };
+        let ping = self
+            .run(|connection| Ok(connection.ping().map_err(|error| error.to_string())))
+            .await
+            .unwrap_or_else(|error| Err(error.to_string()));
+        Ok(match ping {
+            Ok(true) => CacheConnectionResult {
+                status: CacheConnectionStatus::Success,
+                message: format!("{label} connection test successful"),
+                error: None,
+            },
+            Ok(false) => CacheConnectionResult {
+                status: CacheConnectionStatus::Failed,
+                message: format!("{label} ping returned False"),
+                error: None,
+            },
+            Err(error) => CacheConnectionResult {
+                status: CacheConnectionStatus::Failed,
+                message: format!("{label} connection failed: {error}"),
+                error: Some(error),
+            },
+        })
+    }
+}
+
+impl<S, C> DisconnectCache for RedisCache<S, C>
+where
+    S: CacheCodec,
+    C: redis::ConnectionLike + Send + 'static,
+{
+    async fn disconnect(&self) -> Result<(), Error> {
+        self.connections.disconnect();
+        Ok(())
+    }
+}
+
+impl<S, C> ClientInfoCache for RedisCache<S, C>
+where
+    S: CacheCodec,
+    C: redis::ConnectionLike + Send + 'static,
+{
+    type ClientList = String;
+    type Info = String;
+
+    fn client_list(&self) -> Result<Self::ClientList, Error> {
+        self.execute(|connection| connection.node_text(redis::cmd("CLIENT").arg("LIST")))
+    }
+
+    fn info(&self) -> Result<Self::Info, Error> {
+        self.execute(|connection| connection.node_text(&redis::cmd("INFO")))
+    }
+}

litellm-rust/crates/cache-redis/src/queue.rs

@@ -0,0 +1,226 @@
+use std::time::Duration;
+
+use litellm_cache::{CacheCodec, Error, PopOperation, PushOperation, QueueCache, SetCache};
+
+use crate::{cache::RedisCache, script::RedisArg};
+
+pub type RedisRpushOperation = PushOperation<RedisArg>;
+pub type RedisLpopOperation = PopOperation;
+
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub enum RedisLpopResult {
+    Missing,
+    Value(Vec<u8>),
+    Values(Vec<Vec<u8>>),
+}
+
+impl<S, C> SetCache for RedisCache<S, C>
+where
+    S: CacheCodec,
+    C: redis::ConnectionLike + Send + 'static,
+{
+    type SetValue = RedisArg;
+    type SetResult = usize;
+
+    async fn async_set_cache_sadd(
+        &self,
+        key: &str,
+        values: Vec<Self::SetValue>,
+        ttl: Option<Duration>,
+    ) -> Result<Self::SetResult, Error> {
+        if values.is_empty() {
+            return Err(Error::InvalidEntry);
+        }
+        let key = self.namespaced_key(key);
+        let ttl = self.ttl_or_default(ttl);
+        let mut pipeline = redis::pipe();
+        pipeline
+            .cmd("SADD")
+            .arg(&key)
+            .arg(values)
+            .cmd("EXPIRE")
+            .arg(&key)
+            .arg(ttl)
+            .ignore();
+        self.run(move |connection| connection.query_pipeline(&pipeline))
+            .await
+            .map(|(added,)| added)
+    }
+}
+
+impl<S, C> QueueCache for RedisCache<S, C>
+where
+    S: CacheCodec,
+    C: redis::ConnectionLike + Send + 'static,
+{
+    type QueueValue = RedisArg;
+    type PopResult = RedisLpopResult;
+
+    async fn async_rpush(&self, key: &str, values: Vec<Self::QueueValue>) -> Result<usize, Error> {
+        if values.is_empty() {
+            return Err(Error::InvalidEntry);
+        }
+        let key = self.namespaced_key(key);
+        self.run(move |connection| {
+            redis::cmd("RPUSH")
+                .arg(key)
+                .arg(values)
+                .query(connection)
+                .map_err(|_| Error::Unavailable)
+        })
+        .await
+    }
+
+    async fn async_rpush_and_trim(
+        &self,
+        key: &str,
+        values: Vec<Self::QueueValue>,
+        max_len: usize,
+    ) -> Result<usize, Error> {
+        if values.is_empty() {
+            return Err(Error::InvalidEntry);
+        }
+        let key = self.namespaced_key(key);
+        let start = i64::try_from(max_len).map_or(i64::MIN, |max_len| -max_len);
+        let mut pipeline = redis::pipe();
+        pipeline
+            .atomic()
+            .cmd("RPUSH")
+            .arg(&key)
+            .arg(values)
+            .cmd("LTRIM")
+            .arg(&key)
+            .arg(start)
+            .arg(-1)
+            .ignore();
+        self.run(move |connection| connection.query_pipeline(&pipeline))
+            .await
+            .map(|(length,)| length)
+    }
+
+    async fn async_rpush_pipeline(
+        &self,
+        operations: Vec<RedisRpushOperation>,
+    ) -> Result<Vec<usize>, Error> {
+        if operations.is_empty() {
+            return Ok(Vec::new());
+        }
+        let mut pipeline = redis::pipe();
+        for operation in operations {
+            if operation.values.is_empty() {
+                return Err(Error::InvalidEntry);
+            }
+            pipeline
+                .cmd("RPUSH")
+                .arg(self.namespaced_key(&operation.key))
+                .arg(operation.values);
+        }
+        self.run(move |connection| connection.query_pipeline(&pipeline))
+            .await
+    }
+
+    async fn async_lpop(&self, key: &str, count: Option<usize>) -> Result<Self::PopResult, Error> {
+        if let Some(count) = count
+            && self.major_version().await < 7
+        {
+            return self.lpop_one_at_a_time(key, count).await;
+        }
+        let command = lpop(self.namespaced_key(key), count);
+        let value = self
+            .run(move |connection| {
+                command
+                    .query::<redis::Value>(connection)
+                    .map_err(|_| Error::Unavailable)
+            })
+            .await?;
+        lpop_result(value, count.is_some())
+    }
+
+    async fn async_lpop_pipeline(
+        &self,
+        operations: Vec<RedisLpopOperation>,
+    ) -> Result<Vec<Self::PopResult>, Error> {
+        if operations.is_empty() {
+            return Ok(Vec::new());
+        }
+        if operations.iter().any(|operation| operation.count.is_some())
+            && self.major_version().await < 7
+        {
+            let mut results = Vec::with_capacity(operations.len());
+            for operation in &operations {
+                results.push(self.async_lpop(&operation.key, operation.count).await?);
+            }
+            return Ok(results);
+        }
+        let multiple = operations
+            .iter()
+            .map(|operation| operation.count.is_some())
+            .collect::<Vec<_>>();
+        let mut pipeline = redis::pipe();
+        for operation in operations {
+            pipeline.add_command(lpop(self.namespaced_key(&operation.key), operation.count));
+        }
+        self.run(move |connection| connection.query_pipeline::<Vec<redis::Value>>(&pipeline))
+            .await?
+            .into_iter()
+            .zip(multiple)
+            .map(|(value, multiple)| lpop_result(value, multiple))
+            .collect()
+    }
+}
+
+fn lpop(key: String, count: Option<usize>) -> redis::Cmd {
+    let mut command = redis::cmd("LPOP");
+    command.arg(key);
+    if let Some(count) = count {
+        command.arg(count);
+    }
+    command
+}
+
+fn redis_bytes(value: redis::Value) -> Result<Vec<u8>, Error> {
+    match value {
+        redis::Value::BulkString(bytes) => Ok(bytes),
+        redis::Value::SimpleString(text) => Ok(text.into_bytes()),
+        _ => Err(Error::InvalidEntry),
+    }
+}
+
+impl<S, C> RedisCache<S, C>
+where
+    S: CacheCodec,
+    C: redis::ConnectionLike + Send + 'static,
+{
+    /// `handle_lpop_count_for_older_redis_versions`: `count` single-`LPOP` pipelines, keeping
+    /// only the values actually popped.
+    async fn lpop_one_at_a_time(&self, key: &str, count: usize) -> Result<RedisLpopResult, Error> {
+        let key = self.namespaced_key(key);
+        let mut values = Vec::new();
+        for _ in 0..count {
+            let mut pipeline = redis::pipe();
+            pipeline.add_command(lpop(key.clone(), None));
+            let replies = self
+                .run(move |connection| connection.query_pipeline::<Vec<redis::Value>>(&pipeline))
+                .await?;
+            for reply in replies {
+                if reply != redis::Value::Nil {
+                    values.push(redis_bytes(reply)?);
+                }
+            }
+        }
+        Ok(RedisLpopResult::Values(values))
+    }
+}
+
+fn lpop_result(value: redis::Value, multiple: bool) -> Result<RedisLpopResult, Error> {
+    match value {
+        redis::Value::Nil => Ok(RedisLpopResult::Missing),
+        redis::Value::Array(values) if multiple => values
+            .into_iter()
+            .map(redis_bytes)
+            .collect::<Result<Vec<_>, _>>()
+            .map(RedisLpopResult::Values),
+        value if !multiple => redis_bytes(value).map(RedisLpopResult::Value),
+        _ => Err(Error::InvalidEntry),
+    }
+}

litellm-rust/crates/cache-redis/src/script.rs

@@ -0,0 +1,136 @@
+use std::sync::Arc;
+
+use litellm_cache::{CacheCodec, CacheScript, Error, ScriptCache};
+
+use crate::{
+    cache::{RedisCache, namespaced_key},
+    connection::Connections,
+};
+
+#[derive(Clone, Debug, PartialEq)]
+pub enum RedisArg {
+    Bytes(Vec<u8>),
+    Integer(i64),
+    Float(f64),
+}
+
+impl From<&str> for RedisArg {
+    fn from(value: &str) -> Self {
+        Self::Bytes(value.as_bytes().to_vec())
+    }
+}
+
+impl From<String> for RedisArg {
+    fn from(value: String) -> Self {
+        Self::Bytes(value.into_bytes())
+    }
+}
+
+impl From<Vec<u8>> for RedisArg {
+    fn from(value: Vec<u8>) -> Self {
+        Self::Bytes(value)
+    }
+}
+
+impl From<i64> for RedisArg {
+    fn from(value: i64) -> Self {
+        Self::Integer(value)
+    }
+}
+
+impl From<f64> for RedisArg {
+    fn from(value: f64) -> Self {
+        Self::Float(value)
+    }
+}
+
+impl redis::ToRedisArgs for RedisArg {
+    fn write_redis_args<W>(&self, out: &mut W)
+    where
+        W: ?Sized + redis::RedisWrite,
+    {
+        match self {
+            Self::Bytes(value) => value.write_redis_args(out),
+            Self::Integer(value) => value.write_redis_args(out),
+            Self::Float(value) => value.write_redis_args(out),
+        }
+    }
+}
+
+pub struct RedisScript<C> {
+    connections: Arc<Connections<C>>,
+    namespace: Option<String>,
+    source: String,
+}
+
+impl<C> CacheScript for RedisScript<C>
+where
+    C: redis::ConnectionLike + Send + 'static,
+{
+    type Argument = RedisArg;
+    type Output = redis::Value;
+
+    async fn invoke(
+        &self,
+        keys: Vec<String>,
+        arguments: Vec<Self::Argument>,
+    ) -> Result<Self::Output, Error> {
+        let keys = keys
+            .into_iter()
+            .map(|key| namespaced_key(self.namespace.as_deref(), &key))
+            .collect::<Vec<_>>();
+        let source = self.source.clone();
+        Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
+            eval(connection, &source, keys, arguments)
+        })
+        .await
+    }
+}
+
+impl<S, C> RedisCache<S, C>
+where
+    S: CacheCodec,
+    C: redis::ConnectionLike + Send + 'static,
+{
+    pub async fn async_eval(
+        &self,
+        script: String,
+        keys: Vec<String>,
+        arguments: Vec<RedisArg>,
+    ) -> Result<redis::Value, Error> {
+        let keys = self.namespaced_keys(&keys);
+        self.run(move |connection| eval(connection, &script, keys, arguments))
+            .await
+    }
+}
+
+fn eval(
+    connection: &mut impl redis::ConnectionLike,
+    script: &str,
+    keys: Vec<String>,
+    arguments: Vec<RedisArg>,
+) -> Result<redis::Value, Error> {
+    redis::cmd("EVAL")
+        .arg(script)
+        .arg(keys.len())
+        .arg(keys)
+        .arg(arguments)
+        .query(connection)
+        .map_err(|_| Error::Unavailable)
+}
+
+impl<S, C> ScriptCache for RedisCache<S, C>
+where
+    S: CacheCodec,
+    C: redis::ConnectionLike + Send + 'static,
+{
+    type Script = RedisScript<C>;
+
+    fn async_register_script(&self, source: String) -> Self::Script {
+        RedisScript {
+            connections: Arc::clone(&self.connections),
+            namespace: self.namespace.clone(),
+            source,
+        }
+    }
+}

litellm-rust/crates/cache-redis/src/store.rs

@@ -0,0 +1,232 @@
+use std::time::Duration;
+
+use litellm_cache::{
+    BaseCache, BatchCache, BatchEntry, BulkDeleteCache, CacheCodec, DeleteCache, Error,
+    ExactCacheContext, FlushAllCache, FlushCache, TtlPipelineCache,
+};
+use redis::Commands;
+
+use crate::{cache::RedisCache, connection::ConnectionRef};
+
+impl<S, C> BaseCache for RedisCache<S, C>
+where
+    S: CacheCodec,
+    C: redis::ConnectionLike + Send + 'static,
+{
+    type Value = S::Value;
+    type Context = ExactCacheContext;
+
+    fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
+        context.ttl.or(Some(self.default_ttl))
+    }
+
+    fn set_cache(
+        &self,
+        key: &str,
+        value: Self::Value,
+        context: &ExactCacheContext,
+    ) -> Result<(), Error> {
+        let payload = self.codec.encode(&value)?;
+        let ttl = self.ttl_or_default(context.ttl);
+        let key = self.namespaced_key(key);
+        self.execute(|connection| {
+            connection
+                .set_ex::<_, _, ()>(key, payload, ttl)
+                .map_err(|_| Error::Unavailable)
+        })
+    }
+
+    fn get_cache(&self, key: &str, _: &ExactCacheContext) -> Result<Option<Self::Value>, Error> {
+        let key = self.namespaced_key(key);
+        let value = self.execute(|connection| {
+            connection
+                .get::<_, redis::Value>(key)
+                .map_err(|_| Error::Unavailable)
+        })?;
+        self.decode_response(value)
+    }
+
+    async fn async_set_cache(
+        &self,
+        key: &str,
+        value: Self::Value,
+        context: ExactCacheContext,
+    ) -> Result<(), Error> {
+        let payload = self.codec.encode(&value)?;
+        let key = self.namespaced_key(key);
+        let ttl = self.ttl_or_default(context.ttl);
+        self.run(move |connection| {
+            connection
+                .set_ex::<_, _, ()>(key, payload, ttl)
+                .map_err(|_| Error::Unavailable)
+        })
+        .await
+    }
+
+    async fn async_get_cache(
+        &self,
+        key: &str,
+        _: &ExactCacheContext,
+    ) -> Result<Option<Self::Value>, Error> {
+        let key = self.namespaced_key(key);
+        let value = self
+            .run(move |connection| {
+                connection
+                    .get::<_, redis::Value>(key)
+                    .map_err(|_| Error::Unavailable)
+            })
+            .await?;
+        self.decode_response(value)
+    }
+
+    async fn async_set_cache_pipeline(
+        &self,
+        cache_list: Vec<(String, Self::Value)>,
+        context: ExactCacheContext,
+    ) -> Result<(), Error> {
+        self.async_set_cache_pipeline_with_ttls(
+            cache_list
+                .into_iter()
+                .map(|(key, value)| (key, value, context.ttl))
+                .collect(),
+        )
+        .await
+    }
+}
+
+impl<S, C> TtlPipelineCache for RedisCache<S, C>
+where
+    S: CacheCodec,
+    C: redis::ConnectionLike + Send + 'static,
+{
+    async fn async_set_cache_pipeline_with_ttls(
+        &self,
+        entries: Vec<(String, Self::Value, Option<Duration>)>,
+    ) -> Result<(), Error> {
+        if entries.is_empty() {
+            return Ok(());
+        }
+        let mut pipeline = redis::pipe();
+        for (key, value, ttl) in entries {
+            pipeline
+                .cmd("SETEX")
+                .arg(self.namespaced_key(&key))
+                .arg(self.ttl_or_default(ttl))
+                .arg(self.codec.encode(&value)?)
+                .ignore();
+        }
+        self.run(move |connection| connection.query_pipeline(&pipeline))
+            .await
+    }
+}
+
+impl<S, C> BatchCache for RedisCache<S, C>
+where
+    S: CacheCodec,
+    C: redis::ConnectionLike + Send + 'static,
+{
+    fn batch_get_cache(
+        &self,
+        keys: &[String],
+        _: &ExactCacheContext,
+    ) -> Result<Vec<BatchEntry<Self::Value>>, Error> {
+        let keys = self.namespaced_keys(keys);
+        self.execute(|connection| mget(connection, keys))?
+            .into_iter()
+            .map(|value| self.decode_batch_response(value))
+            .collect()
+    }
+
+    async fn async_batch_get_cache(
+        &self,
+        keys: Vec<String>,
+        _: ExactCacheContext,
+    ) -> Result<Vec<BatchEntry<Self::Value>>, Error> {
+        let keys = self.namespaced_keys(&keys);
+        self.run(move |connection| mget(connection, keys))
+            .await?
+            .into_iter()
+            .map(|value| self.decode_batch_response(value))
+            .collect()
+    }
+}
+
+pub(crate) fn mget(
+    connection: &mut ConnectionRef<'_>,
+    keys: Vec<String>,
+) -> Result<Vec<redis::Value>, Error> {
+    redis::cmd("MGET")
+        .arg(keys)
+        .query(connection)
+        .map_err(|_| Error::Unavailable)
+}
+
+impl<S, C> DeleteCache for RedisCache<S, C>
+where
+    S: CacheCodec,
+    C: redis::ConnectionLike + Send + 'static,
+{
+    fn delete_cache(&self, key: &str) -> Result<(), Error> {
+        let key = self.namespaced_key(key);
+        self.execute(|connection| connection.del::<_, ()>(key).map_err(|_| Error::Unavailable))
+    }
+
+    async fn async_delete_cache(&self, key: &str) -> Result<(), Error> {
+        let key = self.namespaced_key(key);
+        self.run(move |connection| connection.del::<_, ()>(key).map_err(|_| Error::Unavailable))
+            .await
+    }
+}
+
+impl<S, C> BulkDeleteCache for RedisCache<S, C>
+where
+    S: CacheCodec,
+    C: redis::ConnectionLike + Send + 'static,
+{
+    async fn delete_cache_keys(&self, keys: Vec<String>) -> Result<usize, Error> {
+        if keys.is_empty() {
+            return Ok(0);
+        }
+        let keys = self.namespaced_keys(&keys);
+        self.run(move |connection| connection.del(keys).map_err(|_| Error::Unavailable))
+            .await
+    }
+}
+
+impl<S, C> FlushCache for RedisCache<S, C>
+where
+    S: CacheCodec,
+    C: redis::ConnectionLike + Send + 'static,
+{
+    fn flush_cache(&self) -> Result<(), Error> {
+        let pattern = self.namespaced_pattern()?;
+        self.execute(|connection| flush_matching(connection, &pattern))
+    }
+
+    async fn async_flush_cache(&self) -> Result<(), Error> {
+        let pattern = self.namespaced_pattern()?;
+        self.run(move |connection| flush_matching(connection, &pattern))
+            .await
+    }
+}
+
+impl<S, C> FlushAllCache for RedisCache<S, C>
+where
+    S: CacheCodec,
+    C: redis::ConnectionLike + Send + 'static,
+{
+    fn flushall(&self) -> Result<(), Error> {
+        self.execute(|connection| connection.flushall())
+    }
+}
+
+fn flush_matching(connection: &mut ConnectionRef<'_>, pattern: &str) -> Result<(), Error> {
+    connection.scan(pattern, 1000, |connection, keys| {
+        if !keys.is_empty() {
+            connection
+                .del::<_, usize>(keys)
+                .map_err(|_| Error::Unavailable)?;
+        }
+        Ok(true)
+    })
+}

litellm-rust/crates/cache-redis/tests/cluster.rs

@@ -1,100 +1,67 @@
-//! Contract tests against a real Redis Cluster. Set `LITELLM_TEST_REDIS_CLUSTER_NODES` to a
-//! comma separated `host:port` list (for example `127.0.0.1:7000,127.0.0.1:7001`) to run them.
+//! Tests against a real Redis Cluster. Set `LITELLM_TEST_REDIS_CLUSTER_NODES` to a comma
+//! separated `host:port` list (for example `127.0.0.1:7000,127.0.0.1:7001`) to run them.
 
-use std::time::{Duration, SystemTime, UNIX_EPOCH};
+mod support;
+
+use std::{collections::HashSet, time::Duration};
 
 use litellm_cache::{
-    BaseCache, BatchCache, BatchEntry, CacheConnectionStatus, CacheScript, ClaimCache,
-    CounterCache, DeleteCache, Error, ExactCacheContext, FlushCache, IncrementOperation, JsonCodec,
-    ScriptCache,
+    BaseCache, BatchCache, BatchEntry, BoundedCounterCache, BulkDeleteCache, CacheConnectionStatus,
+    CacheScript, ClaimCache, ClientInfoCache, ConnectionCache, CounterCache, DeleteCache,
+    DisconnectCache, Error, ExactCacheContext, FlushCache, IncrementOperation, JsonCodec,
+    PingCache, QueueCache, RefreshTtlCache, ScanCache, ScriptCache, SetCache, TtlCache,
+    TtlPipelineCache,
 };
 use litellm_cache_redis::{
     RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisNode, RedisRpushOperation,
     RedisTopology,
 };
 use redis::cluster_routing::Slot;
+use rstest::{fixture, rstest};
+use serde_json::json;
+use support::{JsonCache, cluster_cache, cluster_url};
 
-type Cache = RedisCache<JsonCodec<serde_json::Value>>;
+type Counter = RedisCache<JsonCodec<f64>>;
 
-fn topology() -> Option<RedisTopology> {
-    let nodes = std::env::var("LITELLM_TEST_REDIS_CLUSTER_NODES").ok()?;
-    let startup_nodes = nodes
-        .split(',')
-        .map(|node| {
-            let (host, port) = node.trim().rsplit_once(':').expect("host:port");
-            RedisNode {
-                host: host.to_string(),
-                port: port.parse().expect("port"),
-            }
-        })
-        .collect();
-    Some(RedisTopology::Cluster { startup_nodes })
+#[fixture]
+fn cache(#[default("cache")] label: &str) -> Option<JsonCache> {
+    cluster_cache(label, Duration::from_secs(120), JsonCodec::new())
 }
 
-fn namespace(label: &str) -> String {
-    let nanos = SystemTime::now()
-        .duration_since(UNIX_EPOCH)
-        .unwrap()
-        .as_nanos();
-    format!("cluster-test:{label}:{nanos}")
+#[fixture]
+fn counter(#[default("counter")] label: &str) -> Option<Counter> {
+    cluster_cache(label, Duration::from_secs(60), JsonCodec::new())
 }
 
-fn cluster_url() -> String {
-    std::env::var("LITELLM_TEST_REDIS_CLUSTER_URL")
-        .unwrap_or_else(|_| "redis://127.0.0.1:7000".into())
-}
-
-fn cluster_cache(label: &str) -> Option<Cache> {
-    let topology = topology()?;
-    Some(
-        Cache::connect(
-            &cluster_url(),
-            &topology,
-            Some(Duration::from_secs(120)),
-            JsonCodec::new(),
-        )
-        .expect("cluster connection")
-        .with_namespace(Some(namespace(label))),
-    )
-}
-
-fn counter_cache(label: &str) -> Option<RedisCache<JsonCodec<f64>>> {
-    let topology = topology()?;
-    Some(
-        RedisCache::connect(
-            &cluster_url(),
-            &topology,
-            Some(Duration::from_secs(60)),
-            JsonCodec::new(),
-        )
-        .expect("cluster connection")
-        .with_namespace(Some(namespace(label))),
-    )
+#[fixture]
+fn context() -> ExactCacheContext {
+    ExactCacheContext::default()
 }
 
 fn multi_slot_keys(count: usize) -> Vec<String> {
     let keys: Vec<String> = (0..count).map(|index| format!("key-{index}")).collect();
-    let slots: std::collections::HashSet<Slot> = keys.iter().map(Slot::for_key).collect();
+    let slots: HashSet<Slot> = keys.iter().map(Slot::for_key).collect();
     assert!(slots.len() > 1, "keys must span multiple slots");
     keys
 }
 
-macro_rules! cluster_or_skip {
-    ($label:expr) => {
-        match cluster_cache($label) {
-            Some(cache) => cache,
-            None => return,
-        }
-    };
+fn seconds(seconds: u64) -> Option<Duration> {
+    Some(Duration::from_secs(seconds))
 }
 
-#[test]
-fn constructor_rejects_clusters_without_startup_nodes() {
-    let error = Cache::connect(
-        "redis://127.0.0.1:7000",
-        &RedisTopology::Cluster {
-            startup_nodes: Vec::new(),
-        },
+#[rstest]
+#[case::no_startup_nodes("redis://127.0.0.1:7000", Vec::new())]
+#[case::unix_socket_url(
+    "redis+unix:///tmp/redis.sock",
+    vec![RedisNode { host: "127.0.0.1".into(), port: 7000 }]
+)]
+fn constructor_rejects_unusable_cluster_configs(
+    #[case] url: &str,
+    #[case] startup_nodes: Vec<RedisNode>,
+) {
+    let error = JsonCache::connect(
+        url,
+        &RedisTopology::Cluster { startup_nodes },
         None,
         JsonCodec::new(),
     )
@@ -102,60 +69,56 @@ fn constructor_rejects_clusters_without_startup_nodes() {
     assert!(matches!(error, Some(Error::Unavailable)));
 }
 
-#[test]
-fn constructor_rejects_unix_socket_urls_for_clusters() {
-    let error = Cache::connect(
-        "redis+unix:///tmp/redis.sock",
-        &RedisTopology::Cluster {
-            startup_nodes: vec![RedisNode {
-                host: "127.0.0.1".into(),
-                port: 7000,
-            }],
-        },
-        None,
-        JsonCodec::new(),
-    )
-    .err();
-    assert!(matches!(error, Some(Error::Unavailable)));
-}
-
-#[test]
-fn single_key_operations_round_trip_with_ttl_rounding() {
-    let cache = cluster_or_skip!("single");
+#[rstest]
+#[tokio::test]
+async fn single_key_operations_round_trip_with_ttl_rounding(
+    #[with("single")] cache: Option<JsonCache>,
+) {
+    let Some(cache) = cache else { return };
     let context = ExactCacheContext {
         ttl: Some(Duration::from_millis(1500)),
     };
     let keys = multi_slot_keys(12);
     for (index, key) in keys.iter().enumerate() {
         cache
-            .set_cache(key, serde_json::json!({ "index": index }), &context)
+            .set_cache(key, json!({ "index": index }), &context)
             .unwrap();
     }
     for (index, key) in keys.iter().enumerate() {
         assert_eq!(
             cache.get_cache(key, &context).unwrap(),
-            Some(serde_json::json!({ "index": index }))
+            Some(json!({ "index": index }))
         );
     }
-    let runtime = tokio::runtime::Runtime::new().unwrap();
-    let ttl = runtime.block_on(cache.async_get_ttl(&keys[0])).unwrap();
-    assert_eq!(ttl, Some(2));
+    assert_eq!(cache.async_get_ttl(&keys[0]).await.unwrap(), seconds(2));
+    assert!(
+        cache
+            .async_refresh_ttl(&keys[0], seconds(40))
+            .await
+            .unwrap()
+    );
+    assert_eq!(cache.async_get_ttl(&keys[0]).await.unwrap(), seconds(40));
     cache.delete_cache(&keys[0]).unwrap();
     assert_eq!(cache.get_cache(&keys[0], &context).unwrap(), None);
+    assert!(!cache.async_refresh_ttl(&keys[0], None).await.unwrap());
     assert!(cache.sync_ping().unwrap());
+    cache.async_flush_cache().await.unwrap();
 }
 
+#[rstest]
 #[tokio::test]
-async fn batch_reads_span_slots_and_preserve_order_with_malformed_entries() {
-    let cache = cluster_or_skip!("batch");
-    let context = ExactCacheContext::default();
+async fn batch_reads_span_slots_and_preserve_order_with_malformed_entries(
+    #[with("batch")] cache: Option<JsonCache>,
+    context: ExactCacheContext,
+) {
+    let Some(cache) = cache else { return };
     let keys = multi_slot_keys(40);
     for (index, key) in keys.iter().enumerate() {
         if index % 5 == 0 {
             continue;
         }
         cache
-            .async_set_cache(key, serde_json::json!(index), context.clone())
+            .async_set_cache(key, json!(index), context.clone())
             .await
             .unwrap();
     }
@@ -181,41 +144,65 @@ async fn batch_reads_span_slots_and_preserve_order_with_malformed_entries() {
         } else if index % 5 == 0 {
             BatchEntry::Miss
         } else {
-            BatchEntry::Hit(serde_json::json!(index))
+            BatchEntry::Hit(json!(index))
         };
         assert_eq!(*entry, expected, "entry {index}");
     }
-    let sync_entries = cache.batch_get_cache(&keys, &context).unwrap();
-    assert_eq!(sync_entries, entries);
+    assert_eq!(cache.batch_get_cache(&keys, &context).unwrap(), entries);
 
     cache.delete_cache_keys(keys.clone()).await.unwrap();
     let entries = cache.async_batch_get_cache(keys, context).await.unwrap();
     assert!(entries.iter().all(|entry| *entry == BatchEntry::Miss));
 }
 
+#[rstest]
 #[tokio::test]
-async fn pipelines_group_by_slot_and_return_results_in_submission_order() {
-    let cache = cluster_or_skip!("pipeline");
+async fn pipelines_group_by_slot_and_return_results_in_submission_order(
+    #[with("pipeline")] cache: Option<JsonCache>,
+    counter: Option<Counter>,
+    context: ExactCacheContext,
+) {
+    let (Some(cache), Some(counter)) = (cache, counter) else {
+        return;
+    };
     let keys = multi_slot_keys(30);
-    let entries = keys
-        .iter()
-        .enumerate()
-        .map(|(index, key)| (key.clone(), serde_json::json!(index)))
-        .collect();
     cache
-        .async_set_cache_pipeline(entries, ExactCacheContext::default())
+        .async_set_cache_pipeline(
+            keys.iter()
+                .enumerate()
+                .map(|(index, key)| (key.clone(), json!(index)))
+                .collect(),
+            context.clone(),
+        )
         .await
         .unwrap();
     let hits = cache
-        .async_batch_get_cache(keys.clone(), ExactCacheContext::default())
+        .async_batch_get_cache(keys.clone(), context.clone())
         .await
         .unwrap();
     assert!(
         hits.iter()
             .enumerate()
-            .all(|(index, entry)| *entry == BatchEntry::Hit(serde_json::json!(index)))
+            .all(|(index, entry)| *entry == BatchEntry::Hit(json!(index)))
     );
 
+    cache
+        .async_set_cache_pipeline_with_ttls(
+            keys.iter()
+                .enumerate()
+                .map(|(index, key)| (key.clone(), json!(index), seconds(index as u64 + 10)))
+                .collect(),
+        )
+        .await
+        .unwrap();
+    for (index, key) in keys.iter().enumerate() {
+        assert_eq!(
+            cache.async_get_ttl(key).await.unwrap(),
+            seconds(index as u64 + 10),
+            "{key}"
+        );
+    }
+
     let queues: Vec<String> = keys.iter().map(|key| format!("queue:{key}")).collect();
     let pushed = cache
         .async_rpush_pipeline(
@@ -265,9 +252,6 @@ async fn pipelines_group_by_slot_and_return_results_in_submission_order() {
     }
 
     let counters: Vec<String> = keys.iter().map(|key| format!("counter:{key}")).collect();
-    let Some(counter) = counter_cache("counter") else {
-        return;
-    };
     let totals = counter
         .async_increment_pipeline(
             counters
@@ -284,25 +268,63 @@ async fn pipelines_group_by_slot_and_return_results_in_submission_order() {
         .unwrap();
     let expected: Vec<f64> = (0..keys.len()).map(|index| index as f64 + 0.5).collect();
     assert_eq!(totals, expected);
-    assert_eq!(counter.async_get_ttl(&counters[0]).await.unwrap(), Some(30));
+    assert_eq!(
+        counter.async_get_ttl(&counters[0]).await.unwrap(),
+        seconds(30)
+    );
     assert_eq!(counter.async_get_ttl(&counters[1]).await.unwrap(), None);
     counter.async_flush_cache().await.unwrap();
     cache.async_flush_cache().await.unwrap();
 }
 
+#[rstest]
 #[tokio::test]
-async fn scan_and_scoped_flush_cover_every_primary() {
-    let cache = cluster_or_skip!("flush");
-    let other = cluster_or_skip!("other");
-    let context = ExactCacheContext::default();
+async fn rpush_and_trim_is_one_transaction_on_the_key_slot(
+    #[with("trim")] cache: Option<JsonCache>,
+) {
+    let Some(cache) = cache else { return };
+    let values = |values: &[&str]| values.iter().map(|value| RedisArg::from(*value)).collect();
+    assert_eq!(
+        cache
+            .async_rpush_and_trim("buf", values(&["a", "b"]), 3)
+            .await
+            .unwrap(),
+        2
+    );
+    assert_eq!(
+        cache
+            .async_rpush_and_trim("buf", values(&["c", "d"]), 3)
+            .await
+            .unwrap(),
+        4
+    );
+    assert_eq!(
+        cache.async_lpop("buf", Some(10)).await.unwrap(),
+        RedisLpopResult::Values(vec![b"b".to_vec(), b"c".to_vec(), b"d".to_vec()])
+    );
+    cache.async_flush_cache().await.unwrap();
+}
+
+#[rstest]
+#[tokio::test]
+async fn scan_and_scoped_flush_cover_every_primary(
+    #[with("flush")] cache: Option<JsonCache>,
+    #[from(cache)]
+    #[with("other")]
+    other: Option<JsonCache>,
+    context: ExactCacheContext,
+) {
+    let (Some(cache), Some(other)) = (cache, other) else {
+        return;
+    };
     let keys = multi_slot_keys(60);
     for key in &keys {
         cache
-            .async_set_cache(key, serde_json::json!(true), context.clone())
+            .async_set_cache(key, json!(true), context.clone())
             .await
             .unwrap();
         other
-            .async_set_cache(key, serde_json::json!(true), context.clone())
+            .async_set_cache(key, json!(true), context.clone())
             .await
             .unwrap();
     }
@@ -325,7 +347,7 @@ async fn scan_and_scoped_flush_cover_every_primary() {
     let kept = other.async_batch_get_cache(keys, context).await.unwrap();
     assert!(
         kept.iter()
-            .all(|entry| *entry == BatchEntry::Hit(serde_json::json!(true)))
+            .all(|entry| *entry == BatchEntry::Hit(json!(true)))
     );
     other.async_flush_cache().await.unwrap();
 }
@@ -361,9 +383,10 @@ fn ping_calls_per_node(startup: &redis::Client) -> Vec<(String, u64)> {
     counts
 }
 
+#[rstest]
 #[tokio::test]
-async fn ping_reaches_every_node() {
-    let cache = cluster_or_skip!("ping");
+async fn ping_reaches_every_node(#[with("ping")] cache: Option<JsonCache>) {
+    let Some(cache) = cache else { return };
     let startup = redis::Client::open(cluster_url()).unwrap();
     let before = ping_calls_per_node(&startup);
     assert!(before.len() >= 2, "{before:?}");
@@ -375,14 +398,63 @@ async fn ping_reaches_every_node() {
     assert!(cache.sync_ping().unwrap());
     let result = cache.test_connection().await.unwrap();
     assert_eq!(result.status, CacheConnectionStatus::Success);
+    assert_eq!(result.message, "Redis Cluster connection test successful");
 }
 
+#[rstest]
 #[tokio::test]
-async fn counters_claims_scripts_and_sets_work_on_the_cluster() {
-    let Some(counter) = counter_cache("counter") else {
+async fn disconnect_closes_idle_connections_and_reconnects_on_demand(
+    #[with("disconnect")] cache: Option<JsonCache>,
+) {
+    let Some(cache) = cache else { return };
+    assert!(cache.ping().await.unwrap());
+    cache.disconnect().await.unwrap();
+    assert!(cache.ping().await.unwrap());
+}
+
+#[rstest]
+#[case::keep_existing_ttl(false)]
+#[case::refresh_ttl(true)]
+#[tokio::test]
+async fn increments_refresh_the_ttl_only_when_asked(
+    counter: Option<Counter>,
+    #[case] refresh_ttl: bool,
+) {
+    let Some(counter) = counter else { return };
+    let context = ExactCacheContext { ttl: seconds(60) };
+    counter
+        .async_set_cache("spend", 0.0, ExactCacheContext { ttl: seconds(600) })
+        .await
+        .unwrap();
+    assert_eq!(
+        counter
+            .async_increment("spend", 1.5, context.clone(), refresh_ttl)
+            .await
+            .unwrap(),
+        1.5
+    );
+    assert_eq!(
+        counter
+            .async_increment("spend", 2.0, context, refresh_ttl)
+            .await
+            .unwrap(),
+        3.5
+    );
+    let ttl = counter.async_get_ttl("spend").await.unwrap().unwrap();
+    assert_eq!(ttl <= Duration::from_secs(60), refresh_ttl, "{ttl:?}");
+    counter.async_flush_cache().await.unwrap();
+}
+
+#[rstest]
+#[tokio::test]
+async fn counters_claims_scripts_and_sets_work_on_the_cluster(
+    counter: Option<Counter>,
+    #[with("claim")] cache: Option<JsonCache>,
+    context: ExactCacheContext,
+) {
+    let (Some(counter), Some(cache)) = (counter, cache) else {
         return;
     };
-    let context = ExactCacheContext::default();
     assert_eq!(
         counter
             .increment_cache("spend", 1.5, context.clone())
@@ -391,7 +463,7 @@ async fn counters_claims_scripts_and_sets_work_on_the_cluster() {
     );
     assert_eq!(
         counter
-            .async_increment("spend", 2.0, context.clone())
+            .async_increment("spend", 2.0, context.clone(), false)
             .await
             .unwrap(),
         3.5
@@ -413,9 +485,8 @@ async fn counters_claims_scripts_and_sets_work_on_the_cluster() {
     assert_eq!(counter.async_set_max("peak", 2.0, None).await.unwrap(), 4.0);
     counter.flush_cache().unwrap();
 
-    let cache = cluster_or_skip!("claim");
-    let owner = serde_json::json!("owner-a");
-    let rival = serde_json::json!("owner-b");
+    let owner = json!("owner-a");
+    let rival = json!("owner-b");
     assert_eq!(
         cache
             .claim_cache("lock", owner.clone(), &[], context.clone())
@@ -453,8 +524,8 @@ async fn counters_claims_scripts_and_sets_work_on_the_cluster() {
         .await
         .unwrap();
     assert_eq!(reply, redis::Value::Okay);
-    assert_eq!(cache.async_get_ttl("scripted").await.unwrap(), Some(5));
-    let evaluated: redis::Value = cache
+    assert_eq!(cache.async_get_ttl("scripted").await.unwrap(), seconds(5));
+    let evaluated = cache
         .async_eval(
             "return redis.call('GET', KEYS[1])".into(),
             vec!["scripted".into()],
@@ -472,13 +543,13 @@ async fn counters_claims_scripts_and_sets_work_on_the_cluster() {
                     RedisArg::Bytes(b"a".to_vec()),
                     RedisArg::Bytes(b"b".to_vec())
                 ],
-                Some(Duration::from_secs(9)),
+                seconds(9),
             )
             .await
             .unwrap(),
         2
     );
-    assert_eq!(cache.async_get_ttl("members").await.unwrap(), Some(9));
+    assert_eq!(cache.async_get_ttl("members").await.unwrap(), seconds(9));
 
     let result = cache.test_connection().await.unwrap();
     assert_eq!(result.status, CacheConnectionStatus::Success);

litellm-rust/crates/cache-redis/tests/contract.rs

@@ -0,0 +1,95 @@
+//! The shared cache contracts, run against the in-process fake connection and, when
+//! `LITELLM_TEST_REDIS_CLUSTER_NODES` is set, against a live Redis Cluster.
+
+mod support;
+
+use std::time::Duration;
+
+use litellm_cache::{ExactCacheContext, JsonCodec};
+use litellm_cache_testing as contract;
+use rstest::rstest;
+use serde_json::json;
+use support::{JsonCache, cluster_cache, fake_cache};
+
+const PREFIX: &str = "contract:";
+
+#[derive(Clone, Copy, Debug)]
+enum Contract {
+    HitAndMiss,
+    SyncAsyncEquivalence,
+    OverwriteReplaces,
+    PipelineWritesEveryEntry,
+    BatchPreservesOrder,
+    DeleteRemovesKey,
+    FlushClears,
+    CounterAccumulates,
+}
+
+#[derive(Clone, Copy, Debug)]
+enum Server {
+    Fake,
+    Cluster,
+}
+
+async fn check<C>(contract: Contract, cache: &JsonCache<C>)
+where
+    C: redis::ConnectionLike + Send + 'static,
+{
+    let context = ExactCacheContext::default();
+    match contract {
+        Contract::HitAndMiss => {
+            contract::hit_and_miss(cache, context, PREFIX, json!({"answer": 42})).await
+        }
+        Contract::SyncAsyncEquivalence => {
+            contract::sync_async_equivalence(cache, context, PREFIX, json!("first"), json!([2]))
+                .await
+        }
+        Contract::OverwriteReplaces => {
+            contract::overwrite_replaces(cache, context, PREFIX, json!(1), json!({"b": 2})).await
+        }
+        Contract::PipelineWritesEveryEntry => {
+            contract::pipeline_writes_every_entry(
+                cache,
+                context,
+                PREFIX,
+                vec![json!("a"), json!(2), json!({"c": true})],
+            )
+            .await
+        }
+        Contract::BatchPreservesOrder => {
+            contract::batch_preserves_order(cache, context, PREFIX, json!("first"), json!(2)).await
+        }
+        Contract::DeleteRemovesKey => {
+            contract::delete_removes_key(cache, context, PREFIX, json!("value")).await
+        }
+        Contract::FlushClears => {
+            contract::flush_clears(cache, context, PREFIX, json!("value")).await
+        }
+        Contract::CounterAccumulates => contract::counter_accumulates(cache, context, PREFIX).await,
+    }
+}
+
+#[rstest]
+#[case::hit_and_miss(Contract::HitAndMiss)]
+#[case::sync_async_equivalence(Contract::SyncAsyncEquivalence)]
+#[case::overwrite_replaces(Contract::OverwriteReplaces)]
+#[case::pipeline_writes_every_entry(Contract::PipelineWritesEveryEntry)]
+#[case::batch_preserves_order(Contract::BatchPreservesOrder)]
+#[case::delete_removes_key(Contract::DeleteRemovesKey)]
+#[case::flush_clears(Contract::FlushClears)]
+#[case::counter_accumulates(Contract::CounterAccumulates)]
+#[tokio::test]
+async fn redis_satisfies_the_cache_contract(
+    #[case] contract: Contract,
+    #[values(Server::Fake, Server::Cluster)] server: Server,
+) {
+    match server {
+        Server::Fake => check(contract, &fake_cache("contract")).await,
+        Server::Cluster => {
+            let label = format!("{contract:?}");
+            if let Some(cache) = cluster_cache(&label, Duration::from_secs(120), JsonCodec::new()) {
+                check(contract, &cache).await;
+            }
+        }
+    }
+}

litellm-rust/crates/cache-redis/tests/support/mod.rs

@@ -0,0 +1,231 @@
+#![allow(dead_code)]
+
+use std::{
+    collections::BTreeMap,
+    time::{Duration, SystemTime, UNIX_EPOCH},
+};
+
+use litellm_cache::{CacheCodec, Error, JsonCodec};
+use litellm_cache_redis::{RedisCache, RedisNode, RedisTopology};
+
+/// Encodes a byte behind a tag, so a value written with another tag decodes as invalid.
+pub struct TaggedByteCodec(pub u8);
+
+impl CacheCodec for TaggedByteCodec {
+    type Value = u8;
+
+    fn encode(&self, value: &u8) -> Result<Vec<u8>, Error> {
+        if *value > 127 {
+            return Err(Error::InvalidEntry);
+        }
+        Ok(vec![self.0, *value])
+    }
+
+    fn decode(&self, bytes: &[u8]) -> Result<u8, Error> {
+        match bytes {
+            [tag, value] if *tag == self.0 => Ok(*value),
+            _ => Err(Error::InvalidEntry),
+        }
+    }
+}
+
+/// A stateful in-process stand-in for a Redis server that understands the string commands the
+/// shared contracts exercise, so they run without a live server. TTLs are accepted and ignored.
+#[derive(Default)]
+pub struct FakeRedis {
+    strings: BTreeMap<Vec<u8>, Vec<u8>>,
+}
+
+impl FakeRedis {
+    fn run(&mut self, command: Vec<Vec<u8>>) -> redis::RedisResult<redis::Value> {
+        let name = String::from_utf8_lossy(&command[0]).to_ascii_uppercase();
+        let args = &command[1..];
+        Ok(match name.as_str() {
+            "PING" => redis::Value::SimpleString("PONG".into()),
+            "SET" | "SETEX" => {
+                let value = if name == "SET" { &args[1] } else { &args[2] };
+                self.strings.insert(args[0].clone(), value.clone());
+                redis::Value::Okay
+            }
+            "GET" => self.get(&args[0]),
+            "MGET" => redis::Value::Array(args.iter().map(|key| self.get(key)).collect()),
+            "DEL" => {
+                let removed = args
+                    .iter()
+                    .filter(|key| self.strings.remove(*key).is_some())
+                    .count();
+                redis::Value::Int(removed as i64)
+            }
+            "SCAN" => {
+                let pattern = &args[2];
+                let keys = self
+                    .strings
+                    .keys()
+                    .filter(|key| glob(pattern, key))
+                    .map(|key| redis::Value::BulkString(key.clone()))
+                    .collect();
+                redis::Value::Array(vec![
+                    redis::Value::BulkString(b"0".to_vec()),
+                    redis::Value::Array(keys),
+                ])
+            }
+            "EVAL" if args[0].windows(11).any(|window| window == b"INCRBYFLOAT") => {
+                self.increment_by_float(&args[2], &args[3])
+            }
+            "INCRBYFLOAT" => self.increment_by_float(&args[0], &args[1]),
+            _ => {
+                return Err(redis::RedisError::from((
+                    redis::ErrorKind::Client,
+                    "unsupported command",
+                    name,
+                )));
+            }
+        })
+    }
+
+    fn get(&self, key: &[u8]) -> redis::Value {
+        self.strings.get(key).map_or(redis::Value::Nil, |value| {
+            redis::Value::BulkString(value.clone())
+        })
+    }
+
+    fn increment_by_float(&mut self, key: &[u8], amount: &[u8]) -> redis::Value {
+        let current = self
+            .strings
+            .get(key)
+            .map_or(0.0, |value| parse_float(value));
+        let total = format!("{}", current + parse_float(amount));
+        self.strings
+            .insert(key.to_vec(), total.clone().into_bytes());
+        redis::Value::BulkString(total.into_bytes())
+    }
+}
+
+fn parse_float(bytes: &[u8]) -> f64 {
+    std::str::from_utf8(bytes).unwrap().parse().unwrap()
+}
+
+/// Redis `MATCH` globbing for `*`, `?` and backslash escapes.
+fn glob(pattern: &[u8], key: &[u8]) -> bool {
+    match pattern.split_first() {
+        None => key.is_empty(),
+        Some((b'*', rest)) => (0..=key.len()).any(|skip| glob(rest, &key[skip..])),
+        Some((b'?', rest)) => !key.is_empty() && glob(rest, &key[1..]),
+        Some((b'\\', [escaped, rest @ ..])) => {
+            key.first() == Some(escaped) && glob(rest, &key[1..])
+        }
+        Some((literal, rest)) => key.first() == Some(literal) && glob(rest, &key[1..]),
+    }
+}
+
+/// Splits RESP request bytes into the commands they carry.
+fn commands(mut bytes: &[u8]) -> Vec<Vec<Vec<u8>>> {
+    fn line<'a>(bytes: &mut &'a [u8]) -> &'a [u8] {
+        let end = bytes
+            .windows(2)
+            .position(|window| window == b"\r\n")
+            .unwrap();
+        let (line, rest) = bytes.split_at(end);
+        *bytes = &rest[2..];
+        line
+    }
+    fn length(line: &[u8]) -> usize {
+        std::str::from_utf8(&line[1..]).unwrap().parse().unwrap()
+    }
+    let mut commands = Vec::new();
+    while !bytes.is_empty() {
+        let count = length(line(&mut bytes));
+        let command = (0..count)
+            .map(|_| {
+                let size = length(line(&mut bytes));
+                let (argument, rest) = bytes.split_at(size);
+                bytes = &rest[2..];
+                argument.to_vec()
+            })
+            .collect();
+        commands.push(command);
+    }
+    commands
+}
+
+impl redis::ConnectionLike for FakeRedis {
+    fn req_packed_command(&mut self, cmd: &[u8]) -> redis::RedisResult<redis::Value> {
+        let command = commands(cmd).into_iter().next().unwrap();
+        self.run(command)
+    }
+
+    fn req_packed_commands(
+        &mut self,
+        cmd: &[u8],
+        offset: usize,
+        count: usize,
+    ) -> redis::RedisResult<Vec<redis::Value>> {
+        let replies = commands(cmd)
+            .into_iter()
+            .map(|command| self.run(command))
+            .collect::<redis::RedisResult<Vec<_>>>()?;
+        Ok(replies.into_iter().skip(offset).take(count).collect())
+    }
+
+    fn get_db(&self) -> i64 {
+        0
+    }
+
+    fn check_connection(&mut self) -> bool {
+        true
+    }
+
+    fn is_open(&self) -> bool {
+        true
+    }
+}
+
+pub type JsonCache<C = redis::Connection> = RedisCache<JsonCodec<serde_json::Value>, C>;
+
+pub fn fake_cache(namespace: &str) -> JsonCache<FakeRedis> {
+    RedisCache::with_connection(FakeRedis::default(), None, JsonCodec::new())
+        .with_namespace(Some(namespace.into()))
+}
+
+/// Startup nodes from `LITELLM_TEST_REDIS_CLUSTER_NODES` (`host:port,host:port`); tests that
+/// need a live cluster skip when it is unset.
+pub fn cluster_topology() -> Option<RedisTopology> {
+    let nodes = std::env::var("LITELLM_TEST_REDIS_CLUSTER_NODES").ok()?;
+    let startup_nodes = nodes
+        .split(',')
+        .map(|node| {
+            let (host, port) = node.trim().rsplit_once(':').expect("host:port");
+            RedisNode {
+                host: host.to_string(),
+                port: port.parse().expect("port"),
+            }
+        })
+        .collect();
+    Some(RedisTopology::Cluster { startup_nodes })
+}
+
+pub fn cluster_url() -> String {
+    std::env::var("LITELLM_TEST_REDIS_CLUSTER_URL")
+        .unwrap_or_else(|_| "redis://127.0.0.1:7000".into())
+}
+
+pub fn unique_namespace(label: &str) -> String {
+    let nanos = SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .unwrap()
+        .as_nanos();
+    format!("cluster-test:{label}:{nanos}")
+}
+
+pub fn cluster_cache<S: CacheCodec>(
+    label: &str,
+    default_ttl: Duration,
+    codec: S,
+) -> Option<RedisCache<S>> {
+    let topology = cluster_topology()?;
+    Some(
+        RedisCache::connect(&cluster_url(), &topology, Some(default_ttl), codec)
+            .expect("cluster connection")
+            .with_namespace(Some(unique_namespace(label))),
+    )
+}

litellm-rust/crates/cache-response/Cargo.toml

@@ -17,4 +17,5 @@ litellm-cache-memory.workspace = true
 litellm-cache-redis.workspace = true
 redis = "1.7.0"
 redis-test = "1.0.4"
+rstest.workspace = true
 tokio.workspace = true

litellm-rust/crates/cache-response/README.md

@@ -1,14 +1,16 @@
-# Response cache foundation
+# Response cache
 
 `ResponseCache<B>` adds request keys, independent read/write controls, response envelopes, and freshness checks to any `B: BaseCache<Value = CacheEntry>`
 
 ## Ownership
 
-`litellm-cache` defines typed storage and codec traits. Memory and Redis implement those traits without depending on response policy. Other consumers can store their own value types using the same backend implementations
+`litellm-cache` defines typed storage, codec, and capability traits. `BaseCache` is only get, set, TTL, and pipeline writes. Everything else is an optional capability a backend implements only where its Python class defines the method: `DisconnectCache`, `ConnectionCache` (`test_connection`), `PingCache`, `BatchCache`, `DeleteCache`, `FlushCache`, counters, queues, TTL, scan, and scripts. Memory, Redis, disk, S3, GCS, and Azure Blob implement those traits without depending on response policy, so other consumers can store their own value types in the same backends
+
+Semantic backends (Redis, Valkey, Qdrant) are generic over their embedder and codec, and share one prompt and embedding contract from `litellm_cache::semantic`. They take a `SemanticCacheContext`, so `ResponseCache` drives them the same way it drives exact backends
 
 `litellm-cache-response` owns response keys, controls, entries, the Python-compatible response codec, and `WriteBuffer`, the backend-neutral deferred-write policy. It has no runtime dependency on a specific cache backend or Python
 
-The Python bridge constructs backends and selects them through its private `NativeResponseCache` enum, which only dispatches. Generic Rust callers inject their backend directly. A native gateway can construct the same generic response service in its own host
+`ExactResponseCache` is the object-safe view of a `ResponseCache` over an exact backend. `ConnectionProbe` is the object-safe `test_connection`, implemented only when the backend implements `ConnectionCache`, so a host holds one next to its `ExactResponseCache` and reports the operation as unsupported otherwise, as Python's `BaseCache` does. Lookup, store, batch, and flush never require it
 
 ## Native Rust use
 
@@ -28,36 +30,22 @@ cache.store(&request, json!({"answer": 7}), now)?;
 assert_eq!(cache.async_lookup(&request, now).await?, Some(json!({"answer": 7})));
 ```
 
-For Redis, inject `RedisCache::new(url, ttl, ResponseCacheCodec)` instead. Namespaces are optional and existing namespace prefixes are preserved. Sync operations check out independent connections from a bounded pool, while async callers, including counters and claims, move that blocking work off the executor. The pool skips the checkout PING and instead discards any connection whose command failed
+For Redis, inject `RedisCache::new(url, ttl, ResponseCacheCodec)` instead. Namespaces are optional and existing namespace prefixes are preserved
 
 Callers supply Unix time for response freshness. Backend TTL uses its own clock. A read can reject an entry through `max_age` even while the backend still retains it
 
-## Python integration boundary
+## Python integration
 
-The extension keeps a private test harness for memory and Redis single and batch response lookup and storage. Batch lookup returns ordered values plus missing indices for embedding partial-hit wiring. No bridge-only cache type is part of the public API
+The bridge activates backends through the Rust catalog in `litellm/rust_bridge/catalog.py`. Every cache rule ships as `PYTHON_ONLY`, so SDK, Router, and proxy calls stay on Python and construct no native cache resources until a rule is changed
 
-The bridge also exposes a production-shaped response cache runtime selected through the Rust catalog. Its shipped rule set is empty, so current SDK, Router, and proxy calls stay on Python and do not construct native cache resources. Tests can inject a rule and build the native memory runtime from an ordinary Python `Cache` configuration without changing the legacy cache classes
+When a rule selects a backend, the Python `Cache` facade builds the native runtime from its own configuration and routes its storage calls (sync and async lookup and store, and pipelined batch store) to it. Stream replay, embedding partial-hit merging, response reconstruction, and callbacks stay in Python on top of that native store. The Python backend object remains for its direct API
 
 Object responses are written as they are, and every other response shape is written as a serialized string, which is the pair of shapes Python reads. A string on the wire is therefore always a serialized response, so string-valued responses round trip. Typed backends such as memory never pass through the codec
 
-The resolver reads the namespace's `cache` attribute each time it resolves. A captured binding retains the selected service for its operation, including background writes. `None` disables caching. Custom Python cache objects keep their original methods, arguments, returned awaitables, exceptions, and caller-task execution
-
-Python callbacks use the built-in `Cache` API, so a `Cache` subclass works unchanged. A batch lookup takes one original kwargs mapping per request and returns the list of `get_cache` or gathered `async_get_cache` results, while native bindings return `{values, missing_indices}`. A batch store hands the caller's original result to `async_add_cache_pipeline`. `ping` calls `ping`, and a flush goes to the facade's backend
-
-The private facade test harness checks object identity, method overrides, effective TTL, Redis namespace, memory capacity, and later configuration changes before selecting native execution. Its snapshot includes Redis connection settings, so a later `redis_kwargs` change, including an SSL option, selects Python callback execution. Buffered async writes honor `redis_flush_size`. Public activation must construct the shared native service from the initial Python Redis settings, including `litellm.default_redis_ttl` and SSL options. A buffered entry keeps the time it was produced, and a failed flush drops its batch instead of growing the buffer during an outage. The harness does not migrate entries or replace Python methods. Until activation configures one shared service, the Python facade and native test service can hold separate data. Existing public cache constructors remain on Python
-
-Native cache handles must be recreated after fork. The bridge releases the GIL around native operations, and Redis runs blocking connection operations off the async executor. Native errors propagate to the host, which owns the existing fail-open and logging policy
-
-The Redis backend also provides the primitives needed to preserve its direct Python surface later: TLS URLs, ping, bulk delete, counter batches, TTL, scan, set membership, raw queue push and pop, queue and counter pipelines, counter floor and maximum operations, script evaluation, client information, namespaced flush, and full flush. These are backend operations only and are not exported to Python by this PR. Memory provides TTL, oldest-key, and counter-pipeline operations
+Native cache handles must be recreated after fork. Native errors propagate to the host, which owns the existing fail-open and logging policy
 
 ## Adding another backend
 
-Implement `BaseCache` for the backend with its associated value type, and accept a `CacheCodec` when wire serialization is needed. `ResponseCache<B>` then works without another response implementation. Add a concrete bridge enum variant and constructor only when exposing that backend to Python
+Implement `BaseCache` for the backend with its associated value type and the capability traits its Python class supports, and accept a `CacheCodec` when wire serialization is needed. `ResponseCache<B>` then works without another response implementation
 
-Verify typed values, TTL precedence, missing entries, serialization failures, namespaces, batch ordering, and sync/async behavior. Run response fixtures with `ResponseCacheCodec`, including both Python envelope encodings, before enabling a public facade
-
-## Follow-up scope
-
-Public SDK, Router, and proxy activation still need constructor parity, stream replay, embedding partial-batch integration, response reconstruction, callback scheduling, and failure-policy integration. This foundation does not switch those request paths
-
-Redis cluster, disk, and cloud stores remain follow-ups. Semantic backends plug in through `SemanticCacheContext`, which carries the prompt inputs and metadata alongside the cache TTL. The generic dual cache takes read, write, and remote-failure policies, runs its async operations through the async L2 methods, and provides L2-first counters and atomic affinity claims. Errors propagate by default, and `RemoteFailurePolicy::UseLocal` opts key-value operations and claims into the local tier when L2 is unavailable. Claims compare decoded values, so a pin written by Python still matches. Public Router integration remains follow-up work. Reservations and pubsub still need explicit capabilities owned by their consuming features. Adding a cache backend does not establish those guarantees
+Run the `litellm-cache-testing` contract checks the backend's capabilities allow, and run response fixtures with `ResponseCacheCodec`, including both Python envelope encodings, before adding a catalog rule

litellm-rust/crates/cache-response/src/exact.rs

@@ -1,7 +1,8 @@
 use std::{future::Future, pin::Pin, time::Duration};
 
 use litellm_cache::{
-    BaseCache, BatchCache, CacheConnectionResult, Error, ExactCacheContext, FlushCache,
+    BaseCache, BatchCache, CacheConnectionResult, ConnectionCache, Error, ExactCacheContext,
+    FlushCache,
 };
 use serde_json::Value;
 
@@ -61,10 +62,25 @@ pub trait ExactResponseCache: Send + Sync {
     ) -> BoxFuture<'a, Result<(), Error>>;
 
     fn async_flush<'a>(&'a self) -> BoxFuture<'a, Result<(), Error>>;
+}
 
+/// Object-safe `test_connection` for the exact backends whose Python class defines it. Hosts hold
+/// one next to their `ExactResponseCache` when the backend has it, and report the operation as
+/// unsupported otherwise, as Python's `BaseCache.test_connection` does.
+pub trait ConnectionProbe: Send + Sync {
     fn test_connection<'a>(&'a self) -> BoxFuture<'a, Result<CacheConnectionResult, Error>>;
 }
 
+impl<B> ConnectionProbe for ResponseCache<B>
+where
+    B: ConnectionCache<Value = CacheEntry>,
+    B::Context: Default + PartialEq,
+{
+    fn test_connection<'a>(&'a self) -> BoxFuture<'a, Result<CacheConnectionResult, Error>> {
+        Box::pin(ResponseCache::test_connection(self))
+    }
+}
+
 impl<B> ExactResponseCache for ResponseCache<B>
 where
     B: BaseCache<Value = CacheEntry, Context = ExactCacheContext> + BatchCache + FlushCache,
@@ -141,8 +157,4 @@ where
     fn async_flush<'a>(&'a self) -> BoxFuture<'a, Result<(), Error>> {
         Box::pin(ResponseCache::async_flush(self))
     }
-
-    fn test_connection<'a>(&'a self) -> BoxFuture<'a, Result<CacheConnectionResult, Error>> {
-        Box::pin(ResponseCache::test_connection(self))
-    }
 }

litellm-rust/crates/cache-response/src/lib.rs

@@ -12,5 +12,5 @@ pub use caching::{
 };
 pub use codec::ResponseCacheCodec;
 pub use embedding::PartialHits;
-pub use exact::ExactResponseCache;
+pub use exact::{ConnectionProbe, ExactResponseCache};
 pub use response::{ResponseCache, ResponseCacheRequest};

litellm-rust/crates/cache-response/src/response.rs

@@ -1,7 +1,9 @@
 use std::{sync::Arc, time::Duration};
 
 use litellm_cache::{
-    BaseCache, BatchCache, BatchEntry, CacheConnectionResult, CacheContext, Error, FlushCache,
+    BaseCache, BatchCache, BatchEntry, CacheConnectionResult, CacheContext, ConnectionCache, Error,
+    FlushCache,
+    semantic::{SemanticCache, SemanticLookup},
 };
 use serde_json::Value;
 
@@ -78,7 +80,10 @@ where
         self.backend.async_flush_cache().await
     }
 
-    pub async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
+    pub async fn test_connection(&self) -> Result<CacheConnectionResult, Error>
+    where
+        B: ConnectionCache,
+    {
         self.backend.test_connection().await
     }
 
@@ -121,6 +126,43 @@ where
         Ok(Self::fresh_or_miss(entry, now, request.max_age))
     }
 
+    /// `lookup` plus the similarity the semantic backend reports. Freshness applies to the
+    /// value only: Python stamps the similarity before its max-age check.
+    pub fn lookup_semantic(
+        &self,
+        request: &ResponseCacheRequest<B::Context>,
+        now: Duration,
+    ) -> Result<SemanticLookup<Value>, Error>
+    where
+        B: SemanticCache,
+    {
+        if !request.controls.reads() {
+            return Ok(SemanticLookup::miss(None));
+        }
+        let lookup = self
+            .backend
+            .get_cache_with_similarity(&cache_key(&request.key), &request.context);
+        Self::fresh_semantic(lookup, now, request.max_age)
+    }
+
+    pub async fn async_lookup_semantic(
+        &self,
+        request: &ResponseCacheRequest<B::Context>,
+        now: Duration,
+    ) -> Result<SemanticLookup<Value>, Error>
+    where
+        B: SemanticCache,
+    {
+        if !request.controls.reads() {
+            return Ok(SemanticLookup::miss(None));
+        }
+        let lookup = self
+            .backend
+            .async_get_cache_with_similarity(&cache_key(&request.key), &request.context)
+            .await;
+        Self::fresh_semantic(lookup, now, request.max_age)
+    }
+
     pub fn lookup_batch(
         &self,
         requests: &[ResponseCacheRequest<B::Context>],
@@ -290,6 +332,21 @@ where
         Ok(PartialHits::new(values))
     }
 
+    fn fresh_semantic(
+        lookup: Result<SemanticLookup<CacheEntry>, Error>,
+        now: Duration,
+        max_age: Option<Duration>,
+    ) -> Result<SemanticLookup<Value>, Error> {
+        match lookup {
+            Ok(lookup) => Ok(SemanticLookup {
+                value: Self::fresh_or_miss(lookup.value, now, max_age),
+                similarity: lookup.similarity,
+            }),
+            Err(Error::InvalidEntry) => Ok(SemanticLookup::miss(None)),
+            Err(error) => Err(error),
+        }
+    }
+
     fn fresh_or_miss(
         entry: Option<CacheEntry>,
         now: Duration,

litellm-rust/crates/cache-response/tests/caching.rs

@@ -1,90 +1,174 @@
 use litellm_cache_response::{
     CacheControls, CacheKeyContext, CacheKeyField, CacheKeyInput, cache_key, get_cache_key,
+    should_use_cache,
 };
+use rstest::rstest;
 use sha2::{Digest, Sha256};
 
-#[test]
-fn keys_match_python_order_groups_files_presets_and_namespaces() {
-    let mut input = CacheKeyInput {
-        fields: vec![
-            CacheKeyField {
-                name: "model".into(),
-                value: Some("deployment".into()),
-                api_parameter: true,
-                internal_parameter: false,
-            },
-            CacheKeyField {
-                name: "file".into(),
-                value: None,
-                api_parameter: true,
-                internal_parameter: false,
-            },
-        ],
-        namespace: Some("team".into()),
-        ..Default::default()
-    };
+fn field(name: &str, value: Option<&str>) -> CacheKeyField {
+    CacheKeyField {
+        name: name.into(),
+        value: value.map(str::to_owned),
+        api_parameter: true,
+        internal_parameter: false,
+    }
+}
+
+fn hash(preimage: &[u8]) -> String {
+    format!("{:x}", Sha256::digest(preimage))
+}
+
+#[rstest]
+#[case::caching_group_and_checksum(
     CacheKeyContext {
         model_group: Some("group".into()),
         caching_groups: vec![(vec!["group".into()], "['group']".into())],
         file_checksum: Some("checksum".into()),
         ..Default::default()
-    }
-    .apply(&mut input);
-    assert_eq!(
-        cache_key(&input),
-        format!(
-            "team:{:x}",
-            Sha256::digest(b"model: ['group']file: checksum")
-        )
-    );
-    input.preset = Some("preset".into());
+    },
+    Some("team"),
+    "team:",
+    b"model: ['group']file: checksum".as_slice(),
+)]
+#[case::model_group_outside_caching_groups(
+    CacheKeyContext {
+        model_group: Some("group".into()),
+        caching_groups: vec![(vec!["other".into()], "['other']".into())],
+        file_object_name: Some("object".into()),
+        ..Default::default()
+    },
+    None,
+    "",
+    b"model: groupfile: object".as_slice(),
+)]
+#[case::metadata_file_name_before_parameters(
+    CacheKeyContext {
+        metadata_file_name: Some("metadata".into()),
+        parameters_file_name: Some("parameters".into()),
+        ..Default::default()
+    },
+    Some(""),
+    "",
+    b"model: deploymentfile: metadata".as_slice(),
+)]
+#[case::parameters_file_name_last(
+    CacheKeyContext {
+        parameters_file_name: Some("parameters".into()),
+        ..Default::default()
+    },
+    None,
+    "",
+    b"model: deploymentfile: parameters".as_slice(),
+)]
+#[case::no_context_keeps_the_request_model(
+    CacheKeyContext::default(),
+    Some("team"),
+    "team:",
+    b"model: deployment".as_slice(),
+)]
+fn keys_match_python_order_groups_files_and_namespaces(
+    #[case] context: CacheKeyContext,
+    #[case] namespace: Option<&str>,
+    #[case] prefix: &str,
+    #[case] preimage: &[u8],
+) {
+    let mut input = CacheKeyInput {
+        fields: vec![field("model", Some("deployment")), field("file", None)],
+        namespace: namespace.map(str::to_owned),
+        ..Default::default()
+    };
+    context.apply(&mut input);
+    let expected = format!("{prefix}{}", hash(preimage));
+    assert_eq!(cache_key(&input), expected);
+    assert_eq!(get_cache_key(&input), expected);
+}
+
+#[rstest]
+#[case::api_parameter(true, false, false, true)]
+#[case::provider_parameter_when_included(false, false, true, true)]
+#[case::provider_parameter_when_excluded(false, false, false, false)]
+#[case::internal_parameter_never(false, true, true, false)]
+fn keys_hash_api_and_opted_in_provider_parameters(
+    #[case] api_parameter: bool,
+    #[case] internal_parameter: bool,
+    #[case] include_provider_parameters: bool,
+    #[case] hashed: bool,
+) {
+    let input = CacheKeyInput {
+        fields: vec![
+            field("model", Some("a")),
+            CacheKeyField {
+                name: "extra".into(),
+                value: Some("x".into()),
+                api_parameter,
+                internal_parameter,
+            },
+        ],
+        include_provider_parameters,
+        ..Default::default()
+    };
+    let preimage: &[u8] = if hashed {
+        b"model: aextra: x"
+    } else {
+        b"model: a"
+    };
+    assert_eq!(cache_key(&input), hash(preimage));
+}
+
+#[rstest]
+#[case::without_namespace(None)]
+#[case::with_namespace(Some("team"))]
+fn preset_keys_are_used_verbatim(#[case] namespace: Option<&str>) {
+    let input = CacheKeyInput {
+        fields: vec![field("model", Some("a"))],
+        preset: Some("preset".into()),
+        namespace: namespace.map(str::to_owned),
+        ..Default::default()
+    };
+    assert_eq!(cache_key(&input), "preset");
     assert_eq!(get_cache_key(&input), "preset");
 }
 
-#[test]
-fn cache_controls_honor_default_modes_and_directives() {
-    let enabled = CacheControls {
-        supported_call_type: true,
-        configured: true,
-        default_on: true,
-        ..Default::default()
-    };
-    assert!(enabled.reads());
-    assert!(enabled.writes());
-    assert!(
-        !CacheControls {
-            default_on: false,
-            ..enabled
-        }
-        .reads()
-    );
-    assert!(
-        CacheControls {
-            default_on: false,
-            use_cache: true,
-            ..enabled
-        }
-        .reads()
-    );
-    assert!(
-        !CacheControls {
-            no_cache: true,
-            ..enabled
-        }
-        .reads()
-    );
-    assert!(
-        !CacheControls {
-            no_store: true,
-            ..enabled
-        }
-        .writes()
-    );
-    assert!(
-        !CacheControls {
-            caching: Some(false),
-            ..enabled
-        }
-        .writes()
-    );
+const ENABLED: CacheControls = CacheControls {
+    supported_call_type: true,
+    configured: true,
+    native_backend: false,
+    default_on: true,
+    caching: None,
+    no_cache: false,
+    no_store: false,
+    use_cache: false,
+};
+
+#[rstest]
+#[case::enabled(ENABLED, true, true)]
+#[case::default_off(CacheControls { default_on: false, ..ENABLED }, false, false)]
+#[case::default_off_with_use_cache(
+    CacheControls { default_on: false, use_cache: true, ..ENABLED },
+    true,
+    true
+)]
+#[case::no_cache(CacheControls { no_cache: true, ..ENABLED }, false, true)]
+#[case::no_store(CacheControls { no_store: true, ..ENABLED }, true, false)]
+#[case::no_cache_and_no_store(
+    CacheControls { no_cache: true, no_store: true, ..ENABLED },
+    false,
+    false
+)]
+#[case::caching_disabled(CacheControls { caching: Some(false), ..ENABLED }, false, false)]
+#[case::caching_enabled(CacheControls { caching: Some(true), ..ENABLED }, true, true)]
+#[case::unsupported_call_type(
+    CacheControls { supported_call_type: false, ..ENABLED },
+    false,
+    false
+)]
+#[case::unconfigured(CacheControls { configured: false, ..ENABLED }, false, false)]
+fn cache_controls_honor_default_modes_and_directives(
+    #[case] controls: CacheControls,
+    #[case] reads: bool,
+    #[case] writes: bool,
+) {
+    assert_eq!(controls.reads(), reads);
+    assert_eq!(controls.writes(), writes);
+    assert_eq!(should_use_cache(controls), reads || writes);
 }

litellm-rust/crates/cache-response/tests/codec.rs

@@ -0,0 +1,110 @@
+use litellm_cache::{CacheCodec, Error};
+use litellm_cache_response::{CacheEntry, ResponseCacheCodec};
+use rstest::rstest;
+use serde_json::{Value, json};
+
+fn entry(response: Value) -> CacheEntry {
+    CacheEntry {
+        timestamp: Some(100.0),
+        response,
+    }
+}
+
+#[rstest]
+#[case::python_literal_object(
+    br#"{'timestamp': 100.0, 'response': {'text': 'hello \\ world', 'flag': True, 'empty': None, 'list': [1, 2.5]}}"#.as_slice(),
+    json!({"text": "hello \\ world", "flag": true, "empty": null, "list": [1, 2.5]}),
+)]
+#[case::python_sync_string_response(
+    br#"{'timestamp': 100.0, 'response': '{"ok": true, "text": "cached"}'}"#.as_slice(),
+    json!({"ok": true, "text": "cached"}),
+)]
+#[case::python_literal_string_response(
+    br#"{'timestamp': 100.0, 'response': "{'ok': True, 'items': (1, 2)}"}"#.as_slice(),
+    json!({"ok": true, "items": [1, 2]}),
+)]
+#[case::json_object_response(
+    br#"{"timestamp":100.0,"response":{"ok":true,"text":"cached"}}"#.as_slice(),
+    json!({"ok": true, "text": "cached"}),
+)]
+#[case::json_string_response(
+    br#"{"timestamp": 100.0, "response": "[1,2]"}"#.as_slice(),
+    json!([1, 2]),
+)]
+fn decode_reads_every_python_envelope(#[case] bytes: &[u8], #[case] response: Value) {
+    assert_eq!(ResponseCacheCodec.decode(bytes).unwrap(), entry(response));
+}
+
+#[rstest]
+#[case::code_is_not_executed(b"__import__('os').system('false')".to_vec())]
+#[case::non_numeric_timestamp(b"{'timestamp': 'invalid', 'response': {}}".to_vec())]
+#[case::infinite_timestamp(b"{'timestamp': 1e9999, 'response': {}}".to_vec())]
+#[case::missing_response(br#"{"timestamp": 100.0}"#.to_vec())]
+#[case::unserialized_string_response(br#"{"timestamp": 100.0, "response": "not serialized"}"#.to_vec())]
+#[case::non_utf8(vec![0xff, 0xfe])]
+#[case::deep_nesting(format!("{}None{}", "[".repeat(1000), "]".repeat(1000)).into_bytes())]
+fn decode_rejects_invalid_entries(#[case] bytes: Vec<u8>) {
+    assert_eq!(
+        ResponseCacheCodec.decode(&bytes).unwrap_err(),
+        Error::InvalidEntry
+    );
+}
+
+#[rstest]
+#[case::nan(f64::NAN)]
+#[case::infinity(f64::INFINITY)]
+fn encode_rejects_non_finite_timestamps(#[case] timestamp: f64) {
+    assert_eq!(
+        ResponseCacheCodec
+            .encode(&CacheEntry {
+                timestamp: Some(timestamp),
+                response: json!({}),
+            })
+            .unwrap_err(),
+        Error::InvalidEntry
+    );
+}
+
+#[rstest]
+#[case::object(json!({"choices": [{"text": "cached"}]}), json!({"choices": [{"text": "cached"}]}))]
+#[case::array(json!([1, 2]), json!("[1,2]"))]
+#[case::number(json!(7), json!("7"))]
+#[case::null(json!(null), json!("null"))]
+#[case::string(json!("hello world"), json!("\"hello world\""))]
+#[case::numeric_string(json!("123"), json!("\"123\""))]
+#[case::null_string(json!("null"), json!("\"null\""))]
+fn encode_writes_python_readable_envelopes_that_round_trip(
+    #[case] response: Value,
+    #[case] wire_response: Value,
+) {
+    let wire = ResponseCacheCodec.encode(&entry(response.clone())).unwrap();
+    assert_eq!(
+        serde_json::from_slice::<Value>(&wire).unwrap(),
+        json!({"timestamp": 100.0, "response": wire_response})
+    );
+    assert_eq!(ResponseCacheCodec.decode(&wire).unwrap(), entry(response));
+}
+
+#[rstest]
+fn object_entries_preserve_the_existing_json_representation() {
+    let entry = CacheEntry {
+        timestamp: Some(123.0),
+        response: json!({"choices": [{"text": "cached"}]}),
+    };
+    let bytes = ResponseCacheCodec.encode(&entry).unwrap();
+    assert_eq!(bytes, serde_json::to_vec(&entry).unwrap());
+    assert_eq!(ResponseCacheCodec.decode(&bytes).unwrap(), entry);
+}
+
+#[rstest]
+#[case::json(br#"{"choices": [{"text": "legacy"}]}"#.as_slice())]
+#[case::python_literal(br#"{'choices': [{'text': 'legacy'}]}"#.as_slice())]
+fn values_without_timestamps_decode_as_bare_responses(#[case] bytes: &[u8]) {
+    assert_eq!(
+        ResponseCacheCodec.decode(bytes).unwrap(),
+        CacheEntry {
+            timestamp: None,
+            response: json!({"choices": [{"text": "legacy"}]}),
+        }
+    );
+}

litellm-rust/crates/cache-response/tests/connection.rs

@@ -0,0 +1,123 @@
+mod support;
+
+use std::{sync::Arc, time::Duration};
+
+use litellm_cache::CacheConnectionStatus;
+use litellm_cache_memory::InMemoryCache;
+use litellm_cache_response::{
+    CacheEntry, ConnectionProbe, ExactResponseCache, ResponseCache, ResponseCacheRequest,
+};
+use redis_test::MockCmd;
+use rstest::rstest;
+use serde_json::json;
+use support::{keyed, memory, redis, request};
+
+#[rstest]
+#[case::reachable(
+    Ok("PONG"),
+    CacheConnectionStatus::Success,
+    "Redis connection test successful",
+    false
+)]
+#[case::unexpected_reply(
+    Ok("NOPE"),
+    CacheConnectionStatus::Failed,
+    "Redis ping returned False",
+    false
+)]
+#[case::connection_refused(
+    Err(redis::RedisError::from((redis::ErrorKind::Io, "connection refused"))),
+    CacheConnectionStatus::Failed,
+    "Redis connection failed:",
+    true
+)]
+#[tokio::test]
+async fn connection_backends_are_reachable_as_a_probe(
+    #[case] reply: redis::RedisResult<&'static str>,
+    #[case] status: CacheConnectionStatus,
+    #[case] message: &str,
+    #[case] has_error: bool,
+) {
+    let probe: Arc<dyn ConnectionProbe> =
+        Arc::new(redis(vec![MockCmd::new(redis::cmd("PING"), reply)], None));
+
+    let result = probe.test_connection().await.unwrap();
+    assert_eq!(result.status, status);
+    assert!(result.message.starts_with(message), "{}", result.message);
+    assert_eq!(result.error.is_some(), has_error);
+}
+
+#[rstest]
+#[tokio::test]
+async fn one_service_serves_both_the_exact_cache_and_its_probe(request: ResponseCacheRequest) {
+    let service = Arc::new(redis(
+        vec![
+            MockCmd::new(redis::cmd("PING"), Ok("PONG")),
+            MockCmd::new(
+                redis::cmd("GET").arg("tenant:key"),
+                Ok(br#"{"timestamp":100.0,"response":{"ok":true}}"#.to_vec()),
+            ),
+        ],
+        Some("tenant"),
+    ));
+    let probe: Arc<dyn ConnectionProbe> = service.clone();
+    let exact: Arc<dyn ExactResponseCache> = service;
+
+    assert_eq!(
+        probe.test_connection().await.unwrap().status,
+        CacheConnectionStatus::Success
+    );
+    assert_eq!(
+        exact
+            .async_lookup(&request, Duration::from_secs(100))
+            .await
+            .unwrap(),
+        Some(json!({"ok": true}))
+    );
+}
+
+/// The in-memory backend has no `test_connection`, as in Python, and still serves every response
+/// operation.
+#[rstest]
+#[tokio::test]
+async fn backends_without_a_connection_test_serve_every_response_operation(
+    #[from(memory)] service: Arc<ResponseCache<InMemoryCache<CacheEntry>>>,
+    request: ResponseCacheRequest,
+) {
+    let cache: Arc<dyn ExactResponseCache> = service;
+    let now = Duration::from_secs(100);
+    let other = keyed("tenant:other");
+    let missing = keyed("tenant:missing");
+
+    assert_eq!(cache.default_ttl(), Some(Duration::from_secs(600)));
+    cache.store(&request, json!({"v": 1}), now).unwrap();
+    assert_eq!(cache.lookup(&request, now).unwrap(), Some(json!({"v": 1})));
+    cache
+        .async_store(&other, json!({"v": 2}), now)
+        .await
+        .unwrap();
+    assert_eq!(
+        cache.async_lookup(&other, now).await.unwrap(),
+        Some(json!({"v": 2}))
+    );
+
+    let requests = [request.clone(), missing.clone(), other.clone()];
+    let partial = cache.lookup_batch(&requests, now).unwrap();
+    assert_eq!(
+        partial.values,
+        vec![Some(json!({"v": 1})), None, Some(json!({"v": 2}))]
+    );
+    assert_eq!(partial.missing_indices, vec![1]);
+
+    cache
+        .async_store_batch(vec![(missing.clone(), json!({"v": 3}))], now)
+        .await
+        .unwrap();
+    let partial = cache.async_lookup_batch(&requests, now).await.unwrap();
+    assert!(partial.missing_indices.is_empty());
+    assert_eq!(partial.values[1], Some(json!({"v": 3})));
+
+    cache.async_flush().await.unwrap();
+    let partial = cache.async_lookup_batch(&requests, now).await.unwrap();
+    assert_eq!(partial.missing_indices, vec![0, 1, 2]);
+}

litellm-rust/crates/cache-response/tests/response.rs

@@ -1,3 +1,5 @@
+mod support;
+
 use std::{
     sync::{
         Arc, Mutex,
@@ -7,32 +9,22 @@ use std::{
 };
 
 use litellm_cache::{
-    BaseCache, CacheCodec, CacheConnectionResult, CacheConnectionStatus, Error,
-    SemanticCacheContext,
+    BaseCache, Error, SemanticCacheContext,
+    semantic::{SemanticCache, SemanticLookup},
 };
 use litellm_cache_memory::InMemoryCache;
-use litellm_cache_redis::RedisCache;
 use litellm_cache_response::{
-    CacheEntry, CacheKeyField, CacheKeyInput, ResponseCache, ResponseCacheCodec,
-    ResponseCacheRequest, WriteBuffer,
+    CacheControls, CacheEntry, CacheKeyField, CacheKeyInput, ResponseCache, ResponseCacheRequest,
+    WriteBuffer, cache_key,
 };
-use redis_test::{MockCmd, MockRedisConnection};
-use serde_json::json;
+use redis_test::MockCmd;
+use rstest::rstest;
+use serde_json::{Value, json};
+use support::{keyed, memory, redis, request};
 
-fn memory() -> Arc<ResponseCache<InMemoryCache<CacheEntry>>> {
-    Arc::new(ResponseCache::new(Arc::new(InMemoryCache::new(
-        Some(8),
-        Some(Duration::from_secs(600)),
-    ))))
-}
-
-fn request() -> ResponseCacheRequest {
-    ResponseCacheRequest::new(CacheKeyInput {
-        preset: Some("tenant:key".into()),
-        ..Default::default()
-    })
-}
+type Memory = Arc<ResponseCache<InMemoryCache<CacheEntry>>>;
 
+#[derive(Default)]
 struct SemanticBackend {
     entries: Mutex<Vec<(String, CacheEntry)>>,
     contexts: Mutex<Vec<SemanticCacheContext>>,
@@ -67,50 +59,133 @@ impl BaseCache for SemanticBackend {
             .find(|(entry_key, _)| entry_key == key)
             .map(|(_, entry)| entry.clone()))
     }
-
-    async fn disconnect(&self) -> Result<(), Error> {
-        Ok(())
-    }
-
-    async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
-        Ok(CacheConnectionResult {
-            status: CacheConnectionStatus::Success,
-            message: "ok".into(),
-            error: None,
-        })
-    }
 }
 
-#[test]
-fn semantic_context_reaches_backend_for_store_and_lookup() {
-    let backend = Arc::new(SemanticBackend {
-        entries: Mutex::new(Vec::new()),
-        contexts: Mutex::new(Vec::new()),
-    });
+#[rstest]
+#[tokio::test]
+async fn semantic_context_reaches_backend_for_store_and_lookup(
+    request: ResponseCacheRequest,
+    #[values(false, true)] asynchronous: bool,
+) {
+    let backend = Arc::new(SemanticBackend::default());
     let cache = ResponseCache::new(backend.clone());
     let context = SemanticCacheContext {
         messages: Some(json!([{"role": "user", "content": "hello"}])),
         ..Default::default()
     };
-    let request = request().with_context(context.clone());
+    let request = request.with_context(context.clone());
     let response = json!({"answer": 42});
+    let now = Duration::from_secs(100);
 
-    cache
-        .store(&request, response.clone(), Duration::from_secs(100))
-        .unwrap();
+    let hit = if asynchronous {
+        cache
+            .async_store(&request, response.clone(), now)
+            .await
+            .unwrap();
+        cache.async_lookup(&request, now).await.unwrap()
+    } else {
+        cache.store(&request, response.clone(), now).unwrap();
+        cache.lookup(&request, now).unwrap()
+    };
 
-    assert_eq!(
-        cache.lookup(&request, Duration::from_secs(100)).unwrap(),
-        Some(response)
-    );
+    assert_eq!(hit, Some(response));
     assert_eq!(
         backend.contexts.lock().unwrap().as_slice(),
         &[context.clone(), context]
     );
 }
 
+/// A semantic backend that answers every read with one fixed lookup.
+struct ScoredBackend(Result<SemanticLookup<CacheEntry>, Error>);
+
+impl BaseCache for ScoredBackend {
+    type Value = CacheEntry;
+    type Context = SemanticCacheContext;
+
+    fn get_ttl(&self, _: &Self::Context) -> Option<Duration> {
+        None
+    }
+
+    fn set_cache(&self, _: &str, _: Self::Value, _: &Self::Context) -> Result<(), Error> {
+        Ok(())
+    }
+
+    fn get_cache(&self, key: &str, context: &Self::Context) -> Result<Option<Self::Value>, Error> {
+        self.get_cache_with_similarity(key, context)
+            .map(|lookup| lookup.value)
+    }
+}
+
+impl SemanticCache for ScoredBackend {
+    fn get_cache_with_similarity(
+        &self,
+        _: &str,
+        _: &Self::Context,
+    ) -> Result<SemanticLookup<Self::Value>, Error> {
+        self.0.clone()
+    }
+
+    async fn async_get_cache_with_similarity(
+        &self,
+        key: &str,
+        context: &Self::Context,
+    ) -> Result<SemanticLookup<Self::Value>, Error> {
+        self.get_cache_with_similarity(key, context)
+    }
+}
+
+fn scored(timestamp: f64, similarity: f64) -> Result<SemanticLookup<CacheEntry>, Error> {
+    Ok(SemanticLookup {
+        value: Some(CacheEntry {
+            timestamp: Some(timestamp),
+            response: json!({"answer": 42}),
+        }),
+        similarity: Some(similarity),
+    })
+}
+
+#[rstest]
+#[case::fresh_hit(scored(95.0, 0.95), true, Ok(SemanticLookup { value: Some(json!({"answer": 42})), similarity: Some(0.95) }))]
+#[case::stale_hit_keeps_the_similarity(
+    scored(50.0, 0.95),
+    true,
+    Ok(SemanticLookup::miss(Some(0.95)))
+)]
+#[case::miss_keeps_the_similarity(
+    Ok(SemanticLookup::miss(Some(0.4))),
+    true,
+    Ok(SemanticLookup::miss(Some(0.4)))
+)]
+#[case::no_search(Ok(SemanticLookup::miss(None)), true, Ok(SemanticLookup::miss(None)))]
+#[case::disabled_reads_skip_the_backend(scored(95.0, 0.95), false, Ok(SemanticLookup::miss(None)))]
+#[case::invalid_entry_is_a_miss(Err(Error::InvalidEntry), true, Ok(SemanticLookup::miss(None)))]
+#[case::backend_errors_propagate(Err(Error::Unavailable), true, Err(Error::Unavailable))]
 #[tokio::test]
-async fn sync_and_async_consumers_share_keys_ttls_and_freshness() {
+async fn semantic_lookup_applies_freshness_to_the_value_only(
+    #[case] backend: Result<SemanticLookup<CacheEntry>, Error>,
+    #[case] reads: bool,
+    #[case] expected: Result<SemanticLookup<Value>, Error>,
+    #[values(false, true)] asynchronous: bool,
+    request: ResponseCacheRequest,
+) {
+    let cache = ResponseCache::new(Arc::new(ScoredBackend(backend)));
+    let mut request = request.with_context(SemanticCacheContext::default());
+    request.max_age = Some(Duration::from_secs(10));
+    request.controls.no_cache = !reads;
+    let now = Duration::from_secs(100);
+
+    let lookup = if asynchronous {
+        cache.async_lookup_semantic(&request, now).await
+    } else {
+        cache.lookup_semantic(&request, now)
+    };
+
+    assert_eq!(lookup, expected);
+}
+
+#[rstest]
+#[tokio::test]
+async fn sync_and_async_consumers_share_keys_ttls_and_freshness(mut request: ResponseCacheRequest) {
     let clock = Arc::new(AtomicU64::new(100));
     let backend = Arc::new(InMemoryCache::with_clock(
         Some(8),
@@ -121,7 +196,6 @@ async fn sync_and_async_consumers_share_keys_ttls_and_freshness() {
         },
     ));
     let cache = ResponseCache::new(backend.clone());
-    let mut request = request();
     request.context.ttl = Some(Duration::from_secs(10));
     request.max_age = Some(Duration::from_secs(5));
     cache
@@ -172,80 +246,164 @@ async fn sync_and_async_consumers_share_keys_ttls_and_freshness() {
     );
 }
 
+#[derive(Clone, Copy, Debug)]
+enum Directive {
+    Plain,
+    NoCache,
+    NoStore,
+    DefaultOff,
+    UseCache,
+    CachingOff,
+    Unsupported,
+}
+
+impl Directive {
+    fn apply(self, controls: &mut CacheControls) {
+        match self {
+            Self::Plain => {}
+            Self::NoCache => controls.no_cache = true,
+            Self::NoStore => controls.no_store = true,
+            Self::DefaultOff => controls.default_on = false,
+            Self::UseCache => {
+                controls.default_on = false;
+                controls.use_cache = true;
+            }
+            Self::CachingOff => controls.caching = Some(false),
+            Self::Unsupported => controls.supported_call_type = false,
+        }
+    }
+}
+
+#[rstest]
+#[case::plain(Directive::Plain, Directive::Plain, true)]
+#[case::no_store_skips_the_write(Directive::NoStore, Directive::Plain, false)]
+#[case::no_store_keeps_reads(Directive::Plain, Directive::NoStore, true)]
+#[case::no_cache_keeps_writes(Directive::NoCache, Directive::Plain, true)]
+#[case::no_cache_skips_the_read(Directive::Plain, Directive::NoCache, false)]
+#[case::default_off_skips_the_write(Directive::DefaultOff, Directive::Plain, false)]
+#[case::default_off_skips_the_read(Directive::Plain, Directive::DefaultOff, false)]
+#[case::use_cache_opts_in_under_default_off(Directive::UseCache, Directive::UseCache, true)]
+#[case::caching_off_skips_the_write(Directive::CachingOff, Directive::Plain, false)]
+#[case::caching_off_skips_the_read(Directive::Plain, Directive::CachingOff, false)]
+#[case::unsupported_call_type_skips_the_write(Directive::Unsupported, Directive::Plain, false)]
+#[case::unsupported_call_type_skips_the_read(Directive::Plain, Directive::Unsupported, false)]
 #[tokio::test]
-async fn directives_skip_io_and_keep_reads_and_writes_independent() {
-    let cache = memory();
-    let mut request = request();
+async fn directives_skip_io_and_keep_reads_and_writes_independent(
+    memory: Memory,
+    request: ResponseCacheRequest,
+    #[case] write: Directive,
+    #[case] read: Directive,
+    #[case] hit: bool,
+    #[values(false, true)] asynchronous: bool,
+) {
     let now = Duration::from_secs(100);
-    request.controls.no_store = true;
-    cache
-        .async_store(&request, json!({"v": 1}), now)
-        .await
-        .unwrap();
-    assert_eq!(cache.lookup(&request, now).unwrap(), None);
-    request.controls.no_store = false;
-    request.controls.no_cache = true;
-    cache.store(&request, json!({"v": 2}), now).unwrap();
-    assert_eq!(cache.async_lookup(&request, now).await.unwrap(), None);
-    request.controls.no_cache = false;
-    assert_eq!(cache.lookup(&request, now).unwrap(), Some(json!({"v": 2})));
-    request.controls.default_on = false;
-    cache.store(&request, json!({"v": 3}), now).unwrap();
-    assert_eq!(cache.lookup(&request, now).unwrap(), None);
-    request.controls.use_cache = true;
-    assert_eq!(cache.lookup(&request, now).unwrap(), Some(json!({"v": 2})));
-    request.controls.supported_call_type = false;
-    assert_eq!(cache.lookup(&request, now).unwrap(), None);
-}
+    let mut writer = request.clone();
+    write.apply(&mut writer.controls);
+    let mut reader = request;
+    read.apply(&mut reader.controls);
 
-#[tokio::test]
-async fn redis_consumer_reads_python_sync_and_async_envelopes_and_writes_compatible_json() {
-    let connection = MockRedisConnection::new([
-        MockCmd::new(
-            redis::cmd("GET").arg("tenant:key"),
-            Ok(br#"{'timestamp': 100.0, 'response': '{"ok": true, "text": "cached"}'}"#.to_vec()),
-        ),
-        MockCmd::new(
-            redis::cmd("GET").arg("tenant:key"),
-            Ok(br#"{"timestamp":100.0,"response":{"ok":true,"text":"cached"}}"#.to_vec()),
-        ),
-        MockCmd::new(
-            redis::cmd("SETEX")
-                .arg("tenant:key")
-                .arg(600)
-                .arg(br#"{"timestamp":100.0,"response":{"ok":true,"text":"cached"}}"#.as_slice()),
-            Ok("OK"),
-        ),
-    ])
-    .assert_all_commands_consumed();
-    let backend = RedisCache::with_connection(connection, None, ResponseCacheCodec)
-        .with_namespace(Some("tenant".into()));
-    let cache = ResponseCache::new(Arc::new(backend));
-    let request = request();
-    let expected = json!({"ok": true, "text": "cached"});
-    assert_eq!(
-        cache.lookup(&request, Duration::from_secs(101)).unwrap(),
-        Some(expected.clone())
-    );
-    assert_eq!(
-        cache
-            .async_lookup(&request, Duration::from_secs(101))
+    if asynchronous {
+        memory
+            .async_store(&writer, json!({"v": 1}), now)
             .await
-            .unwrap(),
-        Some(expected.clone())
+            .unwrap();
+    } else {
+        memory.store(&writer, json!({"v": 1}), now).unwrap();
+    }
+    let found = if asynchronous {
+        memory.async_lookup(&reader, now).await.unwrap()
+    } else {
+        memory.lookup(&reader, now).unwrap()
+    };
+
+    assert_eq!(
+        found,
+        hit.then(|| json!({"v": 1})),
+        "{write:?} then {read:?}"
+    );
+}
+
+#[rstest]
+#[case::python_sync_literal(
+    br#"{'timestamp': 100.0, 'response': '{"ok": true, "text": "cached"}'}"#.as_slice()
+)]
+#[case::python_async_json(br#"{"timestamp":100.0,"response":{"ok":true,"text":"cached"}}"#.as_slice())]
+#[tokio::test]
+async fn redis_consumer_reads_python_sync_and_async_envelopes(
+    request: ResponseCacheRequest,
+    #[case] stored: &[u8],
+    #[values(false, true)] asynchronous: bool,
+) {
+    let cache = redis(
+        vec![MockCmd::new(
+            redis::cmd("GET").arg("tenant:key"),
+            Ok(stored.to_vec()),
+        )],
+        Some("tenant"),
+    );
+    let now = Duration::from_secs(101);
+    let found = if asynchronous {
+        cache.async_lookup(&request, now).await.unwrap()
+    } else {
+        cache.lookup(&request, now).unwrap()
+    };
+    assert_eq!(found, Some(json!({"ok": true, "text": "cached"})));
+}
+
+#[rstest]
+#[case::object(
+    json!({"ok": true, "text": "cached"}),
+    br#"{"timestamp":100.0,"response":{"ok":true,"text":"cached"}}"#.as_slice()
+)]
+#[case::array(json!([1, 2]), br#"{"timestamp":100.0,"response":"[1,2]"}"#.as_slice())]
+#[tokio::test]
+async fn redis_consumer_writes_python_compatible_json(
+    request: ResponseCacheRequest,
+    #[case] response: Value,
+    #[case] wire: &[u8],
+) {
+    let cache = redis(
+        vec![MockCmd::new(
+            redis::cmd("SETEX").arg("tenant:key").arg(600).arg(wire),
+            Ok("OK"),
+        )],
+        Some("tenant"),
     );
     cache
-        .async_store(&request, expected, Duration::from_secs(100))
+        .async_store(&request, response, Duration::from_secs(100))
         .await
         .unwrap();
 }
 
+#[rstest]
 #[tokio::test]
-async fn captured_service_keeps_the_selected_backend_for_background_writes() {
-    let original = memory();
+async fn invalid_entries_are_misses_and_disabled_reads_do_not_touch_redis(
+    mut request: ResponseCacheRequest,
+) {
+    let cache = redis(
+        vec![MockCmd::new(
+            redis::cmd("GET").arg("tenant:key"),
+            Ok(b"invalid".to_vec()),
+        )],
+        None,
+    );
+    request.controls.no_cache = true;
+    assert_eq!(cache.lookup(&request, Duration::ZERO).unwrap(), None);
+    request.controls.no_cache = false;
+    assert_eq!(
+        cache.async_lookup(&request, Duration::ZERO).await.unwrap(),
+        None
+    );
+}
+
+#[rstest]
+#[tokio::test]
+async fn captured_service_keeps_the_selected_backend_for_background_writes(
+    #[from(memory)] original: Memory,
+    #[from(memory)] replacement: Memory,
+    request: ResponseCacheRequest,
+) {
     let captured = original.clone();
-    let replacement = memory();
-    let request = request();
     let writer = tokio::spawn({
         let request = request.clone();
         async move {
@@ -271,9 +429,13 @@ async fn captured_service_keeps_the_selected_backend_for_background_writes() {
     );
 }
 
-#[test]
-fn generated_keys_preserve_namespace_and_explicit_keys() {
-    let cache = memory();
+#[rstest]
+#[case::with_namespace(Some("tenant"))]
+#[case::without_namespace(None)]
+fn generated_keys_preserve_namespace_and_explicit_keys(
+    memory: Memory,
+    #[case] namespace: Option<&str>,
+) {
     let key = CacheKeyInput {
         fields: vec![CacheKeyField {
             name: "model".into(),
@@ -281,201 +443,125 @@ fn generated_keys_preserve_namespace_and_explicit_keys() {
             api_parameter: true,
             internal_parameter: false,
         }],
-        namespace: Some("tenant".into()),
+        namespace: namespace.map(str::to_owned),
         ..Default::default()
     };
     let generated = ResponseCacheRequest::new(key.clone());
-    let explicit = ResponseCacheRequest::new(CacheKeyInput {
-        preset: Some(litellm_cache_response::cache_key(&key)),
-        ..Default::default()
-    });
-    cache
+    let explicit = keyed(&cache_key(&key));
+    memory
         .store(&generated, json!({"value": 7}), Duration::from_secs(100))
         .unwrap();
     assert_eq!(
-        cache.lookup(&explicit, Duration::from_secs(100)).unwrap(),
+        memory.lookup(&explicit, Duration::from_secs(100)).unwrap(),
         Some(json!({"value":7}))
     );
 }
 
-#[test]
-fn response_codec_accepts_python_literals_without_executing_code() {
-    let bytes = br#"{'timestamp': 100.0, 'response': {'text': 'hello \\ world', 'flag': True, 'empty': None, 'list': [1, 2.5]}}"#;
-    let entry = ResponseCacheCodec.decode(bytes).unwrap();
-    assert_eq!(
-        entry.response,
-        json!({"text": "hello \\ world", "flag": true, "empty": null, "list": [1, 2.5]})
-    );
-    for bytes in [
-        b"__import__('os').system('false')".as_slice(),
-        b"{'timestamp': 'invalid', 'response': {}}",
-        b"{'timestamp': 1e9999, 'response': {}}",
-    ] {
-        assert_eq!(
-            ResponseCacheCodec.decode(bytes).unwrap_err(),
-            Error::InvalidEntry
-        );
-    }
-    let deep = format!("{}None{}", "[".repeat(1000), "]".repeat(1000));
-    assert_eq!(
-        ResponseCacheCodec.decode(deep.as_bytes()).unwrap_err(),
-        Error::InvalidEntry
-    );
-    assert_eq!(
-        ResponseCacheCodec
-            .encode(&CacheEntry {
-                timestamp: Some(f64::NAN),
-                response: json!({})
-            })
-            .unwrap_err(),
-        Error::InvalidEntry
-    );
-}
-
-#[tokio::test]
-async fn invalid_entries_are_misses_and_disabled_reads_do_not_touch_redis() {
-    let connection = MockRedisConnection::new([MockCmd::new(
-        redis::cmd("GET").arg("tenant:key"),
-        Ok(b"invalid".to_vec()),
-    )])
-    .assert_all_commands_consumed();
-    let backend = RedisCache::with_connection(connection, None, ResponseCacheCodec);
-    let cache = ResponseCache::new(Arc::new(backend));
-    let mut request = request();
-    request.controls.no_cache = true;
-    assert_eq!(cache.lookup(&request, Duration::ZERO).unwrap(), None);
-    request.controls.no_cache = false;
-    assert_eq!(
-        cache.async_lookup(&request, Duration::ZERO).await.unwrap(),
-        None
-    );
-}
-
-#[test]
-fn string_responses_round_trip_through_typed_and_wire_backends() {
-    let cache = ResponseCache::new(Arc::new(InMemoryCache::default()));
+#[rstest]
+#[case::text(json!("hello world"))]
+#[case::numeric_text(json!("123"))]
+#[case::null_text(json!("null"))]
+#[case::array(json!([1, 2]))]
+fn non_object_responses_round_trip_through_a_typed_backend(
+    memory: Memory,
+    request: ResponseCacheRequest,
+    #[case] response: Value,
+) {
     let now = Duration::from_secs(100);
-    for response in [json!("hello world"), json!("123"), json!("null")] {
-        cache.store(&request(), response.clone(), now).unwrap();
-        assert_eq!(
-            cache.lookup(&request(), now).unwrap(),
-            Some(response.clone())
-        );
-
-        let wire = ResponseCacheCodec
-            .encode(&CacheEntry {
-                timestamp: Some(100.0),
-                response: response.clone(),
-            })
-            .unwrap();
-        assert_eq!(ResponseCacheCodec.decode(&wire).unwrap().response, response);
-    }
+    memory.store(&request, response.clone(), now).unwrap();
+    assert_eq!(memory.lookup(&request, now).unwrap(), Some(response));
 }
 
-#[test]
-fn non_object_responses_are_written_as_python_readable_serialized_strings() {
-    let wire = ResponseCacheCodec
-        .encode(&CacheEntry {
-            timestamp: Some(100.0),
-            response: json!([1, 2]),
-        })
-        .unwrap();
-    assert_eq!(
-        serde_json::from_slice::<serde_json::Value>(&wire).unwrap(),
-        json!({"timestamp": 100.0, "response": "[1,2]"})
-    );
-    assert_eq!(
-        ResponseCacheCodec.decode(&wire).unwrap().response,
-        json!([1, 2])
-    );
-    assert_eq!(
-        ResponseCacheCodec.decode(br#"{"timestamp": 100.0, "response": "not serialized"}"#),
-        Err(Error::InvalidEntry)
-    );
-}
-
-#[test]
-fn response_entries_preserve_the_existing_json_representation() {
-    let codec = ResponseCacheCodec;
-    let entry = CacheEntry {
-        timestamp: Some(123.0),
-        response: json!({"choices": [{"text": "cached"}]}),
-    };
-    let bytes = codec.encode(&entry).unwrap();
-    assert_eq!(bytes, serde_json::to_vec(&entry).unwrap());
-    assert_eq!(codec.decode(&bytes).unwrap(), entry);
-}
-
-#[test]
-fn response_codec_preserves_values_without_timestamps() {
-    let codec = ResponseCacheCodec;
-    let raw = json!({"choices": [{"text": "legacy"}]});
-    let entry = codec.decode(&serde_json::to_vec(&raw).unwrap()).unwrap();
-    assert_eq!(entry.timestamp, None);
-    assert_eq!(entry.response, raw);
-
+#[rstest]
+fn entries_without_timestamps_are_always_fresh(request: ResponseCacheRequest) {
     let backend = Arc::new(InMemoryCache::default());
-    BaseCache::set_cache(backend.as_ref(), "tenant:key", entry, &Default::default()).unwrap();
+    BaseCache::set_cache(
+        backend.as_ref(),
+        "tenant:key",
+        CacheEntry {
+            timestamp: None,
+            response: json!({"choices": [{"text": "legacy"}]}),
+        },
+        &Default::default(),
+    )
+    .unwrap();
     let cache = ResponseCache::new(backend);
+    let mut request = request;
+    request.max_age = Some(Duration::from_secs(1));
     assert_eq!(
-        cache.lookup(&request(), Duration::from_secs(100)).unwrap(),
+        cache.lookup(&request, Duration::from_secs(100)).unwrap(),
         Some(json!({"choices": [{"text": "legacy"}]}))
     );
 }
 
+#[rstest]
 #[tokio::test]
-async fn batch_lookup_reports_partial_hits_and_batch_store_populates_misses() {
-    let cache = memory();
-    let requests = ["hit", "miss", "disabled"].map(|key| {
-        ResponseCacheRequest::new(CacheKeyInput {
-            preset: Some(key.into()),
-            ..Default::default()
-        })
-    });
-    cache
-        .store(&requests[0], json!({"value": 1}), Duration::from_secs(100))
+async fn batch_lookup_reports_partial_hits_and_batch_store_populates_misses(
+    memory: Memory,
+    #[values(false, true)] asynchronous: bool,
+) {
+    let now = Duration::from_secs(100);
+    let mut requests = ["hit", "miss", "disabled"].map(keyed).to_vec();
+    memory
+        .store(&requests[0], json!({"value": 1}), now)
         .unwrap();
-    let mut requests = requests.to_vec();
     requests[2].controls.caching = Some(false);
 
-    let partial = cache
-        .async_lookup_batch(&requests, Duration::from_secs(100))
-        .await
-        .unwrap();
+    let partial = if asynchronous {
+        memory.async_lookup_batch(&requests, now).await.unwrap()
+    } else {
+        memory.lookup_batch(&requests, now).unwrap()
+    };
     assert_eq!(partial.values, vec![Some(json!({"value": 1})), None, None]);
     assert_eq!(partial.missing_indices, vec![1, 2]);
 
-    cache
+    memory
         .async_store_batch(
             vec![
                 (requests[1].clone(), json!({"value": 2})),
                 (requests[2].clone(), json!({"value": 3})),
             ],
-            Duration::from_secs(100),
+            now,
         )
         .await
         .unwrap();
     assert_eq!(
-        cache
-            .lookup(&requests[1], Duration::from_secs(100))
-            .unwrap(),
+        memory.lookup(&requests[1], now).unwrap(),
         Some(json!({"value": 2}))
     );
     requests[2].controls.caching = None;
-    assert_eq!(
-        cache
-            .lookup(&requests[2], Duration::from_secs(100))
-            .unwrap(),
-        None
-    );
+    assert_eq!(memory.lookup(&requests[2], now).unwrap(), None);
 }
 
+#[rstest]
 #[tokio::test]
-async fn deferred_entries_keep_the_time_they_were_produced() {
-    let cache = ResponseCache::new(Arc::new(InMemoryCache::default()));
-    let mut request = request();
+async fn batch_lookup_with_no_readable_request_skips_the_backend(
+    #[values(false, true)] asynchronous: bool,
+) {
+    let cache = redis(Vec::new(), None);
+    let mut request = keyed("key");
+    request.controls.no_cache = true;
+    let requests = [request.clone(), request];
+    let partial = if asynchronous {
+        cache
+            .async_lookup_batch(&requests, Duration::ZERO)
+            .await
+            .unwrap()
+    } else {
+        cache.lookup_batch(&requests, Duration::ZERO).unwrap()
+    };
+    assert_eq!(partial.values, vec![None, None]);
+    assert_eq!(partial.missing_indices, vec![0, 1]);
+}
+
+#[rstest]
+#[tokio::test]
+async fn deferred_entries_keep_the_time_they_were_produced(
+    memory: Memory,
+    mut request: ResponseCacheRequest,
+) {
     request.max_age = Some(Duration::from_secs(10));
-    cache
+    memory
         .async_store_entries(vec![(
             request.clone(),
             json!({"answer": 7}),
@@ -485,27 +571,29 @@ async fn deferred_entries_keep_the_time_they_were_produced() {
         .unwrap();
 
     assert_eq!(
-        cache.lookup(&request, Duration::from_secs(110)).unwrap(),
+        memory.lookup(&request, Duration::from_secs(110)).unwrap(),
         Some(json!({"answer": 7}))
     );
     assert_eq!(
-        cache.lookup(&request, Duration::from_secs(111)).unwrap(),
+        memory.lookup(&request, Duration::from_secs(111)).unwrap(),
         None
     );
 }
 
+#[rstest]
 #[tokio::test]
-async fn write_buffer_flushes_at_its_size_and_keeps_each_produced_time() {
-    let cache = ResponseCache::new(Arc::new(InMemoryCache::default()));
+async fn write_buffer_flushes_at_its_size_and_keeps_each_produced_time(
+    memory: Memory,
+    request: ResponseCacheRequest,
+) {
     let buffer = WriteBuffer::new(2);
-    let mut first = request();
+    let mut first = request;
     first.max_age = Some(Duration::from_secs(10));
-    let mut second = request();
-    second.key.preset = Some("tenant:other".into());
+    let second = keyed("tenant:other");
 
     buffer
         .async_store(
-            &cache,
+            memory.as_ref(),
             &first,
             json!({"answer": 7}),
             Duration::from_secs(100),
@@ -513,13 +601,13 @@ async fn write_buffer_flushes_at_its_size_and_keeps_each_produced_time() {
         .await
         .unwrap();
     assert_eq!(
-        cache.lookup(&first, Duration::from_secs(100)).unwrap(),
+        memory.lookup(&first, Duration::from_secs(100)).unwrap(),
         None
     );
 
     buffer
         .async_store(
-            &cache,
+            memory.as_ref(),
             &second,
             json!({"answer": 8}),
             Duration::from_secs(200),
@@ -527,37 +615,36 @@ async fn write_buffer_flushes_at_its_size_and_keeps_each_produced_time() {
         .await
         .unwrap();
     assert_eq!(
-        cache.lookup(&first, Duration::from_secs(110)).unwrap(),
+        memory.lookup(&first, Duration::from_secs(110)).unwrap(),
         Some(json!({"answer": 7}))
     );
     assert_eq!(
-        cache.lookup(&first, Duration::from_secs(111)).unwrap(),
+        memory.lookup(&first, Duration::from_secs(111)).unwrap(),
         None
     );
     assert_eq!(
-        cache.lookup(&second, Duration::from_secs(200)).unwrap(),
+        memory.lookup(&second, Duration::from_secs(200)).unwrap(),
         Some(json!({"answer": 8}))
     );
 }
 
+#[rstest]
 #[tokio::test]
-async fn write_buffer_clear_drops_pending_entries() {
-    let cache = ResponseCache::new(Arc::new(InMemoryCache::default()));
+async fn write_buffer_clear_drops_pending_entries(memory: Memory, request: ResponseCacheRequest) {
     let buffer = WriteBuffer::new(2);
-    let mut other = request();
-    other.key.preset = Some("tenant:other".into());
+    let other = keyed("tenant:other");
     let now = Duration::from_secs(100);
 
     buffer
-        .async_store(&cache, &request(), json!({"answer": 7}), now)
+        .async_store(memory.as_ref(), &request, json!({"answer": 7}), now)
         .await
         .unwrap();
     buffer.clear().unwrap();
     buffer
-        .async_store(&cache, &other, json!({"answer": 8}), now)
+        .async_store(memory.as_ref(), &other, json!({"answer": 8}), now)
         .await
         .unwrap();
 
-    assert_eq!(cache.lookup(&request(), now).unwrap(), None);
-    assert_eq!(cache.lookup(&other, now).unwrap(), None);
+    assert_eq!(memory.lookup(&request, now).unwrap(), None);
+    assert_eq!(memory.lookup(&other, now).unwrap(), None);
 }

litellm-rust/crates/cache-response/tests/support/mod.rs

@@ -0,0 +1,40 @@
+use std::{sync::Arc, time::Duration};
+
+use litellm_cache_memory::InMemoryCache;
+use litellm_cache_redis::RedisCache;
+use litellm_cache_response::{
+    CacheEntry, CacheKeyInput, ResponseCache, ResponseCacheCodec, ResponseCacheRequest,
+};
+use redis_test::{MockCmd, MockRedisConnection};
+use rstest::fixture;
+
+pub type MockedRedis = RedisCache<ResponseCacheCodec, MockRedisConnection>;
+
+#[fixture]
+pub fn memory() -> Arc<ResponseCache<InMemoryCache<CacheEntry>>> {
+    Arc::new(ResponseCache::new(Arc::new(InMemoryCache::new(
+        Some(8),
+        Some(Duration::from_secs(600)),
+    ))))
+}
+
+#[fixture]
+pub fn request() -> ResponseCacheRequest {
+    keyed("tenant:key")
+}
+
+pub fn keyed(key: &str) -> ResponseCacheRequest {
+    ResponseCacheRequest::new(CacheKeyInput {
+        preset: Some(key.into()),
+        ..Default::default()
+    })
+}
+
+/// A Redis response cache that must receive exactly `commands`, in order.
+pub fn redis(commands: Vec<MockCmd>, namespace: Option<&str>) -> ResponseCache<MockedRedis> {
+    let connection = MockRedisConnection::new(commands).assert_all_commands_consumed();
+    ResponseCache::new(Arc::new(
+        RedisCache::with_connection(connection, None, ResponseCacheCodec)
+            .with_namespace(namespace.map(str::to_owned)),
+    ))
+}

litellm-rust/crates/cache-s3/Cargo.toml

@@ -10,11 +10,17 @@ litellm-cache.workspace = true
 litellm-auth-aws.workspace = true
 aws-sdk-s3 = { version = "1.146.1", default-features = false, features = ["rustls", "rt-tokio"] }
 aws-credential-types = "1.3.0"
-aws-smithy-types = "1.6.0"
+aws-smithy-runtime-api = { version = "1.16.2", features = ["client", "http-1x"] }
+aws-smithy-types = { version = "1.6.0", features = ["http-body-1-x"] }
 aws-types = "1.6.0"
+futures-util.workspace = true
+http.workspace = true
+reqwest.workspace = true
 tokio.workspace = true
 
 [dev-dependencies]
+litellm-cache-testing.workspace = true
+rstest.workspace = true
 wiremock = "0.6.5"
 serde_json.workspace = true
 tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }

litellm-rust/crates/cache-s3/src/auth.rs

@@ -5,22 +5,22 @@ use aws_credential_types::{
 use litellm_auth_aws::{AwsAuthConfig, resolve_credentials};
 
 #[derive(Clone)]
-pub(crate) struct Credentials {
+pub struct S3Credentials {
     config: AwsAuthConfig,
     env: fn(&str) -> Option<String>,
 }
 
-impl Credentials {
-    pub(crate) fn new(config: AwsAuthConfig) -> Self {
+impl S3Credentials {
+    pub fn new(config: AwsAuthConfig) -> Self {
         Self::with_env(config, |name| std::env::var(name).ok())
     }
 
-    pub(crate) fn with_env(config: AwsAuthConfig, env: fn(&str) -> Option<String>) -> Self {
+    pub fn with_env(config: AwsAuthConfig, env: fn(&str) -> Option<String>) -> Self {
         Self { config, env }
     }
 }
 
-impl ProvideCredentials for Credentials {
+impl ProvideCredentials for S3Credentials {
     fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a>
     where
         Self: 'a,
@@ -45,57 +45,8 @@ impl ProvideCredentials for Credentials {
     }
 }
 
-impl std::fmt::Debug for Credentials {
+impl std::fmt::Debug for S3Credentials {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        f.debug_struct("Credentials").finish_non_exhaustive()
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[tokio::test]
-    async fn explicit_keys_ignore_an_ambient_session_token() {
-        let provider = Credentials::with_env(
-            AwsAuthConfig {
-                access_key_id: Some("key".to_string()),
-                secret_access_key: Some("secret".to_string()),
-                region_name: Some("us-east-1".to_string()),
-                ..Default::default()
-            },
-            |name| (name == "AWS_SESSION_TOKEN").then(|| "ambient".to_string()),
-        );
-        let credentials = provider.provide_credentials().await.unwrap();
-        assert_eq!(credentials.access_key_id(), "key");
-        assert_eq!(credentials.secret_access_key(), "secret");
-        assert_eq!(credentials.session_token(), None);
-    }
-
-    #[tokio::test]
-    async fn explicit_keys_keep_their_session_token() {
-        let provider = Credentials::new(AwsAuthConfig {
-            access_key_id: Some("key".to_string()),
-            secret_access_key: Some("secret".to_string()),
-            session_token: Some("t".to_string()),
-            region_name: Some("us-east-1".to_string()),
-            ..Default::default()
-        });
-        let credentials = provider.provide_credentials().await.unwrap();
-        assert_eq!(credentials.session_token(), Some("t"));
-    }
-
-    #[tokio::test]
-    async fn environment_keys_resolve_with_their_session_token() {
-        let provider = Credentials::with_env(AwsAuthConfig::default(), |name| match name {
-            "AWS_ACCESS_KEY_ID" => Some("env-key".to_string()),
-            "AWS_SECRET_ACCESS_KEY" => Some("env-secret".to_string()),
-            "AWS_SESSION_TOKEN" => Some("env-token".to_string()),
-            _ => None,
-        });
-        let credentials = provider.provide_credentials().await.unwrap();
-        assert_eq!(credentials.access_key_id(), "env-key");
-        assert_eq!(credentials.secret_access_key(), "env-secret");
-        assert_eq!(credentials.session_token(), Some("env-token"));
+        f.debug_struct("S3Credentials").finish_non_exhaustive()
     }
 }

litellm-rust/crates/cache-s3/src/cache.rs

@@ -10,13 +10,14 @@ use aws_sdk_s3::{
     primitives::ByteStream,
 };
 use aws_smithy_types::{DateTime, date_time::Format};
+use futures_util::future::try_join_all;
 use litellm_auth_aws::AwsAuthConfig;
 use litellm_cache::{
-    BaseCache, BatchCache, CacheCodec, CacheConnectionResult, Error, ExactCacheContext, FlushCache,
+    BaseCache, BatchCache, CacheCodec, DisconnectCache, Error, ExactCacheContext, FlushCache,
 };
 use tokio::runtime::Handle;
 
-use crate::auth::Credentials;
+use crate::{auth::S3Credentials, transport::ReqwestHttpClient};
 
 pub struct S3Endpoint {
     pub url: String,
@@ -41,12 +42,13 @@ pub struct S3Cache<C: CacheCodec> {
 }
 
 impl<C: CacheCodec> S3Cache<C> {
-    pub fn new(config: S3CacheConfig, codec: C, runtime: Handle) -> Self {
+    pub fn new(config: S3CacheConfig, http: reqwest::Client, codec: C, runtime: Handle) -> Self {
         let endpoint_url: Option<String> = config.endpoint.map(|endpoint| endpoint.url);
         let base = aws_sdk_s3::Config::builder()
             .behavior_version(BehaviorVersion::latest())
             .region(Region::new(config.region.clone()))
-            .credentials_provider(Credentials::new(config.auth))
+            .http_client(ReqwestHttpClient(http))
+            .credentials_provider(S3Credentials::new(config.auth))
             .request_checksum_calculation(RequestChecksumCalculation::WhenRequired)
             .response_checksum_validation(ResponseChecksumValidation::WhenRequired);
         let builder = match &endpoint_url {
@@ -202,13 +204,26 @@ impl<C: CacheCodec> BaseCache for S3Cache<C> {
         self.get(key).await
     }
 
+    async fn async_set_cache_pipeline(
+        &self,
+        entries: Vec<(String, Self::Value)>,
+        context: Self::Context,
+    ) -> Result<(), Error> {
+        let context = &context;
+        try_join_all(
+            entries
+                .into_iter()
+                .map(|(key, value)| async move { self.put(&key, value, context).await }),
+        )
+        .await
+        .map(drop)
+    }
+}
+
+impl<C: CacheCodec> DisconnectCache for S3Cache<C> {
     async fn disconnect(&self) -> Result<(), Error> {
         Ok(())
     }
-
-    async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
-        Err(Error::UnsupportedOperation)
-    }
 }
 
 impl<C: CacheCodec> BatchCache for S3Cache<C> {}

litellm-rust/crates/cache-s3/src/lib.rs

@@ -1,4 +1,6 @@
 mod auth;
 mod cache;
+mod transport;
 
+pub use auth::S3Credentials;
 pub use cache::{S3Cache, S3CacheConfig, S3Endpoint};

litellm-rust/crates/cache-s3/src/transport.rs

@@ -0,0 +1,49 @@
+use aws_smithy_runtime_api::client::{
+    http::{
+        HttpClient, HttpConnector, HttpConnectorFuture, HttpConnectorSettings, SharedHttpConnector,
+    },
+    orchestrator::HttpRequest,
+    result::ConnectorError,
+    runtime_components::RuntimeComponents,
+};
+use aws_smithy_types::body::SdkBody;
+
+#[derive(Clone, Debug)]
+pub(crate) struct ReqwestHttpClient(pub(crate) reqwest::Client);
+
+impl HttpClient for ReqwestHttpClient {
+    fn http_connector(
+        &self,
+        _: &HttpConnectorSettings,
+        _: &RuntimeComponents,
+    ) -> SharedHttpConnector {
+        SharedHttpConnector::new(self.clone())
+    }
+}
+
+impl HttpConnector for ReqwestHttpClient {
+    fn call(&self, request: HttpRequest) -> HttpConnectorFuture {
+        let client = self.0.clone();
+        HttpConnectorFuture::new(async move {
+            let request = request
+                .try_into_http1x()
+                .map_err(|error| ConnectorError::other(error.into(), None))?
+                .map(reqwest::Body::wrap);
+            let request = reqwest::Request::try_from(request)
+                .map_err(|error| ConnectorError::other(error.into(), None))?;
+            let response = client.execute(request).await.map_err(|error| {
+                if error.is_timeout() {
+                    ConnectorError::timeout(error.into())
+                } else {
+                    ConnectorError::io(error.into())
+                }
+            })?;
+            let response = http::Response::from(response).map(SdkBody::from_body_1_x);
+            response
+                .try_into()
+                .map_err(|error: aws_smithy_runtime_api::http::HttpError| {
+                    ConnectorError::other(error.into(), None)
+                })
+        })
+    }
+}

litellm-rust/crates/cache-s3/tests/auth.rs

@@ -0,0 +1,57 @@
+use aws_credential_types::provider::ProvideCredentials;
+use litellm_auth_aws::AwsAuthConfig;
+use litellm_cache_s3::S3Credentials;
+use rstest::rstest;
+
+fn explicit(session_token: Option<&str>) -> AwsAuthConfig {
+    AwsAuthConfig {
+        access_key_id: Some("key".to_string()),
+        secret_access_key: Some("secret".to_string()),
+        session_token: session_token.map(str::to_string),
+        region_name: Some("us-east-1".to_string()),
+        ..Default::default()
+    }
+}
+
+fn ambient_token(name: &str) -> Option<String> {
+    (name == "AWS_SESSION_TOKEN").then(|| "ambient".to_string())
+}
+
+fn environment_keys(name: &str) -> Option<String> {
+    match name {
+        "AWS_ACCESS_KEY_ID" => Some("env-key".to_string()),
+        "AWS_SECRET_ACCESS_KEY" => Some("env-secret".to_string()),
+        "AWS_SESSION_TOKEN" => Some("env-token".to_string()),
+        _ => None,
+    }
+}
+
+#[rstest]
+#[case::explicit_keys_ignore_an_ambient_session_token(
+    explicit(None), ambient_token, ("key", "secret", None)
+)]
+#[case::explicit_keys_keep_their_session_token(
+    explicit(Some("t")), ambient_token, ("key", "secret", Some("t"))
+)]
+#[case::environment_keys_resolve_with_their_session_token(
+    AwsAuthConfig::default(), environment_keys, ("env-key", "env-secret", Some("env-token"))
+)]
+#[tokio::test]
+async fn credentials_resolve(
+    #[case] config: AwsAuthConfig,
+    #[case] env: fn(&str) -> Option<String>,
+    #[case] expected: (&str, &str, Option<&str>),
+) {
+    let credentials = S3Credentials::with_env(config, env)
+        .provide_credentials()
+        .await
+        .unwrap();
+    assert_eq!(
+        (
+            credentials.access_key_id(),
+            credentials.secret_access_key(),
+            credentials.session_token(),
+        ),
+        expected
+    );
+}

litellm-rust/crates/cache-s3/tests/cache.rs

@@ -1,41 +1,23 @@
+mod support;
+
 use std::time::{Duration, SystemTime, UNIX_EPOCH};
 
-use litellm_auth_aws::AwsAuthConfig;
+use aws_smithy_types::{DateTime, date_time::Format};
 use litellm_cache::{
-    BaseCache, BatchCache, BatchEntry, Error, ExactCacheContext, FlushCache, JsonCodec,
+    BaseCache, BatchCache, BatchEntry, DisconnectCache, Error, ExactCacheContext, FlushCache,
 };
-use litellm_cache_s3::{S3Cache, S3CacheConfig, S3Endpoint};
+use litellm_cache_s3::S3CacheConfig;
+use rstest::{fixture, rstest};
 use serde_json::{Value, json};
-use tokio::runtime::Handle;
+use support::FakeBucket;
 use wiremock::{
     Mock, MockServer, ResponseTemplate,
+    http::HeaderMap,
     matchers::{method, path},
 };
 
-fn config(endpoint: String) -> S3CacheConfig {
-    S3CacheConfig {
-        bucket: "cache-bucket".to_string(),
-        key_prefix: "team/".to_string(),
-        region: "us-east-1".to_string(),
-        endpoint: Some(S3Endpoint { url: endpoint }),
-        auth: AwsAuthConfig {
-            access_key_id: Some("key".to_string()),
-            secret_access_key: Some("secret".to_string()),
-            region_name: Some("us-east-1".to_string()),
-            ..Default::default()
-        },
-    }
-}
-
-fn cache(endpoint: &str) -> S3Cache<JsonCodec<Value>> {
-    S3Cache::new(
-        config(endpoint.to_string()),
-        JsonCodec::<Value>::new(),
-        Handle::current(),
-    )
-}
-
-async fn mock_server() -> MockServer {
+#[fixture]
+async fn server() -> MockServer {
     let server = MockServer::start().await;
     Mock::given(method("PUT"))
         .respond_with(ResponseTemplate::new(200).insert_header("etag", "\"etag\""))
@@ -44,23 +26,25 @@ async fn mock_server() -> MockServer {
     server
 }
 
-fn http_date_from(headers: &wiremock::http::HeaderMap, name: &str) -> Option<SystemTime> {
-    use aws_smithy_types::{DateTime, date_time::Format};
+fn http_date_from(headers: &HeaderMap, name: &str) -> Option<SystemTime> {
     headers
         .get(name)
         .and_then(|value| DateTime::from_str(value.to_str().ok()?, Format::HttpDate).ok())
         .map(|date| UNIX_EPOCH + Duration::new(date.secs() as u64, date.subsec_nanos()))
 }
 
+fn ttl(seconds: u64) -> ExactCacheContext {
+    ExactCacheContext {
+        ttl: Some(Duration::from_secs(seconds)),
+    }
+}
+
+#[rstest]
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
-async fn set_writes_python_metadata_with_and_without_ttl() {
-    let server = mock_server().await;
-    let cache = cache(&server.uri());
-    let context = ExactCacheContext {
-        ttl: Some(Duration::from_secs(90)),
-    };
+async fn set_writes_python_metadata_with_and_without_ttl(#[future(awt)] server: MockServer) {
+    let cache = support::cache(&server.uri());
     cache
-        .set_cache("alpha:beta", json!({"answer": 1}), &context)
+        .set_cache("alpha:beta", json!({"answer": 1}), &ttl(90))
         .unwrap();
     cache
         .set_cache("plain", json!({"answer": 2}), &ExactCacheContext::default())
@@ -110,61 +94,84 @@ async fn set_writes_python_metadata_with_and_without_ttl() {
     );
 }
 
+#[rstest]
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
-async fn get_hit_miss_expired_and_invalid_entries() {
-    let server = mock_server().await;
-    Mock::given(method("GET"))
-        .and(path("/cache-bucket/team/hit"))
-        .respond_with(ResponseTemplate::new(200).set_body_json(json!({"answer": 3})))
-        .mount(&server)
-        .await;
-    Mock::given(method("GET"))
-        .and(path("/cache-bucket/team/missing"))
-        .respond_with(
-            ResponseTemplate::new(404).set_body_string("<Error><Code>NoSuchKey</Code></Error>"),
-        )
-        .mount(&server)
-        .await;
-    Mock::given(method("GET"))
-        .and(path("/cache-bucket/team/denied"))
-        .respond_with(
-            ResponseTemplate::new(403).set_body_string("<Error><Code>AccessDenied</Code></Error>"),
-        )
-        .mount(&server)
-        .await;
-    Mock::given(method("GET"))
-        .and(path("/cache-bucket/team/expired"))
-        .respond_with(
-            ResponseTemplate::new(200)
-                .insert_header("expires", "Thu, 01 Jan 1970 00:00:00 GMT")
-                .set_body_json(json!({"answer": 4})),
-        )
-        .mount(&server)
-        .await;
-    Mock::given(method("GET"))
-        .and(path("/cache-bucket/team/malformed"))
-        .respond_with(ResponseTemplate::new(200).set_body_string("not a cache entry"))
-        .mount(&server)
-        .await;
-    let cache = cache(&server.uri());
-    let context = ExactCacheContext::default();
+async fn async_set_signs_with_the_configured_keys(#[future(awt)] server: MockServer) {
+    support::cache(&server.uri())
+        .async_set_cache("key", json!({"answer": 1}), ttl(3600))
+        .await
+        .unwrap();
 
-    assert_eq!(
-        cache.get_cache("hit", &context).unwrap(),
-        Some(json!({"answer": 3}))
+    let requests = server.received_requests().await.unwrap();
+    let request = &requests[0];
+    assert_eq!(request.url.path(), "/cache-bucket/team/key");
+    assert!(
+        request.headers["authorization"]
+            .to_str()
+            .unwrap()
+            .contains("Credential=key/")
     );
-    assert_eq!(cache.get_cache("missing", &context).unwrap(), None);
-    assert_eq!(cache.get_cache("denied", &context).unwrap(), None);
-    assert_eq!(cache.get_cache("expired", &context).unwrap(), None);
+    assert!(request.headers.get("x-amz-security-token").is_none());
     assert_eq!(
-        cache.get_cache("malformed", &context),
-        Err(Error::InvalidEntry)
+        request.headers["cache-control"].to_str().unwrap(),
+        "immutable, max-age=3600, s-maxage=3600"
     );
 }
 
+#[rstest]
+#[case::hit("hit", ResponseTemplate::new(200).set_body_json(json!({"answer": 3})), Ok(Some(json!({"answer": 3}))))]
+#[case::no_such_key(
+    "missing",
+    ResponseTemplate::new(404).set_body_string("<Error><Code>NoSuchKey</Code></Error>"),
+    Ok(None)
+)]
+#[case::access_denied(
+    "denied",
+    ResponseTemplate::new(403).set_body_string("<Error><Code>AccessDenied</Code></Error>"),
+    Ok(None)
+)]
+#[case::expired(
+    "expired",
+    ResponseTemplate::new(200)
+        .insert_header("expires", "Thu, 01 Jan 1970 00:00:00 GMT")
+        .set_body_json(json!({"answer": 4})),
+    Ok(None)
+)]
+#[case::not_yet_expired(
+    "fresh",
+    ResponseTemplate::new(200)
+        .insert_header("expires", "Fri, 01 Jan 2100 00:00:00 GMT")
+        .set_body_json(json!({"answer": 5})),
+    Ok(Some(json!({"answer": 5})))
+)]
+#[case::malformed(
+    "malformed",
+    ResponseTemplate::new(200).set_body_string("not a cache entry"),
+    Err(Error::InvalidEntry)
+)]
+#[case::server_error("broken", ResponseTemplate::new(500), Err(Error::Unavailable))]
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
-async fn batch_get_preserves_order_with_hits_misses_and_invalid() {
-    let server = mock_server().await;
+async fn get_maps_s3_responses(
+    #[future(awt)] server: MockServer,
+    #[case] key: &str,
+    #[case] response: ResponseTemplate,
+    #[case] expected: Result<Option<Value>, Error>,
+) {
+    Mock::given(method("GET"))
+        .and(path(format!("/cache-bucket/team/{key}")))
+        .respond_with(response)
+        .mount(&server)
+        .await;
+    let cache = support::cache(&server.uri());
+    let context = ExactCacheContext::default();
+
+    assert_eq!(cache.get_cache(key, &context), expected);
+    assert_eq!(cache.async_get_cache(key, &context).await, expected);
+}
+
+#[rstest]
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn batch_get_preserves_order_with_hits_misses_and_invalid(#[future(awt)] server: MockServer) {
     for (key, status, body) in [
         ("first", 200, "{\"answer\": 1}"),
         ("invalid", 200, "garbage"),
@@ -180,91 +187,117 @@ async fn batch_get_preserves_order_with_hits_misses_and_invalid() {
         .respond_with(ResponseTemplate::new(404))
         .mount(&server)
         .await;
-    let cache = cache(&server.uri());
+    let cache = support::cache(&server.uri());
     let context = ExactCacheContext::default();
     let keys = vec![
         "first".to_string(),
         "miss".to_string(),
         "invalid".to_string(),
     ];
+    let expected = vec![
+        BatchEntry::Hit(json!({"answer": 1})),
+        BatchEntry::Miss,
+        BatchEntry::Invalid,
+    ];
 
-    let entries = cache.batch_get_cache(&keys, &context).unwrap();
-
+    assert_eq!(cache.batch_get_cache(&keys, &context).unwrap(), expected);
     assert_eq!(
-        entries,
-        vec![
-            BatchEntry::Hit(json!({"answer": 1})),
-            BatchEntry::Miss,
-            BatchEntry::Invalid,
-        ]
+        cache.async_batch_get_cache(keys, context).await.unwrap(),
+        expected
     );
 }
 
+#[rstest]
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
-async fn unsupported_and_noop_capabilities_match_python() {
-    let server = mock_server().await;
-    let cache = cache(&server.uri());
+async fn pipeline_writes_every_entry_with_the_shared_ttl() {
+    let server = FakeBucket::serve().await;
+    let cache = support::cache(&server.uri());
+    cache
+        .async_set_cache_pipeline(
+            vec![
+                ("one".into(), json!({"n": 1})),
+                ("two".into(), json!({"n": 2})),
+            ],
+            ttl(30),
+        )
+        .await
+        .unwrap();
 
+    let requests = server.received_requests().await.unwrap();
+    assert_eq!(requests.len(), 2);
+    assert!(requests.iter().all(|request| {
+        request.headers["cache-control"].to_str().unwrap() == "immutable, max-age=30, s-maxage=30"
+    }));
     assert_eq!(
-        cache.test_connection().await,
-        Err(Error::UnsupportedOperation)
+        cache
+            .get_cache("two", &ExactCacheContext::default())
+            .unwrap(),
+        Some(json!({"n": 2}))
     );
+}
+
+#[rstest]
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn flush_and_disconnect_are_noops_like_python(#[future(awt)] server: MockServer) {
+    let cache = support::cache(&server.uri());
+
     cache.flush_cache().unwrap();
+    cache.async_flush_cache().await.unwrap();
     cache.disconnect().await.unwrap();
-    assert_eq!(cache.get_ttl(&ExactCacheContext::default()), None);
-    assert_eq!(
-        cache.get_ttl(&ExactCacheContext {
-            ttl: Some(Duration::from_secs(45)),
-        }),
-        Some(Duration::from_secs(45))
-    );
     assert!(server.received_requests().await.unwrap().is_empty());
 }
 
-#[test]
-fn key_conversion_prefixes_and_splits_colons() {
+#[rstest]
+#[case::without_ttl(ExactCacheContext::default(), None)]
+#[case::with_ttl(ttl(45), Some(Duration::from_secs(45)))]
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn get_ttl_reports_the_request_ttl(
+    #[case] context: ExactCacheContext,
+    #[case] expected: Option<Duration>,
+) {
+    assert_eq!(
+        support::cache("http://localhost").get_ttl(&context),
+        expected
+    );
+}
+
+#[rstest]
+#[case::prefixed("team/", "a:b:c", "team/a/b/c")]
+#[case::prefixed_plain("team/", "plain", "team/plain")]
+#[case::unprefixed("", "a:b", "a/b")]
+fn key_conversion_prefixes_and_splits_colons(
+    #[case] key_prefix: &str,
+    #[case] key: &str,
+    #[case] expected: &str,
+) {
     let runtime = tokio::runtime::Builder::new_multi_thread()
         .worker_threads(1)
         .enable_all()
         .build()
         .unwrap();
-    let _guard = runtime.enter();
-    let cache = S3Cache::new(
+    let cache = support::cache_with(
         S3CacheConfig {
-            key_prefix: "team/".to_string(),
-            ..config("http://localhost".to_string())
+            key_prefix: key_prefix.to_string(),
+            ..support::config("http://localhost")
         },
-        JsonCodec::<Value>::new(),
         runtime.handle().clone(),
     );
 
     assert_eq!(cache.bucket(), "cache-bucket");
-    assert_eq!(cache.key_prefix(), "team/");
-    assert_eq!(cache.to_s3_key("a:b:c"), "team/a/b/c");
-    assert_eq!(cache.to_s3_key("plain"), "team/plain");
-
-    let unprefixed = S3Cache::new(
-        S3CacheConfig {
-            key_prefix: String::new(),
-            ..config("http://localhost".to_string())
-        },
-        JsonCodec::<Value>::new(),
-        runtime.handle().clone(),
-    );
-    assert_eq!(unprefixed.to_s3_key("a:b"), "a/b");
+    assert_eq!(cache.key_prefix(), key_prefix);
+    assert_eq!(cache.region(), "us-east-1");
+    assert_eq!(cache.endpoint(), Some("http://localhost"));
+    assert_eq!(cache.to_s3_key(key), expected);
 }
 
+#[rstest]
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
-async fn sync_methods_block_inside_and_outside_the_runtime() {
-    let server = mock_server().await;
-    Mock::given(method("GET"))
-        .and(path("/cache-bucket/team/key"))
-        .respond_with(ResponseTemplate::new(200).set_body_json(json!({"answer": 9})))
-        .mount(&server)
-        .await;
+async fn sync_methods_block_outside_the_runtime() {
+    let server = FakeBucket::serve().await;
     let uri = server.uri();
-    let cache = tokio::task::spawn_blocking(move || {
-        let cache = cache(&uri);
+    let handle = tokio::runtime::Handle::current();
+    let cached = tokio::task::spawn_blocking(move || {
+        let cache = support::cache_with(support::config(&uri), handle);
         let context = ExactCacheContext::default();
         cache
             .set_cache("key", json!({"answer": 9}), &context)
@@ -274,5 +307,5 @@ async fn sync_methods_block_inside_and_outside_the_runtime() {
     .await
     .unwrap();
 
-    assert_eq!(cache, Some(json!({"answer": 9})));
+    assert_eq!(cached, Some(json!({"answer": 9})));
 }

litellm-rust/crates/cache-s3/tests/contract.rs

@@ -0,0 +1,65 @@
+mod support;
+
+use litellm_cache::ExactCacheContext;
+use litellm_cache_testing as contract;
+use rstest::{fixture, rstest};
+use serde_json::json;
+use support::{FakeBucket, JsonS3Cache};
+use wiremock::MockServer;
+
+struct S3 {
+    cache: JsonS3Cache,
+    _server: MockServer,
+}
+
+#[fixture]
+async fn s3() -> S3 {
+    let server = FakeBucket::serve().await;
+    S3 {
+        cache: support::cache(&server.uri()),
+        _server: server,
+    }
+}
+
+#[fixture]
+fn context() -> ExactCacheContext {
+    ExactCacheContext::default()
+}
+
+const PREFIX: &str = "contract:";
+
+#[rstest]
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn hit_and_miss(#[future(awt)] s3: S3, context: ExactCacheContext) {
+    contract::hit_and_miss(&s3.cache, context, PREFIX, json!({"answer": 42})).await;
+}
+
+#[rstest]
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn sync_async_equivalence(#[future(awt)] s3: S3, context: ExactCacheContext) {
+    contract::sync_async_equivalence(&s3.cache, context, PREFIX, json!("first"), json!([2])).await;
+}
+
+#[rstest]
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn overwrite_replaces(#[future(awt)] s3: S3, context: ExactCacheContext) {
+    contract::overwrite_replaces(&s3.cache, context, PREFIX, json!(1), json!({"b": 2})).await;
+}
+
+#[rstest]
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn pipeline_writes_every_entry(#[future(awt)] s3: S3, context: ExactCacheContext) {
+    contract::pipeline_writes_every_entry(
+        &s3.cache,
+        context,
+        PREFIX,
+        vec![json!("a"), json!(2), json!({"c": true})],
+    )
+    .await;
+}
+
+#[rstest]
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn batch_preserves_order(#[future(awt)] s3: S3, context: ExactCacheContext) {
+    contract::batch_preserves_order(&s3.cache, context, PREFIX, json!("first"), json!(2)).await;
+}

litellm-rust/crates/cache-s3/tests/support/mod.rs

@@ -0,0 +1,77 @@
+#![allow(dead_code)]
+
+use std::{
+    collections::HashMap,
+    sync::{Arc, Mutex},
+};
+
+use litellm_auth_aws::AwsAuthConfig;
+use litellm_cache::JsonCodec;
+use litellm_cache_s3::{S3Cache, S3CacheConfig, S3Endpoint};
+use serde_json::Value;
+use tokio::runtime::Handle;
+use wiremock::{Mock, MockServer, Request, Respond, ResponseTemplate, http::Method, matchers::any};
+
+pub type JsonS3Cache = S3Cache<JsonCodec<Value>>;
+
+pub fn config(endpoint: &str) -> S3CacheConfig {
+    S3CacheConfig {
+        bucket: "cache-bucket".to_string(),
+        key_prefix: "team/".to_string(),
+        region: "us-east-1".to_string(),
+        endpoint: Some(S3Endpoint {
+            url: endpoint.to_string(),
+        }),
+        auth: AwsAuthConfig {
+            access_key_id: Some("key".to_string()),
+            secret_access_key: Some("secret".to_string()),
+            region_name: Some("us-east-1".to_string()),
+            ..Default::default()
+        },
+    }
+}
+
+pub fn cache_with(config: S3CacheConfig, runtime: Handle) -> JsonS3Cache {
+    S3Cache::new(config, reqwest::Client::new(), JsonCodec::new(), runtime)
+}
+
+pub fn cache(endpoint: &str) -> JsonS3Cache {
+    cache_with(config(endpoint), Handle::current())
+}
+
+/// An in-memory bucket: PUT stores the body under the request path, GET serves it or answers
+/// `NoSuchKey`.
+#[derive(Clone, Default)]
+pub struct FakeBucket {
+    objects: Arc<Mutex<HashMap<String, Vec<u8>>>>,
+}
+
+impl FakeBucket {
+    pub async fn serve() -> MockServer {
+        let server = MockServer::start().await;
+        Mock::given(any())
+            .respond_with(Self::default())
+            .mount(&server)
+            .await;
+        server
+    }
+}
+
+impl Respond for FakeBucket {
+    fn respond(&self, request: &Request) -> ResponseTemplate {
+        let path = request.url.path().to_string();
+        let mut objects = self.objects.lock().unwrap();
+        match request.method {
+            Method::PUT => {
+                objects.insert(path, request.body.clone());
+                ResponseTemplate::new(200).insert_header("etag", "\"etag\"")
+            }
+            Method::GET => match objects.get(&path) {
+                Some(body) => ResponseTemplate::new(200).set_body_bytes(body.clone()),
+                None => ResponseTemplate::new(404)
+                    .set_body_string("<Error><Code>NoSuchKey</Code></Error>"),
+            },
+            _ => ResponseTemplate::new(405),
+        }
+    }
+}