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

.github/workflows/osv-scan.yml

@@ -5,6 +5,9 @@ on:
     branches:
       - main
       - "litellm_**"
+      - "rc/**"
+      - "stable/**"
+      - "lts/**"
   schedule:
     - cron: "23 6 * * *"
   workflow_dispatch:

.github/workflows/test-unit.yml

@@ -5,6 +5,9 @@ on:
     branches:
       - main
       - "litellm_**"
+      - "rc/**"
+      - "stable/**"
+      - "lts/**"
   push:
     branches:
       - main
@@ -280,6 +283,7 @@ jobs:
               tests/unit/llms/e2b
               tests/unit/llms/edenai
               tests/unit/llms/elevenlabs
+              tests/unit/llms/empiriolabs
               tests/unit/llms/exa_ai
               tests/unit/llms/fal_ai
               tests/unit/llms/fastcrw
@@ -973,10 +977,10 @@ jobs:
             dist: loadscope
     steps:
       - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
-        timeout-minutes: ${{ matrix.shard == 'db-and-spend' && 10 || 3 }}
+        timeout-minutes: 3
         with:
           persist-credentials: false
-          fetch-depth: ${{ matrix.shard == 'db-and-spend' && '0' || '1' }}
+          fetch-depth: 1
 
       - name: Detect relevant changes
         id: changes
@@ -1009,13 +1013,11 @@ jobs:
         with:
           artifact: ${{ needs.rust-bridge.outputs.artifact }}
 
-      - name: Install dependencies and check schema
+      - name: Install dependencies
         if: steps.changes.outputs.decision != 'skip'
-        timeout-minutes: 8
+        timeout-minutes: 5
         env:
           RUST_BRIDGE_ARTIFACT: ${{ needs.rust-bridge.outputs.artifact }}
-          UNIT_SHARD: ${{ matrix.shard }}
-          SCHEMA_BASE_REF: origin/${{ github.event.repository.default_branch }}
         run: |
           diff -u model_prices_and_context_window.json litellm/model_prices_and_context_window_backup.json
           if [ -z "$RUST_BRIDGE_ARTIFACT" ]; then
@@ -1027,9 +1029,21 @@ jobs:
             uv run --no-sync python -c "import importlib.metadata; import litellm.rust_bridge._native; print(importlib.metadata.version('litellm'))"
           fi
           uv run --no-sync python -c 'import os, sys; print(sys.version); assert f"{sys.version_info.major}.{sys.version_info.minor}" == os.environ["UV_PYTHON"]'
-          if [ "$UNIT_SHARD" = "db-and-spend" ]; then
-            uv run --no-sync python scripts/check_schema_compatibility.py --base-ref "$SCHEMA_BASE_REF"
-          fi
+
+      - name: Check schema.prisma against the merge base
+        if: steps.changes.outputs.decision != 'skip' && matrix.shard == 'db-and-spend' && github.event_name == 'pull_request'
+        timeout-minutes: 3
+        env:
+          GH_TOKEN: ${{ github.token }}
+          PR_NUMBER: ${{ github.event.pull_request.number }}
+          EVENT_BASE_REF: ${{ github.event.pull_request.base.ref }}
+          HEAD_SHA: ${{ github.event.pull_request.head.sha }}
+        run: |
+          base_ref=$(gh api "repos/${GITHUB_REPOSITORY}/branches/${EVENT_BASE_REF}" --jq .name) || base_ref=$(gh api "repos/${GITHUB_REPOSITORY}/pulls/${PR_NUMBER}" --jq .base.ref)
+          merge_base=$(gh api "repos/${GITHUB_REPOSITORY}/compare/${base_ref}...${HEAD_SHA}?per_page=1" --jq '.merge_base_commit.sha')
+          test -n "$merge_base"
+          git fetch --no-tags --depth=1 origin "$merge_base"
+          uv run --no-sync python scripts/check_schema_compatibility.py --base-commit "$merge_base"
 
       - name: Cache Prisma binaries
         if: steps.changes.outputs.decision != 'skip'

README.md

@@ -398,6 +398,7 @@ Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call th
 | [Deepseek (`deepseek`)](https://docs.litellm.ai/docs/providers/deepseek) | ✅ | ✅ | ✅ |  |  |  |  |  |  |  |
 | [Eden AI (`edenai`)](https://docs.litellm.ai/docs/providers/edenai) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |  |  |  |
 | [ElevenLabs (`elevenlabs`)](https://docs.litellm.ai/docs/providers/elevenlabs) | ✅ | ✅ | ✅ |  |  | ✅ | ✅ |  |  |  |
+| [EmpirioLabs AI (`empiriolabs`)](https://docs.litellm.ai/docs/providers/empiriolabs) | ✅ |  | ✅ |  |  |  |  |  |  |  |
 | [Empower (`empower`)](https://docs.litellm.ai/docs/providers/empower) | ✅ | ✅ | ✅ |  |  |  |  |  |  |  |
 | [eu/jev Decisions API (`eu_jev`)](https://docs.litellm.ai/docs/pass_through/eu_jev) |  |  |  |  |  |  |  |  |  |  |
 | [Fal AI (`fal_ai`)](https://docs.litellm.ai/docs/providers/fal_ai) | ✅ | ✅ | ✅ |  | ✅ |  |  |  |  |  |

enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py

@@ -200,7 +200,6 @@ async def log_billed_background_response(billed: BilledBackgroundResponse) -> Bi
     unpriced: Final = _without_read_price(billed.response)
     cost: Final = logging_obj.response_cost_calculator(result=unpriced)
     end_time: Final = datetime.now()
-    logging_obj.capture_usage_before_logging(unpriced, logging_obj.start_time, end_time)
     GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
         async_coroutine=logging_obj.async_success_handler(
             result=unpriced, start_time=logging_obj.start_time, end_time=end_time

helm/litellm-helm/Chart.yaml

@@ -18,7 +18,7 @@ type: application
 # This is the chart version. This version number should be incremented each time you make changes
 # to the chart and its templates, including the app version.
 # Versions are expected to follow Semantic Versioning (https://semver.org/)
-version: 1.1.3
+version: 1.1.4
 
 # This is the version number of the application being deployed. This version number should be
 # incremented each time you make changes to the application. Versions are not expected to

helm/litellm-helm/README.md

@@ -56,8 +56,8 @@ If `db.useStackgresOperator` is used (not yet implemented):
 
 | `billingMetrics.enabled`    | Enable enterprise billable-request metering. Requires an enterprise license.                                                                                                                                                                  | `false`                   |
 | `billingMetrics.endpoint`   | Collector that the billable-request counter is pushed to.                                                                                                                                                                                    | `https://telemetry.litellm.ai` |
-| `billingMetrics.secretName` | Name of an existing Secret holding the mTLS client certificate, under the keys `tls.crt` and `tls.key`.                                                                                                                                       | `litellm-billing-metrics-mtls` |
-| `billingMetrics.caSecretName` | Name of an existing Secret holding a CA bundle under the key `ca.crt`. Only needed for a private or test collector whose server certificate is not on the public web PKI.                                                                    | `""`                      |
+| `billingMetrics.tokenUrl`   | Cognito OAuth2 client-credentials token endpoint.                                                                                                                                                                                               | `https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token` |
+| `billingMetrics.secretName` | Name of an existing Secret with `client-id` and `client-secret` keys.                                                                                                                                                                           | `litellm-billing-metrics` |
 | `billingMetrics.exportIntervalMs` | How often the counter is pushed, in milliseconds. The proxy defaults to `60000` when unset.                                                                                                                                             | `""`                      |
 
 #### Example `proxy_config` ConfigMap from values (default):
@@ -109,10 +109,10 @@ type: Opaque
 
 #### Enterprise billable-request metering
 
-Enterprise licenses meter billable requests by pushing a counter to LiteLLM's collector over mutual TLS. The chart does not create the client certificate; it mounts one you already hold, read-only, so the private key is never exposed through the environment. Create the Secret under the name the chart expects, then turn the block on:
+Enterprise licenses meter billable requests by pushing a counter to LiteLLM's collector with Cognito bearer authentication. Create a Secret with `client-id` and `client-secret` keys, then enable the chart values:
 
 ```
-kubectl create secret tls litellm-billing-metrics-mtls --cert=client.crt --key=client.key
+kubectl create secret generic litellm-billing-metrics --from-literal=client-id=... --from-literal=client-secret=...
 ```
 
 ```
@@ -120,7 +120,8 @@ billingMetrics:
   enabled: true
 ```
 
-Set `billingMetrics.caSecretName` only when the collector is a private or test one whose server certificate is not on the public web PKI; the production collector needs no CA override. The chart fails the render rather than deploying a proxy that silently never exports, so a missing `secretName` or an emptied `endpoint` surfaces at `helm install` time.
+The endpoint and token URL must use HTTPS. The chart requires both URLs and
+the Secret name when billing metrics are enabled
 
 ### Database Settings
 

helm/litellm-helm/templates/_helpers.tpl

@@ -50,53 +50,27 @@ app.kubernetes.io/name: {{ include "litellm.name" . }}
 app.kubernetes.io/instance: {{ .Release.Name }}
 {{- end }}
 
-{{/*
-Enterprise billable-request metering. The client certificate identifies the
-deployment to LiteLLM's collector, so it is mounted read-only from an existing
-Secret rather than passed through the environment.
-*/}}
-{{- define "litellm.billingMetrics.certDir" -}}/etc/litellm/billing-mtls{{- end -}}
-{{- define "litellm.billingMetrics.caDir" -}}/etc/litellm/billing-mtls-ca{{- end -}}
-
 {{- define "litellm.billingMetricsEnv" -}}
 - name: LITELLM_BILLING_METRICS_ENDPOINT
   value: {{ required "billingMetrics.endpoint is required when billingMetrics.enabled is true" .Values.billingMetrics.endpoint | quote }}
-- name: LITELLM_BILLING_METRICS_CLIENT_CERT
-  value: {{ printf "%s/tls.crt" (include "litellm.billingMetrics.certDir" .) | quote }}
-- name: LITELLM_BILLING_METRICS_CLIENT_KEY
-  value: {{ printf "%s/tls.key" (include "litellm.billingMetrics.certDir" .) | quote }}
-{{- if .Values.billingMetrics.caSecretName }}
-- name: LITELLM_BILLING_METRICS_CA_CERT
-  value: {{ printf "%s/ca.crt" (include "litellm.billingMetrics.caDir" .) | quote }}
-{{- end }}
+- name: LITELLM_BILLING_METRICS_TOKEN_URL
+  value: {{ required "billingMetrics.tokenUrl is required when billingMetrics.enabled is true" .Values.billingMetrics.tokenUrl | quote }}
+- name: LITELLM_BILLING_METRICS_CLIENT_ID
+  valueFrom:
+    secretKeyRef:
+      name: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true" .Values.billingMetrics.secretName }}
+      key: client-id
+- name: LITELLM_BILLING_METRICS_CLIENT_SECRET
+  valueFrom:
+    secretKeyRef:
+      name: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true" .Values.billingMetrics.secretName }}
+      key: client-secret
 {{- with .Values.billingMetrics.exportIntervalMs }}
 - name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS
   value: {{ . | quote }}
 {{- end }}
 {{- end -}}
 
-{{- define "litellm.billingMetricsVolumes" -}}
-- name: billing-metrics-mtls
-  secret:
-    secretName: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key)" .Values.billingMetrics.secretName }}
-{{- if .Values.billingMetrics.caSecretName }}
-- name: billing-metrics-mtls-ca
-  secret:
-    secretName: {{ .Values.billingMetrics.caSecretName }}
-{{- end }}
-{{- end -}}
-
-{{- define "litellm.billingMetricsVolumeMounts" -}}
-- name: billing-metrics-mtls
-  mountPath: {{ include "litellm.billingMetrics.certDir" . }}
-  readOnly: true
-{{- if .Values.billingMetrics.caSecretName }}
-- name: billing-metrics-mtls-ca
-  mountPath: {{ include "litellm.billingMetrics.caDir" . }}
-  readOnly: true
-{{- end }}
-{{- end -}}
-
 {{/*
 Create the name of the service account to use
 */}}

helm/litellm-helm/templates/deployment.yaml

@@ -148,9 +148,6 @@ spec:
             - name: npm
               mountPath: /.npm
           {{- end }}
-          {{- if .Values.billingMetrics.enabled }}
-            {{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }}
-          {{- end }}
           {{- if include "litellm.collector.socketDir" . }}
             - name: collector-socket
               mountPath: {{ include "litellm.collector.socketDir" . }}
@@ -234,9 +231,6 @@ spec:
             items:
               - key: {{ .Values.proxyConfigMap.key | default "config.yaml" }}
                 path: "config.yaml"
-        {{- if .Values.billingMetrics.enabled }}
-        {{- include "litellm.billingMetricsVolumes" . | nindent 8 }}
-        {{- end }}
         {{- if include "litellm.collector.socketDir" . }}
         - name: collector-socket
           emptyDir:

helm/litellm-helm/tests/billing_metrics_tests.yaml

@@ -1,31 +1,21 @@
-suite: test billingMetrics wiring on the proxy deployment
+suite: test bearer-only billingMetrics wiring
 templates:
   - deployment.yaml
   - configmap-litellm.yaml
-  - migrations-job.yaml
 tests:
-  - it: is off by default, adding no env, volume, or mount
+  - it: is disabled by default
     template: deployment.yaml
     asserts:
-      - notContains:
-          path: spec.template.spec.volumes
-          content:
-            name: billing-metrics-mtls
-            secret:
-              secretName: litellm-billing-metrics-mtls
-      - notContains:
-          path: spec.template.spec.containers[0].volumeMounts
-          content:
-            name: billing-metrics-mtls
-            mountPath: /etc/litellm/billing-mtls
-            readOnly: true
       - notContains:
           path: spec.template.spec.containers[0].env
           content:
             name: LITELLM_BILLING_METRICS_ENDPOINT
-            value: https://telemetry.litellm.ai
+      - notContains:
+          path: spec.template.spec.volumes
+          content:
+            name: billing-metrics
 
-  - it: renders the endpoint and the mounted cert paths when enabled
+  - it: renders required URLs and Cognito secret references when enabled
     template: deployment.yaml
     set:
       billingMetrics:
@@ -39,119 +29,34 @@ tests:
       - contains:
           path: spec.template.spec.containers[0].env
           content:
-            name: LITELLM_BILLING_METRICS_CLIENT_CERT
-            value: /etc/litellm/billing-mtls/tls.crt
+            name: LITELLM_BILLING_METRICS_TOKEN_URL
+            value: https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token
       - contains:
           path: spec.template.spec.containers[0].env
           content:
-            name: LITELLM_BILLING_METRICS_CLIENT_KEY
-            value: /etc/litellm/billing-mtls/tls.key
-
-  # The conventional Secret name is the default, so enabling the block is enough.
-  - it: mounts the default cert secret read-only alongside the config volume
-    template: deployment.yaml
-    set:
-      billingMetrics:
-        enabled: true
-    asserts:
-      - contains:
-          path: spec.template.spec.volumes
-          content:
-            name: billing-metrics-mtls
-            secret:
-              secretName: litellm-billing-metrics-mtls
-      - contains:
-          path: spec.template.spec.containers[0].volumeMounts
-          content:
-            name: billing-metrics-mtls
-            mountPath: /etc/litellm/billing-mtls
-            readOnly: true
-
-  - it: honours a secretName override
-    template: deployment.yaml
-    set:
-      billingMetrics:
-        enabled: true
-        secretName: my-billing-mtls
-    asserts:
-      - contains:
-          path: spec.template.spec.volumes
-          content:
-            name: billing-metrics-mtls
-            secret:
-              secretName: my-billing-mtls
-      - notContains:
-          path: spec.template.spec.volumes
-          content:
-            name: billing-metrics-mtls
-            secret:
-              secretName: litellm-billing-metrics-mtls
-
-  - it: honours an endpoint override
-    template: deployment.yaml
-    set:
-      billingMetrics:
-        enabled: true
-        endpoint: https://collector.internal:4318
-    asserts:
+            name: LITELLM_BILLING_METRICS_CLIENT_ID
+            valueFrom:
+              secretKeyRef:
+                name: litellm-billing-metrics
+                key: client-id
       - contains:
           path: spec.template.spec.containers[0].env
           content:
-            name: LITELLM_BILLING_METRICS_ENDPOINT
-            value: https://collector.internal:4318
-
-  # The production collector presents a public web-PKI certificate, so the CA
-  # override must stay absent unless a private collector is configured.
-  - it: omits the CA env, volume, and mount when no caSecretName is set
-    template: deployment.yaml
-    set:
-      billingMetrics:
-        enabled: true
-    asserts:
+            name: LITELLM_BILLING_METRICS_CLIENT_SECRET
+            valueFrom:
+              secretKeyRef:
+                name: litellm-billing-metrics
+                key: client-secret
       - notContains:
           path: spec.template.spec.volumes
           content:
-            name: billing-metrics-mtls-ca
-            secret:
-              secretName: billing-ca
+            name: billing-metrics
       - notContains:
           path: spec.template.spec.containers[0].volumeMounts
           content:
-            name: billing-metrics-mtls-ca
-            mountPath: /etc/litellm/billing-mtls-ca
-            readOnly: true
-      - notContains:
-          path: spec.template.spec.containers[0].env
-          content:
-            name: LITELLM_BILLING_METRICS_CA_CERT
-            value: /etc/litellm/billing-mtls-ca/ca.crt
+            name: billing-metrics
 
-  - it: mounts the CA secret when caSecretName is set
-    template: deployment.yaml
-    set:
-      billingMetrics:
-        enabled: true
-        caSecretName: billing-ca
-    asserts:
-      - contains:
-          path: spec.template.spec.containers[0].env
-          content:
-            name: LITELLM_BILLING_METRICS_CA_CERT
-            value: /etc/litellm/billing-mtls-ca/ca.crt
-      - contains:
-          path: spec.template.spec.volumes
-          content:
-            name: billing-metrics-mtls-ca
-            secret:
-              secretName: billing-ca
-      - contains:
-          path: spec.template.spec.containers[0].volumeMounts
-          content:
-            name: billing-metrics-mtls-ca
-            mountPath: /etc/litellm/billing-mtls-ca
-            readOnly: true
-
-  - it: passes the export interval through only when set
+  - it: renders the export interval only when set
     template: deployment.yaml
     set:
       billingMetrics:
@@ -164,129 +69,7 @@ tests:
             name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS
             value: "5000"
 
-  - it: omits the export interval when unset
-    template: deployment.yaml
-    set:
-      billingMetrics:
-        enabled: true
-    asserts:
-      - notContains:
-          path: spec.template.spec.containers[0].env
-          content:
-            name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS
-            value: "60000"
-
-  # Kubernetes resolves duplicate env names last-wins, so the chart-owned billing
-  # entries must render after .Values.envVars or a user could silently redirect
-  # the metering export. The three billing entries are the last ones emitted here
-  # (migrationJob, which appends DISABLE_SCHEMA_UPDATE, is off for this case).
-  - it: renders the billing endpoint after envVars so it cannot be shadowed
-    template: deployment.yaml
-    set:
-      migrationJob:
-        enabled: false
-      billingMetrics:
-        enabled: true
-      envVars:
-        LITELLM_BILLING_METRICS_ENDPOINT: https://shadowed.example
-    asserts:
-      - contains:
-          path: spec.template.spec.containers[0].env
-          content:
-            name: LITELLM_BILLING_METRICS_ENDPOINT
-            value: https://shadowed.example
-      - equal:
-          path: spec.template.spec.containers[0].env[-3]
-          value:
-            name: LITELLM_BILLING_METRICS_ENDPOINT
-            value: https://telemetry.litellm.ai
-      - equal:
-          path: spec.template.spec.containers[0].env[-2].name
-          value: LITELLM_BILLING_METRICS_CLIENT_CERT
-      - equal:
-          path: spec.template.spec.containers[0].env[-1].name
-          value: LITELLM_BILLING_METRICS_CLIENT_KEY
-
-  - it: keeps user-supplied volumes and mounts alongside the billing secret
-    template: deployment.yaml
-    set:
-      billingMetrics:
-        enabled: true
-      volumes:
-        - name: custom-callbacks
-          configMap:
-            name: my-callbacks
-      volumeMounts:
-        - name: custom-callbacks
-          mountPath: /app/callbacks
-    asserts:
-      - contains:
-          path: spec.template.spec.volumes
-          content:
-            name: custom-callbacks
-            configMap:
-              name: my-callbacks
-      - contains:
-          path: spec.template.spec.volumes
-          content:
-            name: billing-metrics-mtls
-            secret:
-              secretName: litellm-billing-metrics-mtls
-      - contains:
-          path: spec.template.spec.containers[0].volumeMounts
-          content:
-            name: custom-callbacks
-            mountPath: /app/callbacks
-      - contains:
-          path: spec.template.spec.containers[0].volumeMounts
-          content:
-            name: billing-metrics-mtls
-            mountPath: /etc/litellm/billing-mtls
-            readOnly: true
-
-  - it: still mounts the proxy config when enabled
-    template: deployment.yaml
-    set:
-      billingMetrics:
-        enabled: true
-    asserts:
-      - contains:
-          path: spec.template.spec.containers[0].volumeMounts
-          content:
-            name: litellm-config
-            mountPath: /etc/litellm/config.yaml
-            subPath: config.yaml
-
-  # Only the proxy serves billable traffic. The migrations Job must never mount
-  # the client certificate, and it renders its own env and volumes, so nothing
-  # stops a future edit from wiring the billing include into it by mistake.
-  - it: does not touch the migrations job when enabled
-    template: migrations-job.yaml
-    set:
-      billingMetrics:
-        enabled: true
-    asserts:
-      - notContains:
-          path: spec.template.spec.containers[0].env
-          content:
-            name: LITELLM_BILLING_METRICS_ENDPOINT
-            value: https://telemetry.litellm.ai
-      - notExists:
-          path: spec.template.spec.containers[0].volumeMounts
-      - notExists:
-          path: spec.template.spec.volumes
-
-  - it: fails loudly when enabled with an emptied secretName
-    template: deployment.yaml
-    set:
-      billingMetrics:
-        enabled: true
-        secretName: ""
-    asserts:
-      - failedTemplate:
-          errorMessage: billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key)
-
-  - it: fails loudly when enabled without an endpoint
+  - it: requires the endpoint when enabled
     template: deployment.yaml
     set:
       billingMetrics:
@@ -295,3 +78,23 @@ tests:
     asserts:
       - failedTemplate:
           errorMessage: billingMetrics.endpoint is required when billingMetrics.enabled is true
+
+  - it: requires the token URL when enabled
+    template: deployment.yaml
+    set:
+      billingMetrics:
+        enabled: true
+        tokenUrl: ""
+    asserts:
+      - failedTemplate:
+          errorMessage: billingMetrics.tokenUrl is required when billingMetrics.enabled is true
+
+  - it: requires the secret name when enabled
+    template: deployment.yaml
+    set:
+      billingMetrics:
+        enabled: true
+        secretName: ""
+    asserts:
+      - failedTemplate:
+          errorMessage: billingMetrics.secretName is required when billingMetrics.enabled is true

helm/litellm-helm/tests/collector_tests.yaml

@@ -191,7 +191,7 @@ tests:
       metricsServer.port: 9090
       billingMetrics.enabled: true
       billingMetrics.endpoint: https://metering.example.com
-      billingMetrics.secretName: billing-mtls
+      billingMetrics.secretName: billing-metrics
     asserts:
       - contains:
           path: spec.template.spec.containers[0].env
@@ -214,9 +214,9 @@ tests:
             name: LITELLM_BILLING_METRICS_ENDPOINT
           any: true
       - notContains:
-          path: spec.template.spec.containers[1].volumeMounts
+          path: spec.template.spec.containers[1].env
           content:
-            name: billing-metrics-mtls
+            name: LITELLM_BILLING_METRICS_CLIENT_SECRET
           any: true
 
   - it: should give the sidecar the same scratch mounts as the proxy on a read-only root

helm/litellm-helm/values.yaml

@@ -154,19 +154,13 @@ masterkeySecretName: ""
 # if set, use this secret key for the master key; otherwise, use the default key
 masterkeySecretKey: ""
 
-# Optional: enterprise billable-request metering. When enabled, the proxy counts
-# successful requests to inference, MCP, and A2A endpoints and pushes them to
-# LiteLLM's collector over mutual TLS. Requires an enterprise license.
-# The client certificate identifies the deployment, so it is mounted read-only
-# from an existing Secret and never passed through the environment.
+# Optional: enterprise billable-request metering. Requires an enterprise license.
 billingMetrics:
   enabled: false
-  endpoint: https://telemetry.litellm.ai   # collector to push the counter to
-  secretName: litellm-billing-metrics-mtls # existing Secret holding tls.crt and tls.key
-  # Only for private or test collectors whose server certificate is not on the
-  # public web PKI. The production collector needs no CA override.
-  caSecretName: ""                         # existing Secret holding ca.crt
-  exportIntervalMs: ""                     # push cadence; the proxy defaults to 60000
+  endpoint: https://telemetry.litellm.ai
+  tokenUrl: https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token
+  secretName: litellm-billing-metrics
+  exportIntervalMs: ""
 
 proxyConfigMap:
   # when true, creates a new configmap

helm/litellm/Chart.yaml

@@ -2,7 +2,7 @@ apiVersion: v2
 name: litellm
 description: LiteLLM componentized — gateway, UI backend, and UI as separate services
 type: application
-version: 0.1.0
+version: 0.1.1
 appVersion: "0.1.0"
 annotations:
   org.opencontainers.image.source: "https://github.com/BerriAI/litellm"

helm/litellm/README.md

@@ -22,6 +22,23 @@ backend:
   gatewayUrl: https://gateway.example.com
 ```
 
+## Enterprise billing metrics
+
+Billing metrics use Cognito bearer authentication. Create the Secret with
+`client-id` and `client-secret` keys, then enable billing metrics:
+
+```bash
+kubectl create secret generic litellm-billing-metrics --from-literal=client-id=... --from-literal=client-secret=...
+```
+
+```yaml
+billingMetrics:
+  enabled: true
+```
+
+The chart requires an HTTPS endpoint, an HTTPS token URL, and the Secret name
+when billing metrics are enabled
+
 ## Routes to another Service
 
 For a shared hostname, `ingress.extraPaths` can route to a Service in the same namespace using `serviceName` and `servicePort`. The port can be a number or a named Service port. Omit `service` when using these fields

helm/litellm/templates/_helpers.tpl

@@ -58,56 +58,27 @@ app.kubernetes.io/managed-by: {{ .Release.Service }}
 helm.sh/chart: {{ printf "%s-%s" .Chart.Name .Chart.Version | replace "+" "_" }}
 {{- end -}}
 
-{{/*
-Enterprise billable-request metering. Wired into gateway and backend, not the
-migrations job. The gateway serves nearly all billable traffic, but the backend
-keeps the named-server MCP transport (/{mcp_server_name}/mcp), which writes a
-SpendLogs row, so metering only the gateway would silently drop that traffic.
-The client certificate identifies the deployment to LiteLLM's collector, so it is
-mounted read-only from an existing Secret rather than passed through the
-environment.
-*/}}
-{{- define "litellm.billingMetrics.certDir" -}}/etc/litellm/billing-mtls{{- end -}}
-{{- define "litellm.billingMetrics.caDir" -}}/etc/litellm/billing-mtls-ca{{- end -}}
-
 {{- define "litellm.billingMetricsEnv" -}}
 - name: LITELLM_BILLING_METRICS_ENDPOINT
   value: {{ required "billingMetrics.endpoint is required when billingMetrics.enabled is true" .Values.billingMetrics.endpoint | quote }}
-- name: LITELLM_BILLING_METRICS_CLIENT_CERT
-  value: {{ printf "%s/tls.crt" (include "litellm.billingMetrics.certDir" .) | quote }}
-- name: LITELLM_BILLING_METRICS_CLIENT_KEY
-  value: {{ printf "%s/tls.key" (include "litellm.billingMetrics.certDir" .) | quote }}
-{{- if .Values.billingMetrics.caSecretName }}
-- name: LITELLM_BILLING_METRICS_CA_CERT
-  value: {{ printf "%s/ca.crt" (include "litellm.billingMetrics.caDir" .) | quote }}
-{{- end }}
+- name: LITELLM_BILLING_METRICS_TOKEN_URL
+  value: {{ required "billingMetrics.tokenUrl is required when billingMetrics.enabled is true" .Values.billingMetrics.tokenUrl | quote }}
+- name: LITELLM_BILLING_METRICS_CLIENT_ID
+  valueFrom:
+    secretKeyRef:
+      name: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true" .Values.billingMetrics.secretName }}
+      key: client-id
+- name: LITELLM_BILLING_METRICS_CLIENT_SECRET
+  valueFrom:
+    secretKeyRef:
+      name: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true" .Values.billingMetrics.secretName }}
+      key: client-secret
 {{- with .Values.billingMetrics.exportIntervalMs }}
 - name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS
   value: {{ . | quote }}
 {{- end }}
 {{- end -}}
 
-{{- define "litellm.billingMetricsVolumes" -}}
-- name: billing-metrics-mtls
-  secret:
-    secretName: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key)" .Values.billingMetrics.secretName }}
-{{- if .Values.billingMetrics.caSecretName }}
-- name: billing-metrics-mtls-ca
-  secret:
-    secretName: {{ .Values.billingMetrics.caSecretName }}
-{{- end }}
-{{- end -}}
-
-{{- define "litellm.billingMetricsVolumeMounts" -}}
-- name: billing-metrics-mtls
-  mountPath: {{ include "litellm.billingMetrics.certDir" . }}
-  readOnly: true
-{{- if .Values.billingMetrics.caSecretName }}
-- name: billing-metrics-mtls-ca
-  mountPath: {{ include "litellm.billingMetrics.caDir" . }}
-  readOnly: true
-{{- end }}
-{{- end -}}
 
 {{/*
 Per-component selector labels — used in both Service selectors and Deployment matchLabels.

helm/litellm/templates/backend/deployment.yaml

@@ -79,16 +79,13 @@ spec:
             {{- include "litellm.billingMetricsEnv" . | nindent 12 }}
             {{- end }}
           {{- include "litellm.envFrom" .Values.backend | nindent 10 }}
-          {{- if or .Values.gateway.config.create .Values.backend.volumeMounts .Values.billingMetrics.enabled }}
+          {{- if or .Values.gateway.config.create .Values.backend.volumeMounts }}
           volumeMounts:
             {{- if .Values.gateway.config.create }}
             - name: gateway-config
               mountPath: /app/config/config.yaml
               subPath: config.yaml
             {{- end }}
-            {{- if .Values.billingMetrics.enabled }}
-            {{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }}
-            {{- end }}
             {{- with .Values.backend.volumeMounts }}
             {{- toYaml . | nindent 12 }}
             {{- end }}
@@ -114,16 +111,13 @@ spec:
       {{- with .Values.backend.extraContainers }}
         {{- tpl (toYaml .) $ | nindent 8 }}
       {{- end }}
-      {{- if or .Values.gateway.config.create .Values.backend.volumes .Values.billingMetrics.enabled }}
+      {{- if or .Values.gateway.config.create .Values.backend.volumes }}
       volumes:
         {{- if .Values.gateway.config.create }}
         - name: gateway-config
           configMap:
             name: {{ include "litellm.gateway.fullname" . }}-config
         {{- end }}
-        {{- if .Values.billingMetrics.enabled }}
-        {{- include "litellm.billingMetricsVolumes" . | nindent 8 }}
-        {{- end }}
         {{- with .Values.backend.volumes }}
         {{- toYaml . | nindent 8 }}
         {{- end }}

helm/litellm/templates/gateway/deployment.yaml

@@ -86,7 +86,7 @@ spec:
             {{- include "litellm.gateway.collectorEnv" . | nindent 12 }}
             {{- end }}
           {{- include "litellm.envFrom" .Values.gateway | nindent 10 }}
-          {{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled (include "litellm.gateway.collectorSocketDir" .) }}
+          {{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.gateway.metricsServer.enabled (include "litellm.gateway.collectorSocketDir" .) }}
           volumeMounts:
             {{- if .Values.gateway.config.create }}
             - name: gateway-config
@@ -101,9 +101,6 @@ spec:
             - name: collector-socket
               mountPath: {{ include "litellm.gateway.collectorSocketDir" . }}
             {{- end }}
-            {{- if .Values.billingMetrics.enabled }}
-            {{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }}
-            {{- end }}
             {{- with .Values.gateway.volumeMounts }}
             {{- toYaml . | nindent 12 }}
             {{- end }}
@@ -206,7 +203,7 @@ spec:
       {{- with .Values.gateway.extraContainers }}
         {{- tpl (toYaml .) $ | nindent 8 }}
       {{- end }}
-      {{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled (include "litellm.gateway.collectorSocketDir" .) }}
+      {{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.gateway.metricsServer.enabled (include "litellm.gateway.collectorSocketDir" .) }}
       volumes:
         {{- if .Values.gateway.config.create }}
         - name: gateway-config
@@ -222,9 +219,6 @@ spec:
           emptyDir:
             sizeLimit: 1Mi
         {{- end }}
-        {{- if .Values.billingMetrics.enabled }}
-        {{- include "litellm.billingMetricsVolumes" . | nindent 8 }}
-        {{- end }}
         {{- with .Values.gateway.volumes }}
         {{- toYaml . | nindent 8 }}
         {{- end }}

helm/litellm/tests/billing_metrics_tests.yaml

@@ -1,39 +1,30 @@
-suite: test billingMetrics wiring on gateway and backend
+suite: test bearer-only billingMetrics wiring
 templates:
   - gateway/deployment.yaml
   - gateway/configmap.yaml
   - backend/deployment.yaml
+  - backend/configmap.yaml
   - migrations-job.yaml
 values:
   - ./values/required.yaml
 tests:
-  - it: is off by default, adding no env, volume, or mount
+  - it: is disabled by default
     template: gateway/deployment.yaml
     asserts:
-      - notContains:
-          path: spec.template.spec.volumes
-          content:
-            name: billing-metrics-mtls
-            secret:
-              secretName: billing-mtls
-      - equal:
-          path: spec.template.spec.containers[0].volumeMounts
-          value:
-            - name: gateway-config
-              mountPath: /app/config/config.yaml
-              subPath: config.yaml
       - notContains:
           path: spec.template.spec.containers[0].env
           content:
             name: LITELLM_BILLING_METRICS_ENDPOINT
-            value: https://telemetry.litellm.ai
+      - notContains:
+          path: spec.template.spec.volumes
+          content:
+            name: billing-metrics
 
-  - it: renders the endpoint and the mounted cert paths when enabled
+  - it: wires gateway URLs and bearer credentials from Secret keys
     template: gateway/deployment.yaml
     set:
       billingMetrics:
         enabled: true
-        secretName: billing-mtls
     asserts:
       - contains:
           path: spec.template.spec.containers[0].env
@@ -43,87 +34,68 @@ tests:
       - contains:
           path: spec.template.spec.containers[0].env
           content:
-            name: LITELLM_BILLING_METRICS_CLIENT_CERT
-            value: /etc/litellm/billing-mtls/tls.crt
+            name: LITELLM_BILLING_METRICS_TOKEN_URL
+            value: https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token
       - contains:
           path: spec.template.spec.containers[0].env
           content:
-            name: LITELLM_BILLING_METRICS_CLIENT_KEY
-            value: /etc/litellm/billing-mtls/tls.key
-
-  - it: mounts the cert secret read-only alongside the config volume
-    template: gateway/deployment.yaml
-    set:
-      billingMetrics:
-        enabled: true
-        secretName: billing-mtls
-    asserts:
+            name: LITELLM_BILLING_METRICS_CLIENT_ID
+            valueFrom:
+              secretKeyRef:
+                name: litellm-billing-metrics
+                key: client-id
       - contains:
-          path: spec.template.spec.volumes
+          path: spec.template.spec.containers[0].env
           content:
-            name: billing-metrics-mtls
-            secret:
-              secretName: billing-mtls
-      - contains:
-          path: spec.template.spec.containers[0].volumeMounts
-          content:
-            name: billing-metrics-mtls
-            mountPath: /etc/litellm/billing-mtls
-            readOnly: true
-
-  # The production collector presents a public web-PKI certificate, so the CA
-  # override must stay absent unless a private collector is configured.
-  - it: omits the CA env, volume, and mount when no caSecretName is set
-    template: gateway/deployment.yaml
-    set:
-      billingMetrics:
-        enabled: true
-        secretName: billing-mtls
-    asserts:
+            name: LITELLM_BILLING_METRICS_CLIENT_SECRET
+            valueFrom:
+              secretKeyRef:
+                name: litellm-billing-metrics
+                key: client-secret
       - notContains:
           path: spec.template.spec.volumes
           content:
-            name: billing-metrics-mtls-ca
-            secret:
-              secretName: billing-ca
+            name: billing-metrics
       - notContains:
-          path: spec.template.spec.containers[0].env
+          path: spec.template.spec.containers[0].volumeMounts
           content:
-            name: LITELLM_BILLING_METRICS_CA_CERT
-            value: /etc/litellm/billing-mtls-ca/ca.crt
+            name: billing-metrics
 
-  - it: mounts the CA secret when caSecretName is set
-    template: gateway/deployment.yaml
+  - it: wires billing metrics on the backend without volumes
+    template: backend/deployment.yaml
     set:
       billingMetrics:
         enabled: true
-        secretName: billing-mtls
-        caSecretName: billing-ca
     asserts:
       - contains:
           path: spec.template.spec.containers[0].env
           content:
-            name: LITELLM_BILLING_METRICS_CA_CERT
-            value: /etc/litellm/billing-mtls-ca/ca.crt
-      - contains:
+            name: LITELLM_BILLING_METRICS_CLIENT_SECRET
+            valueFrom:
+              secretKeyRef:
+                name: litellm-billing-metrics
+                key: client-secret
+      - notContains:
           path: spec.template.spec.volumes
           content:
-            name: billing-metrics-mtls-ca
-            secret:
-              secretName: billing-ca
-      - contains:
-          path: spec.template.spec.containers[0].volumeMounts
-          content:
-            name: billing-metrics-mtls-ca
-            mountPath: /etc/litellm/billing-mtls-ca
-            readOnly: true
+            name: billing-metrics
 
-  - it: passes the export interval through only when set
+  - it: does not add billing metrics to migrations
+    template: migrations-job.yaml
+    set:
+      billingMetrics:
+        enabled: true
+    asserts:
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: LITELLM_BILLING_METRICS_ENDPOINT
+
+  - it: renders an optional export interval
     template: gateway/deployment.yaml
     set:
       billingMetrics:
         enabled: true
-        secretName: billing-mtls
         exportIntervalMs: 5000
     asserts:
       - contains:
@@ -132,102 +104,27 @@ tests:
             name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS
             value: "5000"
 
-  - it: keeps user-supplied gateway volumes alongside the billing secret
+  - it: requires the endpoint when enabled
     template: gateway/deployment.yaml
     set:
       billingMetrics:
         enabled: true
-        secretName: billing-mtls
-      gateway.volumes:
-        - name: custom-callbacks
-          configMap:
-            name: my-callbacks
-      gateway.volumeMounts:
-        - name: custom-callbacks
-          mountPath: /app/callbacks
+        endpoint: ""
     asserts:
-      - contains:
-          path: spec.template.spec.volumes
-          content:
-            name: custom-callbacks
-            configMap:
-              name: my-callbacks
-      - contains:
-          path: spec.template.spec.volumes
-          content:
-            name: billing-metrics-mtls
-            secret:
-              secretName: billing-mtls
+      - failedTemplate:
+          errorMessage: billingMetrics.endpoint is required when billingMetrics.enabled is true
 
-  # The backend keeps the named-server MCP transport (/{mcp_server_name}/mcp),
-  # which writes a SpendLogs row, so it must meter too or that traffic is lost.
-  - it: meters the backend as well, since it serves the MCP transport
-    template: backend/deployment.yaml
-    set:
-      billingMetrics:
-        enabled: true
-        secretName: billing-mtls
-    asserts:
-      - contains:
-          path: spec.template.spec.containers[0].env
-          content:
-            name: LITELLM_BILLING_METRICS_ENDPOINT
-            value: https://telemetry.litellm.ai
-      - contains:
-          path: spec.template.spec.containers[0].volumeMounts
-          content:
-            name: billing-metrics-mtls
-            mountPath: /etc/litellm/billing-mtls
-            readOnly: true
-      - contains:
-          path: spec.template.spec.volumes
-          content:
-            name: billing-metrics-mtls
-            secret:
-              secretName: billing-mtls
-
-  - it: leaves the backend alone when metering is off
-    template: backend/deployment.yaml
-    asserts:
-      - notContains:
-          path: spec.template.spec.containers[0].env
-          content:
-            name: LITELLM_BILLING_METRICS_ENDPOINT
-            value: https://telemetry.litellm.ai
-
-  # The migrations job runs prisma and serves no traffic; it must never receive
-  # the client key.
-  - it: never mounts the billing cert on the migrations job
-    template: migrations-job.yaml
-    set:
-      billingMetrics:
-        enabled: true
-        secretName: billing-mtls
-    asserts:
-      - notContains:
-          path: spec.template.spec.containers[0].env
-          content:
-            name: LITELLM_BILLING_METRICS_ENDPOINT
-            value: https://telemetry.litellm.ai
-      - isNull:
-          path: spec.template.spec.volumes
-
-  # The conventional Secret name is the default, so enabling metering needs no
-  # secretName at all; the guard below only fires on an explicitly blanked one.
-  - it: uses the conventional secret name by default
+  - it: requires the token URL when enabled
     template: gateway/deployment.yaml
     set:
       billingMetrics:
         enabled: true
+        tokenUrl: ""
     asserts:
-      - contains:
-          path: spec.template.spec.volumes
-          content:
-            name: billing-metrics-mtls
-            secret:
-              secretName: litellm-billing-metrics-mtls
+      - failedTemplate:
+          errorMessage: billingMetrics.tokenUrl is required when billingMetrics.enabled is true
 
-  - it: fails loudly when the secretName is explicitly blanked
+  - it: requires the secret name when enabled
     template: gateway/deployment.yaml
     set:
       billingMetrics:
@@ -235,15 +132,4 @@ tests:
         secretName: ""
     asserts:
       - failedTemplate:
-          errorMessage: billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key)
-
-  - it: fails loudly when enabled without an endpoint
-    template: gateway/deployment.yaml
-    set:
-      billingMetrics:
-        enabled: true
-        endpoint: ""
-        secretName: billing-mtls
-    asserts:
-      - failedTemplate:
-          errorMessage: billingMetrics.endpoint is required when billingMetrics.enabled is true
+          errorMessage: billingMetrics.secretName is required when billingMetrics.enabled is true

helm/litellm/tests/migration_job_tests.yaml

@@ -91,7 +91,7 @@ tests:
             app.kubernetes.io/name: litellm
             app.kubernetes.io/instance: RELEASE-NAME
             app.kubernetes.io/managed-by: Helm
-            helm.sh/chart: litellm-0.1.0
+            helm.sh/chart: litellm-0.1.1
             app.kubernetes.io/component: migrations
 
   - it: renders pod-level and container-level securityContext in their own scopes

helm/litellm/values.yaml

@@ -180,23 +180,14 @@ masterKey:
   secretKey: master-key
 
 # Optional: enterprise billable-request metering. When enabled, the gateway and
-# backend count successful requests to inference, MCP, and A2A endpoints and push
-# them to LiteLLM's collector over mutual TLS. Both components serve billable
-# routes: the backend keeps the named-server MCP transport. Requires an
-# enterprise license. The client certificate identifies the deployment, so it is
-# mounted read-only from an existing Secret and never passed through the env.
+# backend count successful requests to inference, MCP, and A2A endpoints. Requires
+# an enterprise license.
 billingMetrics:
   enabled: false
-  endpoint: https://telemetry.litellm.ai              # collector to push the counter to
-  # An existing Secret holding the client certificate under tls.crt and its key
-  # under tls.key, usually created from the onboarding artifact. The default is
-  # the conventional name, so the common path is to create that Secret and set
-  # enabled: true. Override only if yours is named differently.
-  secretName: litellm-billing-metrics-mtls
-  # Only for private or test collectors whose server certificate is not on the
-  # public web PKI. The production collector needs no CA override.
-  caSecretName: ""                                    # existing Secret holding ca.crt
-  exportIntervalMs: ""                                # push cadence; the proxy defaults to 60000
+  endpoint: https://telemetry.litellm.ai
+  tokenUrl: https://litellm-prod.auth.us-west-1.amazoncognito.com/oauth2/token
+  secretName: litellm-billing-metrics
+  exportIntervalMs: ""
 
 # External Postgres connection. 
 database:

litellm-proxy-extras/litellm_proxy_extras/migrations/20261103010000_spendlogs_accounting_delivery/migration.sql

@@ -0,0 +1,19 @@
+SET LOCAL lock_timeout = '10s';
+CREATE TABLE IF NOT EXISTS "LiteLLM_SpendLogDelivery" (
+    "delivery_id" TEXT PRIMARY KEY DEFAULT gen_random_uuid()::text,
+    "request_id" TEXT NOT NULL,
+    "state" "AccountingDeliveryState" NOT NULL DEFAULT 'pending',
+    "attempts" INTEGER NOT NULL DEFAULT 0,
+    "next_attempt_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    "applied_at" TIMESTAMP(3),
+    "last_error" TEXT,
+    "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    "accounting_payload" JSONB NOT NULL,
+    "is_batch" BOOLEAN NOT NULL
+);
+CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogDelivery_state_next_attempt_at_idx"
+    ON "LiteLLM_SpendLogDelivery" ("state", "next_attempt_at");
+CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogDelivery_state_applied_at_idx"
+    ON "LiteLLM_SpendLogDelivery" ("state", "applied_at");
+CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogDelivery_request_id_idx"
+    ON "LiteLLM_SpendLogDelivery" ("request_id");

litellm-proxy-extras/litellm_proxy_extras/migrations/20261104000000_skill_zip_upload/migration.sql

@@ -0,0 +1,9 @@
+SET LOCAL lock_timeout = '10s';
+
+CREATE TABLE IF NOT EXISTS "LiteLLM_ClaudeCodePluginUpload" (
+    "plugin_name" TEXT NOT NULL,
+    "content" BYTEA NOT NULL,
+    CONSTRAINT "LiteLLM_ClaudeCodePluginUpload_pkey" PRIMARY KEY ("plugin_name"),
+    CONSTRAINT "LiteLLM_ClaudeCodePluginUpload_plugin_name_fkey" FOREIGN KEY ("plugin_name")
+        REFERENCES "LiteLLM_ClaudeCodePluginTable"("name") ON DELETE CASCADE ON UPDATE CASCADE
+);

litellm-proxy-extras/litellm_proxy_extras/schema.prisma

@@ -1615,9 +1615,18 @@ model LiteLLM_ClaudeCodePluginTable {
   updated_at    DateTime? @default(now()) @updatedAt
   created_by    String?
 
+  uploaded_archive LiteLLM_ClaudeCodePluginUpload?
+
   @@map("LiteLLM_ClaudeCodePluginTable")
 }
 
+model LiteLLM_ClaudeCodePluginUpload {
+  plugin_name String @id
+  content     Bytes
+  plugin      LiteLLM_ClaudeCodePluginTable @relation(fields: [plugin_name], references: [name], onDelete: Cascade, onUpdate: Cascade)
+}
+
+
 // User/team-scoped memory store with a GLOBAL unique key.
 // `value` is a string (typically markdown/text meant for LLM context);
 // `metadata` is an optional JSON envelope for structured tags without schema changes.
@@ -2132,6 +2141,23 @@ model LiteLLM_AccountingDelivery {
   @@index([target, state, next_attempt_at])
 }
 
+model LiteLLM_SpendLogDelivery {
+  delivery_id        String                  @id @default(dbgenerated("gen_random_uuid()::text"))
+  request_id         String
+  state              AccountingDeliveryState @default(pending)
+  accounting_payload Json
+  is_batch           Boolean
+  attempts           Int                     @default(0)
+  next_attempt_at    DateTime                @default(now())
+  applied_at         DateTime?
+  last_error         String?
+  created_at         DateTime                @default(now())
+
+  @@index([state, next_attempt_at])
+  @@index([state, applied_at])
+  @@index([request_id])
+}
+
 model LiteLLM_SubscriptionAccountTable {
   subscription_account_id String @id @default(uuid())
   custom_llm_provider String

litellm-rust/Cargo.lock

@@ -3389,15 +3389,19 @@ version = "0.1.0"
 dependencies = [
  "aws-config",
  "aws-credential-types",
+ "aws-runtime",
  "aws-sdk-sts",
  "aws-sigv4",
  "aws-smithy-runtime-api",
+ "aws-smithy-types",
  "aws-types",
  "bytes",
  "futures-util",
  "litellm-auth-types",
  "litellm-error",
  "litellm-http",
+ "litellm-secret-value",
+ "litellm-secrets",
  "moka",
  "reqwest",
  "rstest",
@@ -3406,6 +3410,7 @@ dependencies = [
  "sha2 0.10.9",
  "thiserror 2.0.19",
  "tokio",
+ "url",
  "veil",
 ]
 
@@ -3418,6 +3423,7 @@ dependencies = [
  "litellm-auth-types",
  "litellm-error",
  "litellm-secret-value",
+ "litellm-secrets",
  "moka",
  "rstest",
  "serde_json",
@@ -3431,6 +3437,7 @@ dependencies = [
 name = "litellm-auth-gcp"
 version = "0.1.0"
 dependencies = [
+ "axum",
  "base64 0.22.1",
  "gcp_auth",
  "google-cloud-auth",
@@ -3440,10 +3447,12 @@ dependencies = [
  "litellm-secret-value",
  "moka",
  "rstest",
+ "rustls 0.23.42",
  "serde",
  "serde_json",
  "sha2 0.10.9",
  "tempfile",
+ "time",
  "tokio",
  "veil",
 ]
@@ -3569,37 +3578,6 @@ dependencies = [
  "uuid",
 ]
 
-[[package]]
-name = "litellm-cache-redis"
-version = "0.1.0"
-dependencies = [
- "litellm-cache",
- "litellm-cache-testing",
- "litellm-error",
- "r2d2",
- "redis",
- "redis-test",
- "rstest",
- "serde_json",
- "thiserror 2.0.19",
- "tokio",
-]
-
-[[package]]
-name = "litellm-cache-redis-semantic"
-version = "0.1.0"
-dependencies = [
- "litellm-cache",
- "litellm-cache-redis",
- "litellm-cache-testing",
- "redis",
- "redis-test",
- "rstest",
- "serde_json",
- "sha2 0.10.9",
- "tokio",
-]
-
 [[package]]
 name = "litellm-cache-response"
 version = "0.1.0"
@@ -3607,11 +3585,8 @@ dependencies = [
  "litellm-cache",
  "litellm-cache-gcs",
  "litellm-cache-memory",
- "litellm-cache-redis",
  "litellm-http",
  "py_literal",
- "redis",
- "redis-test",
  "rstest",
  "serde",
  "serde_json",
@@ -3650,24 +3625,6 @@ dependencies = [
  "litellm-cache",
 ]
 
-[[package]]
-name = "litellm-cache-valkey-semantic"
-version = "0.1.0"
-dependencies = [
- "litellm-cache",
- "litellm-cache-redis",
- "litellm-cache-testing",
- "litellm-error",
- "redis",
- "redis-test",
- "rstest",
- "serde_json",
- "sha2 0.10.9",
- "thiserror 2.0.19",
- "tokio",
- "uuid",
-]
-
 [[package]]
 name = "litellm-callbacks-legacy-python"
 version = "0.1.0"
@@ -4339,8 +4296,6 @@ dependencies = [
  "litellm-auth",
  "litellm-auth-aws",
  "litellm-cache",
- "litellm-cache-memory",
- "litellm-cache-redis",
  "litellm-cache-response",
  "litellm-callbacks-legacy-python",
  "litellm-core-utils",
@@ -4369,7 +4324,6 @@ dependencies = [
  "prost",
  "pyo3",
  "pyo3-async-runtimes",
- "redis",
  "reqwest",
  "rstest",
  "serde",
@@ -5782,17 +5736,6 @@ version = "6.0.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf"
 
-[[package]]
-name = "r2d2"
-version = "0.8.10"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "51de85fb3fb6524929c8a2eb85e6b6d363de4e8c48f9e2c2eac4944abc181c93"
-dependencies = [
- "log",
- "parking_lot",
- "scheduled-thread-pool",
-]
-
 [[package]]
 name = "rand"
 version = "0.8.7"
@@ -5966,18 +5909,6 @@ dependencies = [
  "xxhash-rust",
 ]
 
-[[package]]
-name = "redis-test"
-version = "1.0.4"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "804d36862e4323b69f96440cbb13c9894fc90176abdeaf91264e21d5d77f6aca"
-dependencies = [
- "rand 0.9.5",
- "redis",
- "socket2 0.6.5",
- "tempfile",
-]
-
 [[package]]
 name = "redox_syscall"
 version = "0.5.18"
@@ -6454,15 +6385,6 @@ dependencies = [
  "windows-sys 0.61.2",
 ]
 
-[[package]]
-name = "scheduled-thread-pool"
-version = "0.2.7"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3cbc66816425a074528352f5789333ecff06ca41b36b0b0efdfbb29edc391a19"
-dependencies = [
- "parking_lot",
-]
-
 [[package]]
 name = "schemars"
 version = "0.9.0"

litellm-rust/Cargo.toml

@@ -57,11 +57,9 @@ litellm-db-testing = { path = "crates/db-testing" }
 litellm-cache = { path = "crates/cache" }
 litellm-cache-azure-blob = { path = "crates/cache-azure-blob" }
 litellm-cache-memory = { path = "crates/cache-memory" }
-litellm-cache-redis = { path = "crates/cache-redis" }
 litellm-cache-s3 = { path = "crates/cache-s3" }
 litellm-cache-gcs = { path = "crates/cache-gcs" }
 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" }
@@ -97,7 +95,6 @@ macro_rules_attribute = "0.2.3"
 schemars = "1"
 reqwest = { version = "0.13", default-features = false, features = ["json", "multipart", "rustls", "http2", "stream", "query"] }
 redis = "1.8.0"
-redis-test = "1.0.4"
 testcontainers-modules = "0.15.0"
 qdrant-client = { version = "1.19.0", default-features = false }
 uuid = { version = "1", features = ["v4"] }

litellm-rust/crates/auth-aws/Cargo.toml

@@ -11,14 +11,16 @@ litellm-error.workspace = true
 litellm-http.workspace = true
 reqwest.workspace = true
 
-moka = { workspace = true, features = ["sync"] }
+moka.workspace = true
+litellm-secret-value.workspace = true
 serde.workspace = true
 serde_json = { workspace = true, features = ["preserve_order"] }
 sha2.workspace = true
 thiserror.workspace = true
+url.workspace = true
 veil.workspace = true
 
-aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio"] }
+aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio", "credentials-process"] }
 aws-credential-types = { version = "1.3.0", features = ["hardcoded-credentials"] }
 aws-sdk-sts = { version = "1.108.0", default-features = false, features = ["rustls", "rt-tokio"] }
 aws-sigv4 = "1.5.1"
@@ -26,6 +28,9 @@ aws-types = "1.4.0"
 aws-smithy-runtime-api = "1.13.0"
 
 [dev-dependencies]
+litellm-secrets.workspace = true
+aws-smithy-types = "1.6.4"
+aws-runtime = { version = "1.9.4", default-features = false }
 rstest.workspace = true
 litellm-http = { workspace = true, features = ["test-support"] }
 bytes.workspace = true

litellm-rust/crates/auth-aws/src/constants.rs

@@ -11,8 +11,8 @@ pub const AWS_SESSION_NAME: &str = AwsParams::SESSION_NAME.env[0];
 pub const AWS_PROFILE_NAME: &str = AwsParams::PROFILE_NAME.env[0];
 pub const AWS_ROLE_NAME: &str = AwsParams::ROLE_NAME.env[0];
 pub const AWS_WEB_IDENTITY_TOKEN: &str = AwsParams::WEB_IDENTITY_TOKEN.env[0];
-pub const AWS_ROLE_ARN: &str = "AWS_ROLE_ARN";
-pub const AWS_WEB_IDENTITY_TOKEN_FILE: &str = "AWS_WEB_IDENTITY_TOKEN_FILE";
+pub const AWS_ROLE_ARN: &str = AwsParams::ROLE_ARN_ENV;
+pub const AWS_WEB_IDENTITY_TOKEN_FILE: &str = AwsParams::WEB_IDENTITY_TOKEN_FILE_ENV;
 pub const AWS_STS_ENDPOINT: &str = AwsParams::STS_ENDPOINT.env[0];
 pub const AWS_EXTERNAL_ID: &str = AwsParams::EXTERNAL_ID.env[0];
 pub const AWS_BEARER_TOKEN_BEDROCK: &str = "AWS_BEARER_TOKEN_BEDROCK";

litellm-rust/crates/auth-aws/src/error.rs

@@ -2,6 +2,8 @@ use litellm_error::Source;
 
 #[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)]
 pub enum Error {
+    #[error("AWS credential provider returned expired credentials")]
+    ExpiredCredentials,
     #[error("could not load AWS {flow} credentials")]
     Credentials {
         flow: CredentialFlow,

litellm-rust/crates/auth-aws/src/signer.rs

@@ -75,7 +75,8 @@ impl RequestSigner for SigV4Signer {
         )
         .map_err(|error| match error {
             Error::NonAsciiHeader { .. } => litellm_http::Error::unsignable_header(error),
-            Error::Credentials { .. }
+            Error::ExpiredCredentials
+            | Error::Credentials { .. }
             | Error::MissingWebIdentityCredentials
             | Error::WebIdentityExpiration(_)
             | Error::Signing { .. } => litellm_http::Error::signature(error),

litellm-rust/crates/auth-azure/Cargo.toml

@@ -8,17 +8,19 @@ repository.workspace = true
 [dependencies]
 litellm-secret-value.workspace = true
 litellm-auth-types.workspace = true
+litellm-error.workspace = true
 
 moka.workspace = true
 serde_json.workspace = true
 sha2.workspace = true
 strum.workspace = true
+tokio.workspace = true
 url.workspace = true
 
 azure_core = "1.0.0"
 azure_identity = { version = "1.0.0", features = ["tokio"] }
 
 [dev-dependencies]
-litellm-error.workspace = true
+litellm-secrets.workspace = true
 rstest.workspace = true
-tokio.workspace = true
+tokio = { workspace = true, features = ["test-util"] }

litellm-rust/crates/auth-azure/src/lib.rs

@@ -3,5 +3,5 @@ mod native;
 mod resolve;
 mod types;
 
-pub use resolve::{AzureAuthService, SECRET_NAMES};
+pub use resolve::{AzureAuthService, AzureIdentity, SECRET_NAMES};
 pub use types::{AzureAuthInputs, ConfigValue};

litellm-rust/crates/auth-azure/src/native.rs

@@ -59,6 +59,7 @@ pub(crate) enum NativeAzureRequest {
 pub(crate) struct ValidatedAzureRequest {
     request: NativeAzureRequest,
     credential_source: InputSource,
+    probe_imds: bool,
 }
 
 impl ValidatedAzureRequest {
@@ -68,9 +69,18 @@ impl ValidatedAzureRequest {
         Ok(Self {
             request,
             credential_source,
+            probe_imds: false,
         })
     }
 
+    pub(crate) fn with_imds_probe(self, enabled: bool) -> Self {
+        Self {
+            probe_imds: enabled
+                && matches!(self.request, NativeAzureRequest::ManagedIdentity { .. }),
+            ..self
+        }
+    }
+
     pub(crate) fn credential_source(&self) -> InputSource {
         self.credential_source
     }
@@ -87,9 +97,34 @@ impl ValidatedAzureRequest {
     }
 }
 
+pub(crate) enum CredentialAttempt {
+    Resolved(ResolvedCredential),
+    Unavailable(Error),
+    Failed(Error),
+}
+
+impl CredentialAttempt {
+    pub(crate) fn into_result(self) -> Result<ResolvedCredential, Error> {
+        match self {
+            Self::Resolved(credential) => Ok(credential),
+            Self::Unavailable(error) | Self::Failed(error) => Err(error),
+        }
+    }
+}
+
 pub(crate) struct NativeAzureTokenAcquirer {
     cache: AzureCredentialProviderCache,
     transport: Option<azure_core::http::Transport>,
+    imds_available: tokio::sync::Mutex<ImdsAvailability>,
+}
+
+enum ImdsAvailability {
+    Unknown,
+    Available,
+    Unavailable {
+        retry_at: tokio::time::Instant,
+        error: Error,
+    },
 }
 
 impl Default for NativeAzureTokenAcquirer {
@@ -103,6 +138,7 @@ impl NativeAzureTokenAcquirer {
         Self {
             cache: AzureCredentialProviderCache::new(cache_capacity),
             transport: None,
+            imds_available: tokio::sync::Mutex::new(ImdsAvailability::Unknown),
         }
     }
 
@@ -114,10 +150,51 @@ impl NativeAzureTokenAcquirer {
         Self {
             cache: AzureCredentialProviderCache::new(cache_capacity),
             transport: Some(transport),
+            imds_available: tokio::sync::Mutex::new(ImdsAvailability::Unknown),
         }
     }
 
-    pub(crate) async fn acquire(
+    pub(crate) async fn acquire(&self, request: ValidatedAzureRequest) -> CredentialAttempt {
+        let discover_imds = request.probe_imds;
+        if discover_imds && let Err(error) = self.discover_imds().await {
+            return CredentialAttempt::Unavailable(error);
+        }
+        match self.acquire_token(request).await {
+            Ok(credential) => CredentialAttempt::Resolved(credential),
+            Err(error) => {
+                let missing_identity = discover_imds && unassigned_identity(&error);
+                if missing_identity {
+                    CredentialAttempt::Unavailable(error)
+                } else {
+                    CredentialAttempt::Failed(error)
+                }
+            }
+        }
+    }
+
+    async fn discover_imds(&self) -> Result<(), Error> {
+        let mut availability = self.imds_available.lock().await;
+        match &*availability {
+            ImdsAvailability::Available => return Ok(()),
+            ImdsAvailability::Unavailable { retry_at, error }
+                if tokio::time::Instant::now() < *retry_at =>
+            {
+                return Err(error.clone());
+            }
+            ImdsAvailability::Unknown | ImdsAvailability::Unavailable { .. } => {}
+        }
+        let result = probe_imds(self.transport.clone().unwrap_or_default()).await;
+        *availability = match &result {
+            Ok(()) => ImdsAvailability::Available,
+            Err(error) => ImdsAvailability::Unavailable {
+                retry_at: tokio::time::Instant::now() + Duration::from_secs(30),
+                error: error.clone(),
+            },
+        };
+        result
+    }
+
+    async fn acquire_token(
         &self,
         request: ValidatedAzureRequest,
     ) -> Result<ResolvedCredential, Error> {
@@ -148,6 +225,54 @@ impl NativeAzureTokenAcquirer {
     }
 }
 
+fn unassigned_identity(error: &(dyn std::error::Error + 'static)) -> bool {
+    std::iter::successors(Some(error), |cause| cause.source()).any(|cause| {
+        cause
+            .downcast_ref::<azure_core::Error>()
+            .or_else(|| {
+                cause
+                    .downcast_ref::<litellm_error::Source>()?
+                    .downcast_ref::<azure_core::Error>()
+            })
+            .is_some_and(|error| {
+                error.http_status() == Some(azure_core::http::StatusCode::BadRequest)
+            })
+    })
+}
+
+async fn probe_imds(transport: azure_core::http::Transport) -> Result<(), Error> {
+    use azure_core::http::{Context, Method, Request, Url};
+
+    let endpoint =
+        Url::parse("http://169.254.169.254/metadata/identity/oauth2/token?api-version=2019-08-01")
+            .map_err(|error| {
+                Error::InvalidConfiguration(ErrorDetail::failed("Azure IMDS probe endpoint", error))
+            })?;
+    let mut request = Request::new(endpoint, Method::Get);
+    let response = tokio::time::timeout(
+        Duration::from_secs(1),
+        transport.send(&Context::new(), &mut request),
+    )
+    .await
+    .map_err(|error| {
+        Error::CredentialAcquisition(ErrorDetail::failed("Azure IMDS availability probe", error))
+    })?
+    .map_err(|error| {
+        Error::CredentialAcquisition(ErrorDetail::failed("Azure IMDS availability probe", error))
+    })?;
+    // Headerless IMDS discovery expects 400, checked 2026-11-01
+    // https://learn.microsoft.com/en-us/azure/virtual-machines/instance-metadata-service#errors-and-debugging
+    if response.status() != azure_core::http::StatusCode::BadRequest {
+        return Err(Error::CredentialAcquisition(
+            ErrorDetail::UnexpectedStatus {
+                operation: "Azure IMDS availability probe",
+                status: u16::from(response.status()),
+            },
+        ));
+    }
+    Ok(())
+}
+
 impl NativeAzureRequest {
     fn scope(&self) -> &str {
         match self {
@@ -493,6 +618,7 @@ impl ClientAssertion for StaticAssertion {
 #[cfg(test)]
 mod tests {
     use std::sync::{Arc, Mutex};
+    use std::time::Duration;
 
     use azure_core::http::headers::Headers;
     use azure_core::http::{AsyncRawResponse, HttpClient, Request, StatusCode, Transport};
@@ -569,6 +695,296 @@ mod tests {
         }
     }
 
+    #[derive(Debug)]
+    struct AvailableImdsClient {
+        probe_status: StatusCode,
+        probe_delay: Duration,
+        recover_after: Option<usize>,
+        response_status: Option<StatusCode>,
+        probes: std::sync::atomic::AtomicUsize,
+    }
+
+    impl Default for AvailableImdsClient {
+        fn default() -> Self {
+            Self {
+                probe_status: StatusCode::BadRequest,
+                probe_delay: Duration::ZERO,
+                recover_after: None,
+                response_status: None,
+                probes: std::sync::atomic::AtomicUsize::new(0),
+            }
+        }
+    }
+
+    impl HttpClient for AvailableImdsClient {
+        fn execute_request<'life0, 'life1, 'async_trait>(
+            &'life0 self,
+            request: &'life1 Request,
+        ) -> std::pin::Pin<
+            Box<dyn std::future::Future<Output = Result<AsyncRawResponse>> + Send + 'async_trait>,
+        >
+        where
+            'life0: 'async_trait,
+            'life1: 'async_trait,
+            Self: 'async_trait,
+        {
+            Box::pin(async move {
+                if request
+                    .headers()
+                    .get_optional_str(&azure_core::http::headers::HeaderName::from_static(
+                        "metadata",
+                    ))
+                    .is_none()
+                {
+                    let probe = self
+                        .probes
+                        .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+                    if self.recover_after.is_some_and(|failures| probe >= failures) {
+                        return Ok(AsyncRawResponse::from_bytes(
+                            StatusCode::BadRequest,
+                            Headers::new(),
+                            "metadata header required",
+                        ));
+                    }
+                    tokio::time::sleep(self.probe_delay).await;
+                    return Ok(AsyncRawResponse::from_bytes(
+                        self.probe_status,
+                        Headers::new(),
+                        "metadata header required",
+                    ));
+                }
+                if let Some(status) = self.response_status {
+                    return Ok(AsyncRawResponse::from_bytes(
+                        status,
+                        Headers::new(),
+                        "identity endpoint failure",
+                    ));
+                }
+                tokio::time::sleep(std::time::Duration::from_secs(2)).await;
+                let expires = std::time::SystemTime::now()
+                    .duration_since(std::time::UNIX_EPOCH)
+                    .unwrap()
+                    .as_secs()
+                    + 3600;
+                Ok(AsyncRawResponse::from_bytes(
+                    StatusCode::Ok,
+                    Headers::new(),
+                    serde_json::json!({
+                        "access_token": "managed-token", "expires_on": expires.to_string(),
+                        "token_type": "Bearer", "resource": "https://service.test",
+                    })
+                    .to_string(),
+                ))
+            })
+        }
+    }
+
+    #[rstest::rstest]
+    #[tokio::test(start_paused = true)]
+    async fn available_imds_retains_full_sdk_token_acquisition() {
+        let acquirer = NativeAzureTokenAcquirer::with_transport(
+            4,
+            Transport::new(Arc::new(AvailableImdsClient {
+                probe_status: StatusCode::BadRequest,
+                response_status: None,
+                ..Default::default()
+            })),
+        );
+        let request = ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity {
+            client_id: None,
+            scope: deployment("https://service.test/.default".into()),
+            selection_source: InputSource::Deployment,
+        })
+        .unwrap()
+        .with_imds_probe(true);
+        let credential = acquirer.acquire(request).await.into_result().unwrap();
+        assert_eq!(credential.secret().expose(), "managed-token");
+    }
+
+    #[rstest::rstest]
+    #[case::another_metadata_service(StatusCode::NotFound)]
+    #[case::token_required(StatusCode::Unauthorized)]
+    #[case::proxy_rejection(StatusCode::Forbidden)]
+    #[case::unavailable(StatusCode::ServiceUnavailable)]
+    #[tokio::test(start_paused = true)]
+    async fn unrelated_metadata_responses_do_not_start_managed_identity_acquisition(
+        #[case] status: StatusCode,
+    ) {
+        let client = Arc::new(AvailableImdsClient {
+            probe_status: status,
+            ..Default::default()
+        });
+        let acquirer = NativeAzureTokenAcquirer::with_transport(4, Transport::new(client.clone()));
+        let request = ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity {
+            client_id: None,
+            scope: deployment("https://service.test/.default".into()),
+            selection_source: InputSource::Deployment,
+        })
+        .unwrap()
+        .with_imds_probe(true);
+        assert!(matches!(
+            acquirer.acquire(request.clone()).await,
+            super::CredentialAttempt::Unavailable(_)
+        ));
+        assert!(matches!(
+            acquirer.acquire(request).await,
+            super::CredentialAttempt::Unavailable(_)
+        ));
+        assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 1);
+    }
+
+    #[rstest::rstest]
+    #[tokio::test(start_paused = true)]
+    async fn an_imds_timeout_is_not_repeated_on_the_next_request() {
+        let client = Arc::new(AvailableImdsClient {
+            probe_delay: Duration::from_secs(10),
+            ..Default::default()
+        });
+        let acquirer = NativeAzureTokenAcquirer::with_transport(4, Transport::new(client.clone()));
+        let request = ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity {
+            client_id: None,
+            scope: deployment("https://service.test/.default".into()),
+            selection_source: InputSource::Deployment,
+        })
+        .unwrap()
+        .with_imds_probe(true);
+        assert!(matches!(
+            acquirer.acquire(request.clone()).await,
+            super::CredentialAttempt::Unavailable(_)
+        ));
+        let second = tokio::time::timeout(Duration::from_millis(10), acquirer.acquire(request))
+            .await
+            .unwrap();
+        assert!(matches!(second, super::CredentialAttempt::Unavailable(_)));
+        assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 1);
+    }
+
+    #[rstest::fixture]
+    fn imds_request() -> ValidatedAzureRequest {
+        ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity {
+            client_id: None,
+            scope: deployment("https://service.test/.default".into()),
+            selection_source: InputSource::Deployment,
+        })
+        .unwrap()
+        .with_imds_probe(true)
+    }
+
+    #[rstest::rstest]
+    #[case::timeout(StatusCode::BadRequest, Duration::from_secs(10))]
+    #[case::service_unavailable(StatusCode::ServiceUnavailable, Duration::ZERO)]
+    #[tokio::test(start_paused = true)]
+    async fn failed_imds_discovery_recovers_without_restarting_the_acquirer(
+        imds_request: ValidatedAzureRequest,
+        #[case] status: StatusCode,
+        #[case] delay: Duration,
+    ) {
+        let client = Arc::new(AvailableImdsClient {
+            probe_status: status,
+            probe_delay: delay,
+            recover_after: Some(1),
+            ..Default::default()
+        });
+        let acquirer = NativeAzureTokenAcquirer::with_transport(4, Transport::new(client.clone()));
+        let (first, second) = tokio::join!(
+            acquirer.acquire(imds_request.clone()),
+            acquirer.acquire(imds_request.clone()),
+        );
+        assert!(matches!(first, super::CredentialAttempt::Unavailable(_)));
+        assert!(matches!(second, super::CredentialAttempt::Unavailable(_)));
+        assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 1);
+
+        tokio::time::advance(Duration::from_secs(30)).await;
+        let (first, second) = tokio::join!(
+            acquirer.acquire(imds_request.clone()),
+            acquirer.acquire(imds_request.clone()),
+        );
+        assert_eq!(
+            first.into_result().unwrap().secret().expose(),
+            "managed-token"
+        );
+        assert_eq!(
+            second.into_result().unwrap().secret().expose(),
+            "managed-token"
+        );
+        assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 2);
+
+        tokio::time::advance(Duration::from_secs(60)).await;
+        assert_eq!(
+            acquirer
+                .acquire(imds_request)
+                .await
+                .into_result()
+                .unwrap()
+                .secret()
+                .expose(),
+            "managed-token",
+        );
+        assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 2);
+    }
+
+    #[rstest::rstest]
+    #[tokio::test(start_paused = true)]
+    async fn a_persistent_imds_outage_reuses_each_failed_probe(
+        imds_request: ValidatedAzureRequest,
+    ) {
+        let client = Arc::new(AvailableImdsClient {
+            probe_delay: Duration::from_secs(10),
+            ..Default::default()
+        });
+        let acquirer = NativeAzureTokenAcquirer::with_transport(4, Transport::new(client.clone()));
+        assert!(matches!(
+            acquirer.acquire(imds_request.clone()).await,
+            super::CredentialAttempt::Unavailable(_),
+        ));
+        tokio::time::advance(Duration::from_secs(30)).await;
+        let (first, second) = tokio::join!(
+            acquirer.acquire(imds_request.clone()),
+            acquirer.acquire(imds_request),
+        );
+        assert!(matches!(first, super::CredentialAttempt::Unavailable(_)));
+        assert!(matches!(second, super::CredentialAttempt::Unavailable(_)));
+        assert_eq!(client.probes.load(std::sync::atomic::Ordering::SeqCst), 2);
+    }
+
+    #[rstest::rstest]
+    #[case::default_identity_not_assigned(true, StatusCode::BadRequest, true)]
+    #[case::explicit_identity_not_assigned(false, StatusCode::BadRequest, false)]
+    #[case::default_identity_forbidden(true, StatusCode::Forbidden, false)]
+    #[case::default_identity_service_failure(true, StatusCode::InternalServerError, false)]
+    #[tokio::test(start_paused = true)]
+    async fn only_unavailable_default_identities_allow_fallback(
+        #[case] discovery: bool,
+        #[case] status: StatusCode,
+        #[case] unavailable: bool,
+    ) {
+        let acquirer = NativeAzureTokenAcquirer::with_transport(
+            4,
+            Transport::new(Arc::new(AvailableImdsClient {
+                probe_status: StatusCode::BadRequest,
+                response_status: Some(status),
+                ..Default::default()
+            })),
+        );
+        let request = ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity {
+            client_id: None,
+            scope: deployment("https://service.test/.default".into()),
+            selection_source: InputSource::Deployment,
+        })
+        .unwrap()
+        .with_imds_probe(discovery);
+        let outcome = acquirer.acquire(request).await;
+        let is_unavailable = matches!(outcome, super::CredentialAttempt::Unavailable(_));
+        let error = outcome.into_result().unwrap_err();
+        assert_eq!(
+            is_unavailable,
+            unavailable,
+            "{}",
+            litellm_error::Report::new(&error)
+        );
+    }
+
+    #[rstest::rstest]
     #[tokio::test]
     async fn client_secret_uses_sdk_protocol_and_reuses_cached_credential() {
         let transport = Arc::new(RecordingTokenClient::default());
@@ -582,8 +998,12 @@ mod tests {
             "https://login.test",
         );
 
-        let first = acquirer.acquire(request.clone()).await.unwrap();
-        let second = acquirer.acquire(request).await.unwrap();
+        let first = acquirer
+            .acquire(request.clone())
+            .await
+            .into_result()
+            .unwrap();
+        let second = acquirer.acquire(request).await.into_result().unwrap();
 
         assert_eq!(first.secret().expose(), "native-token");
         assert_eq!(second.secret().expose(), "native-token");
@@ -599,6 +1019,7 @@ mod tests {
         );
     }
 
+    #[rstest::rstest]
     #[tokio::test]
     async fn credential_provider_cache_isolates_every_client_secret_identity_field() {
         let transport = Arc::new(RecordingTokenClient::default());
@@ -645,10 +1066,10 @@ mod tests {
             ),
         ];
 
-        acquirer.acquire(base.clone()).await.unwrap();
-        acquirer.acquire(base).await.unwrap();
+        acquirer.acquire(base.clone()).await.into_result().unwrap();
+        acquirer.acquire(base).await.into_result().unwrap();
         for request in variants.into_iter().skip(1) {
-            acquirer.acquire(request).await.unwrap();
+            acquirer.acquire(request).await.into_result().unwrap();
         }
 
         assert_eq!(transport.requests.lock().unwrap().len(), 6);

litellm-rust/crates/auth-azure/src/resolve.rs

@@ -1,3 +1,4 @@
+use crate::native::CredentialAttempt;
 use litellm_auth_types::{
     CredentialFileRef, CredentialLookup, CredentialRef, InputSource, ResolvedCredential, Sourced,
     TokenProviderHandle,
@@ -12,25 +13,31 @@ use std::sync::Arc;
 use super::native::{NativeAzureRequest, NativeAzureTokenAcquirer, ValidatedAzureRequest};
 use super::types::{AzureAuthInputs, AzureCredentialType, ConfigValue, DEFAULT_AZURE_SCOPE};
 
-const AZURE_AD_TOKEN_ENV: &str = "AZURE_AD_TOKEN";
-const AZURE_TENANT_ID_ENV: &str = "AZURE_TENANT_ID";
-const AZURE_CLIENT_ID_ENV: &str = "AZURE_CLIENT_ID";
-const AZURE_CLIENT_SECRET_ENV: &str = "AZURE_CLIENT_SECRET";
-const AZURE_SCOPE_ENV: &str = "AZURE_SCOPE";
-const AZURE_AUTHORITY_HOST_ENV: &str = "AZURE_AUTHORITY_HOST";
-const AZURE_CREDENTIAL_ENV: &str = "AZURE_CREDENTIAL";
-const AZURE_FEDERATED_TOKEN_FILE_ENV: &str = "AZURE_FEDERATED_TOKEN_FILE";
+#[derive(strum::IntoStaticStr, strum::VariantArray, strum::VariantNames)]
+enum AzureSecretName {
+    #[strum(serialize = "AZURE_AD_TOKEN")]
+    AdToken,
+    #[strum(serialize = "AZURE_TENANT_ID")]
+    TenantId,
+    #[strum(serialize = "AZURE_CLIENT_ID")]
+    ClientId,
+    #[strum(serialize = "AZURE_CLIENT_SECRET")]
+    ClientSecret,
+    #[strum(serialize = "AZURE_SCOPE")]
+    Scope,
+    #[strum(serialize = "AZURE_AUTHORITY_HOST")]
+    AuthorityHost,
+    #[strum(serialize = "AZURE_CREDENTIAL")]
+    Credential,
+    #[strum(serialize = "AZURE_FEDERATED_TOKEN_FILE")]
+    FederatedTokenFile,
+    #[strum(serialize = "IDENTITY_ENDPOINT")]
+    IdentityEndpoint,
+    #[strum(serialize = "MSI_ENDPOINT")]
+    MsiEndpoint,
+}
 
-pub const SECRET_NAMES: &[&str] = &[
-    AZURE_AD_TOKEN_ENV,
-    AZURE_TENANT_ID_ENV,
-    AZURE_CLIENT_ID_ENV,
-    AZURE_CLIENT_SECRET_ENV,
-    AZURE_SCOPE_ENV,
-    AZURE_AUTHORITY_HOST_ENV,
-    AZURE_CREDENTIAL_ENV,
-    AZURE_FEDERATED_TOKEN_FILE_ENV,
-];
+pub const SECRET_NAMES: &[&str] = <AzureSecretName as strum::VariantNames>::VARIANTS;
 
 #[derive(Clone, Debug)]
 pub(crate) enum AzureCredentialPlan {
@@ -57,14 +64,14 @@ trait AzureTokenAcquirer: Send + Sync {
     fn acquire(
         &self,
         request: ValidatedAzureRequest,
-    ) -> Pin<Box<dyn Future<Output = Result<ResolvedCredential, Error>> + Send + '_>>;
+    ) -> Pin<Box<dyn Future<Output = CredentialAttempt> + Send + '_>>;
 }
 
 impl AzureTokenAcquirer for NativeAzureTokenAcquirer {
     fn acquire(
         &self,
         request: ValidatedAzureRequest,
-    ) -> Pin<Box<dyn Future<Output = Result<ResolvedCredential, Error>> + Send + '_>> {
+    ) -> Pin<Box<dyn Future<Output = CredentialAttempt> + Send + '_>> {
         Box::pin(NativeAzureTokenAcquirer::acquire(self, request))
     }
 }
@@ -83,12 +90,36 @@ impl AzureAuthService {
         Self { native }
     }
 
+    pub async fn identity_token(
+        &self,
+        identity: &AzureIdentity,
+        scope: &str,
+        env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
+    ) -> Result<Option<ResolvedCredential>, Error> {
+        let plan = identity_plan(identity, scope, env_lookup)?;
+        let inputs = AzureAuthInputs::default_credential_for_scope(scope);
+        Ok(self
+            .run_plan(plan, &inputs, env_lookup)
+            .await?
+            .map(Sourced::into_value))
+    }
+
     pub async fn get_azure_ad_token(
         &self,
         inputs: &AzureAuthInputs,
         env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
     ) -> Result<Option<Sourced<ResolvedCredential>>, Error> {
-        match select_auth_plan(inputs, env_lookup)? {
+        self.run_plan(select_auth_plan(inputs, env_lookup)?, inputs, env_lookup)
+            .await
+    }
+
+    async fn run_plan(
+        &self,
+        plan: AzureCredentialPlan,
+        inputs: &AzureAuthInputs,
+        env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
+    ) -> Result<Option<Sourced<ResolvedCredential>>, Error> {
+        match plan {
             AzureCredentialPlan::Supplied(credential) => Ok(Some(credential)),
             AzureCredentialPlan::Caller(caller) => {
                 let credential = caller.acquire().await?;
@@ -123,6 +154,7 @@ impl AzureAuthService {
                 self.native
                     .acquire(request)
                     .await
+                    .into_result()
                     .map(|credential| Sourced::new(credential, source))
                     .map(Some)
             }
@@ -131,6 +163,7 @@ impl AzureAuthService {
                 self.native
                     .acquire(request)
                     .await
+                    .into_result()
                     .map(|credential| Some(Sourced::new(credential, source)))
             }
             AzureCredentialPlan::Chain(requests) => {
@@ -138,8 +171,16 @@ impl AzureAuthService {
                 for request in requests {
                     let source = request.credential_source();
                     match self.native.acquire(request).await {
-                        Ok(credential) => return Ok(Some(Sourced::new(credential, source))),
-                        Err(error) => failures.push(error),
+                        CredentialAttempt::Resolved(credential) => {
+                            return Ok(Some(Sourced::new(credential, source)));
+                        }
+                        CredentialAttempt::Unavailable(error) => failures.push(error),
+                        CredentialAttempt::Failed(error) => {
+                            return Err(Error::CredentialChain {
+                                last: Box::new(error),
+                                earlier: failures,
+                            });
+                        }
                     }
                 }
                 match failures.pop() {
@@ -155,39 +196,161 @@ impl AzureAuthService {
     }
 }
 
+/// The identity azure-identity builds for one explicit credential class, without LiteLLM's provider settings
+/// (`AZURE_AD_TOKEN`, `AZURE_CREDENTIAL`, `AZURE_SCOPE`) that `get_azure_ad_token` also reads.
+#[derive(Clone, Debug)]
+pub enum AzureIdentity {
+    ClientSecret {
+        tenant_id: String,
+        client_id: String,
+        client_secret: SecretValue,
+    },
+    ManagedIdentity {
+        client_id: String,
+    },
+    Default,
+}
+
+fn identity_plan(
+    identity: &AzureIdentity,
+    scope: &str,
+    env_lookup: &dyn Fn(&str) -> Option<String>,
+) -> Result<AzureCredentialPlan, Error> {
+    let deployment = |value: &str| Sourced::new(value.to_owned(), InputSource::Deployment);
+    let environment = |name: &str| {
+        env_lookup(name)
+            .filter(|value| !value.trim().is_empty())
+            .map(|value| Sourced::new(value, InputSource::Environment))
+    };
+    let scope = deployment(scope);
+    let authority = environment(AzureSecretName::AuthorityHost.into());
+    match identity {
+        AzureIdentity::ClientSecret {
+            tenant_id,
+            client_id,
+            client_secret,
+        } => Ok(AzureCredentialPlan::Native(ValidatedAzureRequest::new(
+            NativeAzureRequest::ClientSecret {
+                tenant_id: deployment(tenant_id),
+                client_id: deployment(client_id),
+                client_secret: Sourced::new(client_secret.clone(), InputSource::Deployment),
+                scope,
+                authority,
+            },
+        )?)),
+        AzureIdentity::ManagedIdentity { client_id } => Ok(AzureCredentialPlan::Native(
+            ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity {
+                client_id: Some(deployment(client_id)),
+                scope,
+                selection_source: InputSource::Deployment,
+            })?,
+        )),
+        AzureIdentity::Default => {
+            let tenant_id = environment(AzureSecretName::TenantId.into());
+            let client_id = environment(AzureSecretName::ClientId.into());
+            let from_environment = match (
+                tenant_id.clone(),
+                client_id.clone(),
+                environment(AzureSecretName::ClientSecret.into()),
+            ) {
+                (Some(tenant_id), Some(client_id), Some(client_secret)) => {
+                    Some(NativeAzureRequest::ClientSecret {
+                        tenant_id,
+                        client_id,
+                        client_secret: client_secret.map(SecretValue::new),
+                        scope: scope.clone(),
+                        authority: authority.clone(),
+                    })
+                }
+                _ => None,
+            };
+            let probe_imds = env_lookup(AzureSecretName::IdentityEndpoint.into()).is_none()
+                && env_lookup(AzureSecretName::MsiEndpoint.into()).is_none();
+            let AzureCredentialPlan::Chain(chain) = select_native_plan(
+                Some(Sourced::new(
+                    AzureCredentialType::DefaultAzureCredential,
+                    InputSource::Deployment,
+                )),
+                tenant_id,
+                client_id,
+                environment(AzureSecretName::FederatedTokenFile.into()),
+                scope,
+                authority,
+                InputSource::Environment,
+            )?
+            else {
+                unreachable!("DefaultAzureCredential always plans a chain")
+            };
+            let environment_request = from_environment
+                .map(ValidatedAzureRequest::new)
+                .transpose()?;
+            Ok(AzureCredentialPlan::Chain(
+                environment_request
+                    .into_iter()
+                    .chain(chain)
+                    .map(|request| request.with_imds_probe(probe_imds))
+                    .collect(),
+            ))
+        }
+    }
+}
+
 pub(crate) fn select_auth_plan(
     inputs: &AzureAuthInputs,
     env_lookup: &dyn Fn(&str) -> Option<String>,
 ) -> Result<AzureCredentialPlan, Error> {
-    let token = configured_secret(&inputs.azure_ad_token, AZURE_AD_TOKEN_ENV, env_lookup);
-    let tenant_id = configured_string(&inputs.tenant_id, AZURE_TENANT_ID_ENV, env_lookup);
-    let client_id = configured_string(&inputs.client_id, AZURE_CLIENT_ID_ENV, env_lookup);
-    let client_secret =
-        configured_secret(&inputs.client_secret, AZURE_CLIENT_SECRET_ENV, env_lookup);
-    let scope = configured_string(&inputs.azure_scope, AZURE_SCOPE_ENV, env_lookup)
-        .unwrap_or_else(|| Sourced::new(DEFAULT_AZURE_SCOPE.to_string(), InputSource::Environment));
-    let authority = configured_string(
-        &inputs.azure_authority_host,
-        AZURE_AUTHORITY_HOST_ENV,
+    let token = configured_secret(
+        &inputs.azure_ad_token,
+        AzureSecretName::AdToken.into(),
         env_lookup,
     );
-    let selector = configured_string(&inputs.azure_credential, AZURE_CREDENTIAL_ENV, env_lookup)
-        .map(|value| {
-            value
-                .value()
-                .parse::<AzureCredentialType>()
-                .map(|selector| Sourced::new(selector, value.source()))
-        })
-        .transpose()
-        .map_err(|error| {
-            Error::InvalidConfiguration(ErrorDetail::failed(
-                "Azure credential selector parsing",
-                error,
-            ))
-        })?;
+    let tenant_id = configured_string(
+        &inputs.tenant_id,
+        AzureSecretName::TenantId.into(),
+        env_lookup,
+    );
+    let client_id = configured_string(
+        &inputs.client_id,
+        AzureSecretName::ClientId.into(),
+        env_lookup,
+    );
+    let client_secret = configured_secret(
+        &inputs.client_secret,
+        AzureSecretName::ClientSecret.into(),
+        env_lookup,
+    );
+    let scope = configured_string(
+        &inputs.azure_scope,
+        AzureSecretName::Scope.into(),
+        env_lookup,
+    )
+    .unwrap_or_else(|| Sourced::new(DEFAULT_AZURE_SCOPE.to_string(), InputSource::Environment));
+    let authority = configured_string(
+        &inputs.azure_authority_host,
+        AzureSecretName::AuthorityHost.into(),
+        env_lookup,
+    );
+    let selector = configured_string(
+        &inputs.azure_credential,
+        AzureSecretName::Credential.into(),
+        env_lookup,
+    )
+    .map(|value| {
+        value
+            .value()
+            .parse::<AzureCredentialType>()
+            .map(|selector| Sourced::new(selector, value.source()))
+    })
+    .transpose()
+    .map_err(|error| {
+        Error::InvalidConfiguration(ErrorDetail::failed(
+            "Azure credential selector parsing",
+            error,
+        ))
+    })?;
     let federated_token_file = configured_string(
         &inputs.federated_token_file,
-        AZURE_FEDERATED_TOKEN_FILE_ENV,
+        AzureSecretName::FederatedTokenFile.into(),
         env_lookup,
     );
 
@@ -237,7 +400,12 @@ pub(crate) fn select_auth_plan(
         return Ok(AzureCredentialPlan::Missing);
     }
 
-    select_native_plan(
+    let probe_imds = selector
+        .as_ref()
+        .is_none_or(|selector| *selector.value() == AzureCredentialType::DefaultAzureCredential)
+        && env_lookup(AzureSecretName::IdentityEndpoint.into()).is_none()
+        && env_lookup(AzureSecretName::MsiEndpoint.into()).is_none();
+    match select_native_plan(
         selector,
         tenant_id,
         client_id,
@@ -245,7 +413,15 @@ pub(crate) fn select_auth_plan(
         scope,
         authority,
         inputs.enable_azure_ad_token_refresh.source(),
-    )
+    )? {
+        AzureCredentialPlan::Chain(requests) => Ok(AzureCredentialPlan::Chain(
+            requests
+                .into_iter()
+                .map(|request| request.with_imds_probe(probe_imds))
+                .collect(),
+        )),
+        plan => Ok(plan),
+    }
 }
 
 fn select_native_plan(
@@ -509,7 +685,10 @@ mod tests {
     #[derive(Debug)]
     struct FileResolver;
 
+    use crate::native::CredentialAttempt;
+
     struct ChainAcquirer {
+        terminal_on: Option<&'static str>,
         requests: Mutex<Vec<&'static str>>,
         succeed_on: Option<&'static str>,
     }
@@ -518,20 +697,25 @@ mod tests {
         fn acquire(
             &self,
             request: ValidatedAzureRequest,
-        ) -> std::pin::Pin<Box<dyn Future<Output = Result<ResolvedCredential, Error>> + Send + '_>>
-        {
+        ) -> std::pin::Pin<Box<dyn Future<Output = CredentialAttempt> + Send + '_>> {
             let kind = request.kind();
             self.requests.lock().unwrap().push(kind);
             Box::pin(async move {
                 if self.succeed_on == Some(kind) {
-                    Ok(ResolvedCredential::AccessToken {
+                    CredentialAttempt::Resolved(ResolvedCredential::AccessToken {
                         token: SecretValue::new("chain-token"),
                         expires_on: None,
                     })
                 } else {
-                    Err(Error::CredentialAcquisition(
-                        litellm_auth_types::ErrorDetail::Unresolved { subject: kind },
-                    ))
+                    let error =
+                        Error::CredentialAcquisition(litellm_auth_types::ErrorDetail::Unresolved {
+                            subject: kind,
+                        });
+                    if self.terminal_on == Some(kind) {
+                        CredentialAttempt::Failed(error)
+                    } else {
+                        CredentialAttempt::Unavailable(error)
+                    }
                 }
             })
         }
@@ -690,9 +874,210 @@ mod tests {
         assert_eq!(resolved, Some(SecretValue::new("rotated-assertion")));
     }
 
+    #[derive(Debug)]
+    struct UnavailableImdsClient;
+
+    impl azure_core::http::HttpClient for UnavailableImdsClient {
+        fn execute_request<'life0, 'life1, 'async_trait>(
+            &'life0 self,
+            _request: &'life1 azure_core::http::Request,
+        ) -> std::pin::Pin<
+            Box<
+                dyn Future<Output = azure_core::Result<azure_core::http::AsyncRawResponse>>
+                    + Send
+                    + 'async_trait,
+            >,
+        >
+        where
+            'life0: 'async_trait,
+            'life1: 'async_trait,
+            Self: 'async_trait,
+        {
+            Box::pin(std::future::pending())
+        }
+    }
+
+    struct OfflineChainAcquirer(super::NativeAzureTokenAcquirer);
+
+    impl AzureTokenAcquirer for OfflineChainAcquirer {
+        fn acquire(
+            &self,
+            request: ValidatedAzureRequest,
+        ) -> std::pin::Pin<Box<dyn Future<Output = CredentialAttempt> + Send + '_>> {
+            Box::pin(async move {
+                if request.kind() == "developer-tools" {
+                    return CredentialAttempt::Resolved(ResolvedCredential::AccessToken {
+                        token: SecretValue::new("cli-token"),
+                        expires_on: None,
+                    });
+                }
+                self.0.acquire(request).await
+            })
+        }
+    }
+
+    #[rstest::rstest]
+    #[case::default_off_azure("DefaultAzureCredential", None, true)]
+    #[case::explicit_managed_identity("ManagedIdentityCredential", None, false)]
+    #[case::app_service("DefaultAzureCredential", Some("http://app-service.test"), false)]
+    #[tokio::test(start_paused = true)]
+    async fn only_default_imds_discovery_is_bounded(
+        #[case] selector: &str,
+        #[case] identity_endpoint: Option<&str>,
+        #[case] expect_cli: bool,
+    ) {
+        let native = super::NativeAzureTokenAcquirer::with_transport(
+            4,
+            azure_core::http::Transport::new(Arc::new(UnavailableImdsClient)),
+        );
+        let service = AzureAuthService::with_acquirer(Arc::new(OfflineChainAcquirer(native)));
+        let inputs = AzureAuthInputs {
+            azure_credential: crate::types::ConfigValue::Value(Sourced::new(
+                selector.into(),
+                InputSource::Deployment,
+            )),
+            enable_azure_ad_token_refresh: Sourced::new(true, InputSource::Deployment),
+            ..Default::default()
+        };
+        let environment = |name: &str| {
+            (name == "IDENTITY_ENDPOINT")
+                .then_some(identity_endpoint)
+                .flatten()
+                .map(str::to_string)
+        };
+        let outcome = tokio::time::timeout(
+            std::time::Duration::from_secs(2),
+            service.get_azure_ad_token(&inputs, &environment),
+        )
+        .await;
+
+        assert_eq!(outcome.is_ok(), expect_cli);
+        if expect_cli {
+            assert_eq!(
+                outcome.unwrap().unwrap().unwrap().value().secret().expose(),
+                "cli-token"
+            );
+        }
+    }
+
+    #[rstest::rstest]
+    #[case::app_service(Some("IDENTITY_ENDPOINT"), false)]
+    #[case::legacy_app_service(Some("MSI_ENDPOINT"), false)]
+    #[case::off_azure(None, true)]
+    #[tokio::test(start_paused = true)]
+    async fn default_chain_uses_endpoints_from_the_secret_snapshot(
+        #[case] endpoint: Option<&str>,
+        #[case] expect_cli: bool,
+        #[values(false, true)] consumer_identity: bool,
+    ) {
+        use litellm_secrets::source::SecretSource;
+
+        struct EndpointSecrets<'a>(Option<&'a str>);
+
+        impl SecretSource for EndpointSecrets<'_> {
+            fn get_secret_str<'a>(
+                &'a self,
+                name: &'a str,
+            ) -> std::pin::Pin<
+                Box<
+                    dyn Future<Output = Result<Option<SecretValue>, litellm_secrets::Error>>
+                        + Send
+                        + 'a,
+                >,
+            > {
+                Box::pin(async move {
+                    Ok((self.0 == Some(name)).then(|| SecretValue::new("http://app-service.test")))
+                })
+            }
+        }
+
+        let snapshot = EndpointSecrets(endpoint)
+            .resolve(SECRET_NAMES)
+            .await
+            .unwrap();
+        let native = super::NativeAzureTokenAcquirer::with_transport(
+            4,
+            azure_core::http::Transport::new(Arc::new(UnavailableImdsClient)),
+        );
+        let service = AzureAuthService::with_acquirer(Arc::new(OfflineChainAcquirer(native)));
+        let lookup = |name: &str| snapshot.get(name);
+        let scope = "https://service.test/.default";
+        let outcome = tokio::time::timeout(std::time::Duration::from_secs(2), async {
+            if consumer_identity {
+                service
+                    .identity_token(&super::AzureIdentity::Default, scope, &lookup)
+                    .await
+            } else {
+                service
+                    .get_azure_ad_token(
+                        &AzureAuthInputs::default_credential_for_scope(scope),
+                        &lookup,
+                    )
+                    .await
+                    .map(|token| token.map(Sourced::into_value))
+            }
+        })
+        .await;
+
+        assert_eq!(outcome.is_ok(), expect_cli);
+        if expect_cli {
+            assert_eq!(
+                outcome.unwrap().unwrap().unwrap().secret().expose(),
+                "cli-token"
+            );
+        }
+    }
+
+    #[rstest::rstest]
+    #[case::litellm_provider_settings_are_ignored(
+        &[("AZURE_AD_TOKEN", "provider-token"), ("AZURE_CREDENTIAL", "ClientSecretCredential")],
+        &["managed-identity", "developer-tools"]
+    )]
+    #[case::environment_client_secret_comes_first(
+        &[("AZURE_TENANT_ID", "tenant"), ("AZURE_CLIENT_ID", "client"), ("AZURE_CLIENT_SECRET", "secret")],
+        &["client-secret", "managed-identity", "developer-tools"]
+    )]
+    #[case::aks_workload_identity(
+        &[("AZURE_TENANT_ID", "tenant"), ("AZURE_CLIENT_ID", "client"), ("AZURE_FEDERATED_TOKEN_FILE", "/var/run/token")],
+        &["workload-identity", "managed-identity", "developer-tools"]
+    )]
+    #[tokio::test]
+    async fn the_default_identity_follows_azure_identity_defaults(
+        #[case] environment: &[(&str, &str)],
+        #[case] expected: &[&str],
+    ) {
+        let acquirer = Arc::new(ChainAcquirer {
+            terminal_on: None,
+            requests: Mutex::new(Vec::new()),
+            succeed_on: Some("developer-tools"),
+        });
+        let service = AzureAuthService::with_acquirer(acquirer.clone());
+        let lookup = |name: &str| {
+            environment
+                .iter()
+                .find(|(key, _)| *key == name)
+                .map(|(_, value)| (*value).to_owned())
+        };
+
+        let credential = service
+            .identity_token(
+                &super::AzureIdentity::Default,
+                "https://redis.azure.com/.default",
+                &lookup,
+            )
+            .await
+            .unwrap()
+            .unwrap();
+
+        assert_eq!(credential.secret().expose(), "chain-token");
+        assert_eq!(*acquirer.requests.lock().unwrap(), expected);
+    }
+
+    #[rstest::rstest]
     #[tokio::test]
     async fn default_chain_uses_declared_order_and_stops_after_success() {
         let acquirer = Arc::new(ChainAcquirer {
+            terminal_on: None,
             requests: Mutex::new(Vec::new()),
             succeed_on: Some("developer-tools"),
         });
@@ -715,9 +1100,13 @@ mod tests {
         );
     }
 
+    #[rstest::rstest]
+    #[case::all_unavailable(false)]
+    #[case::terminal_failure_after_unavailable(true)]
     #[tokio::test]
-    async fn chain_reports_each_acquisition_failure() {
+    async fn chain_reports_each_acquisition_failure(#[case] terminal: bool) {
         let acquirer = Arc::new(ChainAcquirer {
+            terminal_on: terminal.then_some("developer-tools"),
             requests: Mutex::new(Vec::new()),
             succeed_on: None,
         });
@@ -744,6 +1133,27 @@ mod tests {
         );
     }
 
+    #[rstest::rstest]
+    #[tokio::test]
+    async fn a_configured_identity_failure_does_not_fall_back_to_a_developer_identity() {
+        let acquirer = Arc::new(ChainAcquirer {
+            terminal_on: Some("managed-identity"),
+            requests: Mutex::new(Vec::new()),
+            succeed_on: Some("developer-tools"),
+        });
+        let service = AzureAuthService::with_acquirer(acquirer.clone());
+        let inputs = AzureAuthInputs {
+            enable_azure_ad_token_refresh: Sourced::new(true, InputSource::Deployment),
+            ..Default::default()
+        };
+        let error = service
+            .get_azure_ad_token(&inputs, &|_| None)
+            .await
+            .unwrap_err();
+        assert_eq!(*acquirer.requests.lock().unwrap(), ["managed-identity"]);
+        assert!(matches!(error, Error::CredentialChain { earlier, .. } if earlier.is_empty()));
+    }
+
     #[derive(Debug)]
     struct CallerToken(&'static str);
 

litellm-rust/crates/auth-gcp/Cargo.toml

@@ -6,7 +6,7 @@ license.workspace = true
 repository.workspace = true
 
 [features]
-google-sdk = ["dep:google-cloud-auth", "dep:http"]
+google-sdk = ["dep:http"]
 
 [dependencies]
 litellm-secret-value.workspace = true
@@ -21,10 +21,14 @@ tokio.workspace = true
 veil.workspace = true
 
 gcp_auth = "0.12.7"
-google-cloud-auth = { workspace = true, optional = true }
+google-cloud-auth.workspace = true
 http = { workspace = true, optional = true }
 
 [dev-dependencies]
+time = { workspace = true, features = ["formatting"] }
+axum = { workspace = true, features = ["json"] }
 litellm-error.workspace = true
 rstest.workspace = true
+rustls.workspace = true
 tempfile.workspace = true
+tokio = { workspace = true, features = ["test-util"] }

litellm-rust/crates/auth-gcp/src/auth.rs

@@ -1,20 +1,11 @@
-use std::{future::Future, pin::Pin, sync::Arc};
+use std::sync::Arc;
 
-use gcp_auth::{CustomServiceAccount, TokenProvider};
-use litellm_auth_types::{CredentialPlacement, Error, ErrorDetail, http::apply_credential};
-use moka::future::Cache;
-use serde_json::Value;
-use sha2::{Digest, Sha256};
+use litellm_auth_types::{CredentialPlacement, Error, http::apply_credential};
 
 use crate::{
-    config::{
-        CredentialSource, VertexConfig, credential_source, get_vertex_ai_project,
-        is_base64_encoded_json, names_a_credentials_file, non_empty_env,
-    },
-    constants::{
-        CLOUD_PLATFORM_SCOPE, GOOGLE_OAUTH_TOKEN_ENDPOINT, VERTEX_AI_API_KEY_ENV,
-        VERTEXAI_API_KEY_ENV,
-    },
+    config::{VertexConfig, credential_source, get_vertex_ai_project, non_empty_env},
+    constants::{VERTEX_AI_API_KEY_ENV, VERTEXAI_API_KEY_ENV},
+    identity::{GcpAuth, GcpProviderLoader, GcpTokenSource},
 };
 
 pub struct VertexEnvironment {
@@ -22,28 +13,15 @@ pub struct VertexEnvironment {
     pub project_id: String,
 }
 
-struct VertexAccessToken {
-    token: String,
-    project_id: String,
-}
-
-#[derive(Clone)]
+#[derive(Clone, Default)]
 pub struct VertexAuth {
-    providers: Cache<CredentialCacheKey, Arc<dyn VertexTokenSource>>,
-    loader: Arc<dyn VertexProviderLoader>,
-}
-
-impl Default for VertexAuth {
-    fn default() -> Self {
-        Self::new(Arc::new(GcpProviderLoader))
-    }
+    identity: GcpAuth,
 }
 
 impl VertexAuth {
-    pub fn new(loader: Arc<dyn VertexProviderLoader>) -> Self {
+    pub fn new(loader: Arc<dyn GcpProviderLoader>) -> Self {
         Self {
-            providers: Cache::builder().max_capacity(64).build(),
-            loader,
+            identity: GcpAuth::new(loader),
         }
     }
 
@@ -52,7 +30,10 @@ impl VertexAuth {
         config: &VertexConfig,
         env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
     ) -> Result<String, Error> {
-        self.load_provider(config, env_lookup).await?.token().await
+        self.identity
+            .access_token(credential_source(config, env_lookup))
+            .await
+            .map(|token| token.expose().to_owned())
     }
 
     pub async fn validate_environment(
@@ -77,10 +58,18 @@ impl VertexAuth {
             (true, _) => None,
             (false, Some(token)) => Some(token),
             (false, None) => {
-                let access = self.get_access_token(config, env_lookup).await?;
+                let provider = self.load_provider(config, env_lookup).await?;
+                let (token, project_id) = match project_id {
+                    Some(project_id) => (provider.token().await?, project_id),
+                    None => tokio::try_join!(provider.token(), provider.project_id())?,
+                };
                 return Ok(VertexEnvironment {
-                    headers: apply_credential(headers, &access.token, CredentialPlacement::Bearer)?,
-                    project_id: project_id.unwrap_or(access.project_id),
+                    headers: apply_credential(
+                        headers,
+                        token.expose(),
+                        CredentialPlacement::Bearer,
+                    )?,
+                    project_id,
                 });
             }
         };
@@ -104,248 +93,171 @@ impl VertexAuth {
         })
     }
 
-    async fn get_access_token(
-        &self,
-        config: &VertexConfig,
-        env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
-    ) -> Result<VertexAccessToken, Error> {
-        let provider = self.load_provider(config, env_lookup).await?;
-        let (token, project_id) = tokio::try_join!(provider.token(), provider.project_id())?;
-        Ok(VertexAccessToken { token, project_id })
-    }
-
     async fn load_provider(
         &self,
         config: &VertexConfig,
         env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
-    ) -> Result<Arc<dyn VertexTokenSource>, Error> {
-        let source = credential_source(config, env_lookup);
-        let key = source.cache_key();
-        self.providers
-            .try_get_with(key, self.loader.load(source))
+    ) -> Result<Arc<dyn GcpTokenSource>, Error> {
+        self.identity
+            .load_provider(credential_source(config, env_lookup))
             .await
-            .map_err(|error| (*error).clone())
     }
 }
 
-pub trait VertexTokenSource: Send + Sync {
-    fn project_id(&self) -> VertexAuthFuture<'_, String>;
-    fn token(&self) -> VertexAuthFuture<'_, String>;
-}
-
-pub trait VertexProviderLoader: Send + Sync {
-    fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc<dyn VertexTokenSource>>;
-}
-
-pub type VertexAuthFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
-
-struct GcpTokenSource(Arc<dyn TokenProvider>);
-
-impl VertexTokenSource for GcpTokenSource {
-    fn project_id(&self) -> VertexAuthFuture<'_, String> {
-        Box::pin(async move {
-            self.0
-                .project_id()
-                .await
-                .map(|project| project.to_string())
-                .map_err(auth_acquisition_error)
-        })
-    }
-
-    fn token(&self) -> VertexAuthFuture<'_, String> {
-        Box::pin(async move {
-            self.0
-                .token(&[CLOUD_PLATFORM_SCOPE])
-                .await
-                .map(|token| token.as_str().to_string())
-                .map_err(auth_acquisition_error)
-        })
-    }
-}
-
-struct GcpProviderLoader;
-
-impl VertexProviderLoader for GcpProviderLoader {
-    fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc<dyn VertexTokenSource>> {
-        Box::pin(async move {
-            let provider: Arc<dyn TokenProvider> = match source {
-                CredentialSource::Inline(configured) => Arc::new(
-                    CustomServiceAccount::from_json(validate_request_credentials(
-                        configured.expose(),
-                    )?)
-                    .map_err(auth_acquisition_error)?,
-                ),
-                CredentialSource::Trusted(configured) => {
-                    Arc::new(load_trusted_service_account(configured.expose())?)
-                }
-                CredentialSource::ApplicationCredentials(path) => {
-                    Arc::new(CustomServiceAccount::from_file(path).map_err(auth_acquisition_error)?)
-                }
-                CredentialSource::Adc => {
-                    gcp_auth::provider().await.map_err(auth_acquisition_error)?
-                }
-            };
-            Ok(Arc::new(GcpTokenSource(provider)) as Arc<dyn VertexTokenSource>)
-        })
-    }
-}
-
-fn validate_request_credentials(configured: &str) -> Result<&str, Error> {
-    let token_uri = serde_json::from_str::<Value>(configured)
-        .ok()
-        .and_then(|credentials| {
-            credentials
-                .get("token_uri")
-                .and_then(Value::as_str)
-                .map(str::to_string)
-        });
-    if token_uri.as_deref() != Some(GOOGLE_OAUTH_TOKEN_ENDPOINT) {
-        return Err(Error::InvalidConfiguration(ErrorDetail::InvalidType {
-            field: "request-controlled vertex_credentials token_uri".into(),
-            expected: "the canonical Google OAuth token endpoint",
-        }));
-    }
-    Ok(configured)
-}
-
-impl CredentialSource {
-    fn cache_key(&self) -> CredentialCacheKey {
-        match self {
-            Self::Inline(configured) => {
-                CredentialCacheKey::Inline(Sha256::digest(configured.expose()).into())
-            }
-            Self::Trusted(configured) => {
-                CredentialCacheKey::Trusted(Sha256::digest(configured.expose()).into())
-            }
-            Self::ApplicationCredentials(path) => {
-                CredentialCacheKey::ApplicationCredentials(path.clone())
-            }
-            Self::Adc => CredentialCacheKey::Adc,
-        }
-    }
-}
-
-#[derive(Clone, Debug, Hash, PartialEq, Eq)]
-enum CredentialCacheKey {
-    Inline([u8; 32]),
-    Trusted([u8; 32]),
-    ApplicationCredentials(String),
-    Adc,
-}
-
-fn load_trusted_service_account(configured: &str) -> Result<CustomServiceAccount, Error> {
-    if is_base64_encoded_json(configured) {
-        return Err(Error::InvalidConfiguration(
-            ErrorDetail::Base64EncodedJson {
-                field: "vertex_credentials",
-            },
-        ));
-    }
-    if names_a_credentials_file(configured) {
-        CustomServiceAccount::from_file(configured)
-    } else {
-        CustomServiceAccount::from_json(configured)
-    }
-    .map_err(auth_acquisition_error)
-}
-
-fn auth_acquisition_error(error: gcp_auth::Error) -> Error {
-    Error::CredentialAcquisition(litellm_auth_types::ErrorDetail::failed(
-        "Vertex AI credentials",
-        error,
-    ))
-}
-
 #[cfg(test)]
 mod tests {
+    use std::{
+        collections::{BTreeMap, BTreeSet},
+        sync::{
+            Arc,
+            atomic::{AtomicUsize, Ordering},
+        },
+    };
+
+    use crate::{
+        CredentialSource, GcpAuthFuture, GcpProviderLoader, GcpTokenSource, VertexAuth,
+        VertexConfig, get_vertex_ai_location, secret_names,
+    };
     use litellm_secret_value::SecretValue;
+    use rstest::rstest;
+    use serde_json::{Value, json};
 
-    use super::*;
-
-    #[rstest::rstest]
-    #[case::canonical_endpoint(r#"{"token_uri":"https://oauth2.googleapis.com/token"}"#, true)]
-    #[case::noncanonical_endpoint(r#"{"token_uri":"http://127.0.0.1/token"}"#, false)]
-    #[case::missing_endpoint("{}", false)]
-    fn request_credentials_require_canonical_token_endpoint(
-        #[case] credentials: &str,
-        #[case] accepted: bool,
-    ) {
-        assert_eq!(validate_request_credentials(credentials).is_ok(), accepted);
+    struct FakeProvider {
+        project: Option<&'static str>,
+        calls: Arc<AtomicUsize>,
     }
 
-    #[rstest::rstest]
-    #[case::missing_file(
-        "/nonexistent/vertex-key.json",
-        "failed to open application credentials file"
-    )]
-    #[case::malformed_json(
-        " {\"type\": \"service_account\",",
-        "failed to deserialize ApplicationCredentials"
-    )]
-    #[case::quoted_json(
-        "'{\"type\": \"service_account\"}'",
-        "failed to deserialize ApplicationCredentials"
-    )]
-    #[tokio::test]
-    async fn trusted_credentials_are_read_as_a_file_unless_they_are_json(
-        #[case] configured: &str,
-        #[case] expected: &str,
-    ) {
-        let Err(error) = GcpProviderLoader
-            .load(CredentialSource::Trusted(SecretValue::new(configured)))
-            .await
-        else {
-            panic!("{configured:?} should not load");
-        };
-        let report = litellm_error::Report::new(&error).to_string();
-        assert!(report.contains(expected), "{report}");
-    }
+    impl GcpTokenSource for FakeProvider {
+        fn project_id(&self) -> GcpAuthFuture<'_, String> {
+            self.calls.fetch_add(1, Ordering::SeqCst);
+            Box::pin(async move {
+                self.project.map(str::to_owned).ok_or(
+                    litellm_auth_types::Error::CredentialAcquisition(
+                        litellm_auth_types::ErrorDetail::Unresolved {
+                            subject: "credential project",
+                        },
+                    ),
+                )
+            })
+        }
 
-    #[rstest::rstest]
-    #[case::padded("eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCJ9")]
-    #[case::quoted_prefixed_unpadded("\"base64:eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCJ9\"")]
-    #[case::url_safe_wrapped("eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCIsICJr\nZXkiOiAiPz8_In0=")]
-    #[tokio::test]
-    async fn base64_encoded_credentials_are_named_as_base64(#[case] configured: &str) {
-        assert!(
-            trusted_load_error(configured)
-                .await
-                .contains("is base64-encoded JSON, not the JSON itself")
-        );
-    }
-
-    async fn trusted_load_error(configured: &str) -> String {
-        match GcpProviderLoader
-            .load(CredentialSource::Trusted(SecretValue::new(configured)))
-            .await
-        {
-            Ok(_) => panic!("{configured:?} should not load"),
-            Err(error) => litellm_error::Report::new(&error).to_string(),
+        fn token(&self) -> GcpAuthFuture<'_, SecretValue> {
+            self.calls.fetch_add(1, Ordering::SeqCst);
+            Box::pin(async { Ok(SecretValue::new("adc-token")) })
         }
     }
 
-    #[rstest::rstest]
-    #[tokio::test]
-    async fn a_relative_key_file_named_like_json_is_read_as_a_file() {
-        let key = tempfile::Builder::new()
-            .prefix("{vertex")
-            .suffix(".json")
-            .tempfile_in(".")
-            .unwrap();
-        std::fs::write(key.path(), r#"{"project_id":"file-project"}"#).unwrap();
-        let relative_name = key.path().file_name().unwrap().to_str().unwrap();
-        let absolute_path = std::fs::canonicalize(key.path()).unwrap();
-        assert_eq!(
-            trusted_load_error(relative_name).await,
-            trusted_load_error(absolute_path.to_str().unwrap()).await
-        );
+    struct FakeLoader {
+        loads: Arc<AtomicUsize>,
+        provider: Arc<dyn GcpTokenSource>,
     }
 
-    #[rstest::rstest]
-    fn inline_and_trusted_credentials_never_share_a_cache_entry() {
-        assert_ne!(
-            CredentialSource::Inline(SecretValue::new("same-value")).cache_key(),
-            CredentialSource::Trusted(SecretValue::new("same-value")).cache_key()
-        );
+    impl GcpProviderLoader for FakeLoader {
+        fn load(&self, _source: CredentialSource) -> GcpAuthFuture<'_, Arc<dyn GcpTokenSource>> {
+            let loads = self.loads.clone();
+            let provider = self.provider.clone();
+            Box::pin(async move {
+                loads.fetch_add(1, Ordering::SeqCst);
+                Ok(provider)
+            })
+        }
+    }
+
+    fn config(value: Value) -> VertexConfig {
+        VertexConfig::from_sourced_optional_params(value.as_object().unwrap(), &BTreeMap::new())
+            .unwrap()
+    }
+
+    fn auth(calls: Arc<AtomicUsize>, loads: Arc<AtomicUsize>) -> VertexAuth {
+        let provider: Arc<dyn GcpTokenSource> = Arc::new(FakeProvider {
+            calls,
+            project: Some("adc-project"),
+        });
+        VertexAuth::new(Arc::new(FakeLoader { loads, provider }))
+    }
+
+    #[rstest]
+    #[tokio::test]
+    async fn secret_names_cover_environment_reads() {
+        let seen = Arc::new(std::sync::Mutex::new(BTreeSet::<String>::new()));
+        let recorded = seen.clone();
+        let env = |name: &str| {
+            recorded.lock().unwrap().insert(name.to_string());
+            None
+        };
+        let auth = auth(Arc::new(AtomicUsize::new(0)), Arc::new(AtomicUsize::new(0)));
+        auth.validate_environment(Vec::new(), None, &VertexConfig::default(), &env)
+            .await
+            .unwrap();
+        get_vertex_ai_location(&VertexConfig::default(), &env);
+        let allowed = secret_names()
+            .iter()
+            .map(|name| (*name).to_owned())
+            .collect::<BTreeSet<_>>();
+        assert!(seen.lock().unwrap().is_subset(&allowed));
+    }
+
+    #[rstest]
+    #[tokio::test]
+    async fn explicit_token_and_header_do_not_acquire_adc() {
+        let loads = Arc::new(AtomicUsize::new(0));
+        let auth = auth(Arc::new(AtomicUsize::new(0)), loads.clone());
+        let configured = config(json!({"vertex_project":"project-1"}));
+        let explicit = auth
+            .validate_environment(Vec::new(), Some("access-token"), &configured, &|_| None)
+            .await
+            .unwrap();
+        assert_eq!(explicit.headers[0].1, "Bearer access-token");
+        let existing = auth
+            .validate_environment(
+                vec![("authorization".into(), "Bearer existing".into())],
+                None,
+                &configured,
+                &|_| None,
+            )
+            .await
+            .unwrap();
+        assert_eq!(existing.headers[0].1, "Bearer existing");
+        assert_eq!(loads.load(Ordering::SeqCst), 0);
+    }
+
+    #[rstest]
+    #[tokio::test]
+    async fn provider_is_reused_across_authentication_calls() {
+        let calls = Arc::new(AtomicUsize::new(0));
+        let loads = Arc::new(AtomicUsize::new(0));
+        let auth = auth(calls.clone(), loads.clone());
+        let config = VertexConfig::default();
+        let environment = |_: &str| None;
+        let (first, second) = tokio::try_join!(
+            auth.validate_environment(Vec::new(), None, &config, &environment),
+            auth.validate_environment(Vec::new(), None, &config, &environment),
+        )
+        .unwrap();
+        assert_eq!(first.project_id, "adc-project");
+        assert_eq!(first.headers[0].1, "Bearer adc-token");
+        assert_eq!(first.project_id, second.project_id);
+        assert_eq!(first.headers, second.headers);
+        assert_eq!(loads.load(Ordering::SeqCst), 1);
+        assert_eq!(calls.load(Ordering::SeqCst), 4);
+    }
+
+    #[rstest]
+    #[tokio::test]
+    async fn an_explicit_project_does_not_require_project_metadata_in_the_credentials() {
+        let auth = VertexAuth::new(Arc::new(FakeLoader {
+            loads: Arc::new(AtomicUsize::new(0)),
+            provider: Arc::new(FakeProvider {
+                calls: Arc::new(AtomicUsize::new(0)),
+                project: None,
+            }),
+        }));
+        let configured = VertexConfig::new(None, Some("configured-project".into()), None);
+        let environment = auth
+            .validate_environment(Vec::new(), None, &configured, &|_| None)
+            .await
+            .unwrap();
+        assert_eq!(environment.project_id, "configured-project");
+        assert_eq!(environment.headers[0].1, "Bearer adc-token");
     }
 }

litellm-rust/crates/auth-gcp/src/config.rs

@@ -146,6 +146,20 @@ pub enum CredentialSource {
     Adc,
 }
 
+impl CredentialSource {
+    pub fn from_vertex_environment(env_lookup: &dyn Fn(&str) -> Option<String>) -> Self {
+        VertexParams::CREDENTIALS
+            .resolve(&|_| None, env_lookup)
+            .map(|configured| Self::Trusted(SecretValue::new(configured)))
+            .unwrap_or_else(|| Self::application_default(env_lookup))
+    }
+
+    pub fn application_default(env_lookup: &dyn Fn(&str) -> Option<String>) -> Self {
+        non_empty_env(env_lookup, GOOGLE_APPLICATION_CREDENTIALS_ENV)
+            .map_or(Self::Adc, Self::ApplicationCredentials)
+    }
+}
+
 pub(crate) fn credential_source(
     config: &VertexConfig,
     env_lookup: &dyn Fn(&str) -> Option<String>,
@@ -158,12 +172,7 @@ pub(crate) fn credential_source(
             }
         };
     }
-    if let Some(configured) = VertexParams::CREDENTIALS.resolve(&|_| None, env_lookup) {
-        return CredentialSource::Trusted(SecretValue::new(configured));
-    }
-    non_empty_env(env_lookup, GOOGLE_APPLICATION_CREDENTIALS_ENV)
-        .map(CredentialSource::ApplicationCredentials)
-        .unwrap_or(CredentialSource::Adc)
+    CredentialSource::from_vertex_environment(env_lookup)
 }
 
 fn optional_credentials(
@@ -319,6 +328,24 @@ mod tests {
         ));
     }
 
+    #[rstest::rstest]
+    fn application_default_ignores_vertex_credentials() {
+        let env = |name: &str| match name {
+            VERTEXAI_CREDENTIALS_ENV => Some("environment-json".to_string()),
+            GOOGLE_APPLICATION_CREDENTIALS_ENV => Some("adc.json".to_string()),
+            _ => None,
+        };
+        assert!(
+            matches!(CredentialSource::application_default(&env), CredentialSource::ApplicationCredentials(path) if path == "adc.json")
+        );
+        assert!(matches!(
+            CredentialSource::application_default(
+                &|name| (name == VERTEXAI_CREDENTIALS_ENV).then(|| "environment-json".into())
+            ),
+            CredentialSource::Adc
+        ));
+    }
+
     #[rstest::rstest]
     fn params_become_a_deployment_sourced_config() {
         let params = VertexParams {

litellm-rust/crates/auth-gcp/src/identity.rs

@@ -0,0 +1,445 @@
+use crate::{
+    config::{CredentialSource, is_base64_encoded_json, names_a_credentials_file},
+    constants::{CLOUD_PLATFORM_SCOPE, GOOGLE_OAUTH_TOKEN_ENDPOINT},
+};
+use gcp_auth::{CustomServiceAccount, TokenProvider};
+use litellm_auth_types::{Error, ErrorDetail};
+use litellm_secret_value::SecretValue;
+use moka::future::Cache;
+use serde_json::Value;
+use sha2::{Digest, Sha256};
+use std::{future::Future, pin::Pin, sync::Arc};
+
+#[derive(Clone)]
+pub struct GcpAuth {
+    providers: Cache<CredentialCacheKey, Arc<dyn GcpTokenSource>>,
+    loader: Arc<dyn GcpProviderLoader>,
+}
+
+impl Default for GcpAuth {
+    fn default() -> Self {
+        Self::new(Arc::new(NativeProviderLoader))
+    }
+}
+
+impl GcpAuth {
+    pub fn new(loader: Arc<dyn GcpProviderLoader>) -> Self {
+        Self {
+            providers: Cache::builder().max_capacity(64).build(),
+            loader,
+        }
+    }
+
+    pub async fn access_token(&self, source: CredentialSource) -> Result<SecretValue, Error> {
+        self.load_provider(source).await?.token().await
+    }
+
+    pub(crate) async fn load_provider(
+        &self,
+        source: CredentialSource,
+    ) -> Result<Arc<dyn GcpTokenSource>, Error> {
+        self.providers
+            .try_get_with(source.cache_key(), self.loader.load(source))
+            .await
+            .map_err(|error| (*error).clone())
+    }
+}
+
+pub trait GcpTokenSource: Send + Sync {
+    fn project_id(&self) -> GcpAuthFuture<'_, String>;
+    fn token(&self) -> GcpAuthFuture<'_, SecretValue>;
+}
+
+pub trait GcpProviderLoader: Send + Sync {
+    fn load(&self, source: CredentialSource) -> GcpAuthFuture<'_, Arc<dyn GcpTokenSource>>;
+}
+
+pub type GcpAuthFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
+
+struct ServiceAccountTokenSource(Arc<dyn TokenProvider>);
+
+impl GcpTokenSource for ServiceAccountTokenSource {
+    fn project_id(&self) -> GcpAuthFuture<'_, String> {
+        Box::pin(async move {
+            self.0
+                .project_id()
+                .await
+                .map(|project| project.to_string())
+                .map_err(auth_acquisition_error)
+        })
+    }
+
+    fn token(&self) -> GcpAuthFuture<'_, SecretValue> {
+        Box::pin(async move {
+            self.0
+                .token(&[CLOUD_PLATFORM_SCOPE])
+                .await
+                .map(|token| SecretValue::new(token.as_str()))
+                .map_err(auth_acquisition_error)
+        })
+    }
+}
+
+struct SdkTokenSource {
+    credentials: google_cloud_auth::credentials::AccessTokenCredentials,
+    project_id: Option<String>,
+}
+
+impl GcpTokenSource for SdkTokenSource {
+    fn project_id(&self) -> GcpAuthFuture<'_, String> {
+        Box::pin(async move {
+            self.project_id
+                .clone()
+                .ok_or(Error::CredentialAcquisition(ErrorDetail::Unresolved {
+                    subject: "GCP resource project",
+                }))
+        })
+    }
+
+    fn token(&self) -> GcpAuthFuture<'_, SecretValue> {
+        Box::pin(async move {
+            self.credentials
+                .access_token()
+                .await
+                .map(|token| SecretValue::new(token.token))
+                .map_err(auth_acquisition_error)
+        })
+    }
+}
+
+async fn load_application_credentials(path: &str) -> Result<Arc<dyn GcpTokenSource>, Error> {
+    let contents = tokio::fs::read_to_string(path)
+        .await
+        .map_err(auth_acquisition_error)?;
+    load_credentials_json(&contents)
+}
+
+fn load_credentials_json(contents: &str) -> Result<Arc<dyn GcpTokenSource>, Error> {
+    let configured: Value = serde_json::from_str(contents).map_err(auth_acquisition_error)?;
+    let kind = configured.get("type").and_then(Value::as_str);
+    let project_id = configured
+        .get("project_id")
+        .and_then(Value::as_str)
+        .map(str::to_string);
+    let credentials = match kind {
+        Some("authorized_user") => {
+            google_cloud_auth::credentials::user_account::Builder::new(configured)
+                .with_scopes([CLOUD_PLATFORM_SCOPE])
+                .build_access_token_credentials()
+        }
+        Some("impersonated_service_account") => {
+            google_cloud_auth::credentials::impersonated::Builder::new(configured)
+                .with_scopes([CLOUD_PLATFORM_SCOPE])
+                .build_access_token_credentials()
+        }
+        Some("external_account") => {
+            google_cloud_auth::credentials::external_account::Builder::new(configured)
+                .with_scopes([CLOUD_PLATFORM_SCOPE])
+                .build_access_token_credentials()
+        }
+        Some("service_account") | None => {
+            let provider =
+                CustomServiceAccount::from_json(contents).map_err(auth_acquisition_error)?;
+            return Ok(Arc::new(ServiceAccountTokenSource(Arc::new(provider))));
+        }
+        Some(_) => {
+            return Err(Error::InvalidConfiguration(ErrorDetail::InvalidType {
+                field: "Google application credentials type".into(),
+                expected: "a supported Google credential type",
+            }));
+        }
+    }
+    .map_err(auth_acquisition_error)?;
+    Ok(Arc::new(SdkTokenSource {
+        credentials,
+        project_id,
+    }))
+}
+
+struct NativeProviderLoader;
+
+impl GcpProviderLoader for NativeProviderLoader {
+    fn load(&self, source: CredentialSource) -> GcpAuthFuture<'_, Arc<dyn GcpTokenSource>> {
+        Box::pin(async move {
+            let provider: Arc<dyn TokenProvider> = match source {
+                CredentialSource::Inline(configured) => Arc::new(
+                    CustomServiceAccount::from_json(validate_request_credentials(
+                        configured.expose(),
+                    )?)
+                    .map_err(auth_acquisition_error)?,
+                ),
+                CredentialSource::Trusted(configured) => {
+                    return load_trusted_credentials(configured.expose()).await;
+                }
+                CredentialSource::ApplicationCredentials(path) => {
+                    return load_application_credentials(&path).await;
+                }
+                CredentialSource::Adc => {
+                    gcp_auth::provider().await.map_err(auth_acquisition_error)?
+                }
+            };
+            Ok(Arc::new(ServiceAccountTokenSource(provider)) as Arc<dyn GcpTokenSource>)
+        })
+    }
+}
+
+fn validate_request_credentials(configured: &str) -> Result<&str, Error> {
+    let token_uri = serde_json::from_str::<Value>(configured)
+        .ok()
+        .and_then(|credentials| {
+            credentials
+                .get("token_uri")
+                .and_then(Value::as_str)
+                .map(str::to_string)
+        });
+    if token_uri.as_deref() != Some(GOOGLE_OAUTH_TOKEN_ENDPOINT) {
+        return Err(Error::InvalidConfiguration(ErrorDetail::InvalidType {
+            field: "request-controlled vertex_credentials token_uri".into(),
+            expected: "the canonical Google OAuth token endpoint",
+        }));
+    }
+    Ok(configured)
+}
+
+impl CredentialSource {
+    fn cache_key(&self) -> CredentialCacheKey {
+        match self {
+            Self::Inline(configured) => {
+                CredentialCacheKey::Inline(Sha256::digest(configured.expose()).into())
+            }
+            Self::Trusted(configured) => {
+                CredentialCacheKey::Trusted(Sha256::digest(configured.expose()).into())
+            }
+            Self::ApplicationCredentials(path) => {
+                CredentialCacheKey::ApplicationCredentials(path.clone())
+            }
+            Self::Adc => CredentialCacheKey::Adc,
+        }
+    }
+}
+
+#[derive(Clone, Debug, Hash, PartialEq, Eq)]
+enum CredentialCacheKey {
+    Inline([u8; 32]),
+    Trusted([u8; 32]),
+    ApplicationCredentials(String),
+    Adc,
+}
+
+async fn load_trusted_credentials(configured: &str) -> Result<Arc<dyn GcpTokenSource>, Error> {
+    if is_base64_encoded_json(configured) {
+        return Err(Error::InvalidConfiguration(
+            ErrorDetail::Base64EncodedJson {
+                field: "vertex_credentials",
+            },
+        ));
+    }
+    if names_a_credentials_file(configured) {
+        load_application_credentials(configured).await
+    } else {
+        load_credentials_json(configured)
+    }
+}
+
+fn auth_acquisition_error(error: impl std::error::Error + Send + Sync + 'static) -> Error {
+    Error::CredentialAcquisition(litellm_auth_types::ErrorDetail::failed(
+        "Google credentials",
+        error,
+    ))
+}
+
+#[cfg(test)]
+mod tests {
+    use litellm_secret_value::SecretValue;
+
+    use super::*;
+
+    #[rstest::rstest]
+    #[case::authorized_user("authorized_user", 2)]
+    #[case::external_account("external_account", 2)]
+    #[case::impersonated_service_account("impersonated_service_account", 4)]
+    #[tokio::test]
+    async fn explicit_application_credentials_refresh_expired_tokens(
+        #[case] kind: &str,
+        #[case] expected_requests: usize,
+        #[values(false, true)] inline: bool,
+    ) {
+        use axum::{Json, Router, routing::post};
+        use std::sync::atomic::{AtomicUsize, Ordering};
+
+        let _ = rustls::crypto::ring::default_provider().install_default();
+        let requests = Arc::new(AtomicUsize::new(0));
+        let recorded = requests.clone();
+        let impersonations = requests.clone();
+        let app = Router::new().route("/token", post(move || {
+            let recorded = recorded.clone();
+            async move {
+                Json(serde_json::json!({
+                    "access_token": format!("token-{}", recorded.fetch_add(1, Ordering::SeqCst)),
+                    "expires_in": 3600,
+                    "token_type": "Bearer",
+                    "issued_token_type": "urn:ietf:params:oauth:token-type:access_token",
+                }))
+            }
+        }));
+        let app = app.route("/v1/projects/-/serviceAccounts/test-principal:generateAccessToken", post(move || {
+            let recorded = impersonations.clone();
+            async move {
+                let expiration = (time::OffsetDateTime::now_utc() + time::Duration::hours(1)).format(&time::format_description::well_known::Rfc3339).unwrap();
+                Json(serde_json::json!({
+                    "accessToken": format!("impersonated-{}", recorded.fetch_add(1, Ordering::SeqCst)),
+                    "expireTime": expiration,
+                }))
+            }
+        }));
+        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let address = listener.local_addr().unwrap();
+        let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
+        let file = tempfile::NamedTempFile::new().unwrap();
+        let subject_token = tempfile::NamedTempFile::new().unwrap();
+        std::fs::write(subject_token.path(), "subject-token").unwrap();
+        let configured = if kind == "external_account" {
+            serde_json::json!({
+                "type": "external_account", "audience": "//iam.googleapis.com/projects/test/locations/global/workloadIdentityPools/test/providers/test",
+                "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", "token_url": format!("http://{address}/token"),
+                "credential_source": { "file": subject_token.path().to_str().unwrap() }, "quota_project_id": "quota-project"
+            })
+        } else {
+            let source = serde_json::json!({
+                "type": "authorized_user", "client_id": "user-client", "client_secret": "user-secret",
+                "refresh_token": "user-refresh", "quota_project_id": "quota-project", "token_uri": format!("http://{address}/token")
+            });
+            if kind == "impersonated_service_account" {
+                serde_json::json!({
+                    "type": kind, "source_credentials": source, "quota_project_id": "quota-project",
+                    "service_account_impersonation_url": format!("http://{address}/v1/projects/-/serviceAccounts/test-principal:generateAccessToken")
+                })
+            } else {
+                source
+            }
+        };
+        std::fs::write(file.path(), configured.to_string()).unwrap();
+        let source = if inline {
+            CredentialSource::Trusted(SecretValue::new(configured.to_string()))
+        } else {
+            CredentialSource::ApplicationCredentials(file.path().to_str().unwrap().into())
+        };
+        let provider = NativeProviderLoader.load(source).await.unwrap();
+
+        assert!(matches!(
+            provider.project_id().await,
+            Err(Error::CredentialAcquisition(ErrorDetail::Unresolved {
+                subject: "GCP resource project"
+            }))
+        ));
+        let original = provider.token().await.unwrap();
+        tokio::time::pause();
+        tokio::time::advance(std::time::Duration::from_secs(3601)).await;
+        tokio::time::resume();
+        let refreshed = provider.token().await.unwrap();
+        assert_ne!(original, refreshed);
+        assert_eq!(requests.load(Ordering::SeqCst), expected_requests);
+        server.abort();
+    }
+
+    #[rstest::rstest]
+    #[case::canonical_endpoint(r#"{"token_uri":"https://oauth2.googleapis.com/token"}"#, true)]
+    #[case::noncanonical_endpoint(r#"{"token_uri":"http://127.0.0.1/token"}"#, false)]
+    #[case::missing_endpoint("{}", false)]
+    fn request_credentials_require_canonical_token_endpoint(
+        #[case] credentials: &str,
+        #[case] accepted: bool,
+    ) {
+        assert_eq!(validate_request_credentials(credentials).is_ok(), accepted);
+    }
+
+    #[rstest::rstest]
+    #[case::missing_file("/nonexistent/vertex-key.json", "Google credentials")]
+    #[case::malformed_json(" {\"type\": \"service_account\",", "Google credentials")]
+    #[case::quoted_json("'{\"type\": \"service_account\"}'", "Google credentials")]
+    #[tokio::test]
+    async fn trusted_credentials_are_read_as_a_file_unless_they_are_json(
+        #[case] configured: &str,
+        #[case] expected: &str,
+    ) {
+        let Err(error) = NativeProviderLoader
+            .load(CredentialSource::Trusted(SecretValue::new(configured)))
+            .await
+        else {
+            panic!("{configured:?} should not load");
+        };
+        let report = litellm_error::Report::new(&error).to_string();
+        assert!(report.contains(expected), "{report}");
+    }
+
+    #[rstest::rstest]
+    #[case::padded("eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCJ9")]
+    #[case::quoted_prefixed_unpadded("\"base64:eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCJ9\"")]
+    #[case::url_safe_wrapped("eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCIsICJr\nZXkiOiAiPz8_In0=")]
+    #[tokio::test]
+    async fn base64_encoded_credentials_are_named_as_base64(#[case] configured: &str) {
+        assert!(
+            trusted_load_error(configured)
+                .await
+                .contains("is base64-encoded JSON, not the JSON itself")
+        );
+    }
+
+    async fn trusted_load_error(configured: &str) -> String {
+        match NativeProviderLoader
+            .load(CredentialSource::Trusted(SecretValue::new(configured)))
+            .await
+        {
+            Ok(_) => panic!("{configured:?} should not load"),
+            Err(error) => litellm_error::Report::new(&error).to_string(),
+        }
+    }
+
+    #[rstest::rstest]
+    #[tokio::test]
+    async fn a_relative_key_file_named_like_json_is_read_as_a_file() {
+        let key = tempfile::Builder::new()
+            .prefix("{vertex")
+            .suffix(".json")
+            .tempfile_in(".")
+            .unwrap();
+        std::fs::write(key.path(), r#"{"project_id":"file-project"}"#).unwrap();
+        let relative_name = key.path().file_name().unwrap().to_str().unwrap();
+        let absolute_path = std::fs::canonicalize(key.path()).unwrap();
+        assert_eq!(
+            trusted_load_error(relative_name).await,
+            trusted_load_error(absolute_path.to_str().unwrap()).await
+        );
+    }
+
+    #[rstest::rstest]
+    #[tokio::test]
+    async fn unknown_application_credential_types_are_rejected_without_service_account_fallback() {
+        let file = tempfile::NamedTempFile::new().unwrap();
+        std::fs::write(
+            file.path(),
+            r#"{"type":"unsupported","private_key":"credential-value"}"#,
+        )
+        .unwrap();
+        let result = GcpAuth::default()
+            .access_token(CredentialSource::ApplicationCredentials(
+                file.path().to_str().unwrap().into(),
+            ))
+            .await;
+        let error = result.unwrap_err();
+        assert!(matches!(error, Error::InvalidConfiguration(_)));
+        assert!(
+            !litellm_error::Report::new(&error)
+                .to_string()
+                .contains("credential-value")
+        );
+    }
+
+    #[rstest::rstest]
+    fn inline_and_trusted_credentials_never_share_a_cache_entry() {
+        assert_ne!(
+            CredentialSource::Inline(SecretValue::new("same-value")).cache_key(),
+            CredentialSource::Trusted(SecretValue::new("same-value")).cache_key()
+        );
+    }
+}

litellm-rust/crates/auth-gcp/src/lib.rs

@@ -1,16 +1,16 @@
 mod auth;
 mod config;
 pub mod constants;
+mod identity;
 #[cfg(feature = "google-sdk")]
 mod sdk;
 
-pub use auth::{
-    VertexAuth, VertexAuthFuture, VertexEnvironment, VertexProviderLoader, VertexTokenSource,
-};
+pub use auth::{VertexAuth, VertexEnvironment};
 pub use config::{
     CredentialSource, VertexConfig, get_vertex_ai_location, get_vertex_ai_project,
     get_vertex_ai_project_from_credentials,
 };
 pub use constants::secret_names;
+pub use identity::{GcpAuth, GcpAuthFuture, GcpProviderLoader, GcpTokenSource};
 #[cfg(feature = "google-sdk")]
 pub use sdk::GoogleCredentials;

litellm-rust/crates/auth-gcp/tests/vertex_auth.rs

@@ -1,120 +0,0 @@
-use std::{
-    collections::{BTreeMap, BTreeSet},
-    sync::{
-        Arc,
-        atomic::{AtomicUsize, Ordering},
-    },
-};
-
-use litellm_auth_gcp::{
-    CredentialSource, VertexAuth, VertexAuthFuture, VertexConfig, VertexProviderLoader,
-    VertexTokenSource, get_vertex_ai_location, secret_names,
-};
-use rstest::rstest;
-use serde_json::{Value, json};
-
-struct FakeProvider {
-    calls: Arc<AtomicUsize>,
-}
-
-impl VertexTokenSource for FakeProvider {
-    fn project_id(&self) -> VertexAuthFuture<'_, String> {
-        self.calls.fetch_add(1, Ordering::SeqCst);
-        Box::pin(async { Ok("adc-project".into()) })
-    }
-
-    fn token(&self) -> VertexAuthFuture<'_, String> {
-        self.calls.fetch_add(1, Ordering::SeqCst);
-        Box::pin(async { Ok("adc-token".into()) })
-    }
-}
-
-struct FakeLoader {
-    loads: Arc<AtomicUsize>,
-    provider: Arc<dyn VertexTokenSource>,
-}
-
-impl VertexProviderLoader for FakeLoader {
-    fn load(&self, _source: CredentialSource) -> VertexAuthFuture<'_, Arc<dyn VertexTokenSource>> {
-        let loads = self.loads.clone();
-        let provider = self.provider.clone();
-        Box::pin(async move {
-            loads.fetch_add(1, Ordering::SeqCst);
-            Ok(provider)
-        })
-    }
-}
-
-fn config(value: Value) -> VertexConfig {
-    VertexConfig::from_sourced_optional_params(value.as_object().unwrap(), &BTreeMap::new())
-        .unwrap()
-}
-
-fn auth(calls: Arc<AtomicUsize>, loads: Arc<AtomicUsize>) -> VertexAuth {
-    let provider: Arc<dyn VertexTokenSource> = Arc::new(FakeProvider { calls });
-    VertexAuth::new(Arc::new(FakeLoader { loads, provider }))
-}
-
-#[rstest]
-#[tokio::test]
-async fn secret_names_cover_environment_reads() {
-    let seen = Arc::new(std::sync::Mutex::new(BTreeSet::<String>::new()));
-    let recorded = seen.clone();
-    let env = |name: &str| {
-        recorded.lock().unwrap().insert(name.to_string());
-        None
-    };
-    let auth = auth(Arc::new(AtomicUsize::new(0)), Arc::new(AtomicUsize::new(0)));
-    auth.validate_environment(Vec::new(), None, &VertexConfig::default(), &env)
-        .await
-        .unwrap();
-    get_vertex_ai_location(&VertexConfig::default(), &env);
-    assert!(
-        seen.lock()
-            .unwrap()
-            .iter()
-            .all(|name| secret_names().contains(&name.as_str()))
-    );
-}
-
-#[rstest]
-#[tokio::test]
-async fn explicit_token_and_header_do_not_acquire_adc() {
-    let loads = Arc::new(AtomicUsize::new(0));
-    let auth = auth(Arc::new(AtomicUsize::new(0)), loads.clone());
-    let configured = config(json!({"vertex_project":"project-1"}));
-    let explicit = auth
-        .validate_environment(Vec::new(), Some("access-token"), &configured, &|_| None)
-        .await
-        .unwrap();
-    assert_eq!(explicit.headers[0].1, "Bearer access-token");
-    let existing = auth
-        .validate_environment(
-            vec![("authorization".into(), "Bearer existing".into())],
-            None,
-            &configured,
-            &|_| None,
-        )
-        .await
-        .unwrap();
-    assert_eq!(existing.headers[0].1, "Bearer existing");
-    assert_eq!(loads.load(Ordering::SeqCst), 0);
-}
-
-#[rstest]
-#[tokio::test]
-async fn provider_is_reused_across_authentication_calls() {
-    let calls = Arc::new(AtomicUsize::new(0));
-    let loads = Arc::new(AtomicUsize::new(0));
-    let auth = auth(calls.clone(), loads.clone());
-    for _ in 0..2 {
-        let environment = auth
-            .validate_environment(Vec::new(), None, &VertexConfig::default(), &|_| None)
-            .await
-            .unwrap();
-        assert_eq!(environment.project_id, "adc-project");
-        assert_eq!(environment.headers[0].1, "Bearer adc-token");
-    }
-    assert_eq!(loads.load(Ordering::SeqCst), 1);
-    assert_eq!(calls.load(Ordering::SeqCst), 4);
-}

litellm-rust/crates/auth-types/src/error.rs

@@ -92,6 +92,11 @@ impl Error {
 
 #[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
 pub enum ErrorDetail {
+    #[error("{operation} returned HTTP {status}")]
+    UnexpectedStatus {
+        operation: &'static str,
+        status: u16,
+    },
     #[error("{field} must be {expected}")]
     InvalidType {
         field: String,

litellm-rust/crates/auth-types/src/params/aws.rs

@@ -53,6 +53,9 @@ pub struct AwsParams {
 }
 
 impl AwsParams {
+    pub const ROLE_ARN_ENV: &str = "AWS_ROLE_ARN";
+    pub const WEB_IDENTITY_TOKEN_FILE_ENV: &str = "AWS_WEB_IDENTITY_TOKEN_FILE";
+
     pub const ACCESS_KEY_ID: ParamSpec = spec(
         "access key id",
         ParamKind::Credential,
@@ -149,6 +152,10 @@ impl AwsParams {
             AwsParams::SPECS
                 .iter()
                 .flat_map(|spec| spec.env.iter().copied())
+                .chain([
+                    AwsParams::ROLE_ARN_ENV,
+                    AwsParams::WEB_IDENTITY_TOKEN_FILE_ENV,
+                ])
                 .fold(Vec::new(), |mut names, name| {
                     if !names.contains(&name) {
                         names.push(name);

litellm-rust/crates/auth-types/tests/aws_params.rs

@@ -3,21 +3,13 @@ use rstest::rstest;
 use serde_json::{Map, Value};
 
 #[rstest]
-fn fields_name_every_param_once_and_in_declaration_order() {
+fn every_named_param_roundtrips_without_loss() {
     let filled: Map<String, Value> = AwsParams::fields()
         .map(|name| (name.to_string(), Value::from(format!("value-of-{name}"))))
         .collect();
     let typed = AwsParams::from_optional_params(&filled);
     let serialized = serde_json::to_value(&typed).unwrap();
-    assert_eq!(
-        serialized
-            .as_object()
-            .unwrap()
-            .keys()
-            .map(String::as_str)
-            .collect::<Vec<_>>(),
-        AwsParams::fields().collect::<Vec<_>>()
-    );
+    assert_eq!(filled.len(), AwsParams::fields().count());
     assert_eq!(serialized, Value::Object(filled.clone()));
     assert_eq!(
         serde_json::from_value::<AwsParams>(Value::Object(filled)).unwrap(),

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

@@ -8,7 +8,7 @@ use litellm_cache::{
 use litellm_http::Client;
 use percent_encoding::{AsciiSet, NON_ALPHANUMERIC, percent_encode};
 
-use crate::{GcpTokenSource, TokenSource};
+use crate::{GcsTokenSource, TokenSource};
 
 pub const DEFAULT_ENDPOINT: &str = "https://storage.googleapis.com";
 
@@ -62,7 +62,7 @@ pub struct GcsCache<S: CacheCodec> {
 
 impl<S: CacheCodec> GcsCache<S> {
     pub fn new(config: GcsConfig, client: Client, codec: S) -> Self {
-        let token = Arc::new(GcpTokenSource::new(config.path_service_account.clone()));
+        let token = Arc::new(GcsTokenSource::new(config.path_service_account.clone()));
         Self::with_token_source(config, client, codec, token)
     }
 

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

@@ -2,4 +2,4 @@ mod cache;
 mod token;
 
 pub use cache::{DEFAULT_ENDPOINT, GcsCache, GcsConfig, key_prefix};
-pub use token::{GcpTokenSource, StaticTokenSource, TokenSource};
+pub use token::{GcsTokenSource, StaticTokenSource, TokenSource};

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

@@ -1,7 +1,6 @@
 use std::{future::Future, pin::Pin};
 
-use litellm_auth_gcp::{VertexAuth, VertexConfig};
-use litellm_auth_types::{InputSource, Sourced};
+use litellm_auth_gcp::{CredentialSource, GcpAuth};
 use litellm_secret_value::SecretValue;
 
 use litellm_cache::Error;
@@ -10,28 +9,30 @@ pub trait TokenSource: Send + Sync + 'static {
     fn bearer_token(&self) -> Pin<Box<dyn Future<Output = Result<String, Error>> + Send + '_>>;
 }
 
-pub struct GcpTokenSource {
-    auth: VertexAuth,
-    config: VertexConfig,
+pub struct GcsTokenSource {
+    auth: GcpAuth,
+    credentials: Option<String>,
 }
 
-impl GcpTokenSource {
+impl GcsTokenSource {
     pub fn new(path_service_account: Option<String>) -> Self {
-        let credentials = path_service_account
-            .map(|path| Sourced::new(SecretValue::new(path), InputSource::Deployment));
         Self {
-            auth: VertexAuth::default(),
-            config: VertexConfig::new(credentials, None, None),
+            auth: GcpAuth::default(),
+            credentials: path_service_account.filter(|path| !path.trim().is_empty()),
         }
     }
 }
 
-impl TokenSource for GcpTokenSource {
+impl TokenSource for GcsTokenSource {
     fn bearer_token(&self) -> Pin<Box<dyn Future<Output = Result<String, Error>> + Send + '_>> {
         Box::pin(async move {
             self.auth
-                .access_token(&self.config, &|name| std::env::var(name).ok())
+                .access_token(match &self.credentials {
+                    Some(path) => CredentialSource::Trusted(SecretValue::new(path.clone())),
+                    None => CredentialSource::application_default(&|name| std::env::var(name).ok()),
+                })
                 .await
+                .map(|token| token.expose().to_owned())
                 .map_err(Error::unavailable)
         })
     }

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

@@ -1,19 +0,0 @@
-[package]
-name = "litellm-cache-redis-semantic"
-version = "0.1.0"
-edition.workspace = true
-license.workspace = true
-repository.workspace = true
-
-[dependencies]
-litellm-cache.workspace = true
-litellm-cache-redis.workspace = true
-redis = { workspace = true, features = ["tls-rustls"] }
-sha2.workspace = true
-
-[dev-dependencies]
-litellm-cache-testing.workspace = true
-redis-test.workspace = true
-rstest.workspace = true
-serde_json.workspace = true
-tokio.workspace = true

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

@@ -1,404 +0,0 @@
-use std::{
-    sync::Arc,
-    time::{Duration, SystemTime, UNIX_EPOCH},
-};
-
-use litellm_cache::{
-    BaseCache, CacheCodec, Error, SemanticCacheContext,
-    semantic::{Embedder, SemanticCache, SemanticLookup, prompt_from_context},
-};
-use litellm_cache_redis::{
-    RedisTopology,
-    connection::{ConnectionRef, Connections},
-};
-use sha2::{Digest, Sha256};
-
-use crate::{
-    RedisSemanticConfig,
-    index::{CACHE_KEY_FIELD, Index, VECTOR_FIELD},
-    reply::{bytes_field, first_document, number_field, string_field},
-};
-
-struct Inner {
-    index: Index,
-    distance_threshold: f64,
-    clock: fn() -> f64,
-}
-
-impl Inner {
-    fn new(config: RedisSemanticConfig, clock: fn() -> f64) -> Self {
-        Self {
-            index: Index::new(config.index_name),
-            distance_threshold: 1.0 - f64::from(config.similarity_threshold),
-            clock,
-        }
-    }
-
-    fn store(
-        &self,
-        connection: &mut ConnectionRef<'_>,
-        tag: &str,
-        response: Vec<u8>,
-        prompt: &str,
-        vector: &[f32],
-        ttl: Option<Duration>,
-    ) -> Result<(), Error> {
-        let index = self.index.ensure(connection, vector.len())?;
-        let entry_id = entry_id(prompt, tag);
-        let hash_key = format!("{index}:{entry_id}");
-        redis::cmd("HSET")
-            .arg(&hash_key)
-            .arg("entry_id")
-            .arg(&entry_id)
-            .arg("prompt")
-            .arg(prompt)
-            .arg("response")
-            .arg(response)
-            .arg(VECTOR_FIELD)
-            .arg(vector_buffer(vector))
-            .arg("inserted_at")
-            .arg(format!("{}", (self.clock)()))
-            .arg("updated_at")
-            .arg(format!("{}", (self.clock)()))
-            .arg(CACHE_KEY_FIELD)
-            .arg(tag)
-            .query::<()>(connection)
-            .map_err(Error::unavailable)?;
-        if let Some(ttl) = ttl {
-            redis::cmd("EXPIRE")
-                .arg(&hash_key)
-                .arg(ttl_seconds(ttl))
-                .query::<()>(connection)
-                .map_err(Error::unavailable)?;
-        }
-        Ok(())
-    }
-
-    fn lookup(
-        &self,
-        connection: &mut ConnectionRef<'_>,
-        tag: &str,
-        vector: &[f32],
-    ) -> 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)
-        );
-        let result = redis::cmd("FT.SEARCH")
-            .arg(&index)
-            .arg(query)
-            .arg("RETURN")
-            .arg(8)
-            .arg("entry_id")
-            .arg("prompt")
-            .arg("response")
-            .arg("inserted_at")
-            .arg("updated_at")
-            .arg("metadata")
-            .arg(CACHE_KEY_FIELD)
-            .arg("vector_distance")
-            .arg("SORTBY")
-            .arg("vector_distance")
-            .arg("ASC")
-            .arg("DIALECT")
-            .arg(2)
-            .arg("LIMIT")
-            .arg(0)
-            .arg(1)
-            .arg("PARAMS")
-            .arg(2)
-            .arg("vector")
-            .arg(vector_buffer(vector))
-            .query::<redis::Value>(connection)
-            .map_err(Error::unavailable)?;
-        let Some(fields) = first_document(&result) else {
-            return Ok(SemanticLookup::miss(Some(0.0)));
-        };
-        if string_field(fields, CACHE_KEY_FIELD).as_deref() != Some(tag) {
-            return Ok(SemanticLookup::miss(Some(0.0)));
-        }
-        // 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)));
-        };
-        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),
-        })
-    }
-}
-
-/// `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,
-            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::new(config, clock)),
-            ..self
-        }
-    }
-
-    pub fn embedder(&self) -> &E {
-        &self.embedder
-    }
-
-    pub fn index_name(&self) -> &str {
-        self.inner.index.name()
-    }
-
-    pub fn similarity_threshold(&self) -> f32 {
-        (1.0 - self.inner.distance_threshold) as f32
-    }
-
-    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, S, C> BaseCache for RedisSemanticCache<E, S, C>
-where
-    E: Embedder,
-    S: CacheCodec,
-    C: redis::ConnectionLike + Send + 'static,
-{
-    type Value = S::Value;
-    type Context = SemanticCacheContext;
-
-    fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
-        context.ttl
-    }
-
-    fn set_cache(
-        &self,
-        key: &str,
-        value: Self::Value,
-        context: &Self::Context,
-    ) -> Result<(), Error> {
-        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);
-        self.connections.execute(|connection| {
-            self.inner
-                .store(connection, tag, response, &prompt, &vector, context.ttl)
-        })
-    }
-
-    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)
-    }
-
-    async fn async_set_cache(
-        &self,
-        key: &str,
-        value: Self::Value,
-        context: Self::Context,
-    ) -> Result<(), Error> {
-        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_owned();
-        let inner = Arc::clone(&self.inner);
-        Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
-            inner.store(connection, &tag, response, &prompt, &vector, context.ttl)
-        })
-        .await
-    }
-
-    async fn async_get_cache(
-        &self,
-        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(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_owned();
-        let inner = Arc::clone(&self.inner);
-        let lookup = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
-            inner.lookup(connection, &tag, &vector)
-        })
-        .await?;
-        self.decode(lookup)
-    }
-}
-
-fn timestamp() -> f64 {
-    SystemTime::now()
-        .duration_since(UNIX_EPOCH)
-        .map(|duration| duration.as_secs_f64())
-        .unwrap_or_default()
-}
-
-fn entry_id(prompt: &str, tag: &str) -> String {
-    let mut digest = Sha256::new();
-    digest.update(prompt.as_bytes());
-    digest.update(CACHE_KEY_FIELD.as_bytes());
-    digest.update(tag.as_bytes());
-    format!("{:x}", digest.finalize())
-}
-
-fn vector_buffer(vector: &[f32]) -> Vec<u8> {
-    vector
-        .iter()
-        .flat_map(|component| component.to_le_bytes())
-        .collect()
-}
-
-fn escape_tag(value: &str) -> String {
-    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 {
-    ttl.as_secs()
-        .saturating_add(u64::from(ttl.subsec_nanos() > 0))
-        .max(1)
-}

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

@@ -1,8 +0,0 @@
-/// `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

@@ -1,209 +0,0 @@
-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(error) => match index_compatible(connection, &self.name, dims)? {
-                    Some(true) => self.name.clone(),
-                    Some(false) => self.isolated(connection, dims)?,
-                    None => return Err(error),
-                },
-            },
-        };
-        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(error) => return Err(Error::unavailable(error)),
-    };
-    Ok(Some(schema_compatible(&info, dims)))
-}
-
-fn unknown_index(error: &redis::RedisError) -> bool {
-    error
-        .code()
-        .is_some_and(|code| code.eq_ignore_ascii_case("unknown"))
-        && error
-            .detail()
-            .is_some_and(|detail| detail.to_ascii_lowercase().starts_with("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,7 +0,0 @@
-mod cache;
-mod config;
-mod index;
-mod reply;
-
-pub use cache::RedisSemanticCache;
-pub use config::{DEFAULT_INDEX_NAME, RedisSemanticConfig};

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

@@ -1,57 +0,0 @@
-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/cache.rs

@@ -1,778 +0,0 @@
-mod support;
-
-use std::time::Duration;
-
-use litellm_cache::{
-    BaseCache, Error, JsonCodec, SemanticCacheContext,
-    semantic::{SemanticCache, SemanticLookup},
-};
-use litellm_cache_redis_semantic::{DEFAULT_INDEX_NAME, RedisSemanticCache, RedisSemanticConfig};
-use redis_test::{MockCmd, MockRedisConnection};
-use rstest::{fixture, rstest};
-use serde_json::{Value, json};
-use sha2::{Digest, Sha256};
-use support::FakeEmbedder;
-
-const INDEX: &str = DEFAULT_INDEX_NAME;
-const PROMPT: &str = "hello prompt";
-const CLOCK: fn() -> f64 = || 1700000000.5;
-const VECTOR: [f32; 3] = [0.1, 0.2, 0.3];
-
-type MockCache = RedisSemanticCache<FakeEmbedder, JsonCodec<Value>, MockRedisConnection>;
-
-#[fixture]
-fn config() -> RedisSemanticConfig {
-    RedisSemanticConfig {
-        index_name: INDEX.into(),
-        similarity_threshold: 0.9,
-    }
-}
-
-#[fixture]
-fn entry() -> Value {
-    json!({"timestamp": 1.0, "response": {"answer": "yes"}})
-}
-
-#[fixture]
-fn context() -> SemanticCacheContext {
-    messages_context(vec![json!({"role": "user", "content": PROMPT})])
-}
-
-fn messages_context(messages: Vec<Value>) -> SemanticCacheContext {
-    SemanticCacheContext {
-        messages: Some(Value::Array(messages)),
-        ..Default::default()
-    }
-}
-
-fn cache(commands: Vec<MockCmd>, embedder: FakeEmbedder) -> MockCache {
-    RedisSemanticCache::with_connection(
-        MockRedisConnection::new(commands).assert_all_commands_consumed(),
-        embedder,
-        JsonCodec::new(),
-        config(),
-    )
-    .with_clock(CLOCK)
-}
-
-fn encoded(value: &Value) -> Vec<u8> {
-    serde_json::to_vec(value).unwrap()
-}
-
-fn vector_bytes(vector: &[f32]) -> Vec<u8> {
-    vector
-        .iter()
-        .flat_map(|component| component.to_le_bytes())
-        .collect()
-}
-
-fn entry_id(prompt: &str, tag: &str) -> String {
-    let mut digest = Sha256::new();
-    digest.update(prompt.as_bytes());
-    digest.update(b"litellm_cache_key");
-    digest.update(tag.as_bytes());
-    format!("{:x}", digest.finalize())
-}
-
-fn s(value: &str) -> redis::Value {
-    redis::Value::BulkString(value.as_bytes().to_vec())
-}
-
-fn unknown_index_error() -> redis::RedisError {
-    support::server_error("Unknown index name")
-}
-
-fn info_missing(index: &str) -> MockCmd {
-    MockCmd::new(
-        redis::cmd("FT.INFO").arg(index),
-        Err::<redis::Value, _>(unknown_index_error()),
-    )
-}
-
-fn info(index: &str, value: redis::Value) -> MockCmd {
-    MockCmd::new(redis::cmd("FT.INFO").arg(index), Ok(value))
-}
-
-fn attribute(name: &str, field_type: &str, extra: Vec<redis::Value>) -> redis::Value {
-    let mut parts = vec![
-        s("identifier"),
-        s(name),
-        s("attribute"),
-        s(name),
-        s("type"),
-        s(field_type),
-    ];
-    parts.extend(extra);
-    redis::Value::Array(parts)
-}
-
-fn index_info(attributes: Vec<redis::Value>) -> redis::Value {
-    redis::Value::Array(vec![
-        s("index_name"),
-        s(INDEX),
-        s("attributes"),
-        redis::Value::Array(attributes),
-    ])
-}
-
-fn vector_attribute_with(dims: i64, data_type: &str, distance_metric: &str) -> redis::Value {
-    attribute(
-        "prompt_vector",
-        "VECTOR",
-        vec![
-            s("algorithm"),
-            s("FLAT"),
-            s("data_type"),
-            s(data_type),
-            s("dim"),
-            redis::Value::Int(dims),
-            s("distance_metric"),
-            s(distance_metric),
-        ],
-    )
-}
-
-fn info_with_vector(vector: redis::Value) -> redis::Value {
-    index_info(vec![
-        attribute("prompt", "TEXT", vec![]),
-        attribute("response", "TEXT", vec![]),
-        attribute("inserted_at", "NUMERIC", vec![]),
-        attribute("updated_at", "NUMERIC", vec![]),
-        vector,
-        attribute("litellm_cache_key", "TAG", vec![]),
-    ])
-}
-
-fn compatible_info(dims: i64) -> redis::Value {
-    info_with_vector(vector_attribute_with(dims, "FLOAT32", "COSINE"))
-}
-
-fn unscoped_info(dims: i64) -> redis::Value {
-    index_info(vec![
-        attribute("prompt", "TEXT", vec![]),
-        attribute("response", "TEXT", vec![]),
-        attribute("inserted_at", "NUMERIC", vec![]),
-        attribute("updated_at", "NUMERIC", vec![]),
-        vector_attribute_with(dims, "FLOAT32", "COSINE"),
-    ])
-}
-
-fn create_index(name: &str, dims: usize) -> MockCmd {
-    MockCmd::new(create_index_command(name, dims), Ok("OK"))
-}
-
-fn create_index_command(name: &str, dims: usize) -> redis::Cmd {
-    let mut command = redis::cmd("FT.CREATE");
-    command
-        .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("prompt_vector")
-        .arg("VECTOR")
-        .arg("FLAT")
-        .arg(6)
-        .arg("TYPE")
-        .arg("FLOAT32")
-        .arg("DIM")
-        .arg(dims)
-        .arg("DISTANCE_METRIC")
-        .arg("COSINE")
-        .arg("litellm_cache_key")
-        .arg("TAG")
-        .arg("SEPARATOR")
-        .arg(",");
-    command
-}
-
-fn hset(index: &str, prompt: &str, tag: &str, vector: &[f32], value: &Value) -> MockCmd {
-    MockCmd::new(
-        redis::cmd("HSET")
-            .arg(format!("{index}:{}", entry_id(prompt, tag)))
-            .arg("entry_id")
-            .arg(entry_id(prompt, tag))
-            .arg("prompt")
-            .arg(prompt)
-            .arg("response")
-            .arg(encoded(value))
-            .arg("prompt_vector")
-            .arg(vector_bytes(vector))
-            .arg("inserted_at")
-            .arg("1700000000.5")
-            .arg("updated_at")
-            .arg("1700000000.5")
-            .arg("litellm_cache_key")
-            .arg(tag),
-        Ok(7),
-    )
-}
-
-fn search(
-    index: &str,
-    tag: &str,
-    vector: &[f32],
-    reply: redis::RedisResult<redis::Value>,
-) -> MockCmd {
-    let mut command = redis::cmd("FT.SEARCH");
-    command
-        .arg(index)
-        .arg(format!(
-            "(@litellm_cache_key:{{{tag}}})=>[KNN 1 @prompt_vector $vector AS vector_distance]"
-        ))
-        .arg("RETURN")
-        .arg(8)
-        .arg("entry_id")
-        .arg("prompt")
-        .arg("response")
-        .arg("inserted_at")
-        .arg("updated_at")
-        .arg("metadata")
-        .arg("litellm_cache_key")
-        .arg("vector_distance")
-        .arg("SORTBY")
-        .arg("vector_distance")
-        .arg("ASC")
-        .arg("DIALECT")
-        .arg(2)
-        .arg("LIMIT")
-        .arg(0)
-        .arg(1)
-        .arg("PARAMS")
-        .arg(2)
-        .arg("vector")
-        .arg(vector_bytes(vector));
-    MockCmd::new(command, reply)
-}
-
-fn hit(tag: &str, distance: &str, response: Vec<u8>) -> redis::Value {
-    redis::Value::Array(vec![
-        redis::Value::Int(1),
-        s("litellm_semantic_cache_index:stored-id"),
-        redis::Value::Array(vec![
-            s("entry_id"),
-            s("stored-id"),
-            s("prompt"),
-            s(PROMPT),
-            s("response"),
-            redis::Value::BulkString(response),
-            s("inserted_at"),
-            s("1700000000.5"),
-            s("updated_at"),
-            s("1700000000.5"),
-            s("litellm_cache_key"),
-            s(tag),
-            s("vector_distance"),
-            s(distance),
-        ]),
-    ])
-}
-
-fn empty_result() -> redis::Value {
-    redis::Value::Array(vec![redis::Value::Int(0)])
-}
-
-#[rstest]
-#[case::creates_index_and_expires(false, Some(Duration::from_secs(5)), Some(5))]
-#[case::existing_index_without_ttl(true, None, None)]
-#[case::fractional_ttl_rounds_up(true, Some(Duration::from_millis(1500)), Some(2))]
-fn store_writes_the_redisvl_hash(
-    #[case] index_exists: bool,
-    #[case] ttl: Option<Duration>,
-    #[case] expire: Option<u64>,
-    entry: Value,
-    context: SemanticCacheContext,
-) {
-    let hash_key = format!("{INDEX}:{}", entry_id(PROMPT, "key1"));
-    let mut commands = if index_exists {
-        vec![info(INDEX, compatible_info(3))]
-    } else {
-        vec![info_missing(INDEX), create_index(INDEX, 3)]
-    };
-    commands.push(hset(INDEX, PROMPT, "key1", &VECTOR, &entry));
-    commands.extend(
-        expire.map(|seconds| MockCmd::new(redis::cmd("EXPIRE").arg(&hash_key).arg(seconds), Ok(1))),
-    );
-    let cache = cache(commands, FakeEmbedder::new(&[]));
-
-    cache
-        .set_cache("key1", entry, &SemanticCacheContext { ttl, ..context })
-        .unwrap();
-}
-
-#[rstest]
-#[case::below_distance_threshold("key1", "0.05", None, Ok(Some(entry())))]
-#[case::above_distance_threshold("key1", "0.5", None, Ok(None))]
-#[case::other_cache_key("other", "0.05", None, Ok(None))]
-#[case::malformed_response("key1", "0.05", Some(b"not json!".as_slice()), Err(Error::InvalidEntry(None)))]
-fn lookup_applies_threshold_scope_and_codec(
-    #[case] stored_tag: &str,
-    #[case] distance: &str,
-    #[case] response: Option<&[u8]>,
-    #[case] expected: Result<Option<Value>, Error>,
-    entry: Value,
-    context: SemanticCacheContext,
-) {
-    let response = response.map_or_else(|| encoded(&entry), <[u8]>::to_vec);
-    let cache = cache(
-        vec![
-            info(INDEX, compatible_info(3)),
-            search(
-                INDEX,
-                "key1",
-                &VECTOR,
-                Ok(hit(stored_tag, distance, response)),
-            ),
-        ],
-        FakeEmbedder::new(&[]),
-    );
-
-    let kind = |result: Result<Option<Value>, Error>| {
-        result.map_err(|error| std::mem::discriminant(&error))
-    };
-    assert_eq!(kind(cache.get_cache("key1", &context)), kind(expected));
-}
-
-#[rstest]
-#[case::hit(context(), Some(hit("key1", "0.05", encoded(&entry()))), Some(entry()), Some(1.0 - 0.05))]
-#[case::beyond_distance_threshold(context(), Some(hit("key1", "0.5", encoded(&entry()))), None, Some(0.0))]
-#[case::no_results(context(), Some(empty_result()), None, Some(0.0))]
-#[case::other_cache_key(context(), Some(hit("other", "0.05", encoded(&entry()))), None, Some(0.0))]
-#[case::no_prompt(SemanticCacheContext::default(), None, None, Some(0.0))]
-#[tokio::test]
-async fn lookup_reports_python_semantic_similarity(
-    #[case] context: SemanticCacheContext,
-    #[case] reply: Option<redis::Value>,
-    #[case] value: Option<Value>,
-    #[case] similarity: Option<f64>,
-    #[values(false, true)] use_async: bool,
-) {
-    let commands = reply.map_or_else(Vec::new, |reply| {
-        vec![
-            info(INDEX, compatible_info(3)),
-            search(INDEX, "key1", &VECTOR, Ok(reply)),
-        ]
-    });
-    let cache = cache(commands, FakeEmbedder::new(&[]));
-
-    let lookup = if use_async {
-        cache
-            .async_get_cache_with_similarity("key1", &context)
-            .await
-    } else {
-        cache.get_cache_with_similarity("key1", &context)
-    };
-
-    assert_eq!(lookup, Ok(SemanticLookup { value, similarity }));
-}
-
-#[rstest]
-#[tokio::test]
-async fn missing_prompt_is_a_noop_that_never_embeds(entry: Value) {
-    let embedder = FakeEmbedder::new(&[]);
-    let calls = embedder.calls.clone();
-    let cache = cache(Vec::new(), embedder);
-    let context = SemanticCacheContext::default();
-
-    cache.set_cache("key1", entry.clone(), &context).unwrap();
-    assert_eq!(cache.get_cache("key1", &context).unwrap(), None);
-    cache
-        .async_set_cache("key1", entry, context.clone())
-        .await
-        .unwrap();
-    assert_eq!(cache.async_get_cache("key1", &context).await.unwrap(), None);
-    assert!(calls.lock().unwrap().is_empty());
-}
-
-#[rstest]
-fn scope_overrides_key_as_filter_tag(entry: Value, context: SemanticCacheContext) {
-    let cache = cache(
-        vec![
-            info(INDEX, compatible_info(3)),
-            hset(INDEX, PROMPT, "scope-a", &VECTOR, &entry),
-            search(
-                INDEX,
-                "scope\\-a",
-                &VECTOR,
-                Ok(hit("scope-a", "0.05", encoded(&entry))),
-            ),
-        ],
-        FakeEmbedder::new(&[]),
-    );
-    let context = SemanticCacheContext {
-        scope: Some("scope-a".into()),
-        ..context
-    };
-
-    cache.set_cache("key1", entry.clone(), &context).unwrap();
-    assert_eq!(cache.get_cache("key1", &context).unwrap(), Some(entry));
-}
-
-#[rstest]
-#[case::unscoped_schema(unscoped_info(3))]
-#[case::wrong_distance_metric(info_with_vector(vector_attribute_with(3, "FLOAT32", "L2")))]
-#[case::wrong_data_type(info_with_vector(vector_attribute_with(3, "FLOAT64", "COSINE")))]
-fn incompatible_schema_falls_back_to_isolated_index(
-    #[case] base_info: redis::Value,
-    entry: Value,
-    context: SemanticCacheContext,
-) {
-    let isolated = format!("{INDEX}_isolated");
-    let cache = cache(
-        vec![
-            info(INDEX, base_info),
-            info_missing(&isolated),
-            create_index(&isolated, 3),
-            hset(&isolated, PROMPT, "key1", &VECTOR, &entry),
-        ],
-        FakeEmbedder::new(&[]),
-    );
-
-    cache.set_cache("key1", entry, &context).unwrap();
-}
-
-#[rstest]
-fn create_index_race_rechecks_schema_and_stores(entry: Value, context: SemanticCacheContext) {
-    let cache = cache(
-        vec![
-            info_missing(INDEX),
-            MockCmd::new(
-                create_index_command(INDEX, 3),
-                Err::<&str, _>(support::server_error("Index already exists")),
-            ),
-            info(INDEX, compatible_info(3)),
-            hset(INDEX, PROMPT, "key1", &VECTOR, &entry),
-        ],
-        FakeEmbedder::new(&[]),
-    );
-
-    cache.set_cache("key1", entry, &context).unwrap();
-}
-
-#[rstest]
-#[case::punctuation_and_spaces("a:b, c|d", "a\\:b\\,\\ c\\|d")]
-#[case::braces_and_dots("{x}.y", "\\{x\\}\\.y")]
-#[case::plain("key1", "key1")]
-fn tag_special_characters_are_escaped_in_search_filter(
-    #[case] tag: &str,
-    #[case] escaped: &str,
-    context: SemanticCacheContext,
-) {
-    let cache = cache(
-        vec![
-            info(INDEX, compatible_info(3)),
-            search(INDEX, escaped, &VECTOR, Ok(empty_result())),
-        ],
-        FakeEmbedder::new(&[]),
-    );
-
-    assert_eq!(cache.get_cache(tag, &context).unwrap(), None);
-}
-
-#[rstest]
-#[case::content_parts(
-    SemanticCacheContext {
-        messages: Some(json!([
-            {"role": "user", "content": [{"type": "text", "text": "first"}, {"type": "text", "text": "second"}]},
-            {"role": "assistant", "content": "reply"},
-        ])),
-        ..Default::default()
-    },
-    "firstsecondreply"
-)]
-#[case::responses_string_input(
-    SemanticCacheContext { input: Some(json!("  plain input  ")), ..Default::default() },
-    "plain input"
-)]
-#[case::responses_nested_input(
-    SemanticCacheContext {
-        input: Some(json!([{"content": [{"type": "input_text", "text": "nested"}]}, "tail"])),
-        ..Default::default()
-    },
-    "nested\ntail"
-)]
-#[case::responses_output_text(
-    SemanticCacheContext { input: Some(json!({"output_text": "  result text  "})), ..Default::default() },
-    "result text"
-)]
-#[case::search_results(
-    messages_context(vec![json!({
-        "role": "user",
-        "content": "question",
-        "search_results": [{"source": "src", "title": "t", "content": [{"text": "found"}], "citations": {"a": 1}}],
-    })]),
-    "questionsrctfound{\"a\":1}"
-)]
-#[case::empty_messages_fall_back_to_input(
-    SemanticCacheContext { messages: Some(json!([])), input: Some(json!("fallback")), ..Default::default() },
-    "fallback"
-)]
-fn prompt_extraction_matches_python_message_and_input_shapes(
-    #[case] context: SemanticCacheContext,
-    #[case] prompt: &str,
-) {
-    let embedder = FakeEmbedder::new(&[]);
-    let calls = embedder.calls.clone();
-    let cache = cache(
-        vec![
-            info(INDEX, compatible_info(3)),
-            search(INDEX, "key1", &VECTOR, Ok(empty_result())),
-        ],
-        embedder,
-    );
-
-    cache.get_cache("key1", &context).unwrap();
-
-    assert_eq!(*calls.lock().unwrap(), vec![(prompt.to_owned(), None)]);
-}
-
-#[rstest]
-#[case(None)]
-#[case(Some(Duration::from_secs(9)))]
-fn ttl_passes_through_context_only(#[case] ttl: Option<Duration>) {
-    let cache = cache(Vec::new(), FakeEmbedder::new(&[]));
-
-    assert_eq!(
-        cache.get_ttl(&SemanticCacheContext {
-            ttl,
-            ..Default::default()
-        }),
-        ttl
-    );
-}
-
-#[rstest]
-#[tokio::test]
-async fn async_paths_embed_with_metadata_then_run_blocking_redis_work(
-    entry: Value,
-    context: SemanticCacheContext,
-) {
-    let embedder = FakeEmbedder::new(&[]);
-    let calls = embedder.calls.clone();
-    let cache = cache(
-        vec![
-            info(INDEX, compatible_info(3)),
-            hset(INDEX, PROMPT, "key1", &VECTOR, &entry),
-            search(
-                INDEX,
-                "key1",
-                &VECTOR,
-                Ok(hit("key1", "0.05", encoded(&entry))),
-            ),
-        ],
-        embedder,
-    );
-    let context = SemanticCacheContext {
-        metadata: Some(json!({"tenant": "team"})),
-        ..context
-    };
-
-    cache
-        .async_set_cache("key1", entry.clone(), context.clone())
-        .await
-        .unwrap();
-    assert_eq!(
-        cache.async_get_cache("key1", &context).await.unwrap(),
-        Some(entry)
-    );
-    assert_eq!(
-        *calls.lock().unwrap(),
-        vec![(PROMPT.to_owned(), context.metadata.clone()); 2]
-    );
-}
-
-#[rstest]
-fn accessors_report_the_config(config: RedisSemanticConfig) {
-    let cache = cache(Vec::new(), FakeEmbedder::new(&[]));
-
-    assert_eq!(cache.index_name(), config.index_name);
-    assert!((cache.similarity_threshold() - config.similarity_threshold).abs() < 1e-6);
-}
-
-#[rstest]
-fn shared_base_index_across_dimensions_replaces_the_isolated_index(entry: Value) {
-    // Pins parity with Python's `_isolated` + overwrite=True flow.
-    let prompt = "shared prompt";
-    let isolated = format!("{INDEX}_isolated");
-    let context = || messages_context(vec![json!({"role": "user", "content": prompt})]);
-
-    let vector_a = vec![0.1f32; 8];
-    let worker_a = cache(
-        vec![
-            info_missing(INDEX),
-            create_index(INDEX, 8),
-            hset(INDEX, prompt, "key1", &vector_a, &entry),
-        ],
-        FakeEmbedder::new(&[(prompt, &vector_a)]),
-    );
-    worker_a
-        .set_cache("key1", entry.clone(), &context())
-        .unwrap();
-
-    let vector_b = vec![0.2f32; 4];
-    let worker_b = cache(
-        vec![
-            info(INDEX, compatible_info(8)),
-            info_missing(&isolated),
-            create_index(&isolated, 4),
-            hset(&isolated, prompt, "key1", &vector_b, &entry),
-            search(
-                &isolated,
-                "key1",
-                &vector_b,
-                Ok(hit("key1", "0.0", encoded(&entry))),
-            ),
-            search(
-                &isolated,
-                "key1",
-                &vector_b,
-                Err(redis::RedisError::from((
-                    redis::ErrorKind::Extension,
-                    "Vector dimension mismatch",
-                ))),
-            ),
-        ],
-        FakeEmbedder::new(&[(prompt, &vector_b)]),
-    );
-    worker_b
-        .set_cache("key1", entry.clone(), &context())
-        .unwrap();
-    assert_eq!(
-        worker_b.get_cache("key1", &context()).unwrap(),
-        Some(entry.clone())
-    );
-
-    let vector_c = vec![0.3f32; 16];
-    let worker_c = cache(
-        vec![
-            info(INDEX, compatible_info(8)),
-            info(&isolated, compatible_info(4)),
-            MockCmd::new(redis::cmd("FT.DROPINDEX").arg(&isolated), Ok("OK")),
-            create_index(&isolated, 16),
-            hset(&isolated, prompt, "key1", &vector_c, &entry),
-        ],
-        FakeEmbedder::new(&[(prompt, &vector_c)]),
-    );
-    worker_c
-        .set_cache("key1", entry.clone(), &context())
-        .unwrap();
-
-    assert!(matches!(
-        worker_b.get_cache("key1", &context()),
-        Err(Error::Unavailable(Some(_)))
-    ));
-}
-
-#[fixture]
-fn redis_stack_url() -> Option<String> {
-    std::env::var("LITELLM_REDIS_STACK_URL").ok()
-}
-
-fn live_cache(
-    url: &str,
-    index_name: &str,
-    prompt: &str,
-    vector: Vec<f32>,
-) -> RedisSemanticCache<FakeEmbedder, JsonCodec<Value>> {
-    RedisSemanticCache::new(
-        url,
-        FakeEmbedder::new(&[(prompt, vector.as_slice())]),
-        JsonCodec::<Value>::new(),
-        RedisSemanticConfig {
-            index_name: index_name.to_owned(),
-            similarity_threshold: 0.9,
-        },
-    )
-    .unwrap()
-}
-
-#[rstest]
-fn live_shared_index_is_replaced_across_dimensions(redis_stack_url: Option<String>, entry: Value) {
-    let Some(url) = redis_stack_url else {
-        return;
-    };
-    // Pins parity with Python's `_isolated` + overwrite=True flow.
-    let base = format!("rust_semantic_shared_{}", std::process::id());
-    let isolated = format!("{base}_isolated");
-    let prompt = "shared live prompt";
-    let context = || messages_context(vec![json!({"role": "user", "content": prompt})]);
-
-    let worker_a = live_cache(&url, &base, prompt, vec![0.1f32; 8]);
-    worker_a
-        .set_cache("key1", entry.clone(), &context())
-        .unwrap();
-
-    let worker_b = live_cache(&url, &base, prompt, vec![0.2f32; 4]);
-    worker_b
-        .set_cache("key1", entry.clone(), &context())
-        .unwrap();
-    assert_eq!(
-        worker_b.get_cache("key1", &context()).unwrap(),
-        Some(entry.clone())
-    );
-
-    let worker_c = live_cache(&url, &base, prompt, vec![0.3f32; 16]);
-    worker_c
-        .set_cache("key1", entry.clone(), &context())
-        .unwrap();
-
-    assert!(matches!(
-        worker_b.get_cache("key1", &context()),
-        Err(Error::Unavailable(Some(_)))
-    ));
-
-    let mut connection = redis::Client::open(url).unwrap().get_connection().unwrap();
-    for index in [&base, &isolated] {
-        let _: Result<(), _> = redis::cmd("FT.DROPINDEX")
-            .arg(index)
-            .arg("DD")
-            .query(&mut connection);
-    }
-}
-
-#[rstest]
-fn live_store_lookup_and_ttl_against_redis_stack(redis_stack_url: Option<String>, entry: Value) {
-    let Some(url) = redis_stack_url else {
-        return;
-    };
-    let prompt = "rust semantic cache live prompt";
-    let index_name = format!("rust_semantic_test_{}", std::process::id());
-    let cache = live_cache(&url, &index_name, prompt, vec![0.1, 0.2, 0.3, 0.4]);
-    let context = SemanticCacheContext {
-        ttl: Some(Duration::from_secs(120)),
-        ..messages_context(vec![json!({"role": "user", "content": prompt})])
-    };
-
-    cache
-        .set_cache("live-key", entry.clone(), &context)
-        .unwrap();
-    assert_eq!(cache.get_cache("live-key", &context).unwrap(), Some(entry));
-    assert_eq!(cache.get_cache("other-key", &context).unwrap(), None);
-
-    let mut connection = redis::Client::open(url).unwrap().get_connection().unwrap();
-    let ttl: i64 = redis::Commands::ttl(
-        &mut connection,
-        format!("{index_name}:{}", entry_id(prompt, "live-key")),
-    )
-    .unwrap();
-    assert!(
-        ttl > 0,
-        "expected stored hash to carry an expiry, got {ttl}"
-    );
-}

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

@@ -1,63 +0,0 @@
-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

@@ -1,302 +0,0 @@
-#![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(server_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(|| server_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(|| server_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(server_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
-    }
-}
-
-pub fn server_error(reply: &str) -> redis::RedisError {
-    redis::parse_redis_value(format!("-{reply}\r\n").as_bytes())
-        .expect("error reply")
-        .extract_error()
-        .expect_err("server error")
-}
-
-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

@@ -1,20 +0,0 @@
-[package]
-name = "litellm-cache-redis"
-version = "0.1.0"
-edition.workspace = true
-license.workspace = true
-repository.workspace = true
-
-[dependencies]
-litellm-cache.workspace = true
-litellm-error.workspace = true
-redis = { workspace = true, features = ["cluster", "tls-rustls"] }
-r2d2 = "0.8.10"
-thiserror.workspace = true
-tokio.workspace = true
-
-[dev-dependencies]
-litellm-cache-testing.workspace = true
-redis-test.workspace = true
-rstest.workspace = true
-serde_json.workspace = true

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

@@ -1,192 +0,0 @@
-use std::{
-    sync::{Arc, OnceLock},
-    time::Duration,
-};
-
-use litellm_cache::{BatchEntry, CacheCodec, Error};
-
-use crate::{
-    connection::{ConnectionRef, Connections},
-    topology::RedisTopology,
-};
-
-const DEFAULT_TTL: Duration = Duration::from_secs(600);
-
-pub struct RedisCache<S, C = redis::Connection> {
-    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> {
-    pub fn new(url: &str, default_ttl: Option<Duration>, codec: S) -> Result<Self, Error> {
-        Self::connect(url, &RedisTopology::Standalone, default_ttl, codec)
-    }
-
-    pub fn connect(
-        url: &str,
-        topology: &RedisTopology,
-        default_ttl: Option<Duration>,
-        codec: S,
-    ) -> Result<Self, Error> {
-        let connections = Connections::open(url, topology)?;
-        Ok(Self {
-            connections: Arc::new(connections),
-            default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
-            codec,
-            namespace: None,
-            topology: topology.clone(),
-            major_version: Arc::default(),
-        })
-    }
-}
-
-impl<S, C> RedisCache<S, C>
-where
-    S: CacheCodec,
-    C: redis::ConnectionLike + Send + 'static,
-{
-    pub fn with_connection(connection: C, default_ttl: Option<Duration>, codec: S) -> Self {
-        Self {
-            connections: Arc::new(Connections::fixed(connection)),
-            default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
-            codec,
-            namespace: None,
-            topology: RedisTopology::Standalone,
-            major_version: Arc::default(),
-        }
-    }
-
-    pub fn with_namespace(self, namespace: Option<String>) -> Self {
-        Self {
-            namespace: namespace.filter(|value| !value.is_empty()),
-            ..self
-        }
-    }
-
-    pub fn namespace(&self) -> Option<&str> {
-        self.namespace.as_deref()
-    }
-
-    pub fn topology(&self) -> &RedisTopology {
-        &self.topology
-    }
-
-    pub(crate) fn namespaced_key(&self, key: &str) -> String {
-        namespaced_key(self.namespace.as_deref(), key)
-    }
-
-    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()
-            .flat_map(|ch| {
-                if matches!(ch, '*' | '?' | '[' | ']' | '\\') {
-                    vec!['\\', ch]
-                } else {
-                    vec![ch]
-                }
-            })
-            .collect();
-        Ok(format!("{escaped}:*"))
-    }
-
-    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),
-            redis::Value::SimpleString(text) => self.codec.decode(text.as_bytes()).map(Some),
-            _ => Err(Error::InvalidEntry(None)),
-        }
-    }
-
-    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),
-            Err(Error::InvalidEntry(_)) => Ok(BatchEntry::Invalid),
-            Err(error) => Err(error),
-        }
-    }
-}
-
-pub(crate) fn namespaced_key(namespace: Option<&str>, key: &str) -> String {
-    match namespace {
-        Some(namespace) if !key.starts_with(&format!("{namespace}:")) => {
-            format!("{namespace}:{key}")
-        }
-        _ => key.into(),
-    }
-}
-
-pub(crate) fn ttl_seconds(ttl: Duration) -> u64 {
-    ttl.as_secs()
-        .saturating_add(u64::from(ttl.subsec_nanos() > 0))
-        .max(1)
-}
-
-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),
-    }
-}
-
-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/claim.rs

@@ -1,109 +0,0 @@
-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(None)),
-    }
-}
-
-/// 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(None));
-    }
-    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(ClaimContended))
-}
-
-#[derive(Debug, thiserror::Error)]
-#[error("claim kept losing the race after {CLAIM_ATTEMPTS} attempts")]
-struct ClaimContended;
-
-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,513 +0,0 @@
-use std::{
-    sync::{Arc, Mutex},
-    time::Duration,
-};
-
-use litellm_cache::Error;
-use redis::{
-    ConnectionAddr, ConnectionInfo, IntoConnectionInfo,
-    cluster::{
-        ClusterClient, ClusterClientBuilder, ClusterConnection, ClusterPipeline, NodeAddress,
-    },
-    cluster_routing::{
-        MultipleNodeRoutingInfo, ResponsePolicy, RoutingInfo, SingleNodeRoutingInfo,
-    },
-};
-
-use crate::topology::{RedisNode, RedisTopology};
-
-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()?;
-                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(crate) struct ConnectionManager(redis::Client);
-
-impl ConnectionManager {
-    fn open(url: &str) -> Result<Self, Error> {
-        redis::Client::open(url)
-            .map(Self)
-            .map_err(Error::unavailable)
-    }
-}
-
-impl r2d2::ManageConnection for ConnectionManager {
-    type Connection = PooledConnection<redis::Connection>;
-    type Error = redis::RedisError;
-
-    fn connect(&self) -> Result<Self::Connection, redis::RedisError> {
-        let connection = self.0.get_connection()?;
-        connection.set_read_timeout(Some(REDIS_TIMEOUT))?;
-        connection.set_write_timeout(Some(REDIS_TIMEOUT))?;
-        Ok(PooledConnection {
-            connection,
-            failed: false,
-        })
-    }
-
-    fn is_valid(&self, connection: &mut Self::Connection) -> Result<(), redis::RedisError> {
-        redis::cmd("PING").query::<String>(&mut connection.connection)?;
-        Ok(())
-    }
-
-    fn has_broken(&self, connection: &mut Self::Connection) -> bool {
-        connection.failed || !redis::ConnectionLike::is_open(&connection.connection)
-    }
-}
-
-#[derive(Debug, thiserror::Error)]
-enum TopologyFault {
-    #[error("Redis Cluster needs at least one startup node")]
-    NoStartupNodes,
-    #[error("Redis Cluster nodes need a TCP address")]
-    NonTcpAddress,
-    #[error("Redis Cluster answered a fan-out command without a per-node map")]
-    FanOutReply,
-    #[error("scan target does not match the connection topology")]
-    ScanTarget,
-}
-
-pub(crate) struct ClusterConnectionManager(ClusterClient);
-
-impl ClusterConnectionManager {
-    fn open(url: &str, startup_nodes: &[RedisNode]) -> Result<Self, Error> {
-        if startup_nodes.is_empty() {
-            return Err(Error::unavailable(TopologyFault::NoStartupNodes));
-        }
-        let info = url.into_connection_info().map_err(Error::unavailable)?;
-        let nodes = startup_nodes
-            .iter()
-            .map(|node| node_info(&info, node))
-            .collect::<Result<Vec<_>, _>>()?;
-        ClusterClientBuilder::new(nodes)
-            .connection_timeout(REDIS_TIMEOUT)
-            .response_timeout(REDIS_TIMEOUT)
-            .build()
-            .map(Self)
-            .map_err(Error::unavailable)
-    }
-}
-
-fn node_info(info: &ConnectionInfo, node: &RedisNode) -> Result<ConnectionInfo, Error> {
-    let addr = match info.addr() {
-        ConnectionAddr::Tcp(..) => ConnectionAddr::Tcp(node.host.clone(), node.port),
-        ConnectionAddr::TcpTls {
-            insecure,
-            tls_params,
-            ..
-        } => ConnectionAddr::TcpTls {
-            host: node.host.clone(),
-            port: node.port,
-            insecure: *insecure,
-            tls_params: tls_params.clone(),
-        },
-        _ => return Err(Error::unavailable(TopologyFault::NonTcpAddress)),
-    };
-    Ok(info.clone().set_addr(addr))
-}
-
-impl r2d2::ManageConnection for ClusterConnectionManager {
-    type Connection = PooledConnection<ClusterConnection>;
-    type Error = redis::RedisError;
-
-    fn connect(&self) -> Result<Self::Connection, redis::RedisError> {
-        let connection = self.0.get_connection()?;
-        connection.set_read_timeout(Some(REDIS_TIMEOUT))?;
-        connection.set_write_timeout(Some(REDIS_TIMEOUT))?;
-        Ok(PooledConnection {
-            connection,
-            failed: false,
-        })
-    }
-
-    fn is_valid(&self, connection: &mut Self::Connection) -> Result<(), redis::RedisError> {
-        redis::cmd("PING").query::<String>(&mut connection.connection)?;
-        Ok(())
-    }
-
-    fn has_broken(&self, connection: &mut Self::Connection) -> bool {
-        connection.failed || !redis::ConnectionLike::is_open(&connection.connection)
-    }
-}
-
-pub enum ConnectionRef<'a> {
-    Node(&'a mut dyn redis::ConnectionLike),
-    Cluster(&'a mut ClusterConnection),
-}
-
-impl redis::ConnectionLike for ConnectionRef<'_> {
-    fn req_packed_command(&mut self, cmd: &[u8]) -> redis::RedisResult<redis::Value> {
-        match self {
-            Self::Node(connection) => connection.req_packed_command(cmd),
-            Self::Cluster(connection) => connection.req_packed_command(cmd),
-        }
-    }
-
-    fn req_packed_commands(
-        &mut self,
-        cmd: &[u8],
-        offset: usize,
-        count: usize,
-    ) -> redis::RedisResult<Vec<redis::Value>> {
-        match self {
-            Self::Node(connection) => connection.req_packed_commands(cmd, offset, count),
-            Self::Cluster(connection) => connection.req_packed_commands(cmd, offset, count),
-        }
-    }
-
-    fn get_db(&self) -> i64 {
-        match self {
-            Self::Node(connection) => connection.get_db(),
-            Self::Cluster(connection) => redis::ConnectionLike::get_db(*connection),
-        }
-    }
-
-    fn supports_pipelining(&self) -> bool {
-        match self {
-            Self::Node(connection) => connection.supports_pipelining(),
-            Self::Cluster(connection) => redis::ConnectionLike::supports_pipelining(*connection),
-        }
-    }
-
-    fn check_connection(&mut self) -> bool {
-        match self {
-            Self::Node(connection) => connection.check_connection(),
-            Self::Cluster(connection) => connection.check_connection(),
-        }
-    }
-
-    fn is_open(&self) -> bool {
-        match self {
-            Self::Node(connection) => connection.is_open(),
-            Self::Cluster(connection) => redis::ConnectionLike::is_open(*connection),
-        }
-    }
-}
-
-impl ConnectionRef<'_> {
-    /// 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,
-        pipeline: &redis::Pipeline,
-    ) -> Result<T, Error> {
-        match self {
-            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 cluster = ClusterPipeline::with_capacity(pipeline.len());
-                for command in pipeline.cmd_iter() {
-                    cluster.add_command(command.clone());
-                }
-                cluster
-                    .query(connection)
-                    .and_then(|replies| pipeline.query(&mut Replies(Some(replies))))
-            }
-        }
-        .map_err(Error::unavailable)
-    }
-
-    pub(crate) fn scan(
-        &mut self,
-        pattern: &str,
-        count: usize,
-        mut visit: impl FnMut(&mut Self, Vec<String>) -> Result<bool, Error>,
-    ) -> Result<(), Error> {
-        let pages = match self {
-            Self::Node(connection) => {
-                let page = scan_command(0, pattern, count)
-                    .query::<ScanPage>(*connection)
-                    .map_err(Error::unavailable)?;
-                vec![(None, page)]
-            }
-            Self::Cluster(connection) => connection
-                .route_command(
-                    &scan_command(0, pattern, count),
-                    RoutingInfo::MultiNode((
-                        MultipleNodeRoutingInfo::AllMasters,
-                        Some(ResponsePolicy::Special),
-                    )),
-                )
-                .map_err(Error::unavailable)
-                .and_then(primary_pages)?
-                .into_iter()
-                .map(|(node, page)| (Some(node), page))
-                .collect(),
-        };
-        for (node, (mut cursor, mut keys)) in pages {
-            loop {
-                if !visit(self, keys)? {
-                    return Ok(());
-                }
-                if cursor == 0 {
-                    break;
-                }
-                (cursor, keys) = self.scan_page(node.as_ref(), cursor, pattern, count)?;
-            }
-        }
-        Ok(())
-    }
-
-    pub(crate) fn ping(&mut self) -> Result<bool, redis::RedisError> {
-        let command = redis::cmd("PING");
-        match self {
-            Self::Node(connection) => command
-                .query::<String>(*connection)
-                .map(|response| response == "PONG"),
-            Self::Cluster(connection) => connection
-                .route_command(
-                    &command,
-                    RoutingInfo::MultiNode((
-                        MultipleNodeRoutingInfo::AllNodes,
-                        Some(ResponsePolicy::AllSucceeded),
-                    )),
-                )
-                .map(|_| true),
-        }
-    }
-
-    pub(crate) fn node_text(&mut self, command: &redis::Cmd) -> Result<String, Error> {
-        match self {
-            Self::Node(connection) => command.query(*connection).map_err(Error::unavailable),
-            Self::Cluster(connection) => {
-                let value = connection
-                    .route_command(
-                        command,
-                        RoutingInfo::MultiNode((
-                            MultipleNodeRoutingInfo::AllNodes,
-                            Some(ResponsePolicy::Special),
-                        )),
-                    )
-                    .map_err(Error::unavailable)?;
-                let redis::Value::Map(entries) = value else {
-                    return Err(Error::unavailable(TopologyFault::FanOutReply));
-                };
-                let mut replies = entries
-                    .into_iter()
-                    .map(|(node, reply)| {
-                        Ok((
-                            redis::from_redis_value::<String>(node).map_err(Error::unavailable)?,
-                            redis::from_redis_value::<String>(reply).map_err(Error::unavailable)?,
-                        ))
-                    })
-                    .collect::<Result<Vec<(String, String)>, Error>>()?;
-                replies.sort();
-                Ok(replies
-                    .into_iter()
-                    .map(|(_, reply)| reply)
-                    .collect::<Vec<_>>()
-                    .join("\n"))
-            }
-        }
-    }
-
-    pub(crate) fn flushall(&mut self) -> Result<(), Error> {
-        let command = redis::cmd("FLUSHALL");
-        match self {
-            Self::Node(connection) => command.query(*connection).map_err(Error::unavailable),
-            Self::Cluster(connection) => connection
-                .route_command(
-                    &command,
-                    RoutingInfo::MultiNode((
-                        MultipleNodeRoutingInfo::AllMasters,
-                        Some(ResponsePolicy::AllSucceeded),
-                    )),
-                )
-                .map(|_| ())
-                .map_err(Error::unavailable),
-        }
-    }
-
-    fn scan_page(
-        &mut self,
-        node: Option<&NodeAddress>,
-        cursor: u64,
-        pattern: &str,
-        count: usize,
-    ) -> Result<ScanPage, Error> {
-        let command = scan_command(cursor, pattern, count);
-        match (self, node) {
-            (Self::Node(connection), None) => {
-                command.query(*connection).map_err(Error::unavailable)
-            }
-            (Self::Cluster(connection), Some(node)) => connection
-                .route_command(
-                    &command,
-                    RoutingInfo::SingleNode(SingleNodeRoutingInfo::ByAddress {
-                        host: node.host().to_string(),
-                        port: node.port(),
-                    }),
-                )
-                .map_err(Error::unavailable)
-                .and_then(|value| redis::from_redis_value(value).map_err(Error::unavailable)),
-            _ => Err(Error::unavailable(TopologyFault::ScanTarget)),
-        }
-    }
-}
-
-type ScanPage = (u64, Vec<String>);
-
-fn primary_pages(value: redis::Value) -> Result<Vec<(NodeAddress, ScanPage)>, Error> {
-    let redis::Value::Map(entries) = value else {
-        return Err(Error::unavailable(TopologyFault::FanOutReply));
-    };
-    entries
-        .into_iter()
-        .map(|(node, page)| {
-            let node = redis::from_redis_value::<String>(node).map_err(Error::unavailable)?;
-            let node = NodeAddress::try_from(node.as_str()).map_err(Error::unavailable)?;
-            let page = redis::from_redis_value::<ScanPage>(page).map_err(Error::unavailable)?;
-            Ok((node, page))
-        })
-        .collect()
-}
-
-fn scan_command(cursor: u64, pattern: &str, count: usize) -> redis::Cmd {
-    let mut command = redis::cmd("SCAN");
-    command
-        .cursor_arg(cursor)
-        .arg("MATCH")
-        .arg(pattern)
-        .arg("COUNT")
-        .arg(count);
-    command
-}
-
-/// 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
-    }
-}

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

@@ -1,205 +0,0 @@
-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::invalid_entry)
-}
-
-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

@@ -1,63 +0,0 @@
-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,15 +0,0 @@
-mod cache;
-mod claim;
-pub mod connection;
-mod counter;
-mod keys;
-mod lifecycle;
-mod queue;
-mod script;
-mod store;
-mod topology;
-
-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

@@ -1,86 +0,0 @@
-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::unavailable)))
-            .await
-            .and_then(|ping| ping)
-            .map_err(|error| litellm_error::Report::new(&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

@@ -1,226 +0,0 @@
-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(None));
-        }
-        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(None));
-        }
-        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(None));
-        }
-        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(None));
-            }
-            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(None)),
-    }
-}
-
-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(None)),
-    }
-}

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

@@ -1,136 +0,0 @@
-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

@@ -1,232 +0,0 @@
-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/src/topology.rs

@@ -1,14 +0,0 @@
-#[derive(Clone, Debug, PartialEq, Eq)]
-pub struct RedisNode {
-    pub host: String,
-    pub port: u16,
-}
-
-#[derive(Clone, Debug, Default, PartialEq, Eq)]
-pub enum RedisTopology {
-    #[default]
-    Standalone,
-    Cluster {
-        startup_nodes: Vec<RedisNode>,
-    },
-}

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

@@ -1,563 +0,0 @@
-//! 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.
-
-mod support;
-
-use std::{collections::HashSet, time::Duration};
-
-use litellm_cache::{
-    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 Counter = RedisCache<JsonCodec<f64>>;
-
-#[fixture]
-fn cache(#[default("cache")] label: &str) -> Option<JsonCache> {
-    cluster_cache(label, Duration::from_secs(120), JsonCodec::new())
-}
-
-#[fixture]
-fn counter(#[default("counter")] label: &str) -> Option<Counter> {
-    cluster_cache(label, Duration::from_secs(60), JsonCodec::new())
-}
-
-#[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: HashSet<Slot> = keys.iter().map(Slot::for_key).collect();
-    assert!(slots.len() > 1, "keys must span multiple slots");
-    keys
-}
-
-fn seconds(seconds: u64) -> Option<Duration> {
-    Some(Duration::from_secs(seconds))
-}
-
-#[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(),
-    )
-    .err();
-    assert!(matches!(error, Some(Error::Unavailable(_))));
-}
-
-#[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, json!({ "index": index }), &context)
-            .unwrap();
-    }
-    for (index, key) in keys.iter().enumerate() {
-        assert_eq!(
-            cache.get_cache(key, &context).unwrap(),
-            Some(json!({ "index": index }))
-        );
-    }
-    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(
-    #[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, json!(index), context.clone())
-            .await
-            .unwrap();
-    }
-    let mut raw = redis::cluster::ClusterClient::new(vec![cluster_url()])
-        .unwrap()
-        .get_connection()
-        .unwrap();
-    let malformed = format!("{}:{}", cache.namespace().unwrap(), keys[1]);
-    redis::cmd("SET")
-        .arg(&malformed)
-        .arg("not json")
-        .exec(&mut raw)
-        .unwrap();
-
-    let entries = cache
-        .async_batch_get_cache(keys.clone(), context.clone())
-        .await
-        .unwrap();
-    assert_eq!(entries.len(), keys.len());
-    for (index, entry) in entries.iter().enumerate() {
-        let expected = if index == 1 {
-            BatchEntry::Invalid
-        } else if index % 5 == 0 {
-            BatchEntry::Miss
-        } else {
-            BatchEntry::Hit(json!(index))
-        };
-        assert_eq!(*entry, expected, "entry {index}");
-    }
-    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(
-    #[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);
-    cache
-        .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(), context.clone())
-        .await
-        .unwrap();
-    assert!(
-        hits.iter()
-            .enumerate()
-            .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(
-            queues
-                .iter()
-                .enumerate()
-                .map(|(index, key)| RedisRpushOperation {
-                    key: key.clone(),
-                    values: (0..=index)
-                        .map(|value| RedisArg::Integer(value as i64))
-                        .collect(),
-                })
-                .collect(),
-        )
-        .await
-        .unwrap();
-    assert_eq!(pushed, (1..=keys.len()).collect::<Vec<_>>());
-    let popped = cache
-        .async_lpop_pipeline(
-            queues
-                .iter()
-                .enumerate()
-                .map(|(index, key)| RedisLpopOperation {
-                    key: key.clone(),
-                    count: (index % 2 == 0).then_some(2),
-                })
-                .collect(),
-        )
-        .await
-        .unwrap();
-    for (index, result) in popped.into_iter().enumerate() {
-        match result {
-            RedisLpopResult::Value(value) => {
-                assert_eq!(index % 2, 1, "queue {index}");
-                assert_eq!(value, b"0");
-            }
-            RedisLpopResult::Values(values) => {
-                assert_eq!(index % 2, 0, "queue {index}");
-                let expected: Vec<Vec<u8>> = (0..=index)
-                    .take(2)
-                    .map(|value| value.to_string().into_bytes())
-                    .collect();
-                assert_eq!(values, expected);
-            }
-            other => panic!("queue {index}: {other:?}"),
-        }
-    }
-
-    let counters: Vec<String> = keys.iter().map(|key| format!("counter:{key}")).collect();
-    let totals = counter
-        .async_increment_pipeline(
-            counters
-                .iter()
-                .enumerate()
-                .map(|(index, key)| IncrementOperation {
-                    key: key.clone(),
-                    amount: index as f64 + 0.5,
-                    ttl: (index % 3 == 0).then_some(Duration::from_secs(30)),
-                })
-                .collect(),
-        )
-        .await
-        .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(),
-        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 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, json!(true), context.clone())
-            .await
-            .unwrap();
-        other
-            .async_set_cache(key, json!(true), context.clone())
-            .await
-            .unwrap();
-    }
-    let mut scanned = cache.async_scan_iter("key-", 1000).await.unwrap();
-    scanned.sort();
-    let mut expected: Vec<String> = keys
-        .iter()
-        .map(|key| format!("{}:{key}", cache.namespace().unwrap()))
-        .collect();
-    expected.sort();
-    assert_eq!(scanned, expected);
-    assert_eq!(cache.async_scan_iter("key-", 7).await.unwrap().len(), 7);
-
-    cache.flush_cache().unwrap();
-    let flushed = cache
-        .async_batch_get_cache(keys.clone(), context.clone())
-        .await
-        .unwrap();
-    assert!(flushed.iter().all(|entry| *entry == BatchEntry::Miss));
-    let kept = other.async_batch_get_cache(keys, context).await.unwrap();
-    assert!(
-        kept.iter()
-            .all(|entry| *entry == BatchEntry::Hit(json!(true)))
-    );
-    other.async_flush_cache().await.unwrap();
-}
-
-fn ping_calls_per_node(startup: &redis::Client) -> Vec<(String, u64)> {
-    let mut connection = startup.get_connection().unwrap();
-    let nodes: String = redis::cmd("CLUSTER")
-        .arg("NODES")
-        .query(&mut connection)
-        .unwrap();
-    let mut counts: Vec<(String, u64)> = nodes
-        .lines()
-        .map(|line| {
-            let address = line.split_whitespace().nth(1).unwrap();
-            let address = address.split('@').next().unwrap();
-            let mut node = redis::Client::open(format!("redis://{address}"))
-                .unwrap()
-                .get_connection()
-                .unwrap();
-            let stats: String = redis::cmd("INFO")
-                .arg("commandstats")
-                .query(&mut node)
-                .unwrap();
-            let calls = stats
-                .lines()
-                .find_map(|stat| stat.strip_prefix("cmdstat_ping:calls="))
-                .and_then(|rest| rest.split(',').next())
-                .map_or(0, |calls| calls.parse().unwrap());
-            (address.to_string(), calls)
-        })
-        .collect();
-    counts.sort();
-    counts
-}
-
-#[rstest]
-#[tokio::test]
-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:?}");
-    assert!(cache.ping().await.unwrap());
-    let after = ping_calls_per_node(&startup);
-    for ((node, calls_before), (_, calls_after)) in before.iter().zip(&after) {
-        assert!(calls_after > calls_before, "{node} was not pinged");
-    }
-    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 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;
-    };
-    assert_eq!(
-        counter
-            .increment_cache("spend", 1.5, context.clone())
-            .unwrap(),
-        1.5
-    );
-    assert_eq!(
-        counter
-            .async_increment("spend", 2.0, context.clone(), false)
-            .await
-            .unwrap(),
-        3.5
-    );
-    assert_eq!(
-        counter
-            .increment_with_floor("budget", -3, Duration::from_secs(30))
-            .unwrap(),
-        0
-    );
-    assert_eq!(
-        counter
-            .async_increment_with_floor("budget", 7, Duration::from_secs(30))
-            .await
-            .unwrap(),
-        7
-    );
-    assert_eq!(counter.async_set_max("peak", 4.0, None).await.unwrap(), 4.0);
-    assert_eq!(counter.async_set_max("peak", 2.0, None).await.unwrap(), 4.0);
-    counter.flush_cache().unwrap();
-
-    let owner = json!("owner-a");
-    let rival = json!("owner-b");
-    assert_eq!(
-        cache
-            .claim_cache("lock", owner.clone(), &[], context.clone())
-            .unwrap(),
-        owner
-    );
-    assert_eq!(
-        cache
-            .async_claim_cache("lock", rival.clone(), vec![owner.clone()], context.clone())
-            .await
-            .unwrap(),
-        owner
-    );
-    assert_eq!(
-        cache
-            .claim_cache("lock", rival.clone(), &[], context.clone())
-            .unwrap(),
-        owner
-    );
-    assert_eq!(
-        cache
-            .async_claim_cache("lock", rival.clone(), vec![rival.clone()], context.clone())
-            .await
-            .unwrap(),
-        rival
-    );
-
-    let script = cache
-        .async_register_script("return redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2])".into());
-    let reply = script
-        .invoke(
-            vec!["scripted".into()],
-            vec![RedisArg::Bytes(b"payload".to_vec()), RedisArg::Integer(5)],
-        )
-        .await
-        .unwrap();
-    assert_eq!(reply, redis::Value::Okay);
-    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()],
-            Vec::new(),
-        )
-        .await
-        .unwrap();
-    assert_eq!(evaluated, redis::Value::BulkString(b"payload".to_vec()));
-
-    assert_eq!(
-        cache
-            .async_set_cache_sadd(
-                "members",
-                vec![
-                    RedisArg::Bytes(b"a".to_vec()),
-                    RedisArg::Bytes(b"b".to_vec())
-                ],
-                seconds(9),
-            )
-            .await
-            .unwrap(),
-        2
-    );
-    assert_eq!(cache.async_get_ttl("members").await.unwrap(), seconds(9));
-
-    let result = cache.test_connection().await.unwrap();
-    assert_eq!(result.status, CacheConnectionStatus::Success);
-    assert!(cache.ping().await.unwrap());
-    let info = cache.info().unwrap();
-    assert!(info.matches("redis_version").count() > 1, "{info}");
-    assert!(cache.client_list().unwrap().contains("id="));
-    cache.async_flush_cache().await.unwrap();
-    assert_eq!(cache.async_get_ttl("members").await.unwrap(), None);
-    assert_eq!(cache.get_cache("lock", &context).unwrap(), None);
-}

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

@@ -1,95 +0,0 @@
-//! 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

@@ -1,231 +0,0 @@
-#![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(None));
-        }
-        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(None)),
-        }
-    }
-}
-
-/// 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

@@ -16,9 +16,6 @@ sha2.workspace = true
 litellm-cache-gcs.workspace = true
 litellm-http = { workspace = true, features = ["test-support"] }
 litellm-cache-memory.workspace = true
-litellm-cache-redis.workspace = true
-redis.workspace = true
-redis-test.workspace = true
 rstest.workspace = true
 tokio.workspace = true
 wiremock.workspace = true

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

@@ -2,79 +2,11 @@ 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 litellm_cache_response::{CacheEntry, ExactResponseCache, ResponseCache, ResponseCacheRequest};
 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}))
-    );
-}
+use support::{keyed, memory, request};
 
 /// The in-memory backend has no `test_connection`, as in Python, and still serves every response
 /// operation.

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

@@ -17,10 +17,9 @@ use litellm_cache_response::{
     CacheControls, CacheEntry, CacheKeyField, CacheKeyInput, ResponseCache, ResponseCacheRequest,
     WriteBuffer, cache_key,
 };
-use redis_test::MockCmd;
 use rstest::{fixture, rstest};
 use serde_json::{Value, json};
-use support::{keyed, memory, redis, request};
+use support::{keyed, memory, request};
 
 type Memory = Arc<ResponseCache<InMemoryCache<CacheEntry>>>;
 
@@ -331,79 +330,6 @@ async fn directives_skip_io_and_keep_reads_and_writes_independent(
     );
 }
 
-#[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, response, Duration::from_secs(100))
-        .await
-        .unwrap();
-}
-
-#[rstest]
-#[tokio::test]
-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(
@@ -541,27 +467,6 @@ async fn batch_lookup_reports_partial_hits_and_batch_store_populates_misses(
     assert_eq!(memory.lookup(&requests[2], now).unwrap(), None);
 }
 
-#[rstest]
-#[tokio::test]
-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(

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

@@ -1,15 +1,9 @@
 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 litellm_cache_response::{CacheEntry, CacheKeyInput, ResponseCache, ResponseCacheRequest};
 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(
@@ -29,12 +23,3 @@ pub fn keyed(key: &str) -> ResponseCacheRequest {
         ..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/src/auth.rs

@@ -36,9 +36,12 @@ impl ProvideCredentials for S3Credentials {
                 self.config.secret_access_key.clone(),
             ) {
                 return Ok(AwsCredentials::new(
-                    access_key_id,
-                    secret_access_key,
-                    self.config.session_token.clone(),
+                    access_key_id.expose(),
+                    secret_access_key.expose(),
+                    self.config
+                        .session_token
+                        .as_ref()
+                        .map(|value| value.expose().to_owned()),
                     None,
                     "litellm-s3-cache",
                 ));

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

@@ -5,9 +5,9 @@ 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),
+        access_key_id: Some("key".into()),
+        secret_access_key: Some("secret".into()),
+        session_token: session_token.map(Into::into),
         region_name: Some("us-east-1".to_string()),
         ..Default::default()
     }

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

@@ -23,8 +23,8 @@ pub fn config(endpoint: &str) -> S3CacheConfig {
             url: endpoint.to_string(),
         }),
         auth: AwsAuthConfig {
-            access_key_id: Some("key".to_string()),
-            secret_access_key: Some("secret".to_string()),
+            access_key_id: Some("key".into()),
+            secret_access_key: Some("secret".into()),
             region_name: Some("us-east-1".to_string()),
             ..Default::default()
         },

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

@@ -1,22 +0,0 @@
-[package]
-name = "litellm-cache-valkey-semantic"
-version = "0.1.0"
-edition.workspace = true
-license.workspace = true
-repository.workspace = true
-
-[dependencies]
-litellm-cache.workspace = true
-litellm-cache-redis.workspace = true
-redis = { workspace = true, features = ["tls-rustls"] }
-sha2.workspace = true
-thiserror.workspace = true
-uuid = { version = "1", features = ["v4"] }
-
-[dev-dependencies]
-litellm-error.workspace = true
-litellm-cache-testing.workspace = true
-redis-test.workspace = true
-rstest.workspace = true
-serde_json.workspace = true
-tokio.workspace = true

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

@@ -1,245 +0,0 @@
-use std::{
-    sync::{Arc, Mutex},
-    time::Duration,
-};
-
-use litellm_cache::{
-    BaseCache, CacheCodec, Error, SemanticCacheContext,
-    semantic::{Embedder, SemanticCache, SemanticLookup, prompt_from_context},
-};
-use litellm_cache_redis::{RedisTopology, connection::Connections};
-
-use crate::{
-    ValkeySemanticConfig,
-    index::IndexState,
-    search::{embedding_bytes, scope_tag, search_document, write_document},
-};
-
-/// `ValkeySemanticCache`: a semantic cache on valkey-search's TAG + VECTOR index. Values go
-/// through the injected codec, so the response layer decides what a cached entry is.
-pub struct ValkeySemanticCache<E, S, C = redis::Connection> {
-    connections: Arc<Connections<C>>,
-    embedder: E,
-    codec: S,
-    config: ValkeySemanticConfig,
-    index_dimension: Arc<Mutex<Option<usize>>>,
-}
-
-impl<E: Embedder, S: CacheCodec> ValkeySemanticCache<E, S> {
-    pub fn new(
-        url: &str,
-        embedder: E,
-        codec: S,
-        config: ValkeySemanticConfig,
-    ) -> Result<Self, Error> {
-        Ok(Self {
-            connections: Arc::new(Connections::open(url, &RedisTopology::Standalone)?),
-            embedder,
-            codec,
-            config,
-            index_dimension: Arc::new(Mutex::new(None)),
-        })
-    }
-}
-
-impl<E, S, C> ValkeySemanticCache<E, S, C>
-where
-    E: Embedder,
-    S: CacheCodec,
-    C: redis::ConnectionLike + Send + 'static,
-{
-    pub fn with_connection(
-        connection: C,
-        embedder: E,
-        codec: S,
-        config: ValkeySemanticConfig,
-    ) -> Self {
-        Self {
-            connections: Arc::new(Connections::fixed(connection)),
-            embedder,
-            codec,
-            config,
-            index_dimension: Arc::new(Mutex::new(None)),
-        }
-    }
-
-    pub fn similarity_threshold(&self) -> f64 {
-        self.config.similarity_threshold
-    }
-
-    pub fn index_name(&self) -> &str {
-        &self.config.index_name
-    }
-
-    fn index_state(&self) -> IndexState {
-        IndexState {
-            name: self.config.index_name.clone(),
-            prefix: format!("{}:", self.config.index_name),
-            dimension: Arc::clone(&self.index_dimension),
-            similarity_threshold: self.config.similarity_threshold,
-        }
-    }
-
-    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, S, C> ValkeySemanticCache<E, S, C>
-where
-    E: Embedder,
-    S: CacheCodec + Clone,
-    C: redis::ConnectionLike + Send + 'static,
-{
-    /// The same index and connections behind a different embedder.
-    pub fn with_embedder<E2: Embedder>(&self, embedder: E2) -> ValkeySemanticCache<E2, S, C> {
-        ValkeySemanticCache {
-            connections: Arc::clone(&self.connections),
-            embedder,
-            codec: self.codec.clone(),
-            config: self.config.clone(),
-            index_dimension: Arc::clone(&self.index_dimension),
-        }
-    }
-}
-
-impl<E, S, C> BaseCache for ValkeySemanticCache<E, S, C>
-where
-    E: Embedder,
-    S: CacheCodec,
-    C: redis::ConnectionLike + Send + 'static,
-{
-    type Value = S::Value;
-    type Context = SemanticCacheContext;
-
-    fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
-        context.ttl
-    }
-
-    fn set_cache(
-        &self,
-        key: &str,
-        value: Self::Value,
-        context: &Self::Context,
-    ) -> Result<(), Error> {
-        let Some(prompt) = prompt_from_context(context) else {
-            return Ok(());
-        };
-        let embedding = self.embedder.embed(&prompt, context.metadata.as_ref())?;
-        let response = self.codec.encode(&value)?;
-        let index = self.index_state();
-        self.connections.execute(|connection| {
-            write_document(
-                connection,
-                &index,
-                &scope_tag(key),
-                &prompt,
-                response,
-                embedding_bytes(&embedding),
-                self.get_ttl(context),
-            )
-        })
-    }
-
-    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)
-    }
-
-    async fn async_set_cache(
-        &self,
-        key: &str,
-        value: Self::Value,
-        context: Self::Context,
-    ) -> Result<(), Error> {
-        let Some(prompt) = prompt_from_context(&context) else {
-            return Ok(());
-        };
-        let embedding = self
-            .embedder
-            .async_embed(&prompt, context.metadata.as_ref())
-            .await?;
-        let response = self.codec.encode(&value)?;
-        let index = self.index_state();
-        let scope = scope_tag(key);
-        Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
-            write_document(
-                connection,
-                &index,
-                &scope,
-                &prompt,
-                response,
-                embedding_bytes(&embedding),
-                context.ttl,
-            )
-        })
-        .await
-    }
-
-    async fn async_get_cache(
-        &self,
-        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 document in the key's
-/// scope, and the closest document's similarity even when it misses the threshold.
-impl<E, S, C> SemanticCache for ValkeySemanticCache<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(SemanticLookup::miss(Some(0.0)));
-        };
-        let embedding = self.embedder.embed(&prompt, context.metadata.as_ref())?;
-        let index = self.index_state();
-        let lookup = self.connections.execute(|connection| {
-            search_document(
-                connection,
-                &index,
-                &scope_tag(key),
-                embedding_bytes(&embedding),
-            )
-        })?;
-        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 embedding = self
-            .embedder
-            .async_embed(&prompt, context.metadata.as_ref())
-            .await?;
-        let index = self.index_state();
-        let scope = scope_tag(key);
-        let lookup = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
-            search_document(connection, &index, &scope, embedding_bytes(&embedding))
-        })
-        .await?;
-        self.decode(lookup)
-    }
-}

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

@@ -1,8 +0,0 @@
-/// `ValkeySemanticCache.DEFAULT_VALKEY_INDEX_NAME`.
-pub const DEFAULT_INDEX_NAME: &str = "litellm_semantic_cache_index";
-
-#[derive(Clone, Debug, PartialEq)]
-pub struct ValkeySemanticConfig {
-    pub similarity_threshold: f64,
-    pub index_name: String,
-}

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

@@ -1,117 +0,0 @@
-use std::sync::{Arc, Mutex};
-
-use litellm_cache::Error;
-use litellm_cache_redis::connection::ConnectionRef;
-
-use crate::search::value_text;
-
-/// The valkey-search index one cache writes to, with the dimension it was last ensured for.
-#[derive(Clone)]
-pub(crate) struct IndexState {
-    pub(crate) name: String,
-    pub(crate) prefix: String,
-    pub(crate) dimension: Arc<Mutex<Option<usize>>>,
-    pub(crate) similarity_threshold: f64,
-}
-
-/// `_ensure_index_sync` / `_ensure_index_async`: create the TAG + HNSW index once per dimension,
-/// and accept an existing index unless it reports a different dimension.
-pub(crate) fn ensure_index(
-    connection: &mut ConnectionRef<'_>,
-    index: &IndexState,
-    dimension: usize,
-) -> Result<(), Error> {
-    if index
-        .dimension
-        .lock()?
-        .is_some_and(|existing| existing == dimension)
-    {
-        return Ok(());
-    }
-    let create = redis::cmd("FT.CREATE")
-        .arg(&index.name)
-        .arg("ON")
-        .arg("HASH")
-        .arg("PREFIX")
-        .arg(1)
-        .arg(&index.prefix)
-        .arg("SCHEMA")
-        .arg("litellm_cache_key")
-        .arg("TAG")
-        .arg("embedding")
-        .arg("VECTOR")
-        .arg("HNSW")
-        .arg(6)
-        .arg("TYPE")
-        .arg("FLOAT32")
-        .arg("DIM")
-        .arg(dimension)
-        .arg("DISTANCE_METRIC")
-        .arg("COSINE")
-        .query::<String>(connection)
-        .map(|_| ());
-    if let Err(error) = create {
-        if !index_exists(&error) {
-            return Err(Error::unavailable(error));
-        }
-        let info = redis::cmd("FT.INFO")
-            .arg(&index.name)
-            .query::<redis::Value>(connection)
-            .map_err(Error::unavailable)?;
-        if let Some(existing) =
-            index_dimension_from_info(&info).filter(|existing| *existing != dimension)
-        {
-            return Err(Error::unavailable(DimensionMismatch {
-                existing,
-                dimension,
-            }));
-        }
-    }
-    *index.dimension.lock()? = Some(dimension);
-    Ok(())
-}
-
-fn index_exists(error: &redis::RedisError) -> bool {
-    error.code().is_some()
-        && error
-            .detail()
-            .is_some_and(|detail| detail.to_ascii_lowercase().contains("already exists"))
-}
-
-#[derive(Debug, thiserror::Error)]
-#[error("index exists with dimension {existing}, not {dimension}")]
-struct DimensionMismatch {
-    existing: usize,
-    dimension: usize,
-}
-
-/// `_extract_index_dim`: flatten each attribute one level and read the value after
-/// `dimensions`.
-fn index_dimension_from_info(value: &redis::Value) -> Option<usize> {
-    let redis::Value::Array(values) = value else {
-        return None;
-    };
-    let attributes = values.windows(2).find_map(|pair| {
-        (value_text(&pair[0]).as_deref() == Some("attributes")).then_some(&pair[1])
-    })?;
-    let redis::Value::Array(fields) = attributes else {
-        return None;
-    };
-    fields.iter().find_map(|field| {
-        let redis::Value::Array(values) = field else {
-            return None;
-        };
-        let values = values
-            .iter()
-            .flat_map(|value| match value {
-                redis::Value::Array(values) => values.as_slice(),
-                _ => std::slice::from_ref(value),
-            })
-            .collect::<Vec<_>>();
-        values.windows(2).find_map(|pair| {
-            (value_text(pair[0]).as_deref() == Some("dimensions"))
-                .then(|| value_text(pair[1]).and_then(|value| value.parse().ok()))
-                .flatten()
-        })
-    })
-}

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

@@ -1,7 +0,0 @@
-mod cache;
-mod config;
-mod index;
-mod search;
-
-pub use cache::ValkeySemanticCache;
-pub use config::{DEFAULT_INDEX_NAME, ValkeySemanticConfig};

litellm-rust/crates/cache-valkey-semantic/src/search.rs

@@ -1,161 +0,0 @@
-use std::time::Duration;
-
-use litellm_cache::{Error, semantic::SemanticLookup};
-use litellm_cache_redis::connection::ConnectionRef;
-use sha2::{Digest, Sha256};
-use uuid::Uuid;
-
-use crate::index::{IndexState, ensure_index};
-
-/// `_scope_tag`: valkey-search TAG fields cannot match arbitrary keys verbatim, so scopes are
-/// the key's lowercase SHA-256.
-pub(crate) fn scope_tag(key: &str) -> String {
-    let digest = Sha256::digest(key.as_bytes());
-    digest.iter().map(|byte| format!("{byte:02x}")).collect()
-}
-
-pub(crate) fn embedding_bytes(embedding: &[f32]) -> Vec<u8> {
-    embedding
-        .iter()
-        .flat_map(|value| value.to_le_bytes())
-        .collect()
-}
-
-/// `HSET` a fresh `<prefix><scope>:<uuid4>` document, then `EXPIRE` it when a TTL is set.
-pub(crate) fn write_document(
-    connection: &mut ConnectionRef<'_>,
-    index: &IndexState,
-    scope: &str,
-    prompt: &str,
-    response: Vec<u8>,
-    vector: Vec<u8>,
-    ttl: Option<Duration>,
-) -> Result<(), Error> {
-    ensure_index(connection, index, vector.len() / size_of::<f32>())?;
-    let document = format!("{}{scope}:{}", index.prefix, Uuid::new_v4());
-    let mut pipeline = redis::pipe();
-    pipeline
-        .cmd("HSET")
-        .arg(&document)
-        .arg("litellm_cache_key")
-        .arg(scope)
-        .arg("prompt")
-        .arg(prompt)
-        .arg("response")
-        .arg(response)
-        .arg("embedding")
-        .arg(vector)
-        .ignore();
-    if let Some(ttl) = ttl {
-        pipeline
-            .cmd("EXPIRE")
-            .arg(&document)
-            .arg(ttl.as_secs())
-            .ignore();
-    }
-    pipeline.query::<()>(connection).map_err(Error::unavailable)
-}
-
-/// The KNN-1 search within `scope`: the closest document's similarity, and its stored response
-/// when that similarity reaches the threshold. No document reads as a similarity of `0.0`.
-pub(crate) fn search_document(
-    connection: &mut ConnectionRef<'_>,
-    index: &IndexState,
-    scope: &str,
-    vector: Vec<u8>,
-) -> Result<SemanticLookup<Vec<u8>>, Error> {
-    ensure_index(connection, index, vector.len() / size_of::<f32>())?;
-    let query =
-        format!("(@litellm_cache_key:{{{scope}}})=>[KNN 1 @embedding $vec AS vector_distance]");
-    let response = redis::cmd("FT.SEARCH")
-        .arg(&index.name)
-        .arg(query)
-        .arg("PARAMS")
-        .arg(2)
-        .arg("vec")
-        .arg(vector)
-        .arg("RETURN")
-        .arg(2)
-        .arg("response")
-        .arg("vector_distance")
-        .arg("DIALECT")
-        .arg(2)
-        .query::<redis::Value>(connection)
-        .map_err(Error::unavailable)?;
-    let Some(fields) = search_fields(response)? else {
-        return Ok(SemanticLookup::miss(Some(0.0)));
-    };
-    let field = |name: &str| {
-        fields
-            .iter()
-            .find_map(|(field, value)| (field == name).then(|| value.clone()))
-            .ok_or(Error::InvalidEntry(None))
-    };
-    let response = field("response")?;
-    let similarity = 1.0 - parse_f64(&field("vector_distance")?)?;
-    Ok(SemanticLookup {
-        value: (similarity >= index.similarity_threshold).then_some(response),
-        similarity: Some(similarity),
-    })
-}
-
-type SearchFields = Vec<(String, Vec<u8>)>;
-
-fn search_fields(value: redis::Value) -> Result<Option<SearchFields>, Error> {
-    let redis::Value::Array(values) = value else {
-        return Err(Error::InvalidEntry(None));
-    };
-    let total = parse_i64(values.first().ok_or(Error::InvalidEntry(None))?)?;
-    if total <= 0 || values.len() < 3 {
-        return Ok(None);
-    }
-    let redis::Value::Array(fields) = &values[2] else {
-        return Err(Error::InvalidEntry(None));
-    };
-    let (pairs, remainder) = fields.as_chunks::<2>();
-    if !remainder.is_empty() {
-        return Err(Error::InvalidEntry(None));
-    }
-    pairs
-        .iter()
-        .map(|pair| {
-            Ok((
-                value_text(&pair[0]).ok_or(Error::InvalidEntry(None))?,
-                value_bytes(&pair[1])?,
-            ))
-        })
-        .collect::<Result<Vec<_>, Error>>()
-        .map(Some)
-}
-
-fn parse_i64(value: &redis::Value) -> Result<i64, Error> {
-    value_text(value)
-        .ok_or(Error::InvalidEntry(None))?
-        .parse()
-        .map_err(Error::invalid_entry)
-}
-
-fn parse_f64(value: &[u8]) -> Result<f64, Error> {
-    std::str::from_utf8(value)
-        .map_err(Error::invalid_entry)?
-        .parse()
-        .map_err(Error::invalid_entry)
-}
-
-pub(crate) fn value_text(value: &redis::Value) -> Option<String> {
-    match value {
-        redis::Value::BulkString(bytes) => String::from_utf8(bytes.clone()).ok(),
-        redis::Value::SimpleString(value) => Some(value.clone()),
-        redis::Value::Int(value) => Some(value.to_string()),
-        _ => None,
-    }
-}
-
-fn value_bytes(value: &redis::Value) -> Result<Vec<u8>, Error> {
-    match value {
-        redis::Value::BulkString(bytes) => Ok(bytes.clone()),
-        redis::Value::SimpleString(value) => Ok(value.as_bytes().to_vec()),
-        redis::Value::Int(value) => Ok(value.to_string().into_bytes()),
-        _ => Err(Error::InvalidEntry(None)),
-    }
-}

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

@@ -1,430 +0,0 @@
-mod support;
-
-use std::{
-    sync::{Arc, Mutex},
-    time::Duration,
-};
-
-use litellm_cache::{
-    BaseCache, Error, JsonCodec, SemanticCacheContext,
-    semantic::{PreparedEmbedding, SemanticCache, SemanticLookup},
-};
-use litellm_cache_valkey_semantic::{
-    DEFAULT_INDEX_NAME, ValkeySemanticCache, ValkeySemanticConfig,
-};
-use redis_test::MockRedisConnection;
-use rstest::{fixture, rstest};
-use serde_json::{Value, json};
-use support::{EmbedCalls, FakeEmbedder, RecordingConnection};
-
-type Requests = Arc<Mutex<Vec<Vec<u8>>>>;
-type RecordingCache = ValkeySemanticCache<FakeEmbedder, JsonCodec<Value>, RecordingConnection>;
-
-const KEY_SCOPE: &str = "2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683";
-
-struct Recording {
-    cache: RecordingCache,
-    requests: Requests,
-    calls: EmbedCalls,
-}
-
-impl Recording {
-    fn text(&self) -> String {
-        self.requests
-            .lock()
-            .unwrap()
-            .iter()
-            .map(|request| String::from_utf8_lossy(request).into_owned())
-            .collect::<Vec<_>>()
-            .join("\n")
-    }
-}
-
-fn config() -> ValkeySemanticConfig {
-    ValkeySemanticConfig {
-        similarity_threshold: 0.8,
-        index_name: "test".into(),
-    }
-}
-
-fn recording(replies: impl IntoIterator<Item = redis::RedisResult<redis::Value>>) -> Recording {
-    let connection = RecordingConnection::new(replies);
-    let requests = connection.requests();
-    let embedder = FakeEmbedder::new(&[]);
-    let calls = Arc::clone(&embedder.calls);
-    Recording {
-        cache: ValkeySemanticCache::with_connection(
-            connection,
-            embedder,
-            JsonCodec::new(),
-            config(),
-        ),
-        requests,
-        calls,
-    }
-}
-
-#[fixture]
-fn entry() -> Value {
-    json!({"timestamp": 1.0, "response": {"answer": "ok"}})
-}
-
-#[fixture]
-fn context() -> SemanticCacheContext {
-    SemanticCacheContext {
-        messages: Some(json!([{"role": "user", "content": "hello"}])),
-        metadata: Some(json!({"source": "test"})),
-        ..Default::default()
-    }
-}
-
-fn ok() -> redis::RedisResult<redis::Value> {
-    Ok(redis::Value::SimpleString("OK".into()))
-}
-
-fn already_exists() -> redis::RedisResult<redis::Value> {
-    Err(support::server_error("Index already exists"))
-}
-
-fn bulk(value: &[u8]) -> redis::Value {
-    redis::Value::BulkString(value.to_vec())
-}
-
-fn search_hit(response: &[u8], distance: &str) -> redis::Value {
-    redis::Value::Array(vec![
-        redis::Value::Int(1),
-        bulk(b"test:document"),
-        redis::Value::Array(vec![
-            bulk(b"response"),
-            bulk(response),
-            bulk(b"vector_distance"),
-            bulk(distance.as_bytes()),
-        ]),
-    ])
-}
-
-fn encoded(value: &Value) -> Vec<u8> {
-    serde_json::to_vec(value).unwrap()
-}
-
-/// `FT.INFO` with the vector field's dimension nested one level down, as valkey-search reports.
-fn nested_dimension_info(dimension: i64) -> redis::Value {
-    redis::Value::Array(vec![
-        redis::Value::SimpleString("attributes".into()),
-        redis::Value::Array(vec![redis::Value::Array(vec![
-            redis::Value::SimpleString("embedding".into()),
-            redis::Value::Array(vec![
-                redis::Value::SimpleString("dimensions".into()),
-                redis::Value::Int(dimension),
-            ]),
-        ])]),
-    ])
-}
-
-/// `FT.INFO` with `dimensions` as a sibling string of the identifier.
-fn flat_dimension_info(dimension: &str) -> redis::Value {
-    redis::Value::Array(vec![
-        redis::Value::SimpleString("attributes".into()),
-        redis::Value::Array(vec![redis::Value::Array(vec![
-            redis::Value::SimpleString("identifier".into()),
-            redis::Value::SimpleString("embedding".into()),
-            redis::Value::Array(vec![
-                redis::Value::SimpleString("dimensions".into()),
-                redis::Value::SimpleString(dimension.into()),
-            ]),
-        ])]),
-    ])
-}
-
-#[rstest]
-#[case::string_content(json!([{"content": "hello"}]), None, Some("hello"))]
-#[case::text_parts(json!([{"content": [{"text": "hello"}, {"text": " world"}]}]), None, Some("hello world"))]
-#[case::non_object_parts_are_skipped(json!([{"content": ["raw", {"text": "hello"}]}]), None, Some("hello"))]
-#[case::search_results(json!([{"search_results": [{"source": "s", "title": "t", "content": [{"text": "c"}], "citations": ["x"]}]}]), None, Some(r#"stc["x"]"#))]
-#[case::responses_string_input(json!([]), Some(json!(" hello ")), Some("hello"))]
-#[case::responses_item_input(json!([]), Some(json!([{"content": "first"}, {"text": "second"}])), Some("first\nsecond"))]
-#[case::blank_input(json!([]), Some(json!("   ")), None)]
-fn prompt_shapes_follow_redis_semantic_extraction(
-    #[case] messages: Value,
-    #[case] input: Option<Value>,
-    #[case] expected: Option<&str>,
-) {
-    let recording = recording([ok(), Ok(redis::Value::Array(vec![redis::Value::Int(0)]))]);
-    let context = SemanticCacheContext {
-        messages: Some(messages),
-        input,
-        ..Default::default()
-    };
-
-    assert_eq!(recording.cache.get_cache("key", &context).unwrap(), None);
-
-    let prompts = recording
-        .calls
-        .lock()
-        .unwrap()
-        .iter()
-        .map(|(prompt, _)| prompt.clone())
-        .collect::<Vec<_>>();
-    assert_eq!(
-        prompts,
-        expected.into_iter().map(str::to_owned).collect::<Vec<_>>()
-    );
-    assert_eq!(
-        recording.requests.lock().unwrap().is_empty(),
-        expected.is_none()
-    );
-}
-
-#[rstest]
-#[case::key("key", KEY_SCOPE)]
-#[case::empty_key("", "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855")]
-fn documents_are_scoped_by_the_keys_sha256(
-    #[case] key: &str,
-    #[case] scope: &str,
-    entry: Value,
-    context: SemanticCacheContext,
-) {
-    let recording = recording([ok()]);
-
-    recording.cache.set_cache(key, entry, &context).unwrap();
-
-    let text = recording.text();
-    assert!(text.contains(&format!("test:{scope}:")));
-    assert!(text.contains(&format!("litellm_cache_key\r\n$64\r\n{scope}\r\n")));
-}
-
-#[rstest]
-#[case::no_ttl(None, None)]
-#[case::whole_seconds(Some(Duration::from_secs(5)), Some("5"))]
-#[case::fractional_seconds_truncate(Some(Duration::from_millis(1900)), Some("1"))]
-fn set_writes_hset_and_expires_only_with_a_ttl(
-    #[case] ttl: Option<Duration>,
-    #[case] expire: Option<&str>,
-    entry: Value,
-    context: SemanticCacheContext,
-) {
-    let recording = recording([ok()]);
-
-    recording
-        .cache
-        .set_cache(
-            "key",
-            entry,
-            &SemanticCacheContext {
-                ttl,
-                ..context.clone()
-            },
-        )
-        .unwrap();
-
-    let text = recording.text();
-    assert!(text.contains("FT.CREATE"));
-    assert!(text.contains("HSET"));
-    let expire_seconds = text
-        .split_once("EXPIRE\r\n")
-        .and_then(|(_, rest)| rest.split("\r\n").nth(3));
-    assert_eq!(expire_seconds, expire);
-    assert_eq!(
-        *recording.calls.lock().unwrap(),
-        vec![("hello".to_owned(), context.metadata)]
-    );
-}
-
-#[rstest]
-fn second_set_skips_create_after_dimension_is_cached(entry: Value, context: SemanticCacheContext) {
-    let recording = recording([ok()]);
-
-    recording
-        .cache
-        .set_cache("key", entry.clone(), &context)
-        .unwrap();
-    recording.cache.set_cache("key", entry, &context).unwrap();
-
-    let text = recording.text();
-    assert_eq!(text.matches("FT.CREATE").count(), 1);
-    assert_eq!(text.matches("HSET").count(), 2);
-}
-
-#[rstest]
-#[case::nested_matching(nested_dimension_info(3), Ok(()))]
-#[case::nested_mismatch(
-    nested_dimension_info(2),
-    Err("cache is unavailable: index exists with dimension 2, not 3")
-)]
-#[case::flat_matching(flat_dimension_info("3"), Ok(()))]
-#[case::flat_mismatch(
-    flat_dimension_info("2"),
-    Err("cache is unavailable: index exists with dimension 2, not 3")
-)]
-#[case::unreported_dimension_is_accepted(redis::Value::Array(vec![]), Ok(()))]
-fn existing_index_dimension_must_match_embedding(
-    #[case] info: redis::Value,
-    #[case] expected: Result<(), &str>,
-    entry: Value,
-    context: SemanticCacheContext,
-) {
-    let recording = recording([already_exists(), Ok(info)]);
-
-    let result = recording.cache.set_cache("key", entry, &context);
-    assert_eq!(
-        result.map_err(|error| litellm_error::Report::new(&error).to_string()),
-        expected.map_err(str::to_owned)
-    );
-}
-
-#[rstest]
-#[case::create_failure(Err(redis::RedisError::from((redis::ErrorKind::Io, "boom"))))]
-fn index_creation_failures_are_unavailable(
-    #[case] reply: redis::RedisResult<redis::Value>,
-    entry: Value,
-    context: SemanticCacheContext,
-) {
-    let recording = recording([reply]);
-
-    assert!(matches!(
-        recording.cache.set_cache("key", entry, &context),
-        Err(Error::Unavailable(Some(_)))
-    ));
-}
-
-#[rstest]
-#[case::within_threshold(search_hit(&encoded(&entry()), "0.1"), Ok(Some(entry())))]
-#[case::at_threshold(search_hit(&encoded(&entry()), "0.2"), Ok(Some(entry())))]
-#[case::beyond_threshold(search_hit(&encoded(&entry()), "0.5"), Ok(None))]
-#[case::zero_documents(redis::Value::Array(vec![redis::Value::Int(0)]), Ok(None))]
-#[case::missing_response(
-    redis::Value::Array(vec![
-        redis::Value::Int(1),
-        bulk(b"document"),
-        redis::Value::Array(vec![bulk(b"vector_distance"), bulk(b"0.1")]),
-    ]),
-    Err(Error::InvalidEntry(None))
-)]
-#[case::unparsable_distance(search_hit(b"not-json", "abc"), Err(Error::InvalidEntry(None)))]
-#[case::undecodable_response(search_hit(b"not-json", "0.1"), Err(Error::InvalidEntry(None)))]
-fn get_applies_threshold_and_decodes_entry(
-    #[case] reply: redis::Value,
-    #[case] expected: Result<Option<Value>, Error>,
-    context: SemanticCacheContext,
-) {
-    let recording = recording([ok(), Ok(reply)]);
-
-    let kind = |result: Result<Option<Value>, Error>| {
-        result.map_err(|error| std::mem::discriminant(&error))
-    };
-    assert_eq!(
-        kind(recording.cache.get_cache("key", &context)),
-        kind(expected)
-    );
-}
-
-#[rstest]
-#[case::hit(context(), Some(search_hit(&encoded(&entry()), "0.1")), Some(entry()), Some(1.0 - 0.1))]
-#[case::below_threshold(context(), Some(search_hit(&encoded(&entry()), "0.5")), None, Some(1.0 - 0.5))]
-#[case::no_results(context(), Some(redis::Value::Array(vec![redis::Value::Int(0)])), None, Some(0.0))]
-#[case::no_prompt(SemanticCacheContext::default(), None, None, Some(0.0))]
-#[tokio::test]
-async fn lookup_reports_python_semantic_similarity(
-    #[case] context: SemanticCacheContext,
-    #[case] reply: Option<redis::Value>,
-    #[case] value: Option<Value>,
-    #[case] similarity: Option<f64>,
-    #[values(false, true)] use_async: bool,
-) {
-    let searched = reply.is_some();
-    let recording = recording(reply.map_or_else(Vec::new, |reply| vec![ok(), Ok(reply)]));
-
-    let lookup = if use_async {
-        recording
-            .cache
-            .async_get_cache_with_similarity("key", &context)
-            .await
-    } else {
-        recording.cache.get_cache_with_similarity("key", &context)
-    };
-
-    assert_eq!(lookup, Ok(SemanticLookup { value, similarity }));
-    assert_eq!(recording.text().contains("FT.SEARCH"), searched);
-}
-
-#[rstest]
-#[tokio::test]
-async fn missing_prompt_does_not_touch_valkey(entry: Value) {
-    let cache = ValkeySemanticCache::with_connection(
-        MockRedisConnection::new([]).assert_all_commands_consumed(),
-        FakeEmbedder::new(&[]),
-        JsonCodec::new(),
-        config(),
-    );
-    let context = SemanticCacheContext::default();
-
-    cache.set_cache("key", entry.clone(), &context).unwrap();
-    assert_eq!(cache.get_cache("key", &context).unwrap(), None);
-    cache
-        .async_set_cache("key", entry, context.clone())
-        .await
-        .unwrap();
-    assert_eq!(cache.async_get_cache("key", &context).await.unwrap(), None);
-    assert_eq!(cache.get_ttl(&context), None);
-}
-
-#[rstest]
-fn with_embedder_shares_index_state_and_connections(entry: Value, context: SemanticCacheContext) {
-    let recording = recording([ok(), ok(), Ok(search_hit(&encoded(&entry), "0.1"))]);
-    recording
-        .cache
-        .set_cache("key", entry.clone(), &context)
-        .unwrap();
-
-    let prepared = recording
-        .cache
-        .with_embedder(PreparedEmbedding(vec![0.1, 0.2, 0.3]));
-
-    assert_eq!(prepared.get_cache("key", &context).unwrap(), Some(entry));
-    assert_eq!(recording.text().matches("FT.CREATE").count(), 1);
-}
-
-#[rstest]
-fn accessors_report_the_config() {
-    let recording = recording([]);
-
-    assert_eq!(recording.cache.index_name(), "test");
-    assert_eq!(recording.cache.similarity_threshold(), 0.8);
-    assert_eq!(DEFAULT_INDEX_NAME, "litellm_semantic_cache_index");
-}
-
-#[rstest]
-#[tokio::test]
-async fn async_set_and_get_use_shared_document_helpers(
-    entry: Value,
-    context: SemanticCacheContext,
-) {
-    let recording = recording([ok(), ok(), ok(), Ok(search_hit(&encoded(&entry), "0.1"))]);
-    let context = SemanticCacheContext {
-        ttl: Some(Duration::from_millis(1900)),
-        ..context
-    };
-
-    recording
-        .cache
-        .async_set_cache("key", entry.clone(), context.clone())
-        .await
-        .unwrap();
-    assert_eq!(
-        recording
-            .cache
-            .async_get_cache("key", &context)
-            .await
-            .unwrap(),
-        Some(entry)
-    );
-
-    let text = recording.text();
-    assert!(text.contains("FT.CREATE"));
-    assert!(text.contains("HSET"));
-    assert!(text.contains("EXPIRE"));
-    assert_eq!(
-        *recording.calls.lock().unwrap(),
-        vec![("hello".to_owned(), context.metadata.clone()); 2]
-    );
-}

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

@@ -1,62 +0,0 @@
-//! `overwrite_replaces` does not apply: like Python, every write is a new `<scope>:<uuid4>`
-//! document, so a second write with the same prompt adds a tie instead of replacing the first.
-
-mod support;
-
-use litellm_cache::{JsonCodec, SemanticCacheContext, semantic::PreparedEmbedding};
-use litellm_cache_testing as contract;
-use litellm_cache_valkey_semantic::{
-    DEFAULT_INDEX_NAME, ValkeySemanticCache, ValkeySemanticConfig,
-};
-use rstest::{fixture, rstest};
-use serde_json::{Value, json};
-use support::FakeSearch;
-
-type Cache = ValkeySemanticCache<PreparedEmbedding, JsonCodec<Value>, FakeSearch>;
-
-const PREFIX: &str = "contract:";
-
-#[fixture]
-fn cache() -> Cache {
-    ValkeySemanticCache::with_connection(
-        FakeSearch::default(),
-        PreparedEmbedding(vec![0.6, 0.8]),
-        JsonCodec::new(),
-        ValkeySemanticConfig {
-            similarity_threshold: 0.9,
-            index_name: DEFAULT_INDEX_NAME.into(),
-        },
-    )
-}
-
-#[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 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-valkey-semantic/tests/support/mod.rs

@@ -1,352 +0,0 @@
-#![allow(dead_code)]
-
-use std::{
-    collections::{BTreeMap, HashMap, VecDeque},
-    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 valkey-search 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(server_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(|| server_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(|| server_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(server_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
-    }
-}
-
-pub fn server_error(reply: &str) -> redis::RedisError {
-    redis::parse_redis_value(format!("-{reply}\r\n").as_bytes())
-        .expect("error reply")
-        .extract_error()
-        .expect_err("server error")
-}
-
-fn bulk(bytes: &[u8]) -> redis::Value {
-    redis::Value::BulkString(bytes.to_vec())
-}
-
-fn index_info(index: &FakeIndex) -> redis::Value {
-    redis::Value::Array(vec![
-        bulk(b"index_name"),
-        bulk(b"fake"),
-        bulk(b"attributes"),
-        redis::Value::Array(vec![
-            redis::Value::Array(vec![
-                bulk(b"identifier"),
-                bulk(b"litellm_cache_key"),
-                bulk(b"type"),
-                bulk(b"TAG"),
-            ]),
-            redis::Value::Array(vec![
-                bulk(b"identifier"),
-                bulk(index.vector_field.as_bytes()),
-                bulk(b"type"),
-                bulk(b"VECTOR"),
-                bulk(b"index"),
-                redis::Value::Array(vec![
-                    bulk(b"dimensions"),
-                    redis::Value::Int(index.dims as i64),
-                ]),
-            ]),
-        ]),
-    ])
-}
-
-/// 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
-}
-
-/// Records every packed request and answers from a script, `OK` once the script runs out.
-pub struct RecordingConnection {
-    requests: Arc<Mutex<Vec<Vec<u8>>>>,
-    replies: Mutex<VecDeque<redis::RedisResult<redis::Value>>>,
-}
-
-impl RecordingConnection {
-    pub fn new(replies: impl IntoIterator<Item = redis::RedisResult<redis::Value>>) -> Self {
-        Self {
-            requests: Arc::default(),
-            replies: Mutex::new(replies.into_iter().collect()),
-        }
-    }
-
-    pub fn requests(&self) -> Arc<Mutex<Vec<Vec<u8>>>> {
-        Arc::clone(&self.requests)
-    }
-
-    fn reply(&self) -> redis::RedisResult<redis::Value> {
-        self.replies
-            .lock()
-            .unwrap()
-            .pop_front()
-            .unwrap_or_else(|| Ok(redis::Value::SimpleString("OK".into())))
-    }
-}
-
-impl redis::ConnectionLike for RecordingConnection {
-    fn req_packed_command(&mut self, command: &[u8]) -> redis::RedisResult<redis::Value> {
-        self.requests.lock().unwrap().push(command.to_vec());
-        self.reply()
-    }
-
-    fn req_packed_commands(
-        &mut self,
-        command: &[u8],
-        _offset: usize,
-        count: usize,
-    ) -> redis::RedisResult<Vec<redis::Value>> {
-        self.requests.lock().unwrap().push(command.to_vec());
-        (0..count).map(|_| self.reply()).collect()
-    }
-
-    fn get_db(&self) -> i64 {
-        0
-    }
-
-    fn check_connection(&mut self) -> bool {
-        true
-    }
-
-    fn is_open(&self) -> bool {
-        true
-    }
-}

litellm-rust/crates/inference-ocr/tests/resources.rs

@@ -7,7 +7,7 @@ use std::sync::{
 
 use litellm_auth::AuthServices;
 use litellm_auth_gcp::{
-    CredentialSource, VertexAuth, VertexAuthFuture, VertexProviderLoader, VertexTokenSource,
+    CredentialSource, GcpAuthFuture, GcpProviderLoader, GcpTokenSource, VertexAuth,
 };
 use litellm_http::{HttpSettings, Resolution};
 use litellm_inference::resources::CoreResources;
@@ -20,28 +20,28 @@ use support::{ReceivedRequest, json_response, upstream};
 
 struct TokenSource(String);
 
-impl VertexTokenSource for TokenSource {
-    fn project_id(&self) -> VertexAuthFuture<'_, String> {
+impl GcpTokenSource for TokenSource {
+    fn project_id(&self) -> GcpAuthFuture<'_, String> {
         Box::pin(async { Ok(self.0.clone()) })
     }
 
-    fn token(&self) -> VertexAuthFuture<'_, String> {
-        Box::pin(async { Ok(self.0.clone()) })
+    fn token(&self) -> GcpAuthFuture<'_, litellm_secret_value::SecretValue> {
+        Box::pin(async { Ok(litellm_secret_value::SecretValue::new(self.0.clone())) })
     }
 }
 
 #[derive(Default)]
 struct Loader(AtomicUsize);
 
-impl VertexProviderLoader for Loader {
-    fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc<dyn VertexTokenSource>> {
+impl GcpProviderLoader for Loader {
+    fn load(&self, source: CredentialSource) -> GcpAuthFuture<'_, Arc<dyn GcpTokenSource>> {
         Box::pin(async move {
             self.0.fetch_add(1, Ordering::SeqCst);
             let identity = match source {
                 CredentialSource::Trusted(secret) => secret.expose().to_string(),
                 other => panic!("unexpected credential source: {other:?}"),
             };
-            Ok(Arc::new(TokenSource(identity)) as Arc<dyn VertexTokenSource>)
+            Ok(Arc::new(TokenSource(identity)) as Arc<dyn GcpTokenSource>)
         })
     }
 }

litellm-rust/crates/python-bridge/Cargo.toml

@@ -29,8 +29,6 @@ bytes.workspace = true
 futures-util.workspace = true
 litellm-cache.workspace = true
 litellm-error.workspace = true
-litellm-cache-memory.workspace = true
-litellm-cache-redis.workspace = true
 litellm-cache-response.workspace = true
 serde.workspace = true
 litellm-auth.workspace = true
@@ -56,7 +54,6 @@ pyo3.workspace = true
 prost.workspace = true
 pyo3-async-runtimes.workspace = true
 reqwest.workspace = true
-redis = { workspace = true, features = ["tls-rustls"] }
 serde_json.workspace = true
 strum.workspace = true
 veil.workspace = true

litellm-rust/crates/python-bridge/src/cache/AGENTS.md

@@ -1,7 +1,7 @@
 # Cache boundary
 
-This folder owns how Rust inference reaches the selected cache: global cache selection, route admission, delegation to a Python cache and the experimental V2 native handles. Cache algorithms, storage protocols and response-cache semantics belong to their cache crates
+This folder owns how Rust inference reaches the selected Python cache. Cache algorithms, storage protocols and response-cache semantics belong to their cache crates
 
-`mod.rs` exposes the cache boundary to routes and module registration. Adapter directories remain private. `selection.rs` asks `litellm.rust_bridge.host.cache` for the live configured cache, and owns route admission and inference protocol composition for both adapters
+`mod.rs` exposes the cache boundary to routes. `selection.rs` asks `litellm.rust_bridge.host.cache` for the live configured cache and owns route admission and inference protocol composition
 
-`python/` delegates operations to the selected Python cache without discovering configuration. `native/` owns the experimental V2 native cache handles. Neither adapter depends on shared selection or the other adapter. Shared composition depends on the adapters, and routes use only the parent module's exports
+`python/` delegates operations to the selected Python cache without discovering configuration. Routes use only the parent module's exports

litellm-rust/crates/python-bridge/src/cache/mod.rs

@@ -1,32 +1,15 @@
-mod native;
 mod python;
 mod selection;
 
-pub(crate) use native::NativeCacheHandle;
 pub(crate) use python::{CacheCall, PythonCache};
-pub(crate) use selection::{Cached, Selection, configure, configured_native};
+pub(crate) use selection::{Cached, Selection, configure, ensure_cache_disabled};
 
-use litellm_cache::Error;
-use pyo3::{
-    exceptions::{PyNotImplementedError, PyRuntimeError, PyValueError},
-    prelude::*,
-};
+use pyo3::{exceptions::PyRuntimeError, prelude::*};
 
-/// A cache call the bridge rejects itself, before or after the cache ran.
 #[derive(Debug, thiserror::Error)]
 pub(crate) enum HandleError {
-    #[error("recreate the v2 cache after fork")]
-    Forked,
-    #[error("cache limits must be positive")]
-    Limits,
-    #[error("namespace and a positive cache limit are required")]
-    RedisLimits,
-    #[error("cache durations must be finite and positive")]
-    Duration(#[source] Option<std::time::TryFromFloatSecsError>),
-    #[error("unsupported v2 cache control: {0}")]
-    Control(String),
-    #[error("native inference needs a v2 cache, the catalog keeps Python caches on Python")]
-    NotNative,
+    #[error("this inference route does not support configured caches")]
+    Unsupported,
     #[error("cache operation without configured cache")]
     Unconfigured,
     #[error("cache arguments unavailable")]
@@ -39,66 +22,10 @@ impl From<HandleError> for PyErr {
     fn from(error: HandleError) -> Self {
         let message = litellm_error::Report::new(&error).to_string();
         match error {
-            HandleError::Limits
-            | HandleError::RedisLimits
-            | HandleError::Duration(_)
-            | HandleError::Control(_) => PyValueError::new_err(message),
-            HandleError::Forked
-            | HandleError::NotNative
+            HandleError::Unsupported
             | HandleError::Unconfigured
             | HandleError::MissingArguments
             | HandleError::NothingPending => PyRuntimeError::new_err(message),
         }
     }
 }
-
-fn cache_error(error: Error) -> PyErr {
-    let message = litellm_error::Report::new(&error).to_string();
-    match error {
-        Error::InvalidEntry(_) => PyValueError::new_err(message),
-        Error::UnsupportedOperation => PyNotImplementedError::new_err(message),
-        Error::Unavailable(_) | Error::UnscopedFlush | Error::MissingPrompt => {
-            PyRuntimeError::new_err(message)
-        }
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use pyo3::{
-        exceptions::{PyRuntimeError, PyValueError},
-        prelude::*,
-    };
-    use rstest::rstest;
-
-    use super::{Error, cache_error};
-
-    fn refused() -> std::io::Error {
-        std::io::Error::other("connection refused")
-    }
-
-    #[rstest]
-    #[case::unavailable(
-        Error::unavailable(refused()),
-        false,
-        "cache is unavailable: connection refused"
-    )]
-    #[case::invalid_entry(
-        Error::invalid_entry(refused()),
-        true,
-        "invalid cache entry: connection refused"
-    )]
-    fn python_sees_the_class_and_the_whole_chain(
-        #[case] error: Error,
-        #[case] caller_fault: bool,
-        #[case] message: &str,
-    ) {
-        Python::initialize();
-        Python::attach(|py| {
-            let failure = cache_error(error);
-            assert_eq!(failure.is_instance_of::<PyValueError>(py), caller_fault);
-            assert_eq!(failure.is_instance_of::<PyRuntimeError>(py), !caller_fault);
-            assert_eq!(failure.value(py).to_string(), message);
-        });
-    }
-}

litellm-rust/crates/python-bridge/src/cache/native/AGENTS.md

@@ -1,9 +0,0 @@
-# Native cache bindings
-
-This directory exposes the experimental V2 native cache handles (`NativeCacheHandle`) to Python. Cache algorithms and storage protocols remain in their cache crates
-
-Accept the cache object selected by the parent module. Do not read global `litellm.cache`, decide route admission, or select the Python cache adapter here
-
-Keep Python-facing cache classes and method signatures stable when reorganizing modules. Native internals stay private to this directory unless the shared cache boundary or Python module registration needs them
-
-Verify changes with the V2 cache tests using a freshly built extension. Test cache behavior, not module paths or file structure

litellm-rust/crates/python-bridge/src/cache/native/mod.rs

@@ -1,3 +0,0 @@
-pub(super) mod v2;
-
-pub(crate) use v2::NativeCacheHandle;

litellm-rust/crates/python-bridge/src/cache/native/v2.rs

@@ -1,336 +0,0 @@
-use crate::cache::{HandleError, cache_error};
-use std::{
-    sync::Arc,
-    time::{Duration, SystemTime, UNIX_EPOCH},
-};
-
-use litellm_cache::{DeleteCache, DisconnectCache, PingCache};
-use litellm_host_python::{from_py, release_gil, to_py};
-use serde_json::Value;
-
-use litellm_cache_memory::InMemoryCache;
-use litellm_cache_redis::{RedisCache, RedisTopology};
-use litellm_cache_response::{
-    CacheEntry, CacheKeyInput, ExactResponseCache, ResponseCache, ResponseCacheCodec,
-    ResponseCacheConfig, ResponseCacheRequest, ResponseCacheService,
-};
-use pyo3::{prelude::*, types::PyDict};
-
-#[pyclass(
-    frozen,
-    name = "NativeCacheHandle",
-    module = "litellm.rust_bridge._native"
-)]
-pub(crate) struct NativeCacheHandle {
-    service: Arc<dyn ResponseCacheService>,
-    backend: Arc<dyn ExactResponseCache>,
-    storage: Storage,
-    pid: u32,
-}
-
-#[derive(Clone)]
-enum Storage {
-    Memory(Arc<InMemoryCache<CacheEntry>>),
-    Redis(Arc<RedisCache<ResponseCacheCodec>>),
-}
-
-impl NativeCacheHandle {
-    fn check_process(&self) -> PyResult<()> {
-        if self.pid != std::process::id() {
-            return Err(HandleError::Forked.into());
-        }
-        Ok(())
-    }
-}
-
-fn request(key: String, ttl: Option<f64>) -> PyResult<ResponseCacheRequest> {
-    let mut request: ResponseCacheRequest = ResponseCacheRequest::new(CacheKeyInput {
-        preset: Some(key),
-        ..Default::default()
-    });
-    request.context.ttl = ttl.map(duration).transpose()?;
-    Ok(request)
-}
-
-#[pymethods]
-impl NativeCacheHandle {
-    #[staticmethod]
-    #[pyo3(signature = (*, ttl=600.0, capacity=200, max_entry_bytes=4194304))]
-    fn memory(ttl: f64, capacity: usize, max_entry_bytes: usize) -> PyResult<Self> {
-        let ttl = duration(ttl)?;
-        if capacity == 0 || max_entry_bytes == 0 {
-            return Err(HandleError::Limits.into());
-        }
-        let storage = Arc::new(InMemoryCache::new(Some(capacity), Some(ttl)));
-        let backend = Arc::new(ResponseCache::new(storage.clone()).with_config(
-            ResponseCacheConfig {
-                namespace: "sdk".into(),
-                max_entry_bytes,
-            },
-        ));
-        Ok(Self {
-            service: backend.clone(),
-            backend,
-            storage: Storage::Memory(storage),
-            pid: std::process::id(),
-        })
-    }
-
-    #[staticmethod]
-    #[pyo3(signature = (url, *, namespace, ttl=600.0, max_entry_bytes=4194304))]
-    fn redis(
-        py: Python<'_>,
-        url: &str,
-        namespace: String,
-        ttl: f64,
-        max_entry_bytes: usize,
-    ) -> PyResult<Self> {
-        let ttl = duration(ttl)?;
-        if namespace.is_empty() || max_entry_bytes == 0 {
-            return Err(HandleError::RedisLimits.into());
-        }
-        let storage = Arc::new(
-            release_gil(py, || {
-                RedisCache::connect(
-                    url,
-                    &RedisTopology::Standalone,
-                    Some(ttl),
-                    ResponseCacheCodec,
-                )
-            })
-            .map_err(cache_error)?
-            .with_namespace(Some(namespace.clone())),
-        );
-        let backend = Arc::new(ResponseCache::new(storage.clone()).with_config(
-            ResponseCacheConfig {
-                namespace,
-                max_entry_bytes,
-            },
-        ));
-        Ok(Self {
-            service: backend.clone(),
-            backend,
-            storage: Storage::Redis(storage),
-            pid: std::process::id(),
-        })
-    }
-    fn get(&self, py: Python<'_>, key: String) -> PyResult<Py<PyAny>> {
-        self.check_process()?;
-        let request = request(key, None)?;
-        let value =
-            release_gil(py, || self.backend.lookup(&request, now())).map_err(cache_error)?;
-        to_py(py, &value)
-    }
-
-    #[pyo3(signature = (key, value, *, ttl=None))]
-    fn set(
-        &self,
-        py: Python<'_>,
-        key: String,
-        value: &Bound<'_, PyAny>,
-        ttl: Option<f64>,
-    ) -> PyResult<()> {
-        self.check_process()?;
-        let request = request(key, ttl)?;
-        let value: Value = from_py(value)?;
-        release_gil(py, || self.backend.store(&request, value, now())).map_err(cache_error)
-    }
-
-    fn async_get<'py>(&self, py: Python<'py>, key: String) -> PyResult<Bound<'py, PyAny>> {
-        self.check_process()?;
-        let request = request(key, None)?;
-        let backend = self.backend.clone();
-        crate::execution::run_async(
-            py,
-            async move { backend.async_lookup(&request, now()).await },
-            cache_error,
-        )
-    }
-
-    #[pyo3(signature = (key, value, *, ttl=None))]
-    fn async_set<'py>(
-        &self,
-        py: Python<'py>,
-        key: String,
-        value: &Bound<'_, PyAny>,
-        ttl: Option<f64>,
-    ) -> PyResult<Bound<'py, PyAny>> {
-        self.check_process()?;
-        let request = request(key, ttl)?;
-        let value: Value = from_py(value)?;
-        let backend = self.backend.clone();
-        crate::execution::run_async(
-            py,
-            async move { backend.async_store(&request, value, now()).await },
-            cache_error,
-        )
-    }
-
-    #[pyo3(signature = (entries, *, ttl=None))]
-    fn async_set_many<'py>(
-        &self,
-        py: Python<'py>,
-        entries: &Bound<'_, PyAny>,
-        ttl: Option<f64>,
-    ) -> PyResult<Bound<'py, PyAny>> {
-        self.check_process()?;
-        let entries: Vec<(String, Value)> = from_py(entries)?;
-        let entries = entries
-            .into_iter()
-            .map(|(key, value)| Ok((request(key, ttl)?, value)))
-            .collect::<PyResult<Vec<_>>>()?;
-        let backend = self.backend.clone();
-        crate::execution::run_async(
-            py,
-            async move { backend.async_store_batch(entries, now()).await },
-            cache_error,
-        )
-    }
-
-    fn flush(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
-        self.check_process()?;
-        let backend = self.backend.clone();
-        crate::execution::run_sync(py, async move { backend.async_flush().await }, cache_error)
-    }
-
-    fn async_flush<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
-        self.check_process()?;
-        let backend = self.backend.clone();
-        crate::execution::run_async(py, async move { backend.async_flush().await }, cache_error)
-    }
-
-    fn ping<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
-        self.check_process()?;
-        let storage = self.storage.clone();
-        crate::execution::run_async(
-            py,
-            async move {
-                match storage {
-                    Storage::Memory(_) => Ok(true),
-                    Storage::Redis(cache) => cache.ping().await,
-                }
-            },
-            cache_error,
-        )
-    }
-
-    fn disconnect<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
-        self.check_process()?;
-        let storage = self.storage.clone();
-        crate::execution::run_async(
-            py,
-            async move {
-                match storage {
-                    Storage::Memory(cache) => cache.disconnect().await,
-                    Storage::Redis(cache) => cache.disconnect().await,
-                }
-            },
-            cache_error,
-        )
-    }
-
-    fn delete<'py>(&self, py: Python<'py>, keys: Vec<String>) -> PyResult<Bound<'py, PyAny>> {
-        self.check_process()?;
-        let storage = self.storage.clone();
-        crate::execution::run_async(
-            py,
-            async move {
-                for key in keys {
-                    match &storage {
-                        Storage::Memory(cache) => cache.async_delete_cache(&key).await?,
-                        Storage::Redis(cache) => cache.async_delete_cache(&key).await?,
-                    }
-                }
-                Ok(())
-            },
-            cache_error,
-        )
-    }
-}
-
-fn duration(seconds: f64) -> Result<Duration, HandleError> {
-    match Duration::try_from_secs_f64(seconds) {
-        Ok(duration) if !duration.is_zero() => Ok(duration),
-        Ok(_) => Err(HandleError::Duration(None)),
-        Err(error) => Err(HandleError::Duration(Some(error))),
-    }
-}
-
-pub(in crate::cache) fn native_handle<'py>(
-    configured: &Bound<'py, PyAny>,
-) -> PyResult<Option<Bound<'py, PyAny>>> {
-    Ok(configured
-        .getattr_opt("cache")?
-        .map(|backend| backend.getattr_opt("native_handle"))
-        .transpose()?
-        .flatten()
-        .filter(|handle| handle.is_instance_of::<NativeCacheHandle>()))
-}
-
-pub(in crate::cache) fn configured(
-    configured: &Bound<'_, PyAny>,
-    kwargs: &Bound<'_, PyDict>,
-) -> PyResult<(
-    Option<Arc<dyn ResponseCacheService>>,
-    litellm_cache_response::CacheOptions,
-)> {
-    let handle = native_handle(configured)?.ok_or(HandleError::NotNative)?;
-    let cache = handle.extract::<PyRef<'_, NativeCacheHandle>>()?;
-    cache.check_process()?;
-    let controls = kwargs.get_item("cache")?.filter(|value| !value.is_none());
-    let controls = controls
-        .as_ref()
-        .map(|value| value.cast::<PyDict>())
-        .transpose()?;
-    if let Some(controls) = controls {
-        for name in controls.keys() {
-            let name = name.extract::<String>()?;
-            if !matches!(
-                name.as_str(),
-                "no-cache" | "no-store" | "ttl" | "s-maxage" | "s-max-age" | "use-cache"
-            ) {
-                return Err(HandleError::Control(name).into());
-            }
-        }
-    }
-    let boolean = |name: &str| -> PyResult<bool> {
-        controls
-            .map(|values| values.get_item(name))
-            .transpose()?
-            .flatten()
-            .map(|value| value.extract())
-            .transpose()
-            .map(|value| value.unwrap_or(false))
-    };
-    let seconds = |name: &str| -> PyResult<Option<Duration>> {
-        controls
-            .map(|values| values.get_item(name))
-            .transpose()?
-            .flatten()
-            .map(|value| Ok(duration(value.extract()?)?))
-            .transpose()
-    };
-    Ok((
-        Some(cache.service.clone()),
-        litellm_cache_response::CacheOptions {
-            policy: litellm_cache_response::CachePolicy {
-                caching: kwargs
-                    .get_item("caching")?
-                    .filter(|value| !value.is_none())
-                    .map(|value| value.extract())
-                    .transpose()?,
-                no_cache: boolean("no-cache")?,
-                no_store: boolean("no-store")?,
-                ttl: seconds("ttl")?,
-                max_age: seconds("s-max-age")?.or(seconds("s-maxage")?),
-            },
-            scope: litellm_cache_response::CacheScope::Shared,
-        },
-    ))
-}
-
-fn now() -> Duration {
-    SystemTime::now()
-        .duration_since(UNIX_EPOCH)
-        .unwrap_or_default()
-}

litellm-rust/crates/python-bridge/src/cache/selection.rs

@@ -1,14 +1,10 @@
-use super::{native, python};
-use litellm_cache_response::{
-    CacheOptions, CachePolicy, CacheScope, ResponseCacheService, ScopedCache,
-};
+use super::{HandleError, python};
+use litellm_cache_response::{CacheOptions, CachePolicy, CacheScope, ScopedCache};
 use litellm_host::{
     machine::{HostServices, MachineFault},
     protocol::Protocol,
 };
 use pyo3::{prelude::*, types::PyDict};
-use std::sync::Arc;
-
 const CACHE_HOST_MODULE: &str = "litellm.rust_bridge.host.cache";
 
 pub(crate) struct Cached<P>(std::marker::PhantomData<P>);
@@ -24,7 +20,6 @@ impl<P: Protocol> Protocol for Cached<P> {
 
 enum Backend {
     Disabled,
-    Native(Arc<dyn ResponseCacheService>),
     Python { namespace: String },
 }
 
@@ -43,7 +38,6 @@ impl Selection {
     {
         let service = match self.backend {
             Backend::Disabled => None,
-            Backend::Native(service) => Some(service),
             Backend::Python { namespace } => Some(python::service(services, namespace)),
         };
         (
@@ -76,37 +70,21 @@ fn selected_cache<'py>(
     Ok(Some(configured))
 }
 
-pub(crate) fn configured_native(
+pub(crate) fn ensure_cache_disabled(
     py: Python<'_>,
     kwargs: &Bound<'_, PyDict>,
     call_type: &str,
-) -> PyResult<(
-    Option<Arc<dyn ResponseCacheService>>,
-    litellm_cache_response::CacheOptions,
-)> {
+) -> PyResult<()> {
     let Some(configured) = selected_cache(py, kwargs, call_type)? else {
-        return Ok((
-            None,
-            litellm_cache_response::CacheOptions::new(litellm_cache_response::CacheScope::Shared),
-        ));
+        return Ok(());
     };
-    native_configuration(&configured, kwargs)
-}
-
-fn native_configuration(
-    configured: &Bound<'_, PyAny>,
-    kwargs: &Bound<'_, PyDict>,
-) -> PyResult<(Option<Arc<dyn ResponseCacheService>>, CacheOptions)> {
-    if !configured
+    if configured
         .call_method("should_use_cache", (), Some(kwargs))?
         .extract::<bool>()?
     {
-        return Ok((
-            None,
-            litellm_cache_response::CacheOptions::new(litellm_cache_response::CacheScope::Shared),
-        ));
+        return Err(HandleError::Unsupported.into());
     }
-    native::v2::configured(configured, kwargs)
+    Ok(())
 }
 
 pub(crate) fn configure(
@@ -122,13 +100,6 @@ pub(crate) fn configure(
             options: CacheOptions::new(CacheScope::Shared),
         });
     };
-    if native::v2::native_handle(&cache)?.is_some() {
-        let (native, options) = native_configuration(&cache, arguments)?;
-        return Ok(Selection {
-            backend: native.map_or(Backend::Disabled, Backend::Native),
-            options,
-        });
-    }
     let enabled = cache
         .call_method("should_use_cache", (), Some(arguments))?
         .extract::<bool>()?;

litellm-rust/crates/python-bridge/src/lib.rs

@@ -65,10 +65,6 @@ mod _native {
     fn init(module: &Bound<'_, PyModule>) -> PyResult<()> {
         let py = module.py();
         let dict = module.dict();
-        dict.set_item(
-            "NativeCacheHandle",
-            py.get_type::<crate::cache::NativeCacheHandle>(),
-        )?;
         dict.set_item(
             "_SecretManagerRuntime",
             py.get_type::<crate::secrets::runtime::NativeSecretManager>(),
@@ -92,7 +88,6 @@ mod tests {
         Python::initialize();
         Python::attach(|py| {
             let mut expected = vec![
-                "NativeCacheHandle",
                 "RustBridgeDeclined",
                 "RustUpstreamError",
                 "ForkedAfterNativeRuntimeStarted",

litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs

@@ -4,7 +4,7 @@ use std::sync::Arc;
 
 use host::ChatCompletionsPythonHost;
 use litellm_auth::AuthServices;
-use litellm_cache_response::{CachePolicy, ScopedCache};
+use litellm_cache_response::CachePolicy;
 use litellm_callbacks_legacy_python::LoggingOperation;
 use litellm_host::{call::HostedMachine, protocol::Protocol};
 use litellm_inference_chat::{ChatCompletionsRoute, route::ChatCompletions};
@@ -57,10 +57,6 @@ impl InferenceRoute for ChatCompletionsRoute {
         Self::new(http, auth, secrets)
     }
 
-    fn with_cache(self, cache: ScopedCache) -> Self {
-        self.with_cache(cache)
-    }
-
     fn machine(
         self,
         call: <ChatCompletions as Protocol>::Request,

litellm-rust/crates/python-bridge/src/routes/inference.rs

@@ -1,7 +1,7 @@
 use std::sync::Arc;
 
 use litellm_auth::AuthServices;
-use litellm_cache_response::{CachePolicy, CacheScope, ScopedCache};
+use litellm_cache_response::CachePolicy;
 use litellm_callbacks_legacy_python::LoggingOperation;
 use litellm_core_utils::get_llm_provider_logic::get_custom_llm_provider;
 use litellm_host::{call::HostedMachine, protocol::Protocol};
@@ -142,7 +142,6 @@ pub(super) trait InferenceRoute: Sized + 'static {
         auth: Arc<AuthServices>,
         secrets: Arc<dyn SecretSource>,
     ) -> Self;
-    fn with_cache(self, cache: ScopedCache) -> Self;
     fn machine(
         self,
         call: <Self::Protocol as Protocol>::Request,
@@ -176,12 +175,8 @@ where
                 crate::http::resources().auth.clone(),
                 crate::secrets::source(py)?,
             );
-            let (cache, cache_options) = crate::cache::configured_native(py, arguments, call_type)?;
-            let route = match cache {
-                Some(cache) => route.with_cache(ScopedCache::new(cache, CacheScope::Shared)),
-                None => route,
-            };
-            Ok(route.machine(request, cache_options.policy))
+            crate::cache::ensure_cache_disabled(py, arguments, call_type)?;
+            Ok(route.machine(request, CachePolicy::default()))
         },
         host,
         hooks,

litellm-rust/crates/python-bridge/src/routes/responses/mod.rs

@@ -5,7 +5,7 @@ use std::sync::Arc;
 
 use host::ResponsesPythonHost;
 use litellm_auth::AuthServices;
-use litellm_cache_response::{CachePolicy, ScopedCache};
+use litellm_cache_response::CachePolicy;
 use litellm_callbacks_legacy_python::LoggingOperation;
 use litellm_host::{call::HostedMachine, protocol::Protocol};
 use litellm_inference_responses::{ResponsesRoute, route::Responses};
@@ -49,10 +49,6 @@ impl InferenceRoute for ResponsesRoute {
         Self::new(http, auth, secrets)
     }
 
-    fn with_cache(self, cache: ScopedCache) -> Self {
-        self.with_cache(cache)
-    }
-
     fn machine(
         self,
         call: <Responses as Protocol>::Request,

litellm-rust/crates/secret-value/src/secret_value.rs

@@ -16,6 +16,18 @@ impl SecretValue {
     }
 }
 
+impl From<String> for SecretValue {
+    fn from(value: String) -> Self {
+        Self::new(value)
+    }
+}
+
+impl From<&str> for SecretValue {
+    fn from(value: &str) -> Self {
+        Self::new(value)
+    }
+}
+
 impl PartialEq for SecretValue {
     fn eq(&self, other: &Self) -> bool {
         subtle::ConstantTimeEq::ct_eq(self.0.as_bytes(), other.0.as_bytes()).into()

litellm-rust/crates/secrets-aws/src/auth.rs

@@ -39,30 +39,15 @@ impl Credentials {
         Self {
             auth,
             config: AwsAuthConfig {
-                access_key_id: context
-                    .access_key_id
-                    .as_ref()
-                    .map(|value| value.expose().to_owned()),
-                secret_access_key: context
-                    .secret_access_key
-                    .as_ref()
-                    .map(|value| value.expose().to_owned()),
-                session_token: context
-                    .session_token
-                    .as_ref()
-                    .map(|value| value.expose().to_owned()),
+                access_key_id: context.access_key_id.as_ref().cloned(),
+                secret_access_key: context.secret_access_key.as_ref().cloned(),
+                session_token: context.session_token.as_ref().cloned(),
                 region_name: region(settings, environment.as_ref()).ok(),
                 role_name: settings.aws_role_name.clone(),
                 session_name: settings.aws_session_name.clone(),
-                external_id: settings
-                    .aws_external_id
-                    .as_ref()
-                    .map(|v| v.expose().to_owned()),
+                external_id: settings.aws_external_id.as_ref().cloned(),
                 profile_name: settings.aws_profile_name.clone(),
-                web_identity_token: settings
-                    .aws_web_identity_token
-                    .as_ref()
-                    .map(|v| v.expose().to_owned()),
+                web_identity_token: settings.aws_web_identity_token.as_ref().cloned(),
                 sts_endpoint: settings.aws_sts_endpoint.clone(),
             },
             environment,