Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_redis_breaker_quiet_open

.github/workflows/test-unit.yml

@@ -200,6 +200,7 @@ jobs:
               tests/test_litellm/proxy/types_utils
               tests/test_litellm/proxy/logging_endpoints
               tests/test_litellm/proxy/test_*.py
+              tests/test_gateway
             workers: 4
             reruns: 2
             timeout-minutes: 20

AGENTS.md

@@ -1 +1,3 @@
 Read @CLAUDE.md for coding guidelines
+
+Before requesting maintainer review, verify the current PR tip passes required CI and code coverage, meets Greptile confidence of at least 4/5, and has acceptable Veria and Bugbot reviews. Inspect warnings and findings, fix actionable issues, and rerun the affected checks and reviewers after changes. Record evidence for any false positive or unavailable review; never treat a pending or missing bot result as a pass. Do not lower coverage thresholds or lint budgets to satisfy a check

Dockerfile

@@ -8,9 +8,25 @@ ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7
 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
 # Pinned by digest like the other base images; bump explicitly on Node upgrades.
 ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43
+# Checksum from https://www.pgbouncer.org/downloads/ (the Wolfi repo only carries 1.24.x)
+ARG PGBOUNCER_VERSION=1.25.2
+ARG PGBOUNCER_SHA256=924ad35113fd0a71c8e2dbe85b5d03445532e2b7b37a9f8a48983beea238b332
 
 FROM $UV_IMAGE AS uvbin
 
+FROM $LITELLM_BUILD_IMAGE AS pgbouncer-builder
+ARG PGBOUNCER_VERSION
+ARG PGBOUNCER_SHA256
+USER root
+RUN apk add --no-cache build-base pkgconf libevent-dev openssl-dev curl
+WORKDIR /build
+RUN curl -fsSL -o pgbouncer.tar.gz "https://www.pgbouncer.org/downloads/files/${PGBOUNCER_VERSION}/pgbouncer-${PGBOUNCER_VERSION}.tar.gz" && \
+    echo "${PGBOUNCER_SHA256}  pgbouncer.tar.gz" | sha256sum -c - && \
+    tar xzf pgbouncer.tar.gz --strip-components=1 && \
+    ./configure --prefix=/usr/local --with-openssl=/usr && \
+    make -j"$(nproc)" pgbouncer && \
+    install -m 0755 pgbouncer /usr/local/bin/pgbouncer
+
 # Admin UI builder. Pinned to the build platform so the architecture-independent
 # Next.js static export compiles once natively even in a multi-arch build,
 # instead of once per target arch under QEMU.
@@ -110,7 +126,8 @@ USER root
 RUN echo "https://packages.wolfi.dev/os" >> /etc/apk/repositories
 
 # node (without npm) is required by the prisma CLI at runtime
-RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile
+RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile libevent
+COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer
 
 WORKDIR /app
 ENV PATH="/app/.venv/bin:${PATH}" \

docker/Dockerfile.database

@@ -8,9 +8,25 @@ ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7
 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
 # Pinned by digest like the other base images; bump explicitly on Node upgrades.
 ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43
+# Checksum from https://www.pgbouncer.org/downloads/ (the Wolfi repo only carries 1.24.x)
+ARG PGBOUNCER_VERSION=1.25.2
+ARG PGBOUNCER_SHA256=924ad35113fd0a71c8e2dbe85b5d03445532e2b7b37a9f8a48983beea238b332
 
 FROM $UV_IMAGE AS uvbin
 
+FROM $LITELLM_BUILD_IMAGE AS pgbouncer-builder
+ARG PGBOUNCER_VERSION
+ARG PGBOUNCER_SHA256
+USER root
+RUN apk add --no-cache build-base pkgconf libevent-dev openssl-dev curl
+WORKDIR /build
+RUN curl -fsSL -o pgbouncer.tar.gz "https://www.pgbouncer.org/downloads/files/${PGBOUNCER_VERSION}/pgbouncer-${PGBOUNCER_VERSION}.tar.gz" && \
+    echo "${PGBOUNCER_SHA256}  pgbouncer.tar.gz" | sha256sum -c - && \
+    tar xzf pgbouncer.tar.gz --strip-components=1 && \
+    ./configure --prefix=/usr/local --with-openssl=/usr && \
+    make -j"$(nproc)" pgbouncer && \
+    install -m 0755 pgbouncer /usr/local/bin/pgbouncer
+
 # Admin UI builder. Pinned to the build platform so the architecture-independent
 # Next.js static export compiles once natively even in a multi-arch build,
 # instead of once per target arch under QEMU.
@@ -101,7 +117,8 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime
 USER root
 
 # node (without npm) is required by the prisma CLI at runtime
-RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile
+RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile libevent
+COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer
 
 WORKDIR /app
 ENV PATH="/app/.venv/bin:${PATH}" \

docker/Dockerfile.non_root

@@ -7,9 +7,25 @@ ARG PROXY_EXTRAS_SOURCE=published
 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
 # Pinned by digest like the other base images; bump explicitly on Node upgrades.
 ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43
+# Checksum from https://www.pgbouncer.org/downloads/ (the Wolfi repo only carries 1.24.x)
+ARG PGBOUNCER_VERSION=1.25.2
+ARG PGBOUNCER_SHA256=924ad35113fd0a71c8e2dbe85b5d03445532e2b7b37a9f8a48983beea238b332
 
 FROM $UV_IMAGE AS uvbin
 
+FROM $LITELLM_BUILD_IMAGE AS pgbouncer-builder
+ARG PGBOUNCER_VERSION
+ARG PGBOUNCER_SHA256
+USER root
+RUN apk add --no-cache build-base pkgconf libevent-dev openssl-dev curl
+WORKDIR /build
+RUN curl -fsSL -o pgbouncer.tar.gz "https://www.pgbouncer.org/downloads/files/${PGBOUNCER_VERSION}/pgbouncer-${PGBOUNCER_VERSION}.tar.gz" && \
+    echo "${PGBOUNCER_SHA256}  pgbouncer.tar.gz" | sha256sum -c - && \
+    tar xzf pgbouncer.tar.gz --strip-components=1 && \
+    ./configure --prefix=/usr/local --with-openssl=/usr && \
+    make -j"$(nproc)" pgbouncer && \
+    install -m 0755 pgbouncer /usr/local/bin/pgbouncer
+
 # Admin UI builder. Pinned to the build platform so the architecture-independent
 # Next.js static export compiles once natively even in a multi-arch build,
 # instead of once per target arch under QEMU.
@@ -128,8 +144,9 @@ RUN for i in 1 2 3; do \
       apk upgrade --no-cache && break || sleep 5; \
     done && \
     for i in 1 2 3; do \
-      apk add --no-cache python-3.13 bash openssl tzdata libsndfile nodejs && break || sleep 5; \
+      apk add --no-cache python-3.13 bash openssl tzdata libsndfile nodejs libevent && break || sleep 5; \
     done
+COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer
 
 # Copy only what runtime needs. The application is installed inside the venv;
 # the rest of the builder's /app is source and build metadata that must not

gateway/Dockerfile

@@ -1,9 +1,25 @@
 ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
 ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
+# Checksum from https://www.pgbouncer.org/downloads/ (the Wolfi repo only carries 1.24.x)
+ARG PGBOUNCER_VERSION=1.25.2
+ARG PGBOUNCER_SHA256=924ad35113fd0a71c8e2dbe85b5d03445532e2b7b37a9f8a48983beea238b332
 
 FROM $UV_IMAGE AS uvbin
 
+FROM $LITELLM_BUILD_IMAGE AS pgbouncer-builder
+ARG PGBOUNCER_VERSION
+ARG PGBOUNCER_SHA256
+USER root
+RUN apk add --no-cache build-base pkgconf libevent-dev openssl-dev curl
+WORKDIR /build
+RUN curl -fsSL -o pgbouncer.tar.gz "https://www.pgbouncer.org/downloads/files/${PGBOUNCER_VERSION}/pgbouncer-${PGBOUNCER_VERSION}.tar.gz" && \
+    echo "${PGBOUNCER_SHA256}  pgbouncer.tar.gz" | sha256sum -c - && \
+    tar xzf pgbouncer.tar.gz --strip-components=1 && \
+    ./configure --prefix=/usr/local --with-openssl=/usr && \
+    make -j"$(nproc)" pgbouncer && \
+    install -m 0755 pgbouncer /usr/local/bin/pgbouncer
+
 # ---------- Builder ----------
 FROM $LITELLM_BUILD_IMAGE AS builder
 
@@ -61,6 +77,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \
         --extra bedrock-realtime \
         --python python3.13
 
+# PYTHONPATH=/app makes the source tree shadow the installed package, so the
+# compiled Rust extension must live next to the source or it is never imported.
+RUN cp "$(python -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')"/litellm/rust_bridge/_native*.so litellm/rust_bridge/
+
 RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
     npm_config_cache=/root/.npm \
     prisma generate --schema=./schema.prisma
@@ -73,7 +93,7 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime
 USER root
 
 RUN for i in 1 2 3; do \
-      apk add --no-cache bash openssl tzdata python-3.13 libsndfile libatomic && break; \
+      apk add --no-cache bash openssl tzdata python-3.13 libsndfile libatomic libevent && break; \
       [ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \
       sleep 5; \
     done
@@ -90,15 +110,17 @@ ENV HOME=/home/nonroot \
 
 COPY --from=builder --chown=nonroot:nonroot /app /app
 COPY --from=builder /opt/prisma /opt/prisma
+COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer
 
 RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \
     find /app/.venv -type d -path "*/tornado/test" -delete && \
     chmod -R a+rX /opt/prisma && \
-    python -c "from prisma.client import BINARY_PATHS; paths = list(BINARY_PATHS.query_engine.values()); assert paths and all(p.startswith('/opt/prisma/') for p in paths), paths"
+    python -c "from prisma.client import BINARY_PATHS; paths = list(BINARY_PATHS.query_engine.values()); assert paths and all(p.startswith('/opt/prisma/') for p in paths), paths" && \
+    python -c "import litellm; from litellm.rust_bridge.loader import native_bridge_available; assert litellm.__file__ == '/app/litellm/__init__.py', litellm.__file__; assert native_bridge_available()"
 
 USER nonroot
 
 EXPOSE 4000/tcp
 
-ENTRYPOINT ["sh", "-c", "exec /app/docker/component_entrypoint.sh uvicorn gateway.main:app --workers \"${NUM_WORKERS:-1}\" \"$@\"", "--"]
+ENTRYPOINT ["sh", "-c", "exec /app/docker/component_entrypoint.sh python -m gateway.launch --workers \"${NUM_WORKERS:-1}\" \"$@\"", "--"]
 CMD ["--host", "0.0.0.0", "--port", "4000"]

gateway/launch.py

@@ -0,0 +1,71 @@
+"""Gateway supervisor: assemble DATABASE_URL, start the in-container PgBouncer, then run uvicorn.
+
+``gateway/main.py`` assembles ``DATABASE_URL`` inside every uvicorn worker, which
+is fine for a plain Postgres URL but not for the pooler: PgBouncer must be
+started exactly once per pod, before the workers fork, and the workers must be
+handed the loopback URL it listens on. A pre-existing ``DATABASE_URL`` wins in
+``DatabaseURLSettings.apply_to_env`` under password auth, so setting it here is
+enough for every worker to pick the pooled URL up unchanged.
+
+Run with:
+    python -m gateway.launch --workers 4 --host 0.0.0.0 --port 4000
+"""
+
+import os
+import sys
+from collections.abc import Callable, Mapping, Sequence
+from typing import Final
+
+from uvicorn.main import main as uvicorn_main
+
+from litellm.proxy.db.db_url_settings import DatabaseURLSettings
+from litellm.proxy.db.pgbouncer import PgBouncerError, PgBouncerSettings, start_in_container_pgbouncer
+
+GATEWAY_APP: Final = "gateway.main:app"
+KEEPALIVE_FLAG: Final = "--timeout-keep-alive"
+
+
+def uvicorn_argv(argv: Sequence[str], environ: Mapping[str, str]) -> tuple[str, ...]:
+    """Honor ``KEEPALIVE_TIMEOUT`` like ``proxy_cli.py`` does, unless the flag was passed explicitly."""
+    keepalive: Final = environ.get("KEEPALIVE_TIMEOUT")
+    if keepalive is None or any(arg == KEEPALIVE_FLAG or arg.startswith(f"{KEEPALIVE_FLAG}=") for arg in argv):
+        return (GATEWAY_APP, *argv)
+    return (GATEWAY_APP, *argv, KEEPALIVE_FLAG, keepalive)
+
+
+def pool_database_url(
+    settings: DatabaseURLSettings,
+    pgbouncer: PgBouncerSettings,
+    environ: Mapping[str, str],
+) -> str | PgBouncerError | None:
+    """Start the in-container PgBouncer and return its loopback URL, or None when ``pgbouncer.enabled`` is off.
+
+    The upstream URL is whatever ``apply_to_env`` assembled from the discrete
+    ``DATABASE_*`` vars (or an operator-pinned ``DATABASE_URL``). Token auth is
+    rejected by the pooler itself, since it holds one password for its lifetime.
+    """
+    if not pgbouncer.enabled:
+        return None
+    upstream_url: Final = environ.get("DATABASE_URL")
+    if upstream_url is None:
+        return PgBouncerError("LITELLM_PGBOUNCER_ENABLED is set but no DATABASE_URL could be assembled")
+    return start_in_container_pgbouncer(pgbouncer, upstream_url, token_auth_enabled=settings.token_auth() is not None)
+
+
+def _serve(argv: Sequence[str]) -> None:
+    uvicorn_main(tuple(argv), prog_name="uvicorn")
+
+
+def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) -> None:
+    settings: Final = DatabaseURLSettings.from_env()
+    settings.apply_to_env()
+    pooled_url: Final = pool_database_url(settings, PgBouncerSettings(), os.environ)
+    if isinstance(pooled_url, PgBouncerError):
+        sys.exit(f"LiteLLM gateway: in-container pgbouncer could not start: {pooled_url.reason}")
+    if pooled_url is not None:
+        os.environ["DATABASE_URL"] = pooled_url
+    serve(uvicorn_argv(argv, os.environ))
+
+
+if __name__ == "__main__":
+    main(sys.argv[1:])

helm/litellm-helm/templates/deployment.yaml

@@ -117,6 +117,14 @@ spec:
             - name: DATABASE_URL_READ_REPLICA
               value: {{ .Values.db.readReplicaUrl | quote }}
             {{- end }}
+            {{- if .Values.db.connectionPool.enabled }}
+            - name: LITELLM_PGBOUNCER_ENABLED
+              value: "true"
+            - name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
+              value: {{ .Values.db.connectionPool.maxDbConnections | quote }}
+            - name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
+              value: {{ .Values.db.connectionPool.maxClientConn | quote }}
+            {{- end }}
             - name: PROXY_MASTER_KEY
               valueFrom:
                 secretKeyRef:

helm/litellm-helm/templates/hpa.yaml

@@ -33,4 +33,22 @@ spec:
           type: Utilization
           averageUtilization: {{ .Values.autoscaling.targetMemoryUtilizationPercentage }}
     {{- end }}
+    {{- with .Values.autoscaling.targetRequestsPerSecond }}
+    - type: Pods
+      pods:
+        metric:
+          name: litellm_requests_per_second
+        target:
+          type: AverageValue
+          averageValue: {{ toJson . | trimAll "\"" | quote }}
+    {{- end }}
+    {{- with .Values.autoscaling.targetTokensPerSecond }}
+    - type: Pods
+      pods:
+        metric:
+          name: litellm_tokens_per_second
+        target:
+          type: AverageValue
+          averageValue: {{ toJson . | trimAll "\"" | quote }}
+    {{- end }}
 {{- end }}

helm/litellm-helm/templates/keda.yaml

@@ -23,6 +23,27 @@ spec:
   triggers:
 {{- with .Values.keda.triggers }}
   {{- toYaml . | nindent 2 }}
+{{- end }}
+{{- $prom := .Values.keda.prometheus }}
+{{- if or $prom.requestsPerSecond $prom.tokensPerSecond }}
+{{- if not $prom.serverAddress }}
+{{- fail "keda.prometheus.serverAddress is required when keda.prometheus.requestsPerSecond or tokensPerSecond is set" }}
+{{- end }}
+{{- $selector := printf "namespace=%q,job=%q" .Release.Namespace (printf "%s%s" (include "litellm.fullname" .) (ternary "-metrics" "" .Values.metricsServer.enabled)) }}
+{{- with $prom.requestsPerSecond }}
+  - type: prometheus
+    metadata:
+      serverAddress: {{ $prom.serverAddress | quote }}
+      threshold: {{ toJson . | trimAll "\"" | quote }}
+      query: {{ printf "sum(rate(litellm_proxy_total_requests_metric_total{%s}[1m]))" $selector | quote }}
+{{- end }}
+{{- with $prom.tokensPerSecond }}
+  - type: prometheus
+    metadata:
+      serverAddress: {{ $prom.serverAddress | quote }}
+      threshold: {{ toJson . | trimAll "\"" | quote }}
+      query: {{ printf "sum(rate(litellm_total_tokens_metric_total{%s}[1m]))" $selector | quote }}
+{{- end }}
 {{- end }}
   advanced:
     restoreToOriginalReplicaCount: {{ .Values.keda.restoreToOriginalReplicaCount }}

helm/litellm-helm/tests/connection_pool_tests.yaml

@@ -0,0 +1,61 @@
+suite: test in-container connection pool
+templates:
+  - deployment.yaml
+  - configmap-litellm.yaml
+tests:
+  - it: should not emit pgbouncer env vars by default
+    template: deployment.yaml
+    asserts:
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: LITELLM_PGBOUNCER_ENABLED
+            value: "true"
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
+            value: "20"
+
+  - it: should enable the pool with the default sizing when connectionPool.enabled is set
+    template: deployment.yaml
+    set:
+      db.connectionPool.enabled: true
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: LITELLM_PGBOUNCER_ENABLED
+            value: "true"
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
+            value: "20"
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
+            value: "1000"
+
+  - it: should pass custom sizing through as strings next to the worker count
+    template: deployment.yaml
+    set:
+      numWorkers: 4
+      db.connectionPool.enabled: true
+      db.connectionPool.maxDbConnections: 8
+      db.connectionPool.maxClientConn: 400
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
+            value: "8"
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
+            value: "400"
+      - contains:
+          path: spec.template.spec.containers[0].args
+          content: "4"

helm/litellm-helm/tests/hpa_tests.yaml

@@ -61,6 +61,84 @@ tests:
       - equal: { path: "spec.metrics[1].resource.name", value: memory }
       - equal: { path: "spec.metrics[1].resource.target.averageUtilization", value: 80 }
 
+  - it: "renders no workload metrics by default"
+    set:
+      autoscaling.enabled: true
+      autoscaling.targetMemoryUtilizationPercentage: 80
+    asserts:
+      - lengthEqual: { path: spec.metrics, count: 2 }
+      - notContains: { path: spec.metrics, content: { type: Pods }, any: true }
+
+  - it: "adds a requests-per-second Pods metric after the cpu metric"
+    set:
+      autoscaling.enabled: true
+      autoscaling.targetRequestsPerSecond: 90
+    asserts:
+      - lengthEqual: { path: spec.metrics, count: 2 }
+      - equal: { path: "spec.metrics[0].resource.name", value: cpu }
+      - equal:
+          path: "spec.metrics[1]"
+          value:
+            type: Pods
+            pods:
+              metric: { name: litellm_requests_per_second }
+              target: { type: AverageValue, averageValue: "90" }
+
+  - it: "adds a tokens-per-second Pods metric on its own"
+    set:
+      autoscaling.enabled: true
+      autoscaling.targetTokensPerSecond: 6M
+    asserts:
+      - lengthEqual: { path: spec.metrics, count: 2 }
+      - equal:
+          path: "spec.metrics[1]"
+          value:
+            type: Pods
+            pods:
+              metric: { name: litellm_tokens_per_second }
+              target: { type: AverageValue, averageValue: "6M" }
+      - notContains:
+          path: spec.metrics
+          content: { type: Pods, pods: { metric: { name: litellm_requests_per_second } } }
+          any: true
+
+  - it: "renders requests, tokens, cpu and memory metrics together"
+    set:
+      autoscaling.enabled: true
+      autoscaling.targetMemoryUtilizationPercentage: 80
+      autoscaling.targetRequestsPerSecond: 90
+      autoscaling.targetTokensPerSecond: 6000000
+    asserts:
+      - lengthEqual: { path: spec.metrics, count: 4 }
+      - equal: { path: "spec.metrics[0].resource.name", value: cpu }
+      - equal: { path: "spec.metrics[1].resource.name", value: memory }
+      - equal: { path: "spec.metrics[2].pods.metric.name", value: litellm_requests_per_second }
+      - equal: { path: "spec.metrics[2].pods.target.averageValue", value: "90" }
+      - equal: { path: "spec.metrics[3].pods.metric.name", value: litellm_tokens_per_second }
+      - equal: { path: "spec.metrics[3].pods.target.averageValue", value: "6000000" }
+
+  - it: "scales on workload metrics alone when the cpu target is cleared"
+    set:
+      autoscaling.enabled: true
+      autoscaling.targetCPUUtilizationPercentage: null
+      autoscaling.targetRequestsPerSecond: 90
+      autoscaling.targetTokensPerSecond: 6000000
+    asserts:
+      - lengthEqual: { path: spec.metrics, count: 2 }
+      - notContains: { path: spec.metrics, content: { type: Resource }, any: true }
+      - equal: { path: "spec.metrics[0].pods.metric.name", value: litellm_requests_per_second }
+      - equal: { path: "spec.metrics[1].pods.metric.name", value: litellm_tokens_per_second }
+      - notMatchRegexRaw: { pattern: per_minute }
+
+  - it: "ignores the per-minute keys, which the chart never shipped"
+    set:
+      autoscaling.enabled: true
+      autoscaling.targetRequestsPerMinute: 5400
+      autoscaling.targetTokensPerMinute: 360000000
+    asserts:
+      - lengthEqual: { path: spec.metrics, count: 1 }
+      - notContains: { path: spec.metrics, content: { type: Pods }, any: true }
+
   - it: "renders no hpa when autoscaling is disabled"
     asserts:
       - hasDocuments: { count: 0 }

helm/litellm-helm/tests/keda_tests.yaml

@@ -0,0 +1,106 @@
+suite: "keda"
+templates:
+  - keda.yaml
+release:
+  name: rel
+  namespace: llm
+tests:
+  - it: "renders no scaled object by default"
+    asserts:
+      - hasDocuments: { count: 0 }
+
+  - it: "passes user triggers through and adds no prometheus triggers by default"
+    set:
+      keda.enabled: true
+      keda.triggers:
+        - type: cpu
+          metricType: Utilization
+          metadata: { value: "60" }
+    asserts:
+      - isKind: { of: ScaledObject }
+      - equal:
+          path: spec.triggers
+          value:
+            - type: cpu
+              metricType: Utilization
+              metadata: { value: "60" }
+
+  - it: "scales on release-wide requests per second divided by the per-replica target"
+    set:
+      keda.enabled: true
+      keda.prometheus.serverAddress: http://prometheus-operated.monitoring.svc:9090
+      keda.prometheus.requestsPerSecond: 90
+    asserts:
+      - lengthEqual: { path: spec.triggers, count: 1 }
+      - equal:
+          path: "spec.triggers[0]"
+          value:
+            type: prometheus
+            metadata:
+              serverAddress: http://prometheus-operated.monitoring.svc:9090
+              threshold: "90"
+              query: sum(rate(litellm_proxy_total_requests_metric_total{namespace="llm",job="rel-litellm"}[1m]))
+
+  - it: "scales on tokens per second on its own"
+    set:
+      keda.enabled: true
+      keda.prometheus.serverAddress: http://prom:9090
+      keda.prometheus.tokensPerSecond: 6000000
+    asserts:
+      - lengthEqual: { path: spec.triggers, count: 1 }
+      - equal: { path: "spec.triggers[0].type", value: prometheus }
+      - equal: { path: "spec.triggers[0].metadata.threshold", value: "6000000" }
+      - equal:
+          path: "spec.triggers[0].metadata.query"
+          value: sum(rate(litellm_total_tokens_metric_total{namespace="llm",job="rel-litellm"}[1m]))
+
+  - it: "appends requests and tokens triggers after user triggers and selects the metrics service job"
+    set:
+      keda.enabled: true
+      metricsServer.enabled: true
+      keda.triggers:
+        - type: cpu
+          metricType: Utilization
+          metadata: { value: "60" }
+      keda.prometheus.serverAddress: http://prom:9090
+      keda.prometheus.requestsPerSecond: 90
+      keda.prometheus.tokensPerSecond: 6000000
+    asserts:
+      - lengthEqual: { path: spec.triggers, count: 3 }
+      - equal: { path: "spec.triggers[0].type", value: cpu }
+      - equal: { path: "spec.triggers[1].metadata.threshold", value: "90" }
+      - equal:
+          path: "spec.triggers[1].metadata.query"
+          value: sum(rate(litellm_proxy_total_requests_metric_total{namespace="llm",job="rel-litellm-metrics"}[1m]))
+      - equal: { path: "spec.triggers[2].metadata.threshold", value: "6000000" }
+      - equal:
+          path: "spec.triggers[2].metadata.query"
+          value: sum(rate(litellm_total_tokens_metric_total{namespace="llm",job="rel-litellm-metrics"}[1m]))
+      - notMatchRegexRaw: { pattern: "\\* *60|per_minute|PerMinute" }
+
+  - it: "ignores the per-minute keys, which the chart never shipped"
+    set:
+      keda.enabled: true
+      keda.prometheus.serverAddress: http://prom:9090
+      keda.prometheus.requestsPerMinute: 5400
+      keda.prometheus.tokensPerMinute: 360000000
+    asserts:
+      - isKind: { of: ScaledObject }
+      - isNullOrEmpty: { path: spec.triggers }
+
+  - it: "refuses a workload target without a prometheus server address"
+    set:
+      keda.enabled: true
+      keda.prometheus.requestsPerSecond: 90
+    asserts:
+      - failedTemplate:
+          errorMessage: keda.prometheus.serverAddress is required when keda.prometheus.requestsPerSecond or tokensPerSecond is set
+
+  - it: "yields to the hpa when both autoscalers are enabled"
+    set:
+      autoscaling.enabled: true
+      keda.enabled: true
+      keda.prometheus.serverAddress: http://prom:9090
+      keda.prometheus.requestsPerSecond: 90
+    asserts:
+      - hasDocuments: { count: 0 }

helm/litellm-helm/values.yaml

@@ -222,6 +222,25 @@ autoscaling:
   # Memory is a floor to provision under 'resources', not a signal to scale on.
   # targetMemoryUtilizationPercentage: 80
   # behavior: {}
+  # Opt-in per-pod workload targets, rendered as autoscaling/v2 `Pods` metrics
+  # named `litellm_requests_per_second` and `litellm_tokens_per_second` with an
+  # AverageValue target, alongside whichever resource targets are set (the HPA
+  # follows the metric asking for the most replicas). A Prometheus Adapter must
+  # serve those two names on custom.metrics.k8s.io from the proxy's counters,
+  # grouped by the scrape target's `pod` label (enable serviceMonitor below so
+  # every pod is scraped on its own):
+  #   litellm_requests_per_second:
+  #     sum(rate(litellm_proxy_total_requests_metric_total{<<.LabelMatchers>>}[1m])) by (<<.GroupBy>>)
+  #   litellm_tokens_per_second:
+  #     sum(rate(litellm_total_tokens_metric_total{<<.LabelMatchers>>}[1m])) by (<<.GroupBy>>)
+  # rate() over [1m] is already per second, so no `* 60`. How fast the HPA
+  # reacts is set by that window, the scrape interval and the HPA sync period
+  # (15s by default), not by the unit: keep serviceMonitor.interval at 15s or
+  # faster so a 1m window holds at least 4 samples. averageValue takes SI
+  # suffixes, so "6M" is six million tokens per second per pod. Tokens are
+  # counted when a response completes, so TPS trails long streams.
+  targetRequestsPerSecond: ""
+  targetTokensPerSecond: ""
 
 # Autoscaling with keda is mutually exclusive with hpa
 keda:
@@ -243,6 +262,23 @@ keda:
   #     metricName: http_requests_total
   #     threshold: '100'
   #     query: sum(rate(http_requests_total{deployment="my-deployment"}[2m]))
+  # First-class Prometheus triggers on the proxy's own request and token
+  # counters, appended to `triggers`. Each target is the per-second load one
+  # replica should carry: KEDA divides the release-wide
+  # `sum(rate(<counter>[1m]))` by it to pick the replica count. Thresholds
+  # are plain numbers (KEDA parses them as floats, no SI suffixes). The
+  # queries select samples by the release namespace and the `job` label the
+  # chart's ServiceMonitor produces (the metrics Service name), so enable
+  # serviceMonitor below together with metricsServer: the http port serves
+  # /metrics/ behind virtual-key auth and answers an unauthenticated scrape
+  # with 401. Reaction time comes from the [1m] window, the scrape interval
+  # and pollingInterval above, so keep both at 15s or faster. Tokens are
+  # counted at completion, so TPS trails long streams. serverAddress is
+  # required once either target is set.
+  prometheus:
+    serverAddress: ""
+    requestsPerSecond: ""
+    tokensPerSecond: ""
   behavior: {}
   # scaleDown:
   #   stabilizationWindowSeconds: 300
@@ -319,6 +355,20 @@ db:
   # only (e.g. when IAM_TOKEN_DB_AUTH supplies the token at runtime).
   readReplicaUrl: ""
 
+  # In-container connection pool (PgBouncer, transaction mode) shared by every
+  # worker in the pod. Without it each --num_workers worker opens its own
+  # connection_limit connections to Postgres, so a pod's footprint against the
+  # database's connection ceiling is workers x connection_limit and grows with
+  # every replica. With it, the pod holds at most maxDbConnections upstream
+  # connections no matter how many workers run; the workers connect to the pool
+  # over loopback, with no extra network hop. Migrations still go straight to
+  # Postgres. Starting profile for numWorkers: 4 is maxDbConnections: 20, so
+  # a database with a 5000-connection ceiling fits roughly 200 replicas.
+  connectionPool:
+    enabled: false
+    maxDbConnections: 20
+    maxClientConn: 1000
+
   # Use the Stackgres Helm chart to deploy an instance of a Stackgres cluster.
   #  The Stackgres Operator must already be installed within the target
   #  Kubernetes cluster.

helm/litellm/templates/_helpers.tpl

@@ -360,6 +360,23 @@ harmless no-op for the Job and authoritative for the app pods.
 {{- end }}
 {{- end -}}
 
+{{/*
+In-container PgBouncer env for the gateway container. Fails at render time under IAM or Entra auth: the pooler holds one static password for the life of the pod.
+*/}}
+{{- define "litellm.connectionPoolEnv" -}}
+{{- if or .Values.database.writer.useIAMAuth .Values.database.writer.useAzureEntraAuth }}
+{{- fail "database.connectionPool.enabled cannot be combined with database.writer.useIAMAuth or database.writer.useAzureEntraAuth: the in-container pgbouncer holds a static database password and cannot follow a rotating token. Disable the pool or use a static database password" }}
+{{- end }}
+{{- with .Values.database.connectionPool -}}
+- name: LITELLM_PGBOUNCER_ENABLED
+  value: "true"
+- name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
+  value: {{ required "database.connectionPool.maxDbConnections is required when the pool is enabled" .maxDbConnections | quote }}
+- name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
+  value: {{ required "database.connectionPool.maxClientConn is required when the pool is enabled" .maxClientConn | quote }}
+{{- end }}
+{{- end -}}
+
 {{/*
 PodDisruptionBudget shared by gateway, backend, and ui.
 

helm/litellm/templates/gateway/deployment.yaml

@@ -61,6 +61,9 @@ spec:
             - name: NUM_WORKERS
               value: {{ .Values.gateway.numWorkers | quote }}
             {{- end }}
+            {{- if .Values.database.connectionPool.enabled }}
+            {{- include "litellm.connectionPoolEnv" $ | nindent 12 }}
+            {{- end }}
             {{- if .Values.billingMetrics.enabled }}
             {{- include "litellm.billingMetricsEnv" . | nindent 12 }}
             {{- end }}

helm/litellm/templates/gateway/hpa.yaml

@@ -30,6 +30,24 @@ spec:
           type: Utilization
           averageUtilization: {{ .Values.gateway.hpa.targetMemoryUtilizationPercentage }}
     {{- end }}
+    {{- with .Values.gateway.hpa.targetRequestsPerSecond }}
+    - type: Pods
+      pods:
+        metric:
+          name: litellm_requests_per_second
+        target:
+          type: AverageValue
+          averageValue: {{ toJson . | trimAll "\"" | quote }}
+    {{- end }}
+    {{- with .Values.gateway.hpa.targetTokensPerSecond }}
+    - type: Pods
+      pods:
+        metric:
+          name: litellm_tokens_per_second
+        target:
+          type: AverageValue
+          averageValue: {{ toJson . | trimAll "\"" | quote }}
+    {{- end }}
   {{- with .Values.gateway.hpa.behavior }}
   behavior:
     {{- toYaml . | nindent 4 }}

helm/litellm/templates/gateway/servicemonitor.yaml

@@ -0,0 +1,28 @@
+{{- if and .Values.gateway.enabled .Values.gateway.serviceMonitor.enabled }}
+{{- if not .Values.gateway.metricsServer.enabled }}
+{{- fail "gateway.serviceMonitor.enabled requires gateway.metricsServer.enabled: the http port serves /metrics/ behind virtual-key auth, so an unauthenticated scrape gets 401" }}
+{{- end }}
+apiVersion: monitoring.coreos.com/v1
+kind: ServiceMonitor
+metadata:
+  name: {{ include "litellm.gateway.fullname" . }}
+  labels:
+    {{- include "litellm.commonLabels" . | nindent 4 }}
+    app.kubernetes.io/component: gateway
+    {{- with .Values.gateway.serviceMonitor.labels }}
+    {{- toYaml . | nindent 4 }}
+    {{- end }}
+spec:
+  selector:
+    matchLabels:
+      {{- include "litellm.gateway.selectorLabels" . | nindent 6 }}
+  namespaceSelector:
+    matchNames:
+      - {{ .Release.Namespace | quote }}
+  endpoints:
+    - port: metrics
+      path: /metrics/
+      interval: {{ .Values.gateway.serviceMonitor.interval }}
+      scrapeTimeout: {{ .Values.gateway.serviceMonitor.scrapeTimeout }}
+      scheme: http
+{{- end }}

helm/litellm/tests/connection_pool_tests.yaml

@@ -0,0 +1,117 @@
+suite: test in-container connection pool env vars
+templates:
+  - gateway/deployment.yaml
+  - gateway/configmap.yaml
+  - backend/deployment.yaml
+  - backend/configmap.yaml
+values:
+  - ./values/required.yaml
+tests:
+  - it: renders no pool env by default
+    template: gateway/deployment.yaml
+    asserts:
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: LITELLM_PGBOUNCER_ENABLED
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
+          any: true
+
+  - it: enabled pool renders the three pgbouncer vars with the configured sizes
+    template: gateway/deployment.yaml
+    set:
+      gateway.numWorkers: 4
+      database.connectionPool.enabled: true
+      database.connectionPool.maxDbConnections: 8
+      database.connectionPool.maxClientConn: 250
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: LITELLM_PGBOUNCER_ENABLED
+            value: "true"
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
+            value: "8"
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
+            value: "250"
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: NUM_WORKERS
+            value: "4"
+
+  - it: enabled pool uses the chart default sizes
+    template: gateway/deployment.yaml
+    set:
+      database.connectionPool.enabled: true
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
+            value: "20"
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
+            value: "1000"
+
+  - it: backend never gets the pool env
+    template: backend/deployment.yaml
+    set:
+      database.connectionPool.enabled: true
+    asserts:
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: LITELLM_PGBOUNCER_ENABLED
+          any: true
+
+  - it: pool with IAM auth fails at render time
+    template: gateway/deployment.yaml
+    set:
+      database.connectionPool.enabled: true
+      database.writer.useIAMAuth: true
+    asserts:
+      - failedTemplate:
+          errorMessage: "database.connectionPool.enabled cannot be combined with database.writer.useIAMAuth or database.writer.useAzureEntraAuth: the in-container pgbouncer holds a static database password and cannot follow a rotating token. Disable the pool or use a static database password"
+
+  - it: pool with Entra auth fails at render time
+    template: gateway/deployment.yaml
+    set:
+      database.connectionPool.enabled: true
+      database.writer.useAzureEntraAuth: true
+    asserts:
+      - failedTemplate:
+          errorMessage: "database.connectionPool.enabled cannot be combined with database.writer.useIAMAuth or database.writer.useAzureEntraAuth: the in-container pgbouncer holds a static database password and cannot follow a rotating token. Disable the pool or use a static database password"
+
+  - it: IAM auth without the pool still renders
+    template: gateway/deployment.yaml
+    set:
+      database.writer.useIAMAuth: true
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: IAM_TOKEN_DB_AUTH
+            value: "true"
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: LITELLM_PGBOUNCER_ENABLED
+          any: true

helm/litellm/tests/hpa_workload_metrics_tests.yaml

@@ -0,0 +1,213 @@
+suite: test gateway HPA per-pod requests-per-second and tokens-per-second targets
+templates:
+  - gateway/hpa.yaml
+  - gateway/servicemonitor.yaml
+values:
+  - ./values/required.yaml
+tests:
+  - it: scales on CPU and memory only by default
+    template: gateway/hpa.yaml
+    asserts:
+      - equal:
+          path: spec.metrics
+          value:
+            - type: Resource
+              resource:
+                name: cpu
+                target:
+                  type: Utilization
+                  averageUtilization: 70
+            - type: Resource
+              resource:
+                name: memory
+                target:
+                  type: Utilization
+                  averageUtilization: 80
+
+  - it: adds a requests-per-second Pods metric next to the resource metrics
+    template: gateway/hpa.yaml
+    set:
+      gateway.hpa.targetRequestsPerSecond: 90
+    asserts:
+      - lengthEqual:
+          path: spec.metrics
+          count: 3
+      - contains:
+          path: spec.metrics
+          content:
+            type: Resource
+            resource:
+              name: cpu
+              target:
+                type: Utilization
+                averageUtilization: 70
+      - equal:
+          path: spec.metrics[2]
+          value:
+            type: Pods
+            pods:
+              metric:
+                name: litellm_requests_per_second
+              target:
+                type: AverageValue
+                averageValue: "90"
+      - notContains:
+          path: spec.metrics
+          content:
+            type: Pods
+            pods:
+              metric:
+                name: litellm_tokens_per_second
+          any: true
+
+  - it: adds a tokens-per-second Pods metric on its own
+    template: gateway/hpa.yaml
+    set:
+      gateway.hpa.targetTokensPerSecond: 6M
+    asserts:
+      - lengthEqual:
+          path: spec.metrics
+          count: 3
+      - equal:
+          path: spec.metrics[2]
+          value:
+            type: Pods
+            pods:
+              metric:
+                name: litellm_tokens_per_second
+              target:
+                type: AverageValue
+                averageValue: "6M"
+      - notContains:
+          path: spec.metrics
+          content:
+            type: Pods
+            pods:
+              metric:
+                name: litellm_requests_per_second
+          any: true
+
+  - it: renders requests and tokens targets together and keeps CPU and memory
+    template: gateway/hpa.yaml
+    set:
+      gateway.hpa.targetRequestsPerSecond: 90
+      gateway.hpa.targetTokensPerSecond: 6000000
+    asserts:
+      - lengthEqual:
+          path: spec.metrics
+          count: 4
+      - equal:
+          path: spec.metrics[0].resource.name
+          value: cpu
+      - equal:
+          path: spec.metrics[1].resource.name
+          value: memory
+      - equal:
+          path: spec.metrics[2].pods.metric.name
+          value: litellm_requests_per_second
+      - equal:
+          path: spec.metrics[3].pods.metric.name
+          value: litellm_tokens_per_second
+      - equal:
+          path: spec.metrics[3].pods.target.averageValue
+          value: "6000000"
+
+  - it: scales on workload metrics alone when the resource targets are cleared
+    template: gateway/hpa.yaml
+    set:
+      gateway.hpa.targetCPUUtilizationPercentage: null
+      gateway.hpa.targetMemoryUtilizationPercentage: null
+      gateway.hpa.targetRequestsPerSecond: 90
+      gateway.hpa.targetTokensPerSecond: 6000000
+    asserts:
+      - lengthEqual:
+          path: spec.metrics
+          count: 2
+      - notContains:
+          path: spec.metrics
+          content:
+            type: Resource
+          any: true
+      - equal:
+          path: spec.metrics[0].pods.metric.name
+          value: litellm_requests_per_second
+      - equal:
+          path: spec.metrics[1].pods.metric.name
+          value: litellm_tokens_per_second
+      - notMatchRegexRaw:
+          pattern: per_minute
+
+  - it: ignores the per-minute keys, which the chart never shipped
+    template: gateway/hpa.yaml
+    set:
+      gateway.hpa.targetRequestsPerMinute: 5400
+      gateway.hpa.targetTokensPerMinute: 360000000
+    asserts:
+      - lengthEqual:
+          path: spec.metrics
+          count: 2
+      - notContains:
+          path: spec.metrics
+          content:
+            type: Pods
+          any: true
+
+  - it: renders no ServiceMonitor by default
+    template: gateway/servicemonitor.yaml
+    asserts:
+      - hasDocuments:
+          count: 0
+
+  - it: refuses a ServiceMonitor without the metrics server, whose http port needs a bearer token
+    template: gateway/servicemonitor.yaml
+    set:
+      gateway.serviceMonitor.enabled: true
+    asserts:
+      - failedTemplate:
+          errorPattern: gateway.serviceMonitor.enabled requires gateway.metricsServer.enabled
+
+  - it: scrapes each gateway pod through the metrics port
+    template: gateway/servicemonitor.yaml
+    release:
+      name: rel
+      namespace: llm
+    set:
+      gateway.serviceMonitor.enabled: true
+      gateway.metricsServer.enabled: true
+      gateway.serviceMonitor.labels:
+        release: kube-prometheus-stack
+    asserts:
+      - isKind:
+          of: ServiceMonitor
+      - equal:
+          path: metadata.labels.release
+          value: kube-prometheus-stack
+      - equal:
+          path: spec.selector.matchLabels
+          value:
+            app.kubernetes.io/name: litellm
+            app.kubernetes.io/instance: rel
+            app.kubernetes.io/component: gateway
+      - equal:
+          path: spec.namespaceSelector.matchNames
+          value:
+            - llm
+      - equal:
+          path: spec.endpoints
+          value:
+            - port: metrics
+              path: /metrics/
+              interval: 15s
+              scrapeTimeout: 10s
+              scheme: http
+
+  - it: honours a custom scrape interval
+    template: gateway/servicemonitor.yaml
+    set:
+      gateway.serviceMonitor.enabled: true
+      gateway.serviceMonitor.interval: 30s
+      gateway.metricsServer.enabled: true
+    asserts:
+      - equal:
+          path: spec.endpoints[0].interval
+          value: 30s

helm/litellm/values.yaml

@@ -225,6 +225,26 @@ database:
       usernameKey: username
       passwordKey: password
 
+  # In-container connection pool (PgBouncer, transaction mode) shared by every
+  # gateway worker in the pod. Without it each of the `gateway.numWorkers`
+  # workers opens its own Prisma pool straight to Postgres, so a pod's
+  # footprint against the database's connection ceiling is
+  # numWorkers x connection_limit and grows with every replica. With it, the
+  # pod holds at most maxDbConnections upstream connections no matter how many
+  # workers run; the workers connect to the pool over loopback, with no extra
+  # network hop. The chart emits LITELLM_PGBOUNCER_ENABLED /
+  # LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS / LITELLM_PGBOUNCER_MAX_CLIENT_CONN on
+  # the gateway container only: the backend runs a single worker and the
+  # migrations Job must keep a direct connection. The pool holds a static
+  # password, so it cannot be combined with `database.writer.useIAMAuth` or
+  # `useAzureEntraAuth` (rendering fails). Starting profile for
+  # `gateway.numWorkers: 4` is maxDbConnections: 20, so a database with a
+  # 5000-connection ceiling fits roughly 200 gateway replicas.
+  connectionPool:
+    enabled: false
+    maxDbConnections: 20
+    maxClientConn: 1000
+
 # Optional Redis. Leave host empty to disable.
 #
 # This is the proxy's coordination store: cross-pod tpm/rpm rate limits, spend
@@ -284,6 +304,16 @@ gateway:
         memory: 128Mi
       limits:
         memory: 512Mi
+  # Prometheus Operator ServiceMonitor for the gateway pods. Scrapes the
+  # `<gateway>-metrics` Service, so it requires metricsServer above (the http
+  # port serves /metrics/ behind virtual-key auth). Every pod is its own scrape
+  # target, so the samples carry the `pod` label the per-pod autoscaling
+  # queries below group by.
+  serviceMonitor:
+    enabled: false
+    labels: {}
+    interval: 15s
+    scrapeTimeout: 10s
   image:
     repository: ghcr.io/berriai/litellm-gateway
     tag: ""              # defaults to .Chart.AppVersion
@@ -340,6 +370,25 @@ gateway:
     #     policies:
     #       - { type: Percent, value: 100, periodSeconds: 30 }
     behavior: {}
+    # Opt-in per-pod workload targets, rendered as autoscaling/v2 `Pods` metrics
+    # named `litellm_requests_per_second` and `litellm_tokens_per_second` with an
+    # AverageValue target. They coexist with the CPU/memory targets above: the
+    # HPA scales on whichever metric asks for the most replicas. Kubernetes has
+    # no idea what a token is, so a Prometheus Adapter must serve those two
+    # names on custom.metrics.k8s.io from the proxy's counters, grouped by the
+    # scrape target's `pod` label (enable serviceMonitor above):
+    #   litellm_requests_per_second:
+    #     sum(rate(litellm_proxy_total_requests_metric_total{<<.LabelMatchers>>}[1m])) by (<<.GroupBy>>)
+    #   litellm_tokens_per_second:
+    #     sum(rate(litellm_total_tokens_metric_total{<<.LabelMatchers>>}[1m])) by (<<.GroupBy>>)
+    # rate() over [1m] is already per second, so no `* 60`. How fast the HPA
+    # reacts is set by that window, the scrape interval and the HPA sync period
+    # (15s by default), not by the unit: keep serviceMonitor.interval at 15s or
+    # faster so a 1m window holds at least 4 samples. averageValue takes SI
+    # suffixes, so "6M" is six million tokens per second per pod. Tokens are
+    # counted when a response completes, so TPS trails long streams.
+    targetRequestsPerSecond: ""
+    targetTokensPerSecond: ""
   # PodDisruptionBudget for the gateway pods. Set exactly one of
   # `minAvailable` / `maxUnavailable` (minAvailable wins if both are set;
   # enabling without either falls back to `maxUnavailable: 1`). Disabled by

litellm-proxy-extras/litellm_proxy_extras/migrations/20260910000000_skills_on_object_permission/migration.sql

@@ -0,0 +1 @@
+ALTER TABLE "LiteLLM_ObjectPermissionTable" ADD COLUMN IF NOT EXISTS "skills" TEXT[] DEFAULT ARRAY[]::TEXT[];

litellm-proxy-extras/litellm_proxy_extras/schema.prisma

@@ -282,6 +282,7 @@ model LiteLLM_ObjectPermissionTable {
   mcp_toolsets          String[]       @default([]) // Toolset IDs granted to this key/team/user
   search_tools          String[]       @default([]) // search_tool_name values this key/team/user may call
   mcp_tool_search_enabled Boolean?
+  skills                String[]       @default([]) // Claude Code plugin names granted to this key/team beyond the public (enabled) set
   teams                 LiteLLM_TeamTable[]
   projects              LiteLLM_ProjectTable[]
   verification_tokens   LiteLLM_VerificationToken[]

litellm-rust/AGENTS.md

@@ -1,18 +1,19 @@
 # AGENTS.md
 
-litellm-rust has five crates. A crate is a layer or shared foundation, not a route. Routes (ocr, realtime, chat) and providers (mistral, openai) are modules inside the layers.
+litellm-rust has six crates. A crate is a layer or shared foundation, not a route. Routes (ocr, realtime, chat) and providers (mistral, openai) are modules inside the layers.
 
 ## Crates
 
 | Crate | Role |
 |-------|------|
 | litellm-core | The LiteLLM SDK in Rust. One public entrypoint per top-level call (`messages::messages()`), owning types, transforms, provider resolution, auth, and the provider HTTP call. Call it, get a typed response. |
+| litellm-token-counter | Standalone input token counting shared by host integrations without pulling in the full SDK. |
 | litellm-config | Config-loading boundary. Returns resolved core deployment data and optionally delegates loading to Python. |
 | litellm-ai-gateway | The axum server (behind the `server` feature) plus the WebSocket hosts. Translates HTTP/WS to core entrypoints; owns no provider logic and no handlers. |
 | litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. |
 | litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. Owns API registration, domain wiring, and Python exception mapping. |
 
-Dependency direction is acyclic: `litellm-config` depends on `litellm-core`, the gateway depends on both, and `litellm-python-bridge` depends on the domain layers and `litellm-python-interop`. The interop foundation depends on no LiteLLM domain crate.
+Dependency direction is acyclic: `litellm-config` depends on `litellm-core`, the gateway depends on both, and `litellm-python-bridge` depends on the domain layers, `litellm-token-counter`, and `litellm-python-interop`. The token counter and interop foundations depend on no LiteLLM domain crate.
 
 ## Where a route lives
 

litellm-rust/Cargo.lock

@@ -2,6 +2,20 @@
 # It is not intended for manual editing.
 version = 4
 
+[[package]]
+name = "ahash"
+version = "0.8.12"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75"
+dependencies = [
+ "cfg-if",
+ "getrandom 0.3.4",
+ "once_cell",
+ "serde",
+ "version_check",
+ "zerocopy",
+]
+
 [[package]]
 name = "aho-corasick"
 version = "1.1.5"
@@ -418,7 +432,7 @@ checksum = "edca88bc138befd0323b20752846e6587272d3b03b0343c8ea28a6f819e6e71f"
 dependencies = [
  "async-trait",
  "axum-core",
- "base64",
+ "base64 0.22.1",
  "bytes",
  "futures-util",
  "http 1.4.2",
@@ -468,6 +482,12 @@ dependencies = [
  "tracing",
 ]
 
+[[package]]
+name = "base64"
+version = "0.13.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8"
+
 [[package]]
 name = "base64"
 version = "0.22.1"
@@ -542,6 +562,15 @@ version = "0.3.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5"
 
+[[package]]
+name = "castaway"
+version = "0.2.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "dec551ab6e7578819132c713a93c022a05d60159dc86e7a7050223577484c55a"
+dependencies = [
+ "rustversion",
+]
+
 [[package]]
 name = "cc"
 version = "1.3.0"
@@ -644,6 +673,21 @@ version = "0.5.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "0c9ea0ac24bc397ab3c98583a3c9ba74fa56b09a4449bbe172b9b1ddb016027a"
 
+[[package]]
+name = "compact_str"
+version = "0.9.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9dfdd1c2274d9aa354115b09dc9a901d6c5576818cdf70d14cae2bdb47df00ab"
+dependencies = [
+ "castaway",
+ "cfg-if",
+ "itoa",
+ "rustversion",
+ "ryu",
+ "serde",
+ "static_assertions",
+]
+
 [[package]]
 name = "const-oid"
 version = "0.10.2"
@@ -696,7 +740,7 @@ dependencies = [
  "ciborium",
  "clap",
  "criterion-plot",
- "itertools",
+ "itertools 0.13.0",
  "num-traits",
  "oorandom",
  "page_size",
@@ -716,7 +760,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d8d80a2f4f5b554395e47b5d8305bc3d27813bacb73493eb1001e8f76dae29ea"
 dependencies = [
  "cast",
- "itertools",
+ "itertools 0.13.0",
 ]
 
 [[package]]
@@ -778,6 +822,56 @@ dependencies = [
  "cmov",
 ]
 
+[[package]]
+name = "daachorse"
+version = "3.0.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5614204febbc33cc07a2806aa6440b904ac012b68eecc37f4493ea4a76455a3d"
+
+[[package]]
+name = "darling"
+version = "0.20.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fc7f46116c46ff9ab3eb1597a45688b6715c6e628b5c133e288e709a29bcb4ee"
+dependencies = [
+ "darling_core",
+ "darling_macro",
+]
+
+[[package]]
+name = "darling_core"
+version = "0.20.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0d00b9596d185e565c2207a0b01f8bd1a135483d02d9b7b0a54b11da8d53412e"
+dependencies = [
+ "fnv",
+ "ident_case",
+ "proc-macro2",
+ "quote",
+ "strsim",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "darling_macro"
+version = "0.20.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead"
+dependencies = [
+ "darling_core",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "dary_heap"
+version = "0.3.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8b1e3a325bc115f096c8b77bbf027a7c2592230e70be2d985be950d3d5e60ebe"
+dependencies = [
+ "serde",
+]
+
 [[package]]
 name = "data-encoding"
 version = "2.11.0"
@@ -790,6 +884,37 @@ version = "0.5.8"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "7cd812cc2bc1d69d4764bd80df88b4317eaef9e773c75226407d9bc0876b211c"
 
+[[package]]
+name = "derive_builder"
+version = "0.20.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "507dfb09ea8b7fa618fcf76e953f4f5e192547945816d5358edffe39f6f94947"
+dependencies = [
+ "derive_builder_macro",
+]
+
+[[package]]
+name = "derive_builder_core"
+version = "0.20.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2d5bcf7b024d6835cfb3d473887cd966994907effbe9227e8c8219824d06c4e8"
+dependencies = [
+ "darling",
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "derive_builder_macro"
+version = "0.20.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ab63b0e2bf4d5928aff72e83a7dace85d7bba5fe12dcc3c5a572d78caffd3f3c"
+dependencies = [
+ "derive_builder_core",
+ "syn 2.0.119",
+]
+
 [[package]]
 name = "digest"
 version = "0.10.7"
@@ -841,6 +966,12 @@ version = "1.0.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f"
 
+[[package]]
+name = "esaxx-rs"
+version = "0.1.10"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d817e038c30374a4bcb22f94d0a8a0e216958d4c3dcde369b1439fec4bdda6e6"
+
 [[package]]
 name = "fastrand"
 version = "2.5.0"
@@ -964,6 +1095,18 @@ dependencies = [
  "wasm-bindgen",
 ]
 
+[[package]]
+name = "getrandom"
+version = "0.3.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd"
+dependencies = [
+ "cfg-if",
+ "libc",
+ "r-efi 5.3.0",
+ "wasip2",
+]
+
 [[package]]
 name = "getrandom"
 version = "0.4.3"
@@ -973,7 +1116,7 @@ dependencies = [
  "cfg-if",
  "js-sys",
  "libc",
- "r-efi",
+ "r-efi 6.0.0",
  "rand_core 0.10.1",
  "wasm-bindgen",
 ]
@@ -1220,7 +1363,7 @@ version = "0.1.20"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0"
 dependencies = [
- "base64",
+ "base64 0.22.1",
  "bytes",
  "futures-channel",
  "futures-util",
@@ -1319,6 +1462,12 @@ dependencies = [
  "zerovec",
 ]
 
+[[package]]
+name = "ident_case"
+version = "1.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39"
+
 [[package]]
 name = "idna"
 version = "1.1.0"
@@ -1348,6 +1497,8 @@ checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9"
 dependencies = [
  "equivalent",
  "hashbrown",
+ "serde",
+ "serde_core",
 ]
 
 [[package]]
@@ -1365,6 +1516,15 @@ dependencies = [
  "either",
 ]
 
+[[package]]
+name = "itertools"
+version = "0.14.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285"
+dependencies = [
+ "either",
+]
+
 [[package]]
 name = "itoa"
 version = "1.0.18"
@@ -1409,7 +1569,7 @@ name = "litellm-ai-gateway"
 version = "0.1.0"
 dependencies = [
  "axum",
- "base64",
+ "base64 0.22.1",
  "futures-channel",
  "futures-util",
  "litellm-config",
@@ -1447,17 +1607,19 @@ dependencies = [
  "aws-sigv4",
  "aws-smithy-runtime-api",
  "aws-types",
- "base64",
+ "base64 0.22.1",
  "rand 0.8.7",
  "reqwest",
  "rstest",
  "serde",
  "serde_json",
+ "serde_path_to_error",
  "sha2 0.10.9",
  "thiserror 2.0.19",
  "tokio",
  "tracing",
  "tracing-subscriber",
+ "url",
 ]
 
 [[package]]
@@ -1469,6 +1631,7 @@ dependencies = [
  "litellm-ai-gateway",
  "litellm-core",
  "litellm-python-interop",
+ "litellm-token-counter",
  "pyo3",
  "pyo3-async-runtimes",
  "serde",
@@ -1489,6 +1652,22 @@ dependencies = [
  "serde_json",
 ]
 
+[[package]]
+name = "litellm-token-counter"
+version = "0.1.0"
+dependencies = [
+ "criterion",
+ "indexmap",
+ "itoa",
+ "rand 0.8.7",
+ "rstest",
+ "serde",
+ "serde_json",
+ "thiserror 2.0.19",
+ "tokenizers",
+ "unicode-normalization-alignments",
+]
+
 [[package]]
 name = "litemap"
 version = "0.8.2"
@@ -1507,6 +1686,22 @@ version = "0.1.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154"
 
+[[package]]
+name = "macro_rules_attribute"
+version = "0.2.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b3ae8f6d608c795738406608304d30a2dfbdc8e58e44f7ba43236da5208ded3c"
+dependencies = [
+ "macro_rules_attribute-proc_macro",
+ "pastey",
+]
+
+[[package]]
+name = "macro_rules_attribute-proc_macro"
+version = "0.2.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fc04a4c58212d57930a24bf47d3fa87485264a3a054e9c10e042eb373573ad3c"
+
 [[package]]
 name = "matchit"
 version = "0.7.3"
@@ -1535,6 +1730,12 @@ dependencies = [
  "unicase",
 ]
 
+[[package]]
+name = "minimal-lexical"
+version = "0.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a"
+
 [[package]]
 name = "mio"
 version = "1.2.2"
@@ -1546,6 +1747,38 @@ dependencies = [
  "windows-sys 0.61.2",
 ]
 
+[[package]]
+name = "monostate"
+version = "0.1.18"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3341a273f6c9d5bef1908f17b7267bbab0e95c9bf69a0d4dcf8e9e1b2c76ef67"
+dependencies = [
+ "monostate-impl",
+ "serde",
+ "serde_core",
+]
+
+[[package]]
+name = "monostate-impl"
+version = "0.1.18"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e4db6d5580af57bf992f59068d4ea26fd518574ff48d7639b255a36f9de6e7e9"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "nom"
+version = "7.1.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a"
+dependencies = [
+ "memchr",
+ "minimal-lexical",
+]
+
 [[package]]
 name = "num-conv"
 version = "0.2.2"
@@ -1576,6 +1809,28 @@ version = "1.21.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50"
 
+[[package]]
+name = "onig"
+version = "6.5.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0cc3cbf698f9438986c11a880c90a6d04b9de27575afd28bbf45b154b6c709e2"
+dependencies = [
+ "bitflags",
+ "libc",
+ "once_cell",
+ "onig_sys",
+]
+
+[[package]]
+name = "onig_sys"
+version = "69.9.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1e68317604e77e53b85896388e1a803c1d21b74c899ec9e5e1112db90735edd7"
+dependencies = [
+ "cc",
+ "pkg-config",
+]
+
 [[package]]
 name = "oorandom"
 version = "11.1.5"
@@ -1604,6 +1859,18 @@ dependencies = [
  "winapi",
 ]
 
+[[package]]
+name = "paste"
+version = "1.0.15"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a"
+
+[[package]]
+name = "pastey"
+version = "0.2.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2ee67f1008b1ba2321834326597b8e186293b049a023cdef258527550b9935b4"
+
 [[package]]
 name = "percent-encoding"
 version = "2.3.2"
@@ -1850,6 +2117,12 @@ dependencies = [
  "proc-macro2",
 ]
 
+[[package]]
+name = "r-efi"
+version = "5.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f"
+
 [[package]]
 name = "r-efi"
 version = "6.0.0"
@@ -1863,10 +2136,20 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "22f6172bdec972074665ed81ed53b71da00bfc44b65a753cfde883ec4c702a1a"
 dependencies = [
  "libc",
- "rand_chacha",
+ "rand_chacha 0.3.1",
  "rand_core 0.6.4",
 ]
 
+[[package]]
+name = "rand"
+version = "0.9.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b9ef1d0d795eb7d84685bca4f72f3649f064e6641543d3a8c415898726a57b41"
+dependencies = [
+ "rand_chacha 0.9.0",
+ "rand_core 0.9.5",
+]
+
 [[package]]
 name = "rand"
 version = "0.10.2"
@@ -1888,6 +2171,16 @@ dependencies = [
  "rand_core 0.6.4",
 ]
 
+[[package]]
+name = "rand_chacha"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb"
+dependencies = [
+ "ppv-lite86",
+ "rand_core 0.9.5",
+]
+
 [[package]]
 name = "rand_core"
 version = "0.6.4"
@@ -1897,6 +2190,15 @@ dependencies = [
  "getrandom 0.2.17",
 ]
 
+[[package]]
+name = "rand_core"
+version = "0.9.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c"
+dependencies = [
+ "getrandom 0.3.4",
+]
+
 [[package]]
 name = "rand_core"
 version = "0.10.1"
@@ -1922,6 +2224,17 @@ dependencies = [
  "rayon-core",
 ]
 
+[[package]]
+name = "rayon-cond"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2964d0cf57a3e7a06e8183d14a8b527195c706b7983549cd5462d5aa3747438f"
+dependencies = [
+ "either",
+ "itertools 0.14.0",
+ "rayon",
+]
+
 [[package]]
 name = "rayon-core"
 version = "1.13.0"
@@ -1979,7 +2292,7 @@ version = "0.12.28"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147"
 dependencies = [
- "base64",
+ "base64 0.22.1",
  "bytes",
  "futures-channel",
  "futures-core",
@@ -2361,12 +2674,36 @@ dependencies = [
  "windows-sys 0.61.2",
 ]
 
+[[package]]
+name = "spm_precompiled"
+version = "0.1.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5851699c4033c63636f7ea4cf7b7c1f1bf06d0cc03cfb42e711de5a5c46cf326"
+dependencies = [
+ "base64 0.13.1",
+ "nom",
+ "serde",
+ "unicode-segmentation",
+]
+
 [[package]]
 name = "stable_deref_trait"
 version = "1.2.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596"
 
+[[package]]
+name = "static_assertions"
+version = "1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f"
+
+[[package]]
+name = "strsim"
+version = "0.11.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f"
+
 [[package]]
 name = "subtle"
 version = "2.6.1"
@@ -2535,6 +2872,39 @@ version = "0.1.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20"
 
+[[package]]
+name = "tokenizers"
+version = "0.23.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7afbf6e88718afcc138bad01d6ccc3051dbbc3b2ce9793d8b8a3aeb610969cfc"
+dependencies = [
+ "ahash",
+ "compact_str",
+ "daachorse",
+ "dary_heap",
+ "derive_builder",
+ "esaxx-rs",
+ "getrandom 0.3.4",
+ "itertools 0.14.0",
+ "log",
+ "macro_rules_attribute",
+ "monostate",
+ "onig",
+ "paste",
+ "rand 0.9.5",
+ "rayon",
+ "rayon-cond",
+ "regex",
+ "regex-syntax",
+ "serde",
+ "serde_json",
+ "spm_precompiled",
+ "thiserror 2.0.19",
+ "unicode-normalization-alignments",
+ "unicode-segmentation",
+ "unicode_categories",
+]
+
 [[package]]
 name = "tokio"
 version = "1.53.0"
@@ -2773,6 +3143,27 @@ version = "1.0.24"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75"
 
+[[package]]
+name = "unicode-normalization-alignments"
+version = "0.1.12"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "43f613e4fa046e69818dd287fdc4bc78175ff20331479dab6e1b0f98d57062de"
+dependencies = [
+ "smallvec",
+]
+
+[[package]]
+name = "unicode-segmentation"
+version = "1.13.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c6f5d3c3b1bf09027a88a6bc961fc00497d651009560b5463668dc81b0fa87a8"
+
+[[package]]
+name = "unicode_categories"
+version = "0.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e"
+
 [[package]]
 name = "untrusted"
 version = "0.9.0"
@@ -2856,6 +3247,15 @@ version = "0.11.1+wasi-snapshot-preview1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b"
 
+[[package]]
+name = "wasip2"
+version = "1.0.4+wasi-0.2.12"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b67efb37e106e55ce722a510d6b5f9c17f083e5fc79afc2badeb12cc313d9487"
+dependencies = [
+ "wit-bindgen",
+]
+
 [[package]]
 name = "wasm-bindgen"
 version = "0.2.126"
@@ -3081,6 +3481,12 @@ dependencies = [
  "memchr",
 ]
 
+[[package]]
+name = "wit-bindgen"
+version = "0.57.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e"
+
 [[package]]
 name = "writeable"
 version = "0.6.3"

litellm-rust/Cargo.toml

@@ -1,6 +1,7 @@
 [workspace]
 members = [
     "crates/core",
+    "crates/token-counter",
     "crates/config",
     "crates/ai-gateway",
     "crates/python-interop",
@@ -18,6 +19,7 @@ repository = "https://github.com/BerriAI/litellm"
 tracing = "0.1"
 tracing-subscriber = { version = "0.3", default-features = false, features = ["registry", "std"] }
 litellm-core = { path = "crates/core" }
+litellm-token-counter = { path = "crates/token-counter" }
 litellm-config = { path = "crates/config" }
 litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false }
 litellm-python-interop = { path = "crates/python-interop" }
@@ -39,6 +41,8 @@ tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"]
 tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] }
 futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] }
 base64 = "0.22"
+url = "2.5.8"
+criterion = "0.8.2"
 
 [profile.release]
 opt-level = 3

litellm-rust/crates/ai-gateway/README.md

@@ -6,17 +6,18 @@ dials OpenAI upstream, and splices the two sockets frame-by-frame.
 
 ## Crates
 
-`litellm-rust` has five crates. A crate is a layer or shared foundation, not a route:
+`litellm-rust` has six crates. A crate is a layer or shared foundation, not a route:
 
 | Crate | Role |
 |-------|------|
 | litellm-core | The LiteLLM SDK in Rust — per-route entrypoints (`messages::messages()`) that resolve the provider, transform, and make the call; plus types, provider transforms, and the router. |
+| litellm-token-counter | Standalone input token counting shared by host integrations without pulling in the full SDK. |
 | litellm-config | Config-loading boundary. Returns resolved deployments and optionally delegates loading to Python. |
 | litellm-ai-gateway | The Axum server (behind the `server` feature) and WebSocket hosts. Translates HTTP/WS to core entrypoints; no provider handlers. |
 | litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. |
 | litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. |
 
-Dependency direction is acyclic: config depends on core, the gateway depends on config and core, and the Python bridge depends on the domain layers and Python interop.
+Dependency direction is acyclic: config depends on core, the gateway depends on config and core, and the Python bridge depends on the domain layers, token counter, and Python interop.
 
 - **Client endpoint:** `wss://<host>/v1/realtime?model=<model>` (WebSocket)
 - **Auth:** `Authorization: Bearer $LITELLM_MASTER_KEY` (fails closed if unset)

litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs

@@ -269,7 +269,7 @@ fn guardrail_error_to_core_error(error: GuardrailError) -> Error {
 
 fn core_error_kind(error: &Error) -> &'static str {
     match error {
-        Error::Auth(_) => "AuthError",
+        Error::Auth(_) | Error::MissingApiKey { .. } => "AuthError",
         Error::InvalidProvider(_) => "InvalidProvider",
         Error::InvalidRequest(_) => "InvalidRequest",
         Error::InvalidType { .. } => "InvalidType",

litellm-rust/crates/ai-gateway/src/ocr/hooks.rs

@@ -386,7 +386,7 @@ fn guardrail_error_to_core_error(error: GuardrailError) -> Error {
 
 fn core_error_kind(error: &Error) -> &'static str {
     match error {
-        Error::Auth(_) => "AuthError",
+        Error::Auth(_) | Error::MissingApiKey { .. } => "AuthError",
         Error::InvalidProvider(_) => "InvalidProvider",
         Error::InvalidRequest(_) => "InvalidRequest",
         Error::InvalidType { .. } => "InvalidType",

litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs

@@ -114,7 +114,8 @@ impl IntoResponse for MessagesRouteError {
             | Error::Connect(_)
             | Error::InvalidResponse(_)
             | Error::InvalidType { .. }
-            | Error::MissingField(_) => (
+            | Error::MissingField(_)
+            | Error::MissingApiKey { .. } => (
                 StatusCode::BAD_GATEWAY,
                 "messages provider request failed".to_string(),
             ),

litellm-rust/crates/ai-gateway/src/trace_parity.rs

@@ -47,10 +47,10 @@ pub async fn messages_request(
         .header(CONTENT_TYPE, "application/json")
         .body(Body::from(body.to_string()))
         .map_err(|error| Error::InvalidRequest(error.to_string()))?;
-    let response = routes::app(state)
-        .oneshot(request)
-        .await
-        .map_err(|error| match error {})?;
+    let response = match routes::app(state).oneshot(request).await {
+        Ok(response) => response,
+        Err(error) => match error {},
+    };
     let status: StatusCode = response.status();
     let bytes = to_bytes(response.into_body(), usize::MAX)
         .await

litellm-rust/crates/core/Cargo.toml

@@ -4,6 +4,11 @@ version = "0.1.0"
 edition.workspace = true
 license.workspace = true
 repository.workspace = true
+autotests = false
+
+[[test]]
+name = "workspace_crate_allowlist"
+path = "tests/workspace_crate_allowlist.rs"
 
 [dependencies]
 base64.workspace = true
@@ -11,10 +16,13 @@ rand.workspace = true
 reqwest.workspace = true
 serde.workspace = true
 serde_json.workspace = true
+serde_path_to_error = "0.1"
+tokio.workspace = true
 thiserror.workspace = true
 tracing.workspace = true
 tracing-subscriber = { workspace = true, optional = true }
 sha2.workspace = true
+url.workspace = true
 aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true }
 aws-credential-types = { version = "1.3.0", features = ["hardcoded-credentials"], optional = true }
 aws-sdk-sts = { version = "1.108.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true }

litellm-rust/crates/core/src/constants.rs

@@ -43,3 +43,6 @@ pub const EMPTY_TEXT_PLACEHOLDER: &str =
     "[System: Empty message content sanitised to satisfy protocol]";
 
 pub const FUNCTION_TRACE_TARGET: &str = "litellm::function_trace";
+pub(crate) const OCR_HTTP_TIMEOUT_SECS: u64 = 600;
+pub(crate) const OCR_CONNECT_TIMEOUT_SECS: u64 = 10;
+pub(crate) const MISTRAL_OCR_API_BASE: &str = "https://api.mistral.ai/v1";

litellm-rust/crates/core/src/error.rs

@@ -17,6 +17,10 @@ pub enum Error {
     InvalidRequest(String),
     #[error("{0}")]
     Auth(String),
+    #[error(
+        "Missing {provider} API Key - A call is being made to {provider} but no key is set either in the environment variables or via params"
+    )]
+    MissingApiKey { provider: &'static str },
     #[error("upstream request failed with status {status}: {body}")]
     Http { status: u16, body: String },
     #[error("upstream network error: {0}")]
@@ -36,6 +40,59 @@ pub enum Error {
     Unsupported(&'static str),
 }
 
+#[derive(Clone, Debug, ThisError, PartialEq, Eq)]
+pub enum TransportError {
+    #[error("upstream request failed with status {status}: {body}")]
+    Http { status: u16, body: String },
+    #[error("upstream network error: {0}")]
+    Network(String),
+    #[error("could not reach the provider: {0}")]
+    Connect(String),
+}
+
+impl TransportError {
+    pub fn from_reqwest_before_dispatch(error: reqwest::Error) -> Self {
+        let before_dispatch = !error.is_timeout() && (error.is_connect() || error.is_builder());
+        let message = error.without_url().to_string();
+        if before_dispatch {
+            Self::Connect(message)
+        } else {
+            Self::Network(message)
+        }
+    }
+}
+
+impl From<reqwest::Error> for TransportError {
+    fn from(error: reqwest::Error) -> Self {
+        Self::Network(error.without_url().to_string())
+    }
+}
+
+impl From<crate::ocr::error::OcrRequestError> for Error {
+    fn from(error: crate::ocr::error::OcrRequestError) -> Self {
+        match error {
+            crate::ocr::error::OcrRequestError::MissingField(field) => Self::MissingField(field),
+            error => Self::InvalidRequest(error.to_string()),
+        }
+    }
+}
+
+impl From<crate::ocr::error::OcrResponseError> for Error {
+    fn from(error: crate::ocr::error::OcrResponseError) -> Self {
+        Self::InvalidResponse(error.to_string())
+    }
+}
+
+impl From<TransportError> for Error {
+    fn from(error: TransportError) -> Self {
+        match error {
+            TransportError::Http { status, body } => Self::Http { status, body },
+            TransportError::Network(message) => Self::Network(message),
+            TransportError::Connect(message) => Self::Connect(message),
+        }
+    }
+}
+
 pub fn json_type_name(value: &serde_json::Value) -> &'static str {
     match value {
         serde_json::Value::Null => "null",
@@ -46,3 +103,53 @@ pub fn json_type_name(value: &serde_json::Value) -> &'static str {
         serde_json::Value::Object(_) => "object",
     }
 }
+
+#[cfg(test)]
+mod transport_tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn transport_errors_remove_urls_and_keep_dispatch_context() {
+        let error = reqwest::Client::builder()
+            .no_proxy()
+            .build()
+            .expect("client")
+            .get("http://localhost:invalid/private?api_key=secret")
+            .send()
+            .await
+            .expect_err("invalid port");
+        let error = TransportError::from_reqwest_before_dispatch(error);
+        assert!(matches!(error, TransportError::Connect(_)));
+        assert!(!error.to_string().contains("secret"));
+        assert!(!error.to_string().contains("private"));
+    }
+
+    #[tokio::test]
+    async fn request_timeout_is_not_safe_to_retry_as_a_connect_failure() {
+        use std::time::Duration;
+        let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
+            .await
+            .expect("bind");
+        let address = listener.local_addr().expect("address");
+        let request = reqwest::Client::builder()
+            .no_proxy()
+            .build()
+            .expect("client")
+            .get(format!("http://{address}"))
+            .timeout(Duration::from_millis(200))
+            .send();
+        let (response, accepted) = tokio::join!(
+            request,
+            tokio::time::timeout(Duration::from_secs(2), listener.accept())
+        );
+        let _connection = accepted
+            .expect("accept deadline")
+            .expect("accepted connection");
+        let error = response.expect_err("server does not respond");
+        assert!(error.is_timeout());
+        assert!(matches!(
+            TransportError::from_reqwest_before_dispatch(error),
+            TransportError::Network(_)
+        ));
+    }
+}

litellm-rust/crates/core/src/http_utils.rs

@@ -1,10 +1,43 @@
-//! Header and upstream-body helpers shared by every route module.
-
 use serde_json::{Map, Value};
 
 use crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS;
 use crate::error::{Error, json_type_name};
 
+#[allow(
+    dead_code,
+    reason = "used by the OCR architecture in the next stacked PR"
+)]
+pub(crate) enum HeaderPolicy<'a> {
+    All,
+    Only(&'a [&'a str]),
+    Except(&'a [&'a str]),
+}
+
+#[allow(
+    dead_code,
+    reason = "used by the OCR architecture in the next stacked PR"
+)]
+pub(crate) fn with_headers(
+    builder: reqwest::RequestBuilder,
+    headers: &[(String, String)],
+    policy: HeaderPolicy<'_>,
+) -> reqwest::RequestBuilder {
+    headers
+        .iter()
+        .filter(|(name, _)| match policy {
+            HeaderPolicy::All => true,
+            HeaderPolicy::Only(names) => names
+                .iter()
+                .any(|allowed| name.eq_ignore_ascii_case(allowed)),
+            HeaderPolicy::Except(names) => !names
+                .iter()
+                .any(|excluded| name.eq_ignore_ascii_case(excluded)),
+        })
+        .fold(builder, |builder, (name, value)| {
+            builder.header(name, value)
+        })
+}
+
 #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
 pub async fn http_request(
     request: reqwest::RequestBuilder,
@@ -12,8 +45,6 @@ pub async fn http_request(
     request.send().await
 }
 
-/// Bound an upstream error body before it crosses a host boundary, so provider
-/// bodies stay data-minimized.
 pub fn truncate_error_body(body: &str) -> String {
     if body.chars().count() <= UPSTREAM_ERROR_BODY_MAX_CHARS {
         return body.to_string();
@@ -61,11 +92,77 @@ pub fn has_bearer_auth(headers: &[(String, String)]) -> bool {
     })
 }
 
+#[allow(
+    dead_code,
+    reason = "used by the OCR architecture in the next stacked PR"
+)]
+pub(crate) fn deserialize_optional_param<'de, D, T>(
+    deserializer: D,
+) -> Result<Option<Option<T>>, D::Error>
+where
+    D: serde::Deserializer<'de>,
+    T: serde::Deserialize<'de>,
+{
+    <Option<T> as serde::Deserialize>::deserialize(deserializer).map(Some)
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
     use serde_json::json;
 
+    #[rstest::rstest]
+    #[case(HeaderPolicy::All, true, true)]
+    #[case(HeaderPolicy::Only(&["authorization"]), true, false)]
+    #[case(HeaderPolicy::Except(&["authorization"]), false, true)]
+    fn forwarding_policy_preserves_matching_headers_and_duplicates(
+        #[case] policy: HeaderPolicy<'_>,
+        #[case] auth: bool,
+        #[case] trace: bool,
+    ) {
+        let request = with_headers(
+            reqwest::Client::new().get("https://example.com"),
+            &[
+                ("AuThOrIzAtIoN".into(), "Bearer token".into()),
+                ("X-Trace".into(), "first".into()),
+                ("x-trace".into(), "second".into()),
+            ],
+            policy,
+        )
+        .build()
+        .unwrap();
+        assert_eq!(request.headers().contains_key("authorization"), auth);
+        let traces: Vec<_> = request.headers().get_all("x-trace").iter().collect();
+        if trace {
+            assert_eq!(traces, ["first", "second"]);
+        } else {
+            assert!(traces.is_empty());
+        }
+    }
+
+    #[test]
+    fn multipart_policy_leaves_content_headers_to_reqwest() {
+        let request = with_headers(
+            reqwest::Client::new()
+                .post("https://example.com")
+                .multipart(reqwest::multipart::Form::new().text("file", "abc")),
+            &[
+                ("Content-Type".into(), "application/json".into()),
+                ("CONTENT-LENGTH".into(), "0".into()),
+            ],
+            HeaderPolicy::Except(&["content-type", "content-length"]),
+        )
+        .build()
+        .unwrap();
+        assert!(
+            request.headers()["content-type"]
+                .to_str()
+                .unwrap()
+                .starts_with("multipart/form-data; boundary=")
+        );
+        assert_ne!(request.headers()["content-length"], "0");
+    }
+
     #[test]
     fn truncate_leaves_short_bodies_untouched() {
         assert_eq!(truncate_error_body("short"), "short");

litellm-rust/crates/core/src/lib.rs

@@ -14,5 +14,6 @@ pub mod realtime;
 pub mod responses;
 pub mod router;
 pub mod routing_utils;
+mod url_utils;
 
 pub use error::Error;

litellm-rust/crates/core/src/ocr/adapters/mistral.rs

@@ -0,0 +1,147 @@
+use super::OcrAdapter;
+use crate::Error;
+use crate::constants::MISTRAL_OCR_API_BASE;
+use crate::ocr::OcrClient;
+use crate::ocr::codecs::mistral::{self, MistralOcrParams, MistralOcrResponse};
+use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError};
+use crate::ocr::prepare::{
+    _prepare_ocr_request, ParsedProviderParams, credential_env, transform_request_body,
+};
+use crate::ocr::registry::OcrProvider;
+use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection};
+use crate::url_utils::ApiUrl;
+
+const MISTRAL_API_KEY_ENV: &str = "MISTRAL_API_KEY";
+
+#[derive(Clone, Debug)]
+pub(crate) struct MistralAdapter;
+
+impl OcrAdapter for MistralAdapter {
+    type ProviderResponse = MistralOcrResponse;
+    const PROVIDER: OcrProvider = OcrProvider::Mistral;
+
+    async fn prepare_request(
+        &self,
+        request: &LiteLLMOcrRequest,
+        client: &OcrClient,
+    ) -> Result<reqwest::Request, OcrError> {
+        let ParsedProviderParams {
+            known: params,
+            extra_params: _extra_params,
+        } = _prepare_ocr_request::<MistralOcrParams>(request)?;
+        let headers = validate_environment(&request.connection, &credential_env)?;
+        let url = get_complete_url(request.connection.api_base.as_deref())?;
+        let body =
+            mistral::transform_ocr_request(&request.model, request.document.clone(), &params)?;
+        transform_request_body(client, request, &url, &headers, body, |_| Ok(())).await
+    }
+
+    fn transform_ocr_response(
+        &self,
+        request: &LiteLLMOcrRequest,
+        response: Self::ProviderResponse,
+    ) -> Result<LiteLLMOcrResponse, OcrResponseError> {
+        mistral::transform_ocr_response(&request.model, response)
+    }
+}
+
+pub(crate) fn get_complete_url(api_base: Option<&str>) -> Result<String, OcrError> {
+    let base = api_base
+        .map(str::trim)
+        .filter(|base| !base.is_empty())
+        .unwrap_or(MISTRAL_OCR_API_BASE);
+    ApiUrl::parse(base)
+        .and_then(|url| url.complete_path(&["v1", "ocr"]))
+        .map(|url| url.into_string())
+        .map_err(|_| {
+            OcrRequestError::RequestField {
+                path: "api_base".into(),
+            }
+            .into()
+        })
+}
+
+fn validate_environment(
+    connection: &OcrConnection,
+    env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
+) -> Result<Vec<(String, String)>, OcrError> {
+    if crate::http_utils::has_header(&connection.extra_headers, "authorization") {
+        return Ok(connection.extra_headers.clone());
+    }
+    let api_key = connection
+        .api_key
+        .as_deref()
+        .map(str::trim)
+        .filter(|key| !key.is_empty())
+        .map(str::to_string)
+        .or_else(|| env_lookup(MISTRAL_API_KEY_ENV).filter(|key| !key.trim().is_empty()))
+        .ok_or(Error::MissingApiKey {
+            provider: "Mistral",
+        })?;
+    Ok(
+        std::iter::once(("Authorization".into(), format!("Bearer {api_key}")))
+            .chain(connection.extra_headers.clone())
+            .collect(),
+    )
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn complete_url_defaults_and_dedupes_v1() {
+        assert_eq!(
+            get_complete_url(None).unwrap(),
+            "https://api.mistral.ai/v1/ocr"
+        );
+        assert_eq!(
+            get_complete_url(Some("https://example.com/v1?tenant=a")).unwrap(),
+            "https://example.com/v1/ocr?tenant=a"
+        );
+        assert_eq!(
+            get_complete_url(Some("https://example.com/v1/ocr?tenant=a")).unwrap(),
+            "https://example.com/v1/ocr?tenant=a"
+        );
+    }
+
+    #[test]
+    fn environment_prefers_explicit_key_then_environment() {
+        let explicit = OcrConnection {
+            api_key: Some("explicit".into()),
+            ..OcrConnection::default()
+        };
+        assert_eq!(
+            validate_environment(&explicit, &|_| Some("environment".into())).unwrap()[0],
+            ("Authorization".into(), "Bearer explicit".into())
+        );
+
+        assert_eq!(
+            validate_environment(&OcrConnection::default(), &|_| Some("environment".into()))
+                .unwrap()[0],
+            ("Authorization".into(), "Bearer environment".into())
+        );
+    }
+
+    #[test]
+    fn environment_preserves_forwarded_authorization() {
+        let connection = OcrConnection {
+            extra_headers: vec![("authorization".into(), "Bearer forwarded".into())],
+            ..OcrConnection::default()
+        };
+        assert_eq!(
+            validate_environment(&connection, &|_| None).unwrap(),
+            connection.extra_headers
+        );
+    }
+
+    #[test]
+    fn environment_rejects_missing_key() {
+        assert!(matches!(
+            validate_environment(&OcrConnection::default(), &|_| None),
+            Err(OcrError::Public(Error::MissingApiKey {
+                provider: "Mistral"
+            }))
+        ));
+    }
+}

litellm-rust/crates/core/src/ocr/adapters/mod.rs

@@ -0,0 +1,69 @@
+use std::future::Future;
+
+use serde::de::DeserializeOwned;
+
+use super::OcrClient;
+use super::error::{OcrError, OcrResponseError};
+use super::registry::OcrProvider;
+use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrResponseFormat};
+use super::wire::DecodedOcrResponse;
+
+mod mistral;
+
+pub(crate) use mistral::MistralAdapter;
+
+/// Converts a complete LiteLLM OCR call to provider HTTP and normalizes its response.
+pub(crate) trait OcrAdapter: Send + Sync + Sized + 'static {
+    /// Provider JSON schema; direct and Vertex Mistral share `MistralOcrResponse`.
+    type ProviderResponse: DeserializeOwned + Send;
+
+    const PROVIDER: OcrProvider;
+
+    /// Prepares the complete provider HTTP request.
+    /// `request` contains the model, document, connection, and unmapped caller options.
+    /// `client` supplies reusable provider and document HTTP clients.
+    /// Returns the complete HTTP request, whereas Python returns body data.
+    fn prepare_request(
+        &self,
+        request: &LiteLLMOcrRequest,
+        client: &OcrClient,
+    ) -> impl Future<Output = Result<reqwest::Request, OcrError>> + Send;
+
+    /// Python: `transform_ocr_response`.
+    /// `request` supplies caller context, including the fallback model.
+    /// `response` is the decoded provider payload; the output is the shared LiteLLM schema.
+    fn transform_ocr_response(
+        &self,
+        request: &LiteLLMOcrRequest,
+        response: Self::ProviderResponse,
+    ) -> Result<LiteLLMOcrResponse, OcrResponseError>;
+
+    /// Decodes provider HTTP; adapters may override this to poll asynchronous operations.
+    /// Python performs that polling inside `async_transform_ocr_response`.
+    /// `client` is reused for polling; `response` is the initial HTTP response.
+    /// `url` and `headers` describe the submitted call; `request` supplies limits and format.
+    fn read_response(
+        &self,
+        _client: &OcrClient,
+        response: reqwest::Response,
+        _url: &str,
+        _headers: &[(String, String)],
+        request: &LiteLLMOcrRequest,
+    ) -> impl Future<Output = Result<DecodedOcrResponse<Self::ProviderResponse>, OcrError>> + Send
+    {
+        let retain_native = request
+            .response_format()
+            .map(|format| format == OcrResponseFormat::Native);
+        async move { super::client::read_json_response(response, retain_native?).await }
+    }
+}
+
+macro_rules! for_each_ocr_adapter {
+    ($callback:ident) => {
+        $callback! {
+            Mistral, $crate::ocr::adapters::MistralAdapter, $crate::ocr::adapters::MistralAdapter, Mistral;
+        }
+    };
+}
+
+pub(crate) use for_each_ocr_adapter;

litellm-rust/crates/core/src/ocr/client.rs

@@ -0,0 +1,75 @@
+use std::sync::OnceLock;
+use std::time::Duration;
+
+use serde::de::DeserializeOwned;
+
+use super::error::OcrError;
+use super::handler::perform_ocr_request;
+use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse};
+use super::wire::{DecodedOcrResponse, decode_response};
+use crate::Error;
+use crate::constants::OCR_CONNECT_TIMEOUT_SECS;
+use crate::error::TransportError;
+
+#[derive(Clone)]
+pub struct OcrClient {
+    provider_http: reqwest::Client,
+}
+
+impl OcrClient {
+    pub fn new(provider_http: reqwest::Client) -> Result<Self, TransportError> {
+        Ok(Self { provider_http })
+    }
+
+    #[tracing::instrument(
+        name = "ocr",
+        target = "litellm::function_trace",
+        level = "trace",
+        skip_all
+    )]
+    pub async fn perform(&self, request: LiteLLMOcrRequest) -> Result<LiteLLMOcrResponse, Error> {
+        perform_ocr_request(self, request).await
+    }
+
+    pub(crate) fn provider_http(&self) -> &reqwest::Client {
+        &self.provider_http
+    }
+
+    #[cfg(test)]
+    pub(crate) fn for_test(provider_http: reqwest::Client) -> Self {
+        Self { provider_http }
+    }
+}
+
+pub async fn ocr(request: LiteLLMOcrRequest) -> Result<LiteLLMOcrResponse, Error> {
+    static CLIENT: OnceLock<Result<OcrClient, TransportError>> = OnceLock::new();
+    let client = CLIENT
+        .get_or_init(|| {
+            reqwest::Client::builder()
+                .connect_timeout(Duration::from_secs(OCR_CONNECT_TIMEOUT_SECS))
+                .build()
+                .map_err(TransportError::from)
+                .and_then(OcrClient::new)
+        })
+        .clone()?;
+    client.perform(request).await
+}
+
+pub async fn read_json_response<T: DeserializeOwned>(
+    response: reqwest::Response,
+    native: bool,
+) -> Result<DecodedOcrResponse<T>, OcrError> {
+    let status = response.status();
+    let bytes = response
+        .bytes()
+        .await
+        .map_err(crate::error::TransportError::from)?;
+    if !status.is_success() {
+        return Err(crate::error::TransportError::Http {
+            status: status.as_u16(),
+            body: crate::http_utils::truncate_error_body(&String::from_utf8_lossy(&bytes)),
+        }
+        .into());
+    }
+    Ok(decode_response(&bytes, native)?)
+}

litellm-rust/crates/core/src/ocr/codecs/mistral/mod.rs

@@ -0,0 +1,5 @@
+mod transformation;
+mod types;
+
+pub(crate) use transformation::{transform_ocr_request, transform_ocr_response};
+pub(crate) use types::{MistralOcrParams, MistralOcrRequest, MistralOcrResponse};

litellm-rust/crates/core/src/ocr/codecs/mistral/transformation.rs

@@ -0,0 +1,102 @@
+use super::{MistralOcrParams, MistralOcrRequest, MistralOcrResponse};
+use crate::ocr::error::{OcrRequestError, OcrResponseError};
+use crate::ocr::types::{LiteLLMOcrResponse, OcrDocument};
+
+#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
+pub(crate) fn transform_ocr_request(
+    model: &str,
+    document: OcrDocument,
+    params: &MistralOcrParams,
+) -> Result<MistralOcrRequest, OcrRequestError> {
+    Ok(MistralOcrRequest {
+        model: model.to_string(),
+        document,
+        params: params.clone(),
+    })
+}
+
+pub(crate) fn transform_ocr_response(
+    model: &str,
+    response: MistralOcrResponse,
+) -> Result<LiteLLMOcrResponse, OcrResponseError> {
+    Ok(LiteLLMOcrResponse {
+        pages: response.pages,
+        model: response.model.unwrap_or_else(|| model.to_string()),
+        document_annotation: response.document_annotation,
+        usage_info: response.usage_info,
+        object: "ocr".to_string(),
+        extra_fields: response.extra_fields,
+        provider_native_response: None,
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use rstest::rstest;
+    use serde_json::{Value, json};
+
+    #[rstest]
+    #[case("pages", json!([0, 2]))]
+    #[case("include_image_base64", json!(true))]
+    #[case("image_limit", json!(2))]
+    #[case("image_min_size", json!(100))]
+    #[case("bbox_annotation_format", json!({"type":"json_schema"}))]
+    #[case("document_annotation_format", json!({"type":"json_schema"}))]
+    #[case("document_annotation_prompt", json!("extract"))]
+    #[case("extract_header", json!(true))]
+    #[case("extract_footer", json!(false))]
+    #[case("table_format", json!("html"))]
+    #[case("confidence_scores_granularity", json!("word"))]
+    #[case("include_blocks", json!(true))]
+    #[case("id", json!("req-123"))]
+    fn request_mapping_matches_python(#[case] name: &str, #[case] value: Value) {
+        let params: MistralOcrParams =
+            serde_json::from_value(json!({name: value.clone()})).unwrap();
+        let document: OcrDocument = serde_json::from_value(
+            json!({"type":"document_url","document_url":"https://example.com/a.pdf"}),
+        )
+        .unwrap();
+        let result =
+            serde_json::to_value(transform_ocr_request("model", document, &params).unwrap())
+                .unwrap();
+        assert_eq!(result["model"], "model");
+        assert_eq!(result[name], value);
+    }
+
+    #[test]
+    fn request_mapping_filters_unknown_fields() {
+        let params: MistralOcrParams = serde_json::from_value(json!({"unknown": true})).unwrap();
+        let document: OcrDocument = serde_json::from_value(
+            json!({"type":"document_url","document_url":"https://example.com/a.pdf"}),
+        )
+        .unwrap();
+        let result =
+            serde_json::to_value(transform_ocr_request("model", document, &params).unwrap())
+                .unwrap();
+        assert!(result.get("unknown").is_none());
+    }
+
+    #[test]
+    fn response_preserves_provider_fields() {
+        let response: MistralOcrResponse = serde_json::from_value(json!({
+            "pages":[{"index":0,"markdown":"hello","header":"head","confidence_scores":{"mean":0.99}}],
+            "model":"returned-model",
+            "usage_info":{"pages_processed":1,"future_counter":5},
+            "future_response_field":"kept"
+        }))
+        .unwrap();
+        let result = transform_ocr_response("model", response)
+            .unwrap()
+            .into_json();
+        assert_eq!(result["pages"][0]["header"], "head");
+        assert_eq!(result["usage_info"]["future_counter"], 5);
+        assert_eq!(result["future_response_field"], "kept");
+        assert_eq!(result["model"], "returned-model");
+    }
+
+    #[test]
+    fn response_rejects_null_pages() {
+        assert!(serde_json::from_value::<MistralOcrResponse>(json!({"pages":null})).is_err());
+    }
+}

litellm-rust/crates/core/src/ocr/codecs/mistral/types.rs

@@ -0,0 +1,53 @@
+use serde::{Deserialize, Serialize};
+use serde_json::{Map, Value};
+
+use crate::ocr::types::OcrDocument;
+
+#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
+pub(crate) struct MistralOcrParams {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub pages: Option<Vec<i64>>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub include_image_base64: Option<bool>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub image_limit: Option<i64>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub image_min_size: Option<i64>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub bbox_annotation_format: Option<Value>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub document_annotation_format: Option<Value>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub document_annotation_prompt: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub extract_header: Option<bool>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub extract_footer: Option<bool>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub table_format: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub confidence_scores_granularity: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub include_blocks: Option<bool>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub id: Option<String>,
+}
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub(crate) struct MistralOcrRequest {
+    pub model: String,
+    pub document: OcrDocument,
+    #[serde(flatten)]
+    pub params: MistralOcrParams,
+}
+
+#[derive(Clone, Debug, Default, Deserialize)]
+pub(crate) struct MistralOcrResponse {
+    #[serde(default)]
+    pub pages: Vec<Value>,
+    pub model: Option<String>,
+    pub document_annotation: Option<Value>,
+    pub usage_info: Option<Value>,
+    #[serde(flatten)]
+    pub extra_fields: Map<String, Value>,
+}

litellm-rust/crates/core/src/ocr/codecs/mod.rs

@@ -0,0 +1 @@
+pub(crate) mod mistral;

litellm-rust/crates/core/src/ocr/error.rs

@@ -0,0 +1,42 @@
+use thiserror::Error;
+
+use crate::error::TransportError;
+
+#[derive(Debug, Clone, PartialEq, Eq, Error)]
+pub enum OcrRequestError {
+    #[error("Invalid `req_format`. Expected 'native' or 'litellm'.")]
+    RequestFormat,
+    #[error("invalid OCR request field: {path}")]
+    RequestField { path: String },
+    #[error("missing required field: {0}")]
+    MissingField(&'static str),
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Error)]
+pub enum OcrResponseError {
+    #[error("invalid OCR response field: {path}")]
+    ResponseField { path: String },
+}
+
+#[derive(Debug, Error)]
+pub enum OcrError {
+    #[error("{0}")]
+    Request(#[from] OcrRequestError),
+    #[error("{0}")]
+    Response(#[from] OcrResponseError),
+    #[error("{0}")]
+    Transport(#[from] TransportError),
+    #[error("{0}")]
+    Public(#[from] crate::Error),
+}
+
+impl From<OcrError> for crate::Error {
+    fn from(error: OcrError) -> Self {
+        match error {
+            OcrError::Request(error) => error.into(),
+            OcrError::Response(error) => error.into(),
+            OcrError::Transport(error) => error.into(),
+            OcrError::Public(error) => error,
+        }
+    }
+}

litellm-rust/crates/core/src/ocr/handler.rs

@@ -0,0 +1,72 @@
+use super::OcrClient;
+use super::adapters::OcrAdapter;
+use super::hooks::OcrLifecycleHooks;
+use super::registry::OcrAdapterKind;
+use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse};
+use crate::Error;
+use crate::call_lifecycle::{CallLifecycle, CallLifecycleContext};
+
+pub(crate) async fn perform_ocr_request(
+    client: &OcrClient,
+    request: LiteLLMOcrRequest,
+) -> Result<LiteLLMOcrResponse, Error> {
+    let context = CallLifecycleContext::new(
+        "ocr",
+        request.model.clone(),
+        request.adapter.provider().as_str(),
+        request
+            .litellm_call_id
+            .clone()
+            .unwrap_or_else(|| format!("ocr-{:032x}", rand::random::<u128>())),
+    );
+    let hooks = OcrLifecycleHooks {
+        hooks: request.hooks.clone(),
+        provider_name: context.custom_llm_provider.clone(),
+    };
+    CallLifecycle::default().run(context, request, &hooks, |request| async move {
+        macro_rules! execute_selected_adapter {
+            ($( $variant:ident, $adapter:ty, $instance:expr, $provider:ident; )+) => {
+                match request.adapter {
+                    $( OcrAdapterKind::$variant => execute_ocr_provider_call(client, &$instance, request).await, )+
+                }
+            };
+        }
+        super::adapters::for_each_ocr_adapter!(execute_selected_adapter)
+    }).await
+}
+
+#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
+async fn execute_ocr_provider_call<A: OcrAdapter>(
+    client: &OcrClient,
+    adapter: &A,
+    request: LiteLLMOcrRequest,
+) -> Result<LiteLLMOcrResponse, Error> {
+    let provider_request = adapter.prepare_request(&request, client).await?;
+    let url = provider_request.url().to_string();
+    let headers = provider_request
+        .headers()
+        .iter()
+        .map(|(name, value)| {
+            value
+                .to_str()
+                .map(|value| (name.to_string(), value.to_string()))
+                .map_err(|_| super::error::OcrRequestError::RequestField {
+                    path: "headers".into(),
+                })
+        })
+        .collect::<Result<Vec<_>, _>>()?;
+    let response = crate::http_utils::http_request(reqwest::RequestBuilder::from_parts(
+        client.provider_http().clone(),
+        provider_request,
+    ))
+    .await
+    .map_err(crate::error::TransportError::from)?;
+    let decoded = adapter
+        .read_response(client, response, &url, &headers, &request)
+        .await?;
+    let response = adapter.transform_ocr_response(&request, decoded.data)?;
+    Ok(LiteLLMOcrResponse {
+        provider_native_response: decoded.native,
+        ..response
+    })
+}

litellm-rust/crates/core/src/ocr/hooks.rs

@@ -0,0 +1,146 @@
+use std::future::Future;
+use std::pin::Pin;
+use std::sync::Arc;
+
+use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrDocument};
+use crate::Error;
+use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming};
+use serde::Serialize;
+use serde_json::Value;
+
+pub type OcrHookFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
+pub type OcrLogFuture<'a> = Pin<Box<dyn Future<Output = ()> + Send + 'a>>;
+
+#[derive(Clone, Debug, Serialize)]
+pub struct OcrPreCallRequest {
+    pub model: String,
+    pub custom_llm_provider: String,
+    pub document: OcrDocument,
+    pub optional_params: Value,
+}
+
+#[derive(Clone, Debug, Serialize)]
+pub struct OcrDuringCallRequest {
+    pub model: String,
+    pub custom_llm_provider: String,
+    pub url: String,
+    pub body: Value,
+}
+
+pub trait OcrHooks: Send + Sync {
+    fn has_guardrails(&self) -> bool {
+        false
+    }
+    fn pre_call(&self, request: OcrPreCallRequest) -> OcrHookFuture<'_, OcrPreCallRequest> {
+        Box::pin(async move { Ok(request) })
+    }
+    fn during_call(
+        &self,
+        request: OcrDuringCallRequest,
+    ) -> OcrHookFuture<'_, OcrDuringCallRequest> {
+        Box::pin(async move { Ok(request) })
+    }
+    fn success<'a>(
+        &'a self,
+        _context: &'a CallLifecycleContext,
+        _response: &'a LiteLLMOcrResponse,
+        _timing: &'a CallLifecycleTiming,
+    ) -> OcrLogFuture<'a> {
+        Box::pin(async {})
+    }
+    fn failure<'a>(
+        &'a self,
+        _context: &'a CallLifecycleContext,
+        _error: &'a Error,
+        _timing: &'a CallLifecycleTiming,
+    ) -> OcrLogFuture<'a> {
+        Box::pin(async {})
+    }
+}
+
+pub struct NoopOcrHooks;
+impl OcrHooks for NoopOcrHooks {}
+
+pub(crate) struct OcrLifecycleHooks {
+    pub hooks: Arc<dyn OcrHooks>,
+    pub provider_name: String,
+}
+
+impl CallLifecycleHooks<LiteLLMOcrRequest, LiteLLMOcrRequest, LiteLLMOcrResponse>
+    for OcrLifecycleHooks
+{
+    type PreCallFuture<'a> = OcrHookFuture<'a, LiteLLMOcrRequest>;
+    type DuringCallFuture<'a> = OcrHookFuture<'a, LiteLLMOcrRequest>;
+    type SuccessFuture<'a> = OcrLogFuture<'a>;
+    type FailureFuture<'a> = OcrLogFuture<'a>;
+
+    fn async_pre_call_hook<'a>(
+        &'a self,
+        _context: &'a CallLifecycleContext,
+        request: LiteLLMOcrRequest,
+    ) -> Self::PreCallFuture<'a> {
+        Box::pin(async move {
+            if !self.hooks.has_guardrails() {
+                return Ok(request);
+            }
+            let changed = self
+                .hooks
+                .pre_call(OcrPreCallRequest {
+                    model: request.model.clone(),
+                    custom_llm_provider: self.provider_name.clone(),
+                    document: request.document,
+                    optional_params: Value::Object(request.optional_params),
+                })
+                .await?;
+            let Value::Object(optional_params) = changed.optional_params else {
+                return Err(super::error::OcrRequestError::RequestField {
+                    path: "guardrail.optional_params".into(),
+                }
+                .into());
+            };
+            Ok(LiteLLMOcrRequest {
+                document: changed.document,
+                optional_params,
+                ..request
+            })
+        })
+    }
+
+    fn async_during_call_hook<'a>(
+        &'a self,
+        _context: &'a CallLifecycleContext,
+        request: LiteLLMOcrRequest,
+    ) -> Self::DuringCallFuture<'a> {
+        Box::pin(async move { Ok(request) })
+    }
+
+    #[tracing::instrument(
+        name = "success_callback",
+        target = "litellm::function_trace",
+        level = "trace",
+        skip_all
+    )]
+    fn async_log_success_event<'a>(
+        &'a self,
+        context: &'a CallLifecycleContext,
+        response: &'a LiteLLMOcrResponse,
+        timing: &'a CallLifecycleTiming,
+    ) -> Self::SuccessFuture<'a> {
+        self.hooks.success(context, response, timing)
+    }
+
+    #[tracing::instrument(
+        name = "failure_callback",
+        target = "litellm::function_trace",
+        level = "trace",
+        skip_all
+    )]
+    fn async_log_failure_event<'a>(
+        &'a self,
+        context: &'a CallLifecycleContext,
+        error: &'a Error,
+        timing: &'a CallLifecycleTiming,
+    ) -> Self::FailureFuture<'a> {
+        self.hooks.failure(context, error, timing)
+    }
+}

litellm-rust/crates/core/src/ocr/mod.rs

@@ -1,2 +1,21 @@
+mod adapters;
+pub mod client;
+mod codecs;
+pub mod error;
+mod handler;
+pub mod hooks;
+mod prepare;
+mod registry;
 pub mod transformation;
 pub mod types;
+pub mod wire;
+
+pub use client::{OcrClient, ocr};
+pub use types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection, OcrDocument};
+
+#[cfg(test)]
+#[path = "../../tests/ocr/support.rs"]
+pub(crate) mod test_support;
+#[cfg(test)]
+#[path = "../../tests/ocr.rs"]
+pub(crate) mod tests;

litellm-rust/crates/core/src/ocr/prepare.rs

@@ -0,0 +1,142 @@
+use serde::{Deserialize, Serialize, de::DeserializeOwned};
+use serde_json::{Map, Value};
+
+use super::OcrClient;
+use super::error::{OcrError, OcrRequestError};
+use super::hooks::OcrDuringCallRequest;
+use super::types::LiteLLMOcrRequest;
+
+#[derive(Debug, Deserialize)]
+pub(crate) struct ParsedProviderParams<T> {
+    #[serde(flatten)]
+    pub known: T,
+    #[serde(default, flatten)]
+    pub extra_params: Map<String, Value>,
+}
+
+#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
+pub(crate) fn _prepare_ocr_request<T: DeserializeOwned>(
+    request: &LiteLLMOcrRequest,
+) -> Result<ParsedProviderParams<T>, OcrRequestError> {
+    super::wire::decode_request_value(
+        Value::Object(request.optional_params.clone()),
+        "optional_params",
+    )
+}
+
+pub(crate) async fn transform_request_body<B>(
+    client: &OcrClient,
+    request: &LiteLLMOcrRequest,
+    url: &str,
+    headers: &[(String, String)],
+    body: B,
+    validate: impl FnOnce(&B) -> Result<(), OcrRequestError>,
+) -> Result<reqwest::Request, OcrError>
+where
+    B: Serialize + DeserializeOwned,
+{
+    let body = if request.hooks.has_guardrails() {
+        let changed = request
+            .hooks
+            .during_call(OcrDuringCallRequest {
+                model: request.model.clone(),
+                custom_llm_provider: request.adapter.provider().as_str().into(),
+                url: url.into(),
+                body: serde_json::to_value(body).map_err(|_| OcrRequestError::RequestField {
+                    path: "body".into(),
+                })?,
+            })
+            .await?;
+        let body = OcrWireBody::<B>::decode(changed.body)?;
+        validate(&body.body)?;
+        body
+    } else {
+        OcrWireBody {
+            body,
+            extra: Map::new(),
+        }
+    };
+    build_http_request(client, request, url, headers, &body)
+}
+
+pub(crate) fn build_http_request<B: Serialize>(
+    client: &OcrClient,
+    request: &LiteLLMOcrRequest,
+    url: &str,
+    headers: &[(String, String)],
+    body: &B,
+) -> Result<reqwest::Request, OcrError> {
+    let builder = client
+        .provider_http()
+        .post(url)
+        .json(body)
+        .timeout(request.connection.timeout);
+    crate::http_utils::with_headers(builder, headers, crate::http_utils::HeaderPolicy::All)
+        .build()
+        .map_err(crate::error::TransportError::from)
+        .map_err(OcrError::from)
+}
+
+#[derive(Serialize)]
+struct OcrWireBody<B> {
+    #[serde(flatten)]
+    body: B,
+    #[serde(flatten)]
+    extra: Map<String, Value>,
+}
+
+impl<B: Serialize + DeserializeOwned> OcrWireBody<B> {
+    fn decode(value: Value) -> Result<Self, OcrRequestError> {
+        let body: B = super::wire::decode_request_value(value.clone(), "guardrail.body")?;
+        let Value::Object(fields) = value else {
+            return Err(OcrRequestError::RequestField {
+                path: "guardrail.body".into(),
+            });
+        };
+        let known = serde_json::to_value(&body).map_err(|_| OcrRequestError::RequestField {
+            path: "guardrail.body".into(),
+        })?;
+        let extra = fields
+            .into_iter()
+            .filter(|(key, _)| known.get(key).is_none())
+            .collect();
+        Ok(Self { body, extra })
+    }
+}
+
+pub(crate) fn credential_env(name: &str) -> Option<String> {
+    std::env::var(name).ok()
+}
+
+#[cfg(test)]
+mod tests {
+    use serde_json::json;
+
+    use super::*;
+
+    #[derive(Debug, Deserialize, PartialEq)]
+    struct KnownParams {
+        pages: Option<Vec<i64>>,
+    }
+
+    #[test]
+    fn parsed_provider_params_separates_known_and_extra_params() {
+        let parsed: ParsedProviderParams<KnownParams> = super::super::wire::decode_request_value(
+            json!({
+                "pages": [0, 2],
+                "future_ocr_option": true,
+                "extra_body": {"provider_option": "value"}
+            }),
+            "optional_params",
+        )
+        .unwrap();
+
+        assert_eq!(parsed.known.pages, Some(vec![0, 2]));
+        assert_eq!(parsed.extra_params["future_ocr_option"], true);
+        assert_eq!(
+            parsed.extra_params["extra_body"],
+            json!({"provider_option": "value"})
+        );
+        assert_eq!(parsed.extra_params.len(), 2);
+    }
+}

litellm-rust/crates/core/src/ocr/registry.rs

@@ -0,0 +1,53 @@
+use super::adapters::OcrAdapter;
+use crate::Error;
+use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider};
+
+macro_rules! define_adapter_types {
+    ($( $variant:ident, $adapter:ty, $instance:expr, $provider:ident; )+) => {
+        #[derive(Clone, Copy, Debug, PartialEq, Eq)]
+        pub(crate) enum OcrAdapterKind {
+            $( $variant, )+
+        }
+
+        impl OcrAdapterKind {
+            pub(crate) const fn provider(self) -> OcrProvider {
+                match self {
+                    $( Self::$variant => <$adapter>::PROVIDER, )+
+                }
+            }
+        }
+    };
+}
+
+super::adapters::for_each_ocr_adapter!(define_adapter_types);
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub(crate) enum OcrProvider {
+    Mistral,
+}
+
+impl OcrProvider {
+    pub(crate) const fn as_str(self) -> &'static str {
+        match self {
+            Self::Mistral => "mistral",
+        }
+    }
+}
+
+pub(crate) fn resolve_wire_adapter(
+    model: &str,
+    custom_llm_provider: Option<&str>,
+) -> Result<(String, OcrAdapterKind), Error> {
+    let provider =
+        get_custom_llm_provider(model, custom_llm_provider).unwrap_or(CustomLlmProvider {
+            model,
+            custom_llm_provider: OcrProvider::Mistral.as_str(),
+        });
+    let typed_provider = match provider.custom_llm_provider {
+        "mistral" => OcrProvider::Mistral,
+        value => return Err(Error::InvalidProvider(value.to_string())),
+    };
+    match typed_provider {
+        OcrProvider::Mistral => Ok((provider.model.to_string(), OcrAdapterKind::Mistral)),
+    }
+}

litellm-rust/crates/core/src/ocr/transformation.rs

@@ -1,7 +1,7 @@
 use crate::Error;
 use serde_json::{Map, Value};
 
-use super::types::{OcrRequestData, OcrResponseData};
+use super::types::{LiteLLMOcrResponse, OcrRequestData};
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub enum OcrAuthStrategy {
@@ -49,14 +49,14 @@ pub trait OcrProviderConfig: Sync {
         &self,
         model: &str,
         response_json: Value,
-    ) -> Result<OcrResponseData, Error>;
+    ) -> Result<LiteLLMOcrResponse, Error>;
 
     fn transform_ocr_response_with_params(
         &self,
         model: &str,
         response_json: Value,
         _optional_params: &Map<String, Value>,
-    ) -> Result<OcrResponseData, Error> {
+    ) -> Result<LiteLLMOcrResponse, Error> {
         self.transform_ocr_response(model, response_json)
     }
 

litellm-rust/crates/core/src/ocr/types.rs

@@ -1,6 +1,14 @@
+use std::sync::Arc;
+use std::time::Duration;
+
 use serde::{Deserialize, Serialize};
 use serde_json::{Map, Value};
 
+use super::hooks::{NoopOcrHooks, OcrHooks};
+use super::registry::{OcrAdapterKind, resolve_wire_adapter};
+use crate::Error;
+use crate::constants::OCR_HTTP_TIMEOUT_SECS;
+
 #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
 pub struct OcrRequestData {
     pub data: Value,
@@ -8,31 +16,145 @@ pub struct OcrRequestData {
 }
 
 #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
-pub struct OcrResponseData {
+#[serde(tag = "type")]
+pub enum OcrDocument {
+    #[serde(rename = "document_url")]
+    DocumentUrl {
+        document_url: String,
+        #[serde(flatten)]
+        extra_fields: Map<String, Value>,
+    },
+    #[serde(rename = "image_url")]
+    ImageUrl {
+        image_url: String,
+        #[serde(flatten)]
+        extra_fields: Map<String, Value>,
+    },
+}
+
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "lowercase")]
+pub enum OcrResponseFormat {
+    #[default]
+    Litellm,
+    Native,
+}
+
+#[derive(Clone)]
+pub struct OcrConnection {
+    pub api_key: Option<String>,
+    pub api_base: Option<String>,
+    pub extra_headers: Vec<(String, String)>,
+    pub timeout: Duration,
+}
+
+impl Default for OcrConnection {
+    fn default() -> Self {
+        Self {
+            api_key: None,
+            api_base: None,
+            extra_headers: Vec::new(),
+            timeout: Duration::from_secs(OCR_HTTP_TIMEOUT_SECS),
+        }
+    }
+}
+
+pub struct LiteLLMOcrRequest {
+    pub model: String,
+    pub document: OcrDocument,
+    pub connection: OcrConnection,
+    pub hooks: Arc<dyn OcrHooks>,
+    pub litellm_call_id: Option<String>,
+    pub optional_params: Map<String, Value>,
+    pub(crate) adapter: OcrAdapterKind,
+}
+
+impl LiteLLMOcrRequest {
+    pub fn new(
+        model: String,
+        document: OcrDocument,
+        custom_llm_provider: Option<&str>,
+        optional_params: Map<String, Value>,
+    ) -> Result<Self, Error> {
+        let (model, adapter_kind) = resolve_wire_adapter(&model, custom_llm_provider)?;
+
+        Ok(Self {
+            model,
+            document,
+            connection: OcrConnection::default(),
+            hooks: Arc::new(NoopOcrHooks),
+            litellm_call_id: None,
+            optional_params,
+            adapter: adapter_kind,
+        })
+    }
+
+    pub(crate) fn response_format(
+        &self,
+    ) -> Result<OcrResponseFormat, super::error::OcrRequestError> {
+        self.optional_params
+            .get("req_format")
+            .map(|value| {
+                serde_json::from_value(value.clone())
+                    .map_err(|_| super::error::OcrRequestError::RequestFormat)
+            })
+            .transpose()
+            .map(|format| format.unwrap_or_default())
+    }
+
+    pub fn with_host_hooks(
+        self,
+        hooks: Arc<dyn OcrHooks>,
+        litellm_call_id: Option<String>,
+    ) -> Self {
+        Self {
+            hooks,
+            litellm_call_id,
+            ..self
+        }
+    }
+}
+
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+pub struct LiteLLMOcrResponse {
     pub pages: Vec<Value>,
     pub model: String,
     pub document_annotation: Option<Value>,
     pub usage_info: Option<Value>,
     pub object: String,
+    #[serde(flatten)]
     pub extra_fields: Map<String, Value>,
+    #[serde(skip_serializing_if = "Option::is_none")]
     pub provider_native_response: Option<Value>,
 }
 
-impl OcrResponseData {
+impl LiteLLMOcrResponse {
     pub fn into_json(self) -> Value {
-        let mut response = serde_json::json!({
-            "pages": self.pages,
-            "model": self.model,
-            "document_annotation": self.document_annotation,
-            "usage_info": self.usage_info,
-            "object": self.object,
-        });
-        if let Value::Object(object) = &mut response {
-            object.extend(self.extra_fields);
-            if let Some(native_response) = self.provider_native_response {
-                object.insert("provider_native_response".to_string(), native_response);
-            }
-        }
-        response
+        serde_json::to_value(self).expect("OCR response fields are JSON-compatible")
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use serde_json::json;
+
+    #[test]
+    fn response_serialization_flattens_extra_fields_and_omits_absent_native_response() {
+        let response = LiteLLMOcrResponse {
+            pages: vec![],
+            model: "model".into(),
+            document_annotation: None,
+            usage_info: None,
+            object: "ocr".into(),
+            extra_fields: json!({"provider_field":"kept"})
+                .as_object()
+                .unwrap()
+                .clone(),
+            provider_native_response: None,
+        };
+        let serialized = response.into_json();
+        assert_eq!(serialized["provider_field"], "kept");
+        assert!(serialized.get("provider_native_response").is_none());
     }
 }

litellm-rust/crates/core/src/ocr/wire.rs

@@ -0,0 +1,154 @@
+use crate::ocr::error::OcrRequestError;
+use crate::ocr::error::OcrResponseError;
+use std::time::Duration;
+
+use super::hooks::{OcrDuringCallRequest, OcrPreCallRequest};
+use super::types::{LiteLLMOcrRequest, OcrConnection, OcrDocument};
+use crate::Error;
+use serde::{
+    Deserialize,
+    de::{DeserializeOwned, IntoDeserializer},
+};
+use serde_json::{Map, Value};
+
+#[derive(Debug)]
+pub struct DecodedOcrResponse<T> {
+    pub data: T,
+    pub native: Option<Value>,
+}
+
+#[derive(Deserialize)]
+#[serde(deny_unknown_fields)]
+pub struct OcrWireRequest {
+    pub model: String,
+    pub document: Value,
+    pub api_key: Option<String>,
+    pub api_base: Option<String>,
+    pub custom_llm_provider: Option<String>,
+    pub extra_headers: Option<Map<String, Value>>,
+    #[serde(default)]
+    pub optional_params: Map<String, Value>,
+    pub timeout_seconds: Option<f64>,
+}
+
+pub fn is_supported_request(model: &str, custom_llm_provider: Option<&str>) -> bool {
+    super::registry::resolve_wire_adapter(model, custom_llm_provider).is_ok()
+}
+
+pub fn decode_request(wire: OcrWireRequest) -> Result<LiteLLMOcrRequest, Error> {
+    let document = decode_request_value(wire.document, "document")?;
+    let headers = wire
+        .extra_headers
+        .unwrap_or_default()
+        .into_iter()
+        .map(|(name, value)| {
+            let value = value
+                .as_str()
+                .ok_or_else(|| OcrRequestError::RequestField {
+                    path: format!("extra_headers.{name}"),
+                })?;
+            Ok((name, value.to_string()))
+        })
+        .collect::<Result<Vec<_>, OcrRequestError>>()?;
+    let timeout = wire
+        .timeout_seconds
+        .map(|seconds| {
+            Duration::try_from_secs_f64(seconds).map_err(|_| OcrRequestError::RequestField {
+                path: "timeout_seconds".into(),
+            })
+        })
+        .transpose()?;
+    let defaults = OcrConnection::default();
+    let request = LiteLLMOcrRequest::new(
+        wire.model,
+        document,
+        wire.custom_llm_provider.as_deref(),
+        wire.optional_params,
+    )?;
+    let connection = OcrConnection {
+        api_key: nonblank(wire.api_key),
+        api_base: nonblank(wire.api_base),
+        extra_headers: headers,
+        timeout: timeout.unwrap_or(defaults.timeout),
+    };
+    Ok(LiteLLMOcrRequest {
+        connection,
+        ..request
+    })
+}
+
+fn nonblank(value: Option<String>) -> Option<String> {
+    value
+        .map(|s| s.trim().to_string())
+        .filter(|s| !s.is_empty())
+}
+pub fn decode_request_value<T: DeserializeOwned>(
+    value: Value,
+    prefix: &str,
+) -> Result<T, OcrRequestError> {
+    serde_path_to_error::deserialize(value.into_deserializer()).map_err(|error| {
+        OcrRequestError::RequestField {
+            path: format!("{prefix}.{}", error.path()),
+        }
+    })
+}
+
+pub fn decode_response<T: DeserializeOwned>(
+    bytes: &[u8],
+    native: bool,
+) -> Result<DecodedOcrResponse<T>, OcrResponseError> {
+    let mut deserializer = serde_json::Deserializer::from_slice(bytes);
+    let data = serde_path_to_error::deserialize(&mut deserializer).map_err(|error| {
+        OcrResponseError::ResponseField {
+            path: error.path().to_string(),
+        }
+    })?;
+    deserializer
+        .end()
+        .map_err(|_| OcrResponseError::ResponseField {
+            path: "response".into(),
+        })?;
+    let native = if native {
+        Some(
+            serde_json::from_slice(bytes).map_err(|_| OcrResponseError::ResponseField {
+                path: "response".into(),
+            })?,
+        )
+    } else {
+        None
+    };
+    Ok(DecodedOcrResponse { data, native })
+}
+
+pub fn decode_pre_call_result(
+    original: OcrPreCallRequest,
+    value: Value,
+) -> Result<OcrPreCallRequest, OcrRequestError> {
+    #[derive(Deserialize)]
+    struct Changed {
+        document: OcrDocument,
+        #[serde(default)]
+        optional_params: Map<String, Value>,
+    }
+    let changed: Changed = decode_request_value(value, "guardrail")?;
+    Ok(OcrPreCallRequest {
+        document: changed.document,
+        optional_params: Value::Object(changed.optional_params),
+        ..original
+    })
+}
+
+pub fn decode_during_call_result(
+    original: OcrDuringCallRequest,
+    value: Value,
+) -> Result<OcrDuringCallRequest, OcrRequestError> {
+    #[derive(Deserialize)]
+    struct Changed {
+        body: Value,
+    }
+    let changed: Changed = decode_request_value(value, "guardrail")?;
+    Ok(OcrDuringCallRequest {
+        body: changed.body,
+        ..original
+    })
+}

litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs

@@ -2,7 +2,7 @@ use std::collections::BTreeSet;
 
 use crate::error::{Error, json_type_name};
 use crate::ocr::transformation::{OcrAuthStrategy, OcrProviderConfig, OcrResponseHandling};
-use crate::ocr::types::{OcrRequestData, OcrResponseData};
+use crate::ocr::types::{LiteLLMOcrResponse, OcrRequestData};
 use serde_json::{Map, Value, json};
 
 use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG;
@@ -440,7 +440,7 @@ fn transform_document_intelligence_response(
     model: &str,
     response_json: Value,
     preserve_native_response: bool,
-) -> Result<OcrResponseData, Error> {
+) -> Result<LiteLLMOcrResponse, Error> {
     let response = response_json
         .as_object()
         .ok_or_else(|| Error::InvalidType {
@@ -488,7 +488,7 @@ fn transform_document_intelligence_response(
         })
         .collect();
 
-    Ok(OcrResponseData {
+    Ok(LiteLLMOcrResponse {
         usage_info: Some(json!({
             "pages_processed": pages.len(),
             "doc_size_bytes": null,
@@ -521,7 +521,7 @@ impl OcrProviderConfig for AzureAiOcrConfig {
         &self,
         model: &str,
         response_json: Value,
-    ) -> Result<OcrResponseData, Error> {
+    ) -> Result<LiteLLMOcrResponse, Error> {
         MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json)
     }
 
@@ -599,7 +599,7 @@ impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig {
         &self,
         model: &str,
         response_json: Value,
-    ) -> Result<OcrResponseData, Error> {
+    ) -> Result<LiteLLMOcrResponse, Error> {
         transform_document_intelligence_response(model, response_json, false)
     }
 
@@ -608,7 +608,7 @@ impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig {
         model: &str,
         response_json: Value,
         optional_params: &Map<String, Value>,
-    ) -> Result<OcrResponseData, Error> {
+    ) -> Result<LiteLLMOcrResponse, Error> {
         transform_document_intelligence_response(
             model,
             response_json,
@@ -718,7 +718,7 @@ mod tests {
         })
     }
 
-    fn assert_native_fields_preserved(response: &OcrResponseData, operation: &Value) {
+    fn assert_native_fields_preserved(response: &LiteLLMOcrResponse, operation: &Value) {
         let analyze_result = &operation["analyzeResult"];
 
         assert_eq!(response.extra_fields["content"], analyze_result["content"]);

litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs

@@ -1,6 +1,6 @@
 use crate::error::{Error, json_type_name};
 use crate::ocr::transformation::OcrProviderConfig;
-use crate::ocr::types::{OcrRequestData, OcrResponseData};
+use crate::ocr::types::{LiteLLMOcrResponse, OcrRequestData};
 use serde_json::{Map, Value};
 
 const SUPPORTED_OCR_PARAMS: &[&str] = &[
@@ -107,7 +107,7 @@ impl OcrProviderConfig for MistralOcrConfig {
         &self,
         model: &str,
         response_json: Value,
-    ) -> Result<OcrResponseData, Error> {
+    ) -> Result<LiteLLMOcrResponse, Error> {
         let response_object = response_json
             .as_object()
             .ok_or_else(|| Error::InvalidType {
@@ -128,7 +128,7 @@ impl OcrProviderConfig for MistralOcrConfig {
         let document_annotation = response_object.get("document_annotation").cloned();
         let usage_info = response_object.get("usage_info").cloned();
 
-        Ok(OcrResponseData {
+        Ok(LiteLLMOcrResponse {
             pages,
             model,
             document_annotation,
@@ -178,7 +178,10 @@ pub fn transform_ocr_request(
 }
 
 #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
-pub fn transform_ocr_response(model: &str, response_json: Value) -> Result<OcrResponseData, Error> {
+pub fn transform_ocr_response(
+    model: &str,
+    response_json: Value,
+) -> Result<LiteLLMOcrResponse, Error> {
     MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json)
 }
 

litellm-rust/crates/core/src/providers/reducto/ocr/transformation.rs

@@ -6,7 +6,7 @@ use serde_json::{Map, Value, json};
 
 use crate::error::{Error, json_type_name};
 use crate::ocr::transformation::OcrProviderConfig;
-use crate::ocr::types::{OcrRequestData, OcrResponseData};
+use crate::ocr::types::{LiteLLMOcrResponse, OcrRequestData};
 
 pub const REDUCTO_API_BASE: &str = "https://platform.reducto.ai";
 pub const REDUCTO_API_KEY_ENV: &str = "REDUCTO_API_KEY";
@@ -283,7 +283,7 @@ fn build_pages(result: &Map<String, Value>) -> Vec<Value> {
 pub fn transform_reducto_response(
     model: &str,
     response_json: Value,
-) -> Result<OcrResponseData, Error> {
+) -> Result<LiteLLMOcrResponse, Error> {
     let response = response_json
         .as_object()
         .ok_or_else(|| Error::InvalidType {
@@ -311,7 +311,7 @@ pub fn transform_reducto_response(
         "credits": usage.get("credits").cloned().unwrap_or(Value::Null),
     }));
 
-    Ok(OcrResponseData {
+    Ok(LiteLLMOcrResponse {
         pages: build_pages(result),
         model: model.to_string(),
         document_annotation: None,
@@ -341,7 +341,7 @@ impl OcrProviderConfig for ReductoParseV3Config {
         &self,
         model: &str,
         response_json: Value,
-    ) -> Result<OcrResponseData, Error> {
+    ) -> Result<LiteLLMOcrResponse, Error> {
         transform_reducto_response(model, response_json)
     }
 
@@ -383,7 +383,7 @@ impl OcrProviderConfig for ReductoParseLegacyConfig {
         &self,
         model: &str,
         response_json: Value,
-    ) -> Result<OcrResponseData, Error> {
+    ) -> Result<LiteLLMOcrResponse, Error> {
         transform_reducto_response(model, response_json)
     }
 

litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs

@@ -1,6 +1,6 @@
 use crate::error::{Error, json_type_name};
 use crate::ocr::transformation::OcrProviderConfig;
-use crate::ocr::types::{OcrRequestData, OcrResponseData};
+use crate::ocr::types::{LiteLLMOcrResponse, OcrRequestData};
 use serde_json::{Map, Value, json};
 
 use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG;
@@ -226,7 +226,7 @@ impl OcrProviderConfig for VertexAiOcrConfig {
         &self,
         model: &str,
         response_json: Value,
-    ) -> Result<OcrResponseData, Error> {
+    ) -> Result<LiteLLMOcrResponse, Error> {
         MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json)
     }
 
@@ -301,7 +301,7 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig {
         &self,
         model: &str,
         response_json: Value,
-    ) -> Result<OcrResponseData, Error> {
+    ) -> Result<LiteLLMOcrResponse, Error> {
         let response = response_json
             .as_object()
             .ok_or_else(|| Error::InvalidType {
@@ -339,7 +339,7 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig {
             .get("usage_info")
             .cloned()
             .or_else(|| response.get("usage").cloned());
-        Ok(OcrResponseData {
+        Ok(LiteLLMOcrResponse {
             pages,
             model: object
                 .get("model")

litellm-rust/crates/core/src/url_utils.rs

@@ -0,0 +1,95 @@
+use std::marker::PhantomData;
+
+use thiserror::Error;
+use url::Url;
+
+#[derive(Debug, Error)]
+pub(crate) enum ApiUrlError {
+    #[error("invalid URL: {0}")]
+    Parse(#[from] url::ParseError),
+    #[error("URL cannot be used as a base")]
+    CannotBeBase,
+}
+
+pub(crate) struct Base;
+pub(crate) struct Complete;
+
+pub(crate) struct ApiUrl<State> {
+    url: Url,
+    state: PhantomData<State>,
+}
+
+impl ApiUrl<Base> {
+    pub(crate) fn parse(value: &str) -> Result<Self, ApiUrlError> {
+        Ok(Self {
+            url: Url::parse(value.trim())?,
+            state: PhantomData,
+        })
+    }
+
+    pub(crate) fn complete_path(
+        mut self,
+        target: &[&str],
+    ) -> Result<ApiUrl<Complete>, ApiUrlError> {
+        let existing: Vec<String> = self
+            .url
+            .path_segments()
+            .ok_or(ApiUrlError::CannotBeBase)?
+            .filter(|segment| !segment.is_empty())
+            .map(str::to_string)
+            .collect();
+        let overlap = (0..=existing.len().min(target.len()))
+            .rev()
+            .find(|&length| {
+                existing[existing.len() - length..]
+                    .iter()
+                    .map(String::as_str)
+                    .eq(target[..length].iter().copied())
+            })
+            .unwrap_or(0);
+        self.url
+            .path_segments_mut()
+            .map_err(|()| ApiUrlError::CannotBeBase)?
+            .pop_if_empty()
+            .extend(target[overlap..].iter().copied());
+        Ok(ApiUrl {
+            url: self.url,
+            state: PhantomData,
+        })
+    }
+}
+
+impl ApiUrl<Complete> {
+    pub(crate) fn into_string(self) -> String {
+        self.url.into()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn completion_appends_only_the_missing_path_suffix() {
+        for (base, expected) in [
+            ("https://example.test", "https://example.test/v1/ocr"),
+            ("https://example.test/v1", "https://example.test/v1/ocr"),
+            ("https://example.test/v1/ocr", "https://example.test/v1/ocr"),
+        ] {
+            let actual = ApiUrl::parse(base)
+                .and_then(|url| url.complete_path(&["v1", "ocr"]))
+                .map(|url| url.into_string())
+                .expect("url builds");
+            assert_eq!(actual, expected);
+        }
+    }
+
+    #[test]
+    fn completion_places_paths_before_queries() {
+        let actual = ApiUrl::parse("https://example.test/v1?tenant=a")
+            .and_then(|url| url.complete_path(&["v1", "ocr"]))
+            .map(|url| url.into_string())
+            .expect("url builds");
+        assert_eq!(actual, "https://example.test/v1/ocr?tenant=a");
+    }
+}

litellm-rust/crates/core/tests/ocr.rs

@@ -0,0 +1,227 @@
+use std::sync::{Arc, Mutex};
+
+use serde_json::{Value, json};
+
+use super::OcrClient;
+use super::hooks::{OcrHookFuture, OcrHooks, OcrLogFuture, OcrPreCallRequest};
+use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request};
+use super::wire::{OcrWireRequest, decode_request};
+use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleTiming};
+
+#[test]
+fn request_boundary_selects_mistral_and_rejects_unknown_providers() {
+    let request = OcrWireRequest {
+        model: "mistral/model".into(),
+        document: json!({"type":"document_url","document_url":"https://example.com/doc.pdf"}),
+        api_key: Some("key".into()),
+        api_base: None,
+        custom_llm_provider: None,
+        extra_headers: None,
+        optional_params: json!({"extract_header":true,"unknown":42})
+            .as_object()
+            .unwrap()
+            .clone(),
+        timeout_seconds: None,
+    };
+    assert!(decode_request(request).is_ok());
+    assert!(
+        decode_request(OcrWireRequest {
+            model: "model".into(),
+            document: json!({"type":"document_url","document_url":"https://example.com/doc.pdf"}),
+            api_key: Some("key".into()),
+            api_base: None,
+            custom_llm_provider: Some("unknown".into()),
+            extra_headers: None,
+            optional_params: serde_json::Map::new(),
+            timeout_seconds: None,
+        })
+        .is_err()
+    );
+}
+
+#[tokio::test]
+async fn facade_executes_direct_mistral_once() {
+    let (base, seen, server) = mock_server(vec![MockResponse::json(json!({
+        "pages":[{"index":0,"markdown":"hello","custom":"preserved"}],
+        "usage_info":{"pages_processed":1}
+    }))])
+    .await;
+    let result = perform_ocr(wire_request(
+        "mistral/model",
+        &base,
+        json!({"extract_header":true,"unknown":"ignored"}),
+    ))
+    .await
+    .unwrap();
+    server.await.unwrap();
+    assert_eq!(result.pages[0]["markdown"], "hello");
+    assert_eq!(result.pages[0]["custom"], "preserved");
+    let requests = seen.lock().unwrap();
+    assert_eq!(requests.len(), 1);
+    assert!(requests[0].starts_with("POST /v1/ocr "));
+    assert!(
+        requests[0]
+            .to_ascii_lowercase()
+            .contains("authorization: bearer test-key\r\n")
+    );
+    let body: Value = serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap();
+    assert_eq!(
+        body,
+        json!({
+            "model":"model",
+            "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"},
+            "extract_header":true
+        })
+    );
+}
+
+#[tokio::test]
+async fn facade_retains_native_response_when_requested() {
+    let provider_response = json!({
+        "pages":[{"index":0,"markdown":"hello"}],
+        "usage_info":{"pages_processed":1},
+        "provider_only":"preserved"
+    });
+    let (base, _, server) = mock_server(vec![MockResponse::json(provider_response.clone())]).await;
+    let response = perform_ocr(wire_request(
+        "mistral/model",
+        &base,
+        json!({"req_format":"native"}),
+    ))
+    .await
+    .unwrap();
+
+    server.await.unwrap();
+    assert_eq!(response.provider_native_response, Some(provider_response));
+}
+
+#[tokio::test]
+async fn facade_uses_the_injected_http_client() {
+    let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await;
+    let mut default_headers = reqwest::header::HeaderMap::new();
+    default_headers.insert(
+        "x-transport-owner",
+        reqwest::header::HeaderValue::from_static("host"),
+    );
+    let provider_http = reqwest::Client::builder()
+        .default_headers(default_headers)
+        .build()
+        .unwrap();
+    OcrClient::new(provider_http)
+        .unwrap()
+        .perform(wire_request("mistral/model", &base, json!({})))
+        .await
+        .unwrap();
+    server.await.unwrap();
+    assert!(seen.lock().unwrap()[0].contains("x-transport-owner: host"));
+}
+
+struct RecordingHooks {
+    events: Arc<Mutex<Vec<&'static str>>>,
+    block: bool,
+}
+
+impl OcrHooks for RecordingHooks {
+    fn has_guardrails(&self) -> bool {
+        true
+    }
+
+    fn pre_call(&self, request: OcrPreCallRequest) -> OcrHookFuture<'_, OcrPreCallRequest> {
+        Box::pin(async move {
+            self.events.lock().unwrap().push("pre");
+            if self.block {
+                return Err(crate::Error::InvalidRequest("blocked".into()));
+            }
+            Ok(request)
+        })
+    }
+
+    fn during_call(
+        &self,
+        request: super::hooks::OcrDuringCallRequest,
+    ) -> OcrHookFuture<'_, super::hooks::OcrDuringCallRequest> {
+        Box::pin(async move {
+            self.events.lock().unwrap().push("during");
+            Ok(request)
+        })
+    }
+
+    fn success<'a>(
+        &'a self,
+        _context: &'a CallLifecycleContext,
+        _response: &'a super::LiteLLMOcrResponse,
+        _timing: &'a CallLifecycleTiming,
+    ) -> OcrLogFuture<'a> {
+        Box::pin(async move {
+            self.events.lock().unwrap().push("success");
+        })
+    }
+
+    fn failure<'a>(
+        &'a self,
+        _context: &'a CallLifecycleContext,
+        _error: &'a crate::Error,
+        _timing: &'a CallLifecycleTiming,
+    ) -> OcrLogFuture<'a> {
+        Box::pin(async move {
+            self.events.lock().unwrap().push("failure");
+        })
+    }
+}
+
+#[tokio::test]
+async fn lifecycle_orders_hooks_and_emits_one_success() {
+    let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await;
+    let events = Arc::new(Mutex::new(Vec::new()));
+    let request = wire_request("mistral/model", &base, json!({}));
+    let request = super::LiteLLMOcrRequest {
+        hooks: Arc::new(RecordingHooks {
+            events: events.clone(),
+            block: false,
+        }),
+        ..request
+    };
+    perform_ocr(request).await.unwrap();
+    server.await.unwrap();
+    assert_eq!(*events.lock().unwrap(), ["pre", "during", "success"]);
+    assert_eq!(seen.lock().unwrap().len(), 1);
+}
+
+#[tokio::test]
+async fn lifecycle_blocking_prevents_execution_and_emits_one_failure() {
+    let events = Arc::new(Mutex::new(Vec::new()));
+    let request = wire_request("mistral/model", "http://127.0.0.1:1", json!({}));
+    let request = super::LiteLLMOcrRequest {
+        hooks: Arc::new(RecordingHooks {
+            events: events.clone(),
+            block: true,
+        }),
+        ..request
+    };
+    let error = perform_ocr(request).await.unwrap_err();
+    assert!(matches!(error, crate::Error::InvalidRequest(_)));
+    assert_eq!(*events.lock().unwrap(), ["pre", "failure"]);
+}
+
+#[tokio::test]
+async fn upstream_failure_emits_one_terminal_failure() {
+    let (base, seen, server) = mock_server(vec![MockResponse {
+        status: 500,
+        headers: vec![],
+        body: json!({"error":"failed"}),
+    }])
+    .await;
+    let events = Arc::new(Mutex::new(Vec::new()));
+    let request = wire_request("mistral/model", &base, json!({}));
+    let request = super::LiteLLMOcrRequest {
+        hooks: Arc::new(RecordingHooks {
+            events: events.clone(),
+            block: false,
+        }),
+        ..request
+    };
+    assert!(perform_ocr(request).await.is_err());
+    server.await.unwrap();
+    assert_eq!(*events.lock().unwrap(), ["pre", "during", "failure"]);
+    assert_eq!(seen.lock().unwrap().len(), 1);
+}

litellm-rust/crates/core/tests/ocr/support.rs

@@ -0,0 +1,106 @@
+use std::sync::{Arc, Mutex};
+
+use serde_json::{Value, json};
+use tokio::io::{AsyncReadExt, AsyncWriteExt};
+use tokio::net::TcpListener;
+
+use crate::ocr::wire::{OcrWireRequest, decode_request};
+use crate::ocr::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrClient};
+
+pub(crate) fn ocr_client() -> OcrClient {
+    OcrClient::for_test(reqwest::Client::new())
+}
+
+pub(crate) async fn perform_ocr(
+    request: LiteLLMOcrRequest,
+) -> Result<LiteLLMOcrResponse, crate::Error> {
+    ocr_client().perform(request).await
+}
+
+pub(crate) fn wire_request(model: &str, base: &str, options: Value) -> LiteLLMOcrRequest {
+    decode_request(OcrWireRequest {
+        model: model.into(),
+        document: json!({"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}),
+        api_key: Some("test-key".into()),
+        api_base: Some(base.into()),
+        custom_llm_provider: None,
+        extra_headers: None,
+        optional_params: options.as_object().unwrap().clone(),
+        timeout_seconds: Some(2.0),
+    })
+    .unwrap()
+}
+
+pub(crate) struct MockResponse {
+    pub status: u16,
+    pub headers: Vec<(&'static str, String)>,
+    pub body: Value,
+}
+
+impl MockResponse {
+    pub fn json(body: Value) -> Self {
+        Self {
+            status: 200,
+            headers: vec![],
+            body,
+        }
+    }
+}
+
+pub(crate) async fn mock_server(
+    responses: Vec<MockResponse>,
+) -> (String, Arc<Mutex<Vec<String>>>, tokio::task::JoinHandle<()>) {
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let base = format!("http://{}", listener.local_addr().unwrap());
+    let requests = Arc::new(Mutex::new(Vec::new()));
+    let seen = requests.clone();
+    let server_base = base.clone();
+    let task = tokio::spawn(async move {
+        for response in responses {
+            let (mut socket, _) = listener.accept().await.unwrap();
+            let mut bytes = Vec::new();
+            let mut buffer = [0u8; 4096];
+            let header_end = loop {
+                let n = socket.read(&mut buffer).await.unwrap();
+                assert!(n > 0);
+                bytes.extend_from_slice(&buffer[..n]);
+                if let Some(index) = bytes.windows(4).position(|s| s == b"\r\n\r\n") {
+                    break index + 4;
+                }
+            };
+            let length = String::from_utf8_lossy(&bytes[..header_end])
+                .lines()
+                .find_map(|line| {
+                    let (name, value) = line.split_once(':')?;
+                    name.eq_ignore_ascii_case("content-length")
+                        .then(|| value.trim().parse::<usize>().unwrap())
+                })
+                .unwrap_or(0);
+            while bytes.len() < header_end + length {
+                let n = socket.read(&mut buffer).await.unwrap();
+                assert!(n > 0);
+                bytes.extend_from_slice(&buffer[..n]);
+            }
+            seen.lock()
+                .unwrap()
+                .push(String::from_utf8_lossy(&bytes).into_owned());
+            let body = serde_json::to_vec(&response.body).unwrap();
+            let headers = response
+                .headers
+                .into_iter()
+                .map(|(name, value)| {
+                    format!("{name}: {}\r\n", value.replace("{base}", &server_base))
+                })
+                .collect::<String>();
+            let head = format!(
+                "HTTP/1.1 {} OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n{}\r\n",
+                response.status,
+                body.len(),
+                headers
+            );
+            socket.write_all(head.as_bytes()).await.unwrap();
+            socket.write_all(&body).await.unwrap();
+        }
+    });
+    (base, requests, task)
+}

litellm-rust/crates/core/tests/workspace_crate_allowlist.rs

@@ -1,6 +1,7 @@
-//! Enforcement: the litellm-rust workspace has exactly five crates.
+//! Enforcement: the litellm-rust workspace has exactly six crates.
 //!
-//! `core` (the Rust SDK), `config` (the config-loading boundary),
+//! `core` (the Rust SDK), `token-counter` (standalone input token counting),
+//! `config` (the config-loading boundary),
 //! `ai-gateway` (the HTTP/WebSocket host),
 //! `python-interop` (domain-neutral PyO3 primitives), and `python-bridge` (the
 //! PyO3 cdylib). Adding or removing a crate must be a
@@ -20,6 +21,7 @@ use std::path::{Path, PathBuf};
 /// workspace legitimately gains or loses a crate.
 const EXPECTED_MEMBERS: &[&str] = &[
     "crates/core",
+    "crates/token-counter",
     "crates/config",
     "crates/ai-gateway",
     "crates/python-interop",
@@ -29,6 +31,7 @@ const EXPECTED_MEMBERS: &[&str] = &[
 /// The crate subdirectory names that must exist under `crates/`.
 const EXPECTED_CRATE_DIRS: &[&str] = &[
     "core",
+    "token-counter",
     "config",
     "ai-gateway",
     "python-interop",

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

@@ -24,16 +24,17 @@ trace-parity = [
 futures-util.workspace = true
 tracing = { workspace = true, optional = true }
 litellm-core = { workspace = true, features = ["bedrock-auth"] }
+litellm-token-counter.workspace = true
 litellm-ai-gateway = { workspace = true, default-features = false }
 litellm-python-interop.workspace = true
 pyo3.workspace = true
 pyo3-async-runtimes.workspace = true
 serde.workspace = true
 serde_json.workspace = true
-tokio.workspace = true
+tokio = { workspace = true, features = ["sync"] }
 
 [dev-dependencies]
-criterion = "0.8.2"
+criterion.workspace = true
 tokio-tungstenite.workspace = true
 tracing.workspace = true
 

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

@@ -0,0 +1,2 @@
+/// Concurrent token-count encodes allowed when the core count is unavailable.
+pub(crate) const TOKEN_COUNT_FALLBACK_PARALLELISM: usize = 1;

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

@@ -41,6 +41,7 @@ pub(crate) fn chat_completions_error_to_pyerr(err: Error) -> PyErr {
         | Error::InvalidRequest(_)
         | Error::InvalidType { .. }
         | Error::MissingField(_)
+        | Error::MissingApiKey { .. }
         | Error::Routing(_)
         // Nothing reached the provider, so serving it on Python cannot double
         // bill and is the only way the caller gets an answer at all.

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

@@ -3,7 +3,6 @@ use std::panic::AssertUnwindSafe;
 use std::time::Duration;
 
 use futures_util::FutureExt;
-use litellm_core::error::Error;
 use litellm_python_interop::{Pythonized, panic_to_pyerr, release_gil};
 use pyo3::exceptions::PyRuntimeError;
 use pyo3::prelude::*;
@@ -11,14 +10,15 @@ use serde::Serialize;
 use tokio::runtime::{Handle, Runtime};
 use tokio::time::{self, MissedTickBehavior};
 
-pub(crate) fn run_sync<T, F>(
+pub(crate) fn run_sync<T, E, F>(
     py: Python<'_>,
     future: F,
-    map_error: fn(Error) -> PyErr,
+    map_error: fn(E) -> PyErr,
 ) -> PyResult<Py<PyAny>>
 where
     T: Serialize + Send + 'static,
-    F: Future<Output = Result<T, Error>> + Send + 'static,
+    E: Send + 'static,
+    F: Future<Output = Result<T, E>> + Send + 'static,
 {
     run_sync_on(
         py,
@@ -28,15 +28,16 @@ where
     )
 }
 
-fn run_sync_on<T, F>(
+fn run_sync_on<T, E, F>(
     py: Python<'_>,
     runtime: &Runtime,
     future: F,
-    map_error: fn(Error) -> PyErr,
+    map_error: fn(E) -> PyErr,
 ) -> PyResult<Py<PyAny>>
 where
     T: Serialize + Send + 'static,
-    F: Future<Output = Result<T, Error>> + Send + 'static,
+    E: Send + 'static,
+    F: Future<Output = Result<T, E>> + Send + 'static,
 {
     if Handle::try_current().is_ok() {
         return Err(PyRuntimeError::new_err(
@@ -49,14 +50,15 @@ where
     Pythonized(result).into_pyobject(py).map(Bound::unbind)
 }
 
-pub(crate) fn run_async<T, F>(
+pub(crate) fn run_async<T, E, F>(
     py: Python<'_>,
     future: F,
-    map_error: fn(Error) -> PyErr,
+    map_error: fn(E) -> PyErr,
 ) -> PyResult<Bound<'_, PyAny>>
 where
     T: Serialize + Send + 'static,
-    F: Future<Output = Result<T, Error>> + Send + 'static,
+    E: Send + 'static,
+    F: Future<Output = Result<T, E>> + Send + 'static,
 {
     pyo3_async_runtimes::tokio::future_into_py(py, async move {
         let result = catch_future_panic(future).await?;
@@ -65,7 +67,7 @@ where
     })
 }
 
-fn map_core_result<T>(result: Result<T, Error>, map_error: fn(Error) -> PyErr) -> PyResult<T> {
+fn map_core_result<T, E>(result: Result<T, E>, map_error: fn(E) -> PyErr) -> PyResult<T> {
     match result {
         Ok(value) => Ok(value),
         Err(error) => Err(
@@ -75,9 +77,9 @@ fn map_core_result<T>(result: Result<T, Error>, map_error: fn(Error) -> PyErr) -
     }
 }
 
-async fn catch_future_panic<T, F>(future: F) -> PyResult<Result<T, Error>>
+async fn catch_future_panic<T, E, F>(future: F) -> PyResult<Result<T, E>>
 where
-    F: Future<Output = Result<T, Error>>,
+    F: Future<Output = Result<T, E>>,
 {
     AssertUnwindSafe(future)
         .catch_unwind()
@@ -85,9 +87,9 @@ where
         .map_err(panic_to_pyerr)
 }
 
-async fn wait_for_sync_result<T, F>(future: F) -> PyResult<Result<T, Error>>
+async fn wait_for_sync_result<T, E, F>(future: F) -> PyResult<Result<T, E>>
 where
-    F: Future<Output = Result<T, Error>>,
+    F: Future<Output = Result<T, E>>,
 {
     let future = catch_future_panic(future);
     tokio::pin!(future);
@@ -114,6 +116,7 @@ mod tests {
     use std::thread;
     use std::time::Instant;
 
+    use litellm_core::error::Error;
     use pyo3::panic::PanicException;
     use pyo3::types::{PyDict, PyModule};
     use serde::Serializer;
@@ -237,7 +240,7 @@ mod tests {
 
         let error = runtime.block_on(async {
             Python::attach(|py| {
-                run_sync::<bool, _>(py, async { Ok(true) }, runtime_error)
+                run_sync::<bool, Error, _>(py, async { Ok(true) }, runtime_error)
                     .expect_err("sync route should reject a nested Tokio runtime")
             })
         });
@@ -273,7 +276,7 @@ mod tests {
     fn sync_runner_maps_a_panicked_future() {
         Python::initialize();
         Python::attach(|py| {
-            let error = run_sync::<bool, _>(
+            let error = run_sync::<bool, Error, _>(
                 py,
                 poll_fn(|_| -> Poll<Result<bool, Error>> { panic!("route future panicked") }),
                 runtime_error,
@@ -289,7 +292,7 @@ mod tests {
     fn sync_runner_maps_a_panicked_error_mapper() {
         Python::initialize();
         Python::attach(|py| {
-            let error = run_sync::<bool, _>(
+            let error = run_sync::<bool, Error, _>(
                 py,
                 async { Err(Error::InvalidRequest("invalid".to_string())) },
                 panicking_error_mapper,

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

@@ -1,3 +1,4 @@
+mod constants;
 mod diagnostics;
 mod errors;
 mod execution;
@@ -5,6 +6,7 @@ mod execution;
 mod function_trace;
 mod marshal;
 mod routes;
+mod token_counter;
 
 use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection as RustResponsesWebSocketConnection;
 use pyo3::prelude::*;
@@ -71,6 +73,7 @@ mod _native {
         super::errors::register(module)?;
         super::routes::register(module)?;
         module.add_class::<super::ResponsesWebSocketConnection>()?;
+        super::token_counter::register(module)?;
         super::diagnostics::register(module)
     }
 }
@@ -106,6 +109,7 @@ mod tests {
                 "chat_completions",
                 "achat_completions",
                 "ResponsesWebSocketConnection",
+                "TokenCounter",
                 "gil_stats",
             ];
 

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

@@ -0,0 +1,87 @@
+use std::num::NonZero;
+use std::sync::Arc;
+use std::thread::available_parallelism;
+
+use litellm_python_interop::release_gil;
+use litellm_token_counter::{
+    CountableRequest, Error, InputTokenCount, TokenCounter as CoreTokenCounter,
+};
+use pyo3::exceptions::{PyRuntimeError, PyValueError};
+use pyo3::prelude::*;
+use pyo3::types::PyAny;
+use tokio::sync::Semaphore;
+
+use crate::constants::TOKEN_COUNT_FALLBACK_PARALLELISM;
+use crate::errors::RustBridgeDeclined;
+use crate::execution::run_async;
+
+/// Counts the input tokens of a raw request body off the Python event loop with
+/// the GIL released. Python owns which requests get here and what to do with
+/// the count. At most one encode per core runs at a time; the rest wait in the
+/// async task, where a cancelled Python awaiter drops them before any blocking
+/// work is scheduled.
+#[pyclass(frozen)]
+struct TokenCounter {
+    inner: Arc<CoreTokenCounter>,
+    encode_slots: Arc<Semaphore>,
+}
+
+#[pymethods]
+impl TokenCounter {
+    #[new]
+    fn new(py: Python<'_>, tokenizer_json: &str) -> PyResult<Self> {
+        let inner = release_gil(py, || CoreTokenCounter::from_json(tokenizer_json))
+            .map_err(token_count_error_to_pyerr)?;
+        Ok(Self {
+            inner: Arc::new(inner),
+            encode_slots: Arc::new(Semaphore::new(encode_parallelism())),
+        })
+    }
+
+    fn acount_request<'py>(&self, py: Python<'py>, body: &[u8]) -> PyResult<Bound<'py, PyAny>> {
+        let counter = Arc::clone(&self.inner);
+        let encode_slots = Arc::clone(&self.encode_slots);
+        let body = body.to_vec();
+        run_async(
+            py,
+            async move {
+                let _slot = encode_slots
+                    .acquire_owned()
+                    .await
+                    .map_err(|error| Error::Task(error.to_string()))?;
+                tokio::task::spawn_blocking(move || count_body(&counter, &body))
+                    .await
+                    .map_err(|error| Error::Task(error.to_string()))?
+            },
+            token_count_error_to_pyerr,
+        )
+    }
+}
+
+fn encode_parallelism() -> usize {
+    available_parallelism().map_or(TOKEN_COUNT_FALLBACK_PARALLELISM, NonZero::get)
+}
+
+fn count_body(counter: &CoreTokenCounter, body: &[u8]) -> Result<InputTokenCount, Error> {
+    let request = CountableRequest::parse(body)?;
+    counter.count_request(&request)
+}
+
+fn token_count_error_to_pyerr(error: Error) -> PyErr {
+    let message = error.to_string();
+    match error {
+        Error::Load(_) => PyValueError::new_err(message),
+        Error::RequestParse(_)
+        | Error::MissingInput
+        | Error::FloatText
+        | Error::ContentBlock
+        | Error::ArrayItems
+        | Error::JsonSerialization(_)
+        | Error::JsonUtf8(_) => RustBridgeDeclined::new_err(message),
+        Error::Encode(_) | Error::Task(_) => PyRuntimeError::new_err(message),
+    }
+}
+
+pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
+    module.add_class::<TokenCounter>()
+}

litellm-rust/crates/token-counter/Cargo.toml

@@ -0,0 +1,28 @@
+[package]
+name = "litellm-token-counter"
+version = "0.1.0"
+edition.workspace = true
+license.workspace = true
+repository.workspace = true
+
+[dependencies]
+indexmap = { version = "2.14.0", features = ["serde"] }
+itoa = "1.0"
+serde.workspace = true
+serde_json.workspace = true
+thiserror.workspace = true
+tokenizers = { version = "0.23.1", default-features = false, features = ["onig"] }
+unicode-normalization-alignments = "0.1.12"
+
+[dev-dependencies]
+criterion.workspace = true
+rand.workspace = true
+rstest.workspace = true
+
+[[bench]]
+name = "token_counter"
+harness = false
+
+[[bench]]
+name = "allocations"
+harness = false

litellm-rust/crates/token-counter/benches/allocations.rs

@@ -0,0 +1,128 @@
+use std::alloc::{GlobalAlloc, Layout, System};
+use std::hint::black_box;
+use std::sync::atomic::{AtomicUsize, Ordering};
+
+use litellm_token_counter::{CountableRequest, TokenCounter};
+
+struct CountingAllocator;
+
+static ALLOCATIONS: AtomicUsize = AtomicUsize::new(0);
+static BYTES: AtomicUsize = AtomicUsize::new(0);
+
+unsafe impl GlobalAlloc for CountingAllocator {
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        ALLOCATIONS.fetch_add(1, Ordering::Relaxed);
+        BYTES.fetch_add(layout.size(), Ordering::Relaxed);
+        // SAFETY: This allocator delegates the unchanged layout to `System`.
+        unsafe { System.alloc(layout) }
+    }
+
+    unsafe fn dealloc(&self, pointer: *mut u8, layout: Layout) {
+        // SAFETY: The pointer and layout came from the delegated `System` allocation.
+        unsafe { System.dealloc(pointer, layout) }
+    }
+
+    unsafe fn realloc(&self, pointer: *mut u8, layout: Layout, size: usize) -> *mut u8 {
+        ALLOCATIONS.fetch_add(1, Ordering::Relaxed);
+        BYTES.fetch_add(size, Ordering::Relaxed);
+        // SAFETY: The pointer and layout came from `System`; the new size is unchanged.
+        unsafe { System.realloc(pointer, layout, size) }
+    }
+}
+
+#[global_allocator]
+static ALLOCATOR: CountingAllocator = CountingAllocator;
+
+#[derive(Clone, Copy)]
+struct AllocationCount {
+    allocations: usize,
+    bytes: usize,
+}
+
+impl AllocationCount {
+    fn assert_max(self, label: &str, maximum: Self) {
+        eprintln!(
+            "{label}: {} allocations, {} bytes",
+            self.allocations, self.bytes
+        );
+        assert!(
+            self.allocations <= maximum.allocations,
+            "{label} allocation count exceeded {}",
+            maximum.allocations
+        );
+        assert!(
+            self.bytes <= maximum.bytes,
+            "{label} allocated bytes exceeded {}",
+            maximum.bytes
+        );
+    }
+}
+
+const TOKENIZER_JSON: &str = include_str!(concat!(
+    env!("CARGO_MANIFEST_DIR"),
+    "/../../../litellm/litellm_core_utils/tokenizers/anthropic_tokenizer.json"
+));
+const OBJECT_BODY: &[u8] = br#"{"model":"claude-sonnet-4-5","input":{"text":"caf\u00e9","n":3,"ok":true,"list":[1,"a",{"z":[]}]}}"#;
+const INTEGER_BODY: &[u8] = br#"{"input":[-9223372036854775808,0,18446744073709551615]}"#;
+
+fn measure(operation: impl FnOnce()) -> AllocationCount {
+    ALLOCATIONS.store(0, Ordering::Relaxed);
+    BYTES.store(0, Ordering::Relaxed);
+    operation();
+    AllocationCount {
+        allocations: ALLOCATIONS.load(Ordering::Relaxed),
+        bytes: BYTES.load(Ordering::Relaxed),
+    }
+}
+
+fn main() {
+    measure(|| {
+        black_box(CountableRequest::parse(OBJECT_BODY).expect("object request parses"));
+    })
+    .assert_max(
+        "parse object request",
+        AllocationCount {
+            allocations: 16,
+            bytes: 1_900,
+        },
+    );
+
+    let counter = TokenCounter::from_json(TOKENIZER_JSON).expect("tokenizer loads");
+    let object = CountableRequest::parse(OBJECT_BODY).expect("object request parses");
+    counter
+        .count_request(&object)
+        .expect("object warmup succeeds");
+    measure(|| {
+        black_box(
+            counter
+                .count_request(black_box(&object))
+                .expect("object counts"),
+        );
+    })
+    .assert_max(
+        "count object request",
+        AllocationCount {
+            allocations: 74,
+            bytes: 2_200,
+        },
+    );
+
+    let integers = CountableRequest::parse(INTEGER_BODY).expect("integer request parses");
+    counter
+        .count_request(&integers)
+        .expect("integer warmup succeeds");
+    measure(|| {
+        black_box(
+            counter
+                .count_request(black_box(&integers))
+                .expect("integers count"),
+        );
+    })
+    .assert_max(
+        "count integer list",
+        AllocationCount {
+            allocations: 26,
+            bytes: 1_050,
+        },
+    );
+}

litellm-rust/crates/token-counter/benches/token_counter.rs

@@ -0,0 +1,100 @@
+use std::hint::black_box;
+use std::time::Duration;
+
+use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main};
+use litellm_token_counter::TokenCounter;
+use tokenizers::Tokenizer;
+
+const TOKENIZER_JSON: &str = include_str!(concat!(
+    env!("CARGO_MANIFEST_DIR"),
+    "/../../../litellm/litellm_core_utils/tokenizers/anthropic_tokenizer.json"
+));
+const FULL_CONTEXT_TOKENS: usize = 1_000_000;
+const LARGE_PROMPT_UNIT: &str = "The quick brown fox jumps over the lazy dog. 0123456789\n";
+
+fn full_context_input(tokenizer: &Tokenizer) -> String {
+    let unit_tokens = tokenizer
+        .encode_fast(LARGE_PROMPT_UNIT, false)
+        .expect("reference tokenizer should encode")
+        .len();
+    let input = LARGE_PROMPT_UNIT.repeat(FULL_CONTEXT_TOKENS.div_ceil(unit_tokens));
+    let actual_tokens = tokenizer
+        .encode_fast(input.as_str(), true)
+        .expect("reference tokenizer should encode")
+        .len();
+    assert!((FULL_CONTEXT_TOKENS..FULL_CONTEXT_TOKENS + unit_tokens).contains(&actual_tokens));
+    input
+}
+
+fn inputs(tokenizer: &Tokenizer) -> Vec<(&'static str, String)> {
+    vec![
+        (
+            "ascii_chat",
+            "User: Summarize the benefits of statistical benchmarking.\nAssistant:".repeat(8),
+        ),
+        (
+            "unicode_nfkc",
+            "Ａ quick café résumé: مرحبا 世界 🙂 fi Ⅳ.\n".repeat(16),
+        ),
+        ("large_prompt", LARGE_PROMPT_UNIT.repeat(256)),
+        ("full_context_1m_tokens", full_context_input(tokenizer)),
+    ]
+}
+
+fn token_counter(c: &mut Criterion) {
+    let counter = TokenCounter::from_json(TOKENIZER_JSON).expect("token counter should load");
+    let tokenizer = TOKENIZER_JSON
+        .parse::<Tokenizer>()
+        .expect("reference tokenizer should load");
+    let mut group = c.benchmark_group("anthropic_token_counter");
+
+    for (name, input) in inputs(&tokenizer) {
+        let expected = tokenizer
+            .encode_fast(input.as_str(), true)
+            .expect("reference tokenizer should encode")
+            .len();
+        let actual = counter
+            .count_text(input.as_str())
+            .expect("benchmark path should count");
+        assert_eq!(
+            actual, expected,
+            "benchmark paths should produce the same count"
+        );
+        group.throughput(Throughput::Bytes(input.len() as u64));
+        group.bench_with_input(
+            BenchmarkId::new("byte_level_fast_path", name),
+            &input,
+            |b, input| {
+                b.iter(|| {
+                    counter
+                        .count_text(black_box(input.as_str()))
+                        .expect("fast path should count")
+                })
+            },
+        );
+        group.bench_with_input(
+            BenchmarkId::new("full_encoder", name),
+            &input,
+            |b, input| {
+                b.iter(|| {
+                    tokenizer
+                        .encode_fast(black_box(input.as_str()), true)
+                        .expect("reference tokenizer should encode")
+                        .len()
+                })
+            },
+        );
+    }
+
+    group.finish();
+}
+
+criterion_group! {
+    name = benches;
+    config = Criterion::default()
+        .sample_size(20)
+        .warm_up_time(Duration::from_secs(1))
+        .measurement_time(Duration::from_secs(4));
+    targets = token_counter
+}
+criterion_main!(benches);

litellm-rust/crates/token-counter/src/byte_level.rs

@@ -0,0 +1,650 @@
+//! Exact token counting for a supported tokenizer configuration: optional
+//! NFKC normalization, `ByteLevel` pre-tokenization with the GPT-2 split regex,
+//! and no post-processing. A scanner reproduces the regex's piece boundaries
+//! and hands each piece to the tokenizer's model. Unsupported configurations
+//! and added-token inputs fall back to the full encoder.
+
+use std::borrow::Cow;
+use std::iter;
+
+use tokenizers::normalizers::NormalizerWrapper;
+use tokenizers::pre_tokenizers::PreTokenizerWrapper;
+use tokenizers::{Model, Tokenizer};
+use unicode_normalization_alignments::{IsNormalized, UnicodeNormalization, is_nfkc_quick};
+
+use super::unicode_classes::UnicodeClasses;
+
+const CONTRACTIONS: [&str; 7] = ["'s", "'t", "'re", "'ve", "'m", "'ll", "'d"];
+
+pub(super) struct ByteLevelCounter {
+    nfkc: bool,
+    normalized_added_tokens: Vec<String>,
+    unicode_classes: &'static UnicodeClasses,
+}
+
+impl ByteLevelCounter {
+    pub(super) fn detect(tokenizer: &Tokenizer) -> Option<Self> {
+        let nfkc = match tokenizer.get_normalizer() {
+            None => false,
+            Some(NormalizerWrapper::NFKC(_)) => true,
+            Some(_) => return None,
+        };
+        let Some(PreTokenizerWrapper::ByteLevel(byte_level)) = tokenizer.get_pre_tokenizer() else {
+            return None;
+        };
+        let plain = !byte_level.add_prefix_space
+            && byte_level.use_regex
+            && tokenizer.get_post_processor().is_none()
+            && tokenizer.get_truncation().is_none()
+            && tokenizer.get_padding().is_none();
+        if !plain {
+            return None;
+        }
+        let vocabulary = tokenizer.get_added_vocabulary();
+        let normalized_added_tokens = vocabulary
+            .get_vocab()
+            .iter()
+            .filter_map(|(original, id)| {
+                vocabulary
+                    .simple_id_to_token(*id)
+                    .filter(|normalized| normalized != original)
+            })
+            .collect();
+        Some(Self {
+            nfkc,
+            normalized_added_tokens,
+            unicode_classes: UnicodeClasses::get()?,
+        })
+    }
+
+    /// `None` when the text contains an added token or the model rejects a
+    /// piece; the caller then runs the full encoder.
+    pub(super) fn count(&self, tokenizer: &Tokenizer, text: &str) -> Option<usize> {
+        let normalized = self.normalize(text);
+        let added_tokens = tokenizer.get_added_vocabulary().get_vocab();
+        if added_tokens
+            .keys()
+            .chain(self.normalized_added_tokens.iter())
+            .any(|token| text.contains(token.as_str()) || normalized.contains(token.as_str()))
+        {
+            return None;
+        }
+        let model = tokenizer.get_model();
+        let mapped: String = normalized.bytes().map(byte_char).collect();
+        pieces(&normalized, self.unicode_classes)
+            .try_fold((0, 0), |(start, total), piece| {
+                let end = start + mapped_len(piece);
+                let tokens = model.tokenize(&mapped[start..end]).ok()?;
+                Some((end, total + tokens.len()))
+            })
+            .map(|(_, total)| total)
+    }
+
+    /// Same crate and Unicode tables as `NormalizedString::nfkc`, so the
+    /// result is what the full encoder would have tokenized.
+    fn normalize<'a>(&self, text: &'a str) -> Cow<'a, str> {
+        if !self.nfkc || text.is_ascii() || is_nfkc_quick(text.chars()) == IsNormalized::Yes {
+            return Cow::Borrowed(text);
+        }
+        Cow::Owned(text.nfkc().map(|(character, _)| character).collect())
+    }
+}
+
+/// GPT-2 `bytes_to_unicode`: printable Latin-1 bytes map to themselves, the
+/// rest to U+0100 onwards in byte order.
+fn byte_char(byte: u8) -> char {
+    let code = match byte {
+        0x21..=0x7E | 0xA1..=0xAC | 0xAE..=0xFF => u32::from(byte),
+        0x00..=0x20 => 0x100 + u32::from(byte),
+        0x7F..=0xA0 => 0x121 + u32::from(byte - 0x7F),
+        0xAD => 0x143,
+    };
+    char::from_u32(code).unwrap_or(char::REPLACEMENT_CHARACTER)
+}
+
+fn mapped_len(piece: &str) -> usize {
+    piece.len()
+        + piece
+            .bytes()
+            .filter(|byte| !byte.is_ascii_graphic())
+            .count()
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+enum Class {
+    Letter,
+    Number,
+    Space,
+    Other,
+}
+
+fn class(character: char, unicode_classes: &UnicodeClasses) -> Class {
+    match character {
+        'A'..='Z' | 'a'..='z' => Class::Letter,
+        '0'..='9' => Class::Number,
+        '\t'..='\r' | ' ' => Class::Space,
+        _ if character.is_ascii() => Class::Other,
+        _ if unicode_classes.is_letter(character) => Class::Letter,
+        _ if unicode_classes.is_number(character) => Class::Number,
+        _ if unicode_classes.is_space(character) => Class::Space,
+        _ => Class::Other,
+    }
+}
+
+/// The regex matches every character, so the pieces tile the text.
+fn pieces<'a>(
+    text: &'a str,
+    unicode_classes: &'static UnicodeClasses,
+) -> impl Iterator<Item = &'a str> {
+    iter::successors(split_piece(text, unicode_classes), move |(_, rest)| {
+        split_piece(rest, unicode_classes)
+    })
+    .map(|(piece, _)| piece)
+}
+
+fn split_piece<'a>(text: &'a str, unicode_classes: &UnicodeClasses) -> Option<(&'a str, &'a str)> {
+    let first = text.chars().next()?;
+    Some(text.split_at(piece_len(text, first, unicode_classes)))
+}
+
+fn piece_len(text: &str, first: char, unicode_classes: &UnicodeClasses) -> usize {
+    if let Some(contraction) = CONTRACTIONS.iter().find(|word| text.starts_with(**word)) {
+        return contraction.len();
+    }
+    let first_class = class(first, unicode_classes);
+    if first_class != Class::Space {
+        return run_len(text, first_class, unicode_classes);
+    }
+    if first != ' ' {
+        return space_run_len(text, unicode_classes);
+    }
+    let after_space = &text[1..];
+    match after_space
+        .chars()
+        .next()
+        .map(|character| class(character, unicode_classes))
+    {
+        None | Some(Class::Space) => space_run_len(text, unicode_classes),
+        Some(run_class) => 1 + run_len(after_space, run_class, unicode_classes),
+    }
+}
+
+fn run_len(text: &str, run_class: Class, unicode_classes: &UnicodeClasses) -> usize {
+    text.char_indices()
+        .find(|(_, character)| class(*character, unicode_classes) != run_class)
+        .map_or(text.len(), |(index, _)| index)
+}
+
+/// `\s+(?!\S)|\s+`: whitespace followed by a non-space leaves its last
+/// character to start the next piece (` ?` on the following alternatives).
+fn space_run_len(text: &str, unicode_classes: &UnicodeClasses) -> usize {
+    let run = run_len(text, Class::Space, unicode_classes);
+    if run == text.len() {
+        return run;
+    }
+    let last = text[..run].chars().next_back().map_or(0, char::len_utf8);
+    match run - last {
+        0 => run,
+        shorter => shorter,
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use rand::rngs::StdRng;
+    use rand::seq::SliceRandom;
+    use rand::{Rng, SeedableRng};
+    use rstest::{fixture, rstest};
+    use tokenizers::normalizers::NFKC;
+    use tokenizers::pre_tokenizers::byte_level::ByteLevel;
+    use tokenizers::utils::SysRegex;
+    use tokenizers::{
+        NormalizedString, Normalizer, OffsetReferential, OffsetType, PreTokenizedString,
+        PreTokenizer,
+    };
+
+    use super::*;
+
+    #[fixture]
+    fn anthropic_tokenizer() -> Tokenizer {
+        let path = concat!(
+            env!("CARGO_MANIFEST_DIR"),
+            "/../../../litellm/litellm_core_utils/tokenizers/anthropic_tokenizer.json"
+        );
+        std::fs::read_to_string(path)
+            .expect("anthropic tokenizer json is in the repo")
+            .parse()
+            .expect("anthropic tokenizer loads")
+    }
+
+    fn reference_count(tokenizer: &Tokenizer, text: &str) -> usize {
+        tokenizer.encode_fast(text, true).expect("encode").len()
+    }
+
+    fn byte_level_counter(nfkc: bool) -> ByteLevelCounter {
+        ByteLevelCounter {
+            nfkc,
+            normalized_added_tokens: Vec::new(),
+            unicode_classes: UnicodeClasses::get().expect("Oniguruma exposes Unicode classes"),
+        }
+    }
+
+    const ALPHABET: &[&str] = &[
+        "a",
+        "Z",
+        "e",
+        "s",
+        "t",
+        "d",
+        "m",
+        "'",
+        "'s",
+        "'re",
+        "'ll",
+        "'S",
+        "0",
+        "9",
+        " ",
+        "  ",
+        "\t",
+        "\n",
+        "\r\n",
+        "\u{b}",
+        ".",
+        ",",
+        "!",
+        "-",
+        "(",
+        "\"",
+        "\u{a0}",
+        "\u{85}",
+        "\u{2028}",
+        "\u{3000}",
+        "\u{200b}",
+        "\u{200d}",
+        "é",
+        "e\u{301}",
+        "ß",
+        "漢",
+        "字",
+        "ع",
+        "३",
+        "½",
+        "Ⅳ",
+        "🙂",
+        "👍🏽",
+        "Ａ",
+        "fi",
+        "㍿",
+        "㋿",
+        "ꟲ",
+        "𐞁",
+        "a\u{30a}",
+        "\u{1e0b}\u{323}",
+        "<",
+        ">",
+        "EOT",
+        "<EOT>",
+        "<META_START>",
+    ];
+
+    fn random_text(rng: &mut StdRng) -> String {
+        let pieces = rng.gen_range(0..40);
+        (0..pieces)
+            .map(|_| *ALPHABET.choose(rng).expect("alphabet is not empty"))
+            .collect()
+    }
+
+    #[rstest]
+    #[case::plain_text("Hello, how are you today?", true)]
+    #[case::added_token("stop <EOT> here", false)]
+    #[case::normalized_added_token("stop ＜ＥＯＴ＞ here", false)]
+    fn anthropic_tokenizer_takes_the_fast_path(
+        anthropic_tokenizer: Tokenizer,
+        #[case] text: &str,
+        #[case] supported: bool,
+    ) {
+        let fast =
+            ByteLevelCounter::detect(&anthropic_tokenizer).expect("anthropic shape is supported");
+        assert!(fast.nfkc);
+        let count = fast.count(&anthropic_tokenizer, text);
+        if supported {
+            assert_eq!(count, Some(reference_count(&anthropic_tokenizer, text)));
+        } else {
+            assert_eq!(count, None);
+        }
+    }
+
+    #[rstest]
+    fn counts_match_the_full_encoder(anthropic_tokenizer: Tokenizer) {
+        let fast = ByteLevelCounter::detect(&anthropic_tokenizer).expect("supported");
+        let mut rng = StdRng::seed_from_u64(2026);
+        for _ in 0..4000 {
+            let text = random_text(&mut rng).replace('<', "(");
+            let expected = reference_count(&anthropic_tokenizer, &text);
+            assert_eq!(
+                fast.count(&anthropic_tokenizer, &text),
+                Some(expected),
+                "text {text:?}"
+            );
+        }
+    }
+
+    #[rstest]
+    fn nfkc_matches_the_tokenizer_normalizer_for_every_scalar_value() {
+        let fast = byte_level_counter(true);
+        let mut text = String::new();
+        for character in (0..=0x10FFFFu32).filter_map(char::from_u32) {
+            text.clear();
+            text.push(character);
+            let mut expected = NormalizedString::from(text.as_str());
+            NFKC.normalize(&mut expected).expect("nfkc");
+            assert_eq!(
+                fast.normalize(&text),
+                expected.get(),
+                "U+{:04X}",
+                u32::from(character)
+            );
+        }
+    }
+
+    #[rstest]
+    fn nfkc_matches_the_tokenizer_normalizer_on_random_texts() {
+        let fast = byte_level_counter(true);
+        let mut rng = StdRng::seed_from_u64(11);
+        for _ in 0..4000 {
+            let text = random_text(&mut rng);
+            let mut expected = NormalizedString::from(text.as_str());
+            NFKC.normalize(&mut expected).expect("nfkc");
+            assert_eq!(fast.normalize(&text), expected.get(), "text {text:?}");
+        }
+    }
+
+    #[rstest]
+    fn pieces_match_the_byte_level_pre_tokenizer() {
+        let byte_level = ByteLevel::new(false, true, true);
+        let mut rng = StdRng::seed_from_u64(7);
+        for _ in 0..4000 {
+            let text = byte_level_counter(true)
+                .normalize(&random_text(&mut rng))
+                .into_owned();
+            let mut pre_tokenized = PreTokenizedString::from(text.as_str());
+            byte_level
+                .pre_tokenize(&mut pre_tokenized)
+                .expect("pre-tokenize");
+            let expected: Vec<(String, (usize, usize))> = pre_tokenized
+                .get_splits(OffsetReferential::Original, OffsetType::Byte)
+                .into_iter()
+                .map(|(mapped, offsets, _)| (mapped.to_string(), offsets))
+                .collect();
+            let actual: Vec<(String, (usize, usize))> = pieces(
+                &text,
+                UnicodeClasses::get().expect("Oniguruma exposes Unicode classes"),
+            )
+            .map(|piece| {
+                let start = piece.as_ptr() as usize - text.as_ptr() as usize;
+                let mapped: String = piece.bytes().map(byte_char).collect();
+                (mapped, (start, start + piece.len()))
+            })
+            .collect();
+            assert_eq!(actual, expected, "text {text:?}");
+        }
+    }
+
+    #[rstest]
+    fn byte_chars_match_the_byte_level_alphabet() {
+        let byte_level = ByteLevel::new(false, false, false);
+        let characters: Vec<char> = (0..=0x10FFFFu32).filter_map(char::from_u32).collect();
+        for chunk in characters.chunks(1024) {
+            let text: String = chunk.iter().collect();
+            let mut pre_tokenized = PreTokenizedString::from(text.as_str());
+            byte_level
+                .pre_tokenize(&mut pre_tokenized)
+                .expect("pre-tokenize");
+            let expected: String = pre_tokenized
+                .get_splits(OffsetReferential::Original, OffsetType::Byte)
+                .into_iter()
+                .map(|(mapped, _, _)| mapped)
+                .collect();
+            let actual: String = text.bytes().map(byte_char).collect();
+            assert_eq!(actual.len(), mapped_len(&text));
+            assert_eq!(
+                actual,
+                expected,
+                "chunk starting at U+{:04X}",
+                u32::from(chunk[0])
+            );
+        }
+    }
+
+    #[rstest]
+    fn classes_match_oniguruma() {
+        let unicode_classes = UnicodeClasses::get().expect("Oniguruma exposes Unicode classes");
+        let letter = SysRegex::new(r"\p{L}").expect("regex");
+        let number = SysRegex::new(r"\p{N}").expect("regex");
+        let space = SysRegex::new(r"\s").expect("regex");
+        let whole =
+            |regex: &SysRegex, text: &str| regex.find_iter(text).next() == Some((0, text.len()));
+        let mut text = String::new();
+        for character in (0..=0x10FFFFu32).filter_map(char::from_u32) {
+            text.clear();
+            text.push(character);
+            let expected = if whole(&letter, &text) {
+                Class::Letter
+            } else if whole(&number, &text) {
+                Class::Number
+            } else if whole(&space, &text) {
+                Class::Space
+            } else {
+                Class::Other
+            };
+            assert_eq!(
+                class(character, unicode_classes),
+                expected,
+                "U+{:04X}",
+                u32::from(character)
+            );
+        }
+    }
+
+    #[rstest]
+    #[case("prefix")]
+    #[case("regex")]
+    #[case("normalizer")]
+    #[case("no_pre_tokenizer")]
+    #[case("other_pre_tokenizer")]
+    #[case("post_processor")]
+    #[case("truncation")]
+    #[case("padding")]
+    fn other_tokenizer_shapes_are_declined(
+        mut anthropic_tokenizer: Tokenizer,
+        #[case] shape: &str,
+    ) {
+        use tokenizers::{PaddingParams, PaddingStrategy, TruncationParams};
+        match shape {
+            "prefix" => {
+                anthropic_tokenizer.with_pre_tokenizer(Some(ByteLevel::new(true, true, true)));
+            }
+            "regex" => {
+                anthropic_tokenizer.with_pre_tokenizer(Some(ByteLevel::new(false, true, false)));
+            }
+            "normalizer" => {
+                anthropic_tokenizer
+                    .with_normalizer(Some(tokenizers::normalizers::Lowercase))
+                    .expect("normalizer");
+            }
+            "no_pre_tokenizer" => {
+                anthropic_tokenizer.with_pre_tokenizer(None::<PreTokenizerWrapper>);
+            }
+            "other_pre_tokenizer" => {
+                anthropic_tokenizer
+                    .with_pre_tokenizer(Some(tokenizers::pre_tokenizers::whitespace::Whitespace));
+            }
+            "post_processor" => {
+                anthropic_tokenizer.with_post_processor(Some(ByteLevel::default()));
+            }
+            "truncation" => {
+                anthropic_tokenizer
+                    .with_truncation(Some(TruncationParams {
+                        max_length: 2,
+                        ..Default::default()
+                    }))
+                    .expect("truncation");
+            }
+            "padding" => {
+                anthropic_tokenizer.with_padding(Some(PaddingParams {
+                    strategy: PaddingStrategy::Fixed(32),
+                    ..Default::default()
+                }));
+            }
+            _ => unreachable!(),
+        }
+        assert!(ByteLevelCounter::detect(&anthropic_tokenizer).is_none());
+        let counter = crate::TokenCounter::from_json(
+            &anthropic_tokenizer.to_string(false).expect("serialize"),
+        )
+        .expect("load");
+        for text in ["", "Hello WORLD!  ＡＢ fi Ⅳ", "<EOT> stop"] {
+            assert_eq!(
+                counter.count_text(text).expect("count"),
+                reference_count(&anthropic_tokenizer, text)
+            );
+        }
+    }
+
+    #[rstest]
+    #[case(false)]
+    #[case(true)]
+    fn arbitrary_unicode_and_long_inputs_use_fast_path(
+        mut anthropic_tokenizer: Tokenizer,
+        #[case] nfkc: bool,
+    ) {
+        if !nfkc {
+            anthropic_tokenizer
+                .with_normalizer(None::<NormalizerWrapper>)
+                .expect("normalizer");
+        }
+        let fast = ByteLevelCounter::detect(&anthropic_tokenizer).expect("supported");
+        let mut rng = StdRng::seed_from_u64(314159);
+        for _ in 0..1000 {
+            let text: String = (0..64)
+                .filter_map(|_| char::from_u32(rng.gen_range(0..=0x10ffff)))
+                .collect();
+            assert_eq!(
+                fast.count(&anthropic_tokenizer, &text),
+                Some(reference_count(&anthropic_tokenizer, &text)),
+                "text {text:?}"
+            );
+        }
+        for text in [
+            "",
+            "'s't're've'm'll'd'S'RE",
+            "  a \t\r\n b\u{85}\u{a0}c  ",
+            "\0é漢🙂",
+            "a\u{30a}\u{301}",
+            "ＡfiⅣ",
+        ] {
+            let text = text.repeat(2048);
+            assert_eq!(
+                fast.count(&anthropic_tokenizer, &text),
+                Some(reference_count(&anthropic_tokenizer, &text))
+            );
+        }
+    }
+
+    #[rstest]
+    #[case(false, false, false, false)]
+    #[case(true, false, false, false)]
+    #[case(false, true, false, false)]
+    #[case(false, false, true, false)]
+    #[case(false, false, false, true)]
+    fn added_token_options_fall_back(
+        mut anthropic_tokenizer: Tokenizer,
+        #[case] special: bool,
+        #[case] single_word: bool,
+        #[case] lstrip: bool,
+        #[case] rstrip: bool,
+    ) {
+        anthropic_tokenizer
+            .add_tokens([tokenizers::AddedToken::from("custom token", special)
+                .single_word(single_word)
+                .lstrip(lstrip)
+                .rstrip(rstrip)])
+            .expect("add token");
+        let fast = ByteLevelCounter::detect(&anthropic_tokenizer).expect("supported");
+        let counter = crate::TokenCounter::from_json(
+            &anthropic_tokenizer.to_string(false).expect("serialize"),
+        )
+        .expect("load");
+        for text in [
+            "custom token",
+            "a custom token b",
+            "acustom tokenb",
+            "  custom token  ",
+        ] {
+            assert_eq!(fast.count(&anthropic_tokenizer, text), None);
+            assert_eq!(
+                counter.count_text(text).expect("count"),
+                reference_count(&anthropic_tokenizer, text)
+            );
+        }
+    }
+
+    #[test]
+    fn model_errors_reach_public_caller() {
+        let mut tokenizer = Tokenizer::new(tokenizers::models::wordpiece::WordPiece::default());
+        tokenizer.with_pre_tokenizer(Some(ByteLevel::new(false, true, true)));
+        let fast = ByteLevelCounter::detect(&tokenizer).expect("supported");
+        assert_eq!(fast.count(&tokenizer, "hello"), None);
+        assert!(tokenizer.encode_fast("hello", true).is_err());
+        let counter =
+            crate::TokenCounter::from_json(&tokenizer.to_string(false).expect("serialize"))
+                .expect("load");
+        assert!(matches!(
+            counter.count_text("hello"),
+            Err(crate::Error::Encode(_))
+        ));
+    }
+
+    #[rstest]
+    fn shared_counter_matches_encoder_across_threads(anthropic_tokenizer: Tokenizer) {
+        let counter = crate::TokenCounter::from_json(
+            &anthropic_tokenizer.to_string(false).expect("serialize"),
+        )
+        .expect("load");
+        let inputs = [
+            "hello world",
+            "ＡＢ fi\n漢字🙂",
+            " <EOT> stop",
+            "\t  're \r\n",
+        ];
+        let expected = inputs.map(|text| reference_count(&anthropic_tokenizer, text));
+        std::thread::scope(|scope| {
+            for _ in 0..8 {
+                let counter = &counter;
+                scope.spawn(move || {
+                    for _ in 0..100 {
+                        for (text, count) in inputs.iter().zip(expected) {
+                            assert_eq!(counter.count_text(text).expect("count"), count);
+                        }
+                    }
+                });
+            }
+        });
+    }
+
+    #[rstest]
+    fn normalized_added_token_spelling_declines_fast_path(mut anthropic_tokenizer: Tokenizer) {
+        anthropic_tokenizer
+            .add_tokens([tokenizers::AddedToken::from("ＡＢＣＤ　ＥＦＧＨ", false)])
+            .expect("add token");
+        let fast = ByteLevelCounter::detect(&anthropic_tokenizer).expect("supported");
+        assert_eq!(reference_count(&anthropic_tokenizer, "ABCD EFGH"), 1);
+        assert_eq!(fast.count(&anthropic_tokenizer, "ABCD EFGH"), None);
+        let counter = crate::TokenCounter::from_json(
+            &anthropic_tokenizer.to_string(false).expect("serialize"),
+        )
+        .expect("load");
+        assert_eq!(counter.count_text("ABCD EFGH").expect("count"), 1);
+    }
+}

litellm-rust/crates/token-counter/src/counter.rs

@@ -0,0 +1,194 @@
+use serde::Serialize;
+
+use crate::Error;
+use crate::byte_level::ByteLevelCounter;
+use crate::python_json;
+use crate::tools::format_function_definitions;
+use crate::types::{
+    ContentBlock, ContentItem, CountableRequest, Message, MessageContent, TextValue, ToolChoice,
+    ToolDefinition,
+};
+
+const TOKENS_PER_MESSAGE: usize = 3;
+const TOKENS_PER_NAME: usize = 1;
+const REPLY_PRIMING_TOKENS: usize = 3;
+const TOOL_DEFINITIONS_TOKENS: usize = 9;
+const TOOLS_WITH_SYSTEM_MESSAGE_DISCOUNT: usize = 4;
+const TOOL_CHOICE_NONE_TOKENS: usize = 1;
+const NAMED_TOOL_CHOICE_TOKENS: usize = 7;
+
+#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
+pub struct InputTokenCount {
+    pub model: Option<String>,
+    pub input_tokens: usize,
+}
+
+/// A loaded HuggingFace tokenizer plus the message accounting Python applies on
+/// top of it. Encoding is CPU-bound and synchronous; hosts run it off their
+/// event loop.
+pub struct TokenCounter {
+    tokenizer: tokenizers::Tokenizer,
+    byte_level: Option<ByteLevelCounter>,
+}
+
+impl TokenCounter {
+    /// Load a HuggingFace `tokenizer.json` document. The host reads the file.
+    pub fn from_json(tokenizer_json: &str) -> Result<Self, Error> {
+        let tokenizer = tokenizer_json
+            .parse::<tokenizers::Tokenizer>()
+            .map_err(Error::Load)?;
+        let byte_level = ByteLevelCounter::detect(&tokenizer);
+        Ok(Self {
+            tokenizer,
+            byte_level,
+        })
+    }
+
+    pub fn count_text(&self, text: &str) -> Result<usize, Error> {
+        if let Some(count) = self
+            .byte_level
+            .as_ref()
+            .and_then(|counter| counter.count(&self.tokenizer, text))
+        {
+            return Ok(count);
+        }
+        self.tokenizer
+            .encode_fast(text, true)
+            .map(|encoding| encoding.len())
+            .map_err(Error::Encode)
+    }
+
+    /// Mirrors the host's key precedence: `messages`, then `prompt`, then
+    /// `input`, then `query` plus `documents`.
+    pub fn count_request(&self, request: &CountableRequest) -> Result<InputTokenCount, Error> {
+        let input_tokens = if let Some(messages) = &request.messages {
+            self.count_messages(request, messages)?
+        } else if let Some(prompt) = &request.prompt {
+            self.count_text_value(prompt)?
+        } else if let Some(input) = &request.input {
+            self.count_text_value(input)?
+        } else if request.query.is_some() || request.documents.is_some() {
+            self.count_optional_text_value(request.query.as_ref())?
+                + self.count_optional_text_value(request.documents.as_ref())?
+        } else {
+            return Err(Error::MissingInput);
+        };
+        Ok(InputTokenCount {
+            model: request.model.clone(),
+            input_tokens,
+        })
+    }
+
+    fn count_messages(
+        &self,
+        request: &CountableRequest,
+        messages: &[Message],
+    ) -> Result<usize, Error> {
+        let message_tokens = messages
+            .iter()
+            .map(|message| self.count_message(message))
+            .sum::<Result<usize, _>>()?;
+        let includes_system_message = messages
+            .iter()
+            .any(|message| message.role.as_deref() == Some("system"));
+        let extra_tokens = self.count_extra(
+            request.tools.as_deref().unwrap_or_default(),
+            request.tool_choice.as_ref(),
+            includes_system_message,
+        )?;
+        Ok(message_tokens + extra_tokens)
+    }
+
+    fn count_optional_text_value(&self, value: Option<&TextValue>) -> Result<usize, Error> {
+        value.map_or(Ok(0), |value| self.count_text_value(value))
+    }
+
+    /// `str()` for scalars, `json.dumps()` for objects, lists flattened, nulls
+    /// skipped. Floats are declined because Python's `repr` and Rust's float
+    /// formatting disagree on exponents.
+    fn count_text_value(&self, value: &TextValue) -> Result<usize, Error> {
+        match value {
+            TextValue::Null => Ok(0),
+            TextValue::Bool(true) => self.count_text("True"),
+            TextValue::Bool(false) => self.count_text("False"),
+            TextValue::Number(number) => match (number.as_i64(), number.as_u64()) {
+                (Some(number), _) => self.count_text(itoa::Buffer::new().format(number)),
+                (_, Some(number)) => self.count_text(itoa::Buffer::new().format(number)),
+                _ => Err(Error::FloatText),
+            },
+            TextValue::Text(text) => self.count_text(text),
+            TextValue::List(items) => items
+                .iter()
+                .map(|item| self.count_text_value(item))
+                .sum::<Result<usize, _>>(),
+            TextValue::Object(_) => self.count_text(&python_json::dumps(value)?),
+        }
+    }
+
+    fn count_message(&self, message: &Message) -> Result<usize, Error> {
+        let role_tokens = match &message.role {
+            Some(role) => self.count_text(role)?,
+            None => 0,
+        };
+        let name_tokens = match &message.name {
+            Some(name) => self.count_text(name)? + TOKENS_PER_NAME,
+            None => 0,
+        };
+        let content_tokens = match &message.content {
+            Some(MessageContent::Text(text)) => self.count_text(text)?,
+            Some(MessageContent::Blocks(items)) => items
+                .iter()
+                .map(|item| self.count_content_item(item))
+                .sum::<Result<usize, _>>()?,
+            None => 0,
+        };
+        Ok(TOKENS_PER_MESSAGE + role_tokens + name_tokens + content_tokens)
+    }
+
+    fn count_content_item(&self, item: &ContentItem) -> Result<usize, Error> {
+        match item {
+            ContentItem::Text(text) => self.count_text(text),
+            ContentItem::Block(ContentBlock::Text { text }) => self.count_text(text),
+            ContentItem::Block(ContentBlock::Thinking { thinking }) => {
+                if thinking.is_empty() {
+                    return Ok(0);
+                }
+                self.count_text(thinking)
+            }
+            ContentItem::Block(ContentBlock::ToolReference { tool_name }) => {
+                match tool_name.as_deref().filter(|name| !name.is_empty()) {
+                    Some(name) => self.count_text(name),
+                    None => Ok(0),
+                }
+            }
+            ContentItem::Block(ContentBlock::Unsupported) => Err(Error::ContentBlock),
+        }
+    }
+
+    fn count_extra(
+        &self,
+        tools: &[ToolDefinition],
+        tool_choice: Option<&ToolChoice>,
+        includes_system_message: bool,
+    ) -> Result<usize, Error> {
+        let tool_tokens = if tools.is_empty() {
+            0
+        } else {
+            let definitions = self.count_text(&format_function_definitions(tools)?)?;
+            let discount = if includes_system_message {
+                TOOLS_WITH_SYSTEM_MESSAGE_DISCOUNT
+            } else {
+                0
+            };
+            definitions + TOOL_DEFINITIONS_TOKENS - discount
+        };
+        let choice_tokens = match tool_choice {
+            Some(ToolChoice::Mode(mode)) if mode == "none" => TOOL_CHOICE_NONE_TOKENS,
+            Some(ToolChoice::Mode(_)) | None => 0,
+            Some(ToolChoice::Named(named)) => {
+                NAMED_TOOL_CHOICE_TOKENS + self.count_text(&named.function.name)?
+            }
+        };
+        Ok(REPLY_PRIMING_TOKENS + tool_tokens + choice_tokens)
+    }
+}

litellm-rust/crates/token-counter/src/error.rs

@@ -0,0 +1,31 @@
+use std::string::FromUtf8Error;
+
+use thiserror::Error as ThisError;
+
+#[derive(Debug, ThisError)]
+pub enum Error {
+    #[error("failed to load tokenizer: {0}")]
+    Load(#[source] tokenizers::Error),
+    #[error("unsupported by the rust token counter: request body could not be parsed: {0}")]
+    RequestParse(#[source] serde_json::Error),
+    #[error("unsupported by the rust token counter: request has no countable input")]
+    MissingInput,
+    #[error(
+        "unsupported by the rust token counter: float text values are counted by the python path"
+    )]
+    FloatText,
+    #[error(
+        "unsupported by the rust token counter: content block type is counted by the python path"
+    )]
+    ContentBlock,
+    #[error("unsupported by the rust token counter: array parameter without items")]
+    ArrayItems,
+    #[error("unsupported by the rust token counter: text value could not be serialized: {0}")]
+    JsonSerialization(#[source] serde_json::Error),
+    #[error("unsupported by the rust token counter: serialized text value is not UTF-8: {0}")]
+    JsonUtf8(#[source] FromUtf8Error),
+    #[error("tokenization failed: {0}")]
+    Encode(#[source] tokenizers::Error),
+    #[error("token counting task failed: {0}")]
+    Task(String),
+}

litellm-rust/crates/token-counter/src/lib.rs

@@ -0,0 +1,17 @@
+//! Input token counting for a request body, mirroring `litellm.token_counter`
+//! for the shapes it can count exactly. Everything else is declined so the host
+//! keeps its own counter as the reference.
+
+#![forbid(unsafe_code)]
+
+mod byte_level;
+mod counter;
+mod error;
+mod python_json;
+mod tools;
+mod types;
+mod unicode_classes;
+
+pub use counter::{InputTokenCount, TokenCounter};
+pub use error::Error;
+pub use types::CountableRequest;

litellm-rust/crates/token-counter/src/python_json.rs

@@ -0,0 +1,155 @@
+//! `json.dumps(value)` with Python's default arguments: `", "` and `": "`
+//! separators, `ensure_ascii=True`, and keys in insertion order.
+
+use std::io::{self, Write};
+
+use serde::Serialize;
+use serde_json::ser::{Formatter, Serializer};
+
+use super::Error;
+use super::types::TextValue;
+
+pub(super) fn dumps(value: &TextValue) -> Result<String, Error> {
+    let mut output = Vec::with_capacity(serialized_len(value)?);
+    value
+        .serialize(&mut Serializer::with_formatter(
+            &mut output,
+            PythonFormatter,
+        ))
+        .map_err(Error::JsonSerialization)?;
+    debug_assert_eq!(output.len(), output.capacity());
+    String::from_utf8(output).map_err(Error::JsonUtf8)
+}
+
+fn serialized_len(value: &TextValue) -> Result<usize, Error> {
+    match value {
+        TextValue::Null => Ok(4),
+        TextValue::Bool(true) => Ok(4),
+        TextValue::Bool(false) => Ok(5),
+        TextValue::Number(number) => match (number.as_i64(), number.as_u64()) {
+            (Some(number), _) => Ok(unsigned_len(number.unsigned_abs()) + usize::from(number < 0)),
+            (_, Some(number)) => Ok(unsigned_len(number)),
+            _ => Err(Error::FloatText),
+        },
+        TextValue::Text(text) => Ok(quoted_len(text)),
+        TextValue::List(items) => items
+            .iter()
+            .try_fold(2 + items.len().saturating_sub(1) * 2, |len, item| {
+                Ok(len + serialized_len(item)?)
+            }),
+        TextValue::Object(entries) => entries.iter().try_fold(
+            2 + entries.len().saturating_sub(1) * 2,
+            |len, (key, value)| Ok(len + quoted_len(key) + 2 + serialized_len(value)?),
+        ),
+    }
+}
+
+fn unsigned_len(number: u64) -> usize {
+    if number == 0 {
+        1
+    } else {
+        number.ilog10() as usize + 1
+    }
+}
+
+fn quoted_len(value: &str) -> usize {
+    value.chars().fold(2, |len, character| {
+        len + match character {
+            '"' | '\\' | '\u{0008}' | '\u{000c}' | '\n' | '\r' | '\t' => 2,
+            '\u{0000}'..='\u{001f}' => 6,
+            ' '..='~' => 1,
+            _ => character.len_utf16() * 6,
+        }
+    })
+}
+
+struct PythonFormatter;
+
+impl Formatter for PythonFormatter {
+    fn begin_array_value<W>(&mut self, writer: &mut W, first: bool) -> io::Result<()>
+    where
+        W: ?Sized + Write,
+    {
+        if !first {
+            writer.write_all(b", ")?;
+        }
+        Ok(())
+    }
+
+    fn begin_object_key<W>(&mut self, writer: &mut W, first: bool) -> io::Result<()>
+    where
+        W: ?Sized + Write,
+    {
+        if !first {
+            writer.write_all(b", ")?;
+        }
+        Ok(())
+    }
+
+    fn begin_object_value<W>(&mut self, writer: &mut W) -> io::Result<()>
+    where
+        W: ?Sized + Write,
+    {
+        writer.write_all(b": ")
+    }
+
+    fn write_string_fragment<W>(&mut self, writer: &mut W, fragment: &str) -> io::Result<()>
+    where
+        W: ?Sized + Write,
+    {
+        for character in fragment.chars() {
+            if (' '..='~').contains(&character) {
+                write!(writer, "{character}")?;
+                continue;
+            }
+            let mut units = [0u16; 2];
+            for unit in character.encode_utf16(&mut units) {
+                write!(writer, "\\u{unit:04x}")?;
+            }
+        }
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use rstest::rstest;
+
+    use super::dumps;
+
+    #[rstest]
+    #[case::null("null", "null")]
+    #[case::boolean("true", "true")]
+    #[case::signed_integer("-3", "-3")]
+    #[case::unsigned_integer("18446744073709551615", "18446744073709551615")]
+    #[case::empty_array("[]", "[]")]
+    #[case::empty_object("{}", "{}")]
+    #[case::nested(
+        r#"{"first":1,"second":{"ok":true,"none":null},"third":[false,2]}"#,
+        r#"{"first": 1, "second": {"ok": true, "none": null}, "third": [false, 2]}"#
+    )]
+    #[case::string_escaping(
+        r#""caf\u00e9 \u2014 \ud83d\ude00 \"q\" \\ \n\t\u0001\u007f ~ /""#,
+        r#""caf\u00e9 \u2014 \ud83d\ude00 \"q\" \\ \n\t\u0001\u007f ~ /""#
+    )]
+    #[case::short_control_escapes(r#""\b\f\r""#, r#""\b\f\r""#)]
+    fn matches_python_json_dumps(#[case] input: &str, #[case] expected: &str) {
+        let value = serde_json::from_str(input).expect("fixture parses");
+
+        assert_eq!(dumps(&value).expect("fixture dumps"), expected);
+    }
+
+    #[rstest]
+    #[case::top_level("1.5")]
+    #[case::array("[1,2.5]")]
+    #[case::object(r#"{"nested":{"value":-0.25}}"#)]
+    fn rejects_floats(#[case] input: &str) {
+        let value = serde_json::from_str(input).expect("fixture parses");
+        let error = dumps(&value).expect_err("floats are declined");
+
+        assert_eq!(
+            error.to_string(),
+            "unsupported by the rust token counter: float text values are counted by the python path"
+        );
+    }
+}

litellm-rust/crates/token-counter/src/tools.rs

@@ -0,0 +1,304 @@
+//! Renders tool definitions the way `litellm.token_counter` does before
+//! tokenizing them (the TypeScript-like namespace OpenAI appears to use).
+
+use std::fmt::Write;
+
+use super::Error;
+use super::types::{EnumValue, Schema, SchemaType, ToolDefinition};
+
+pub(super) fn format_function_definitions(tools: &[ToolDefinition]) -> Result<String, Error> {
+    ToolFormatter::new(tools.len()).format(tools)
+}
+
+struct ToolFormatter {
+    output: String,
+}
+
+impl ToolFormatter {
+    fn new(tool_count: usize) -> Self {
+        Self {
+            output: String::with_capacity(tool_count.saturating_mul(128).saturating_add(48)),
+        }
+    }
+
+    fn format(mut self, tools: &[ToolDefinition]) -> Result<String, Error> {
+        self.output.push_str("namespace functions {\n\n");
+
+        for tool in tools {
+            self.write_function(tool)?;
+        }
+
+        self.output.push_str("} // namespace functions");
+        Ok(self.output)
+    }
+
+    fn write_function(&mut self, tool: &ToolDefinition) -> Result<(), Error> {
+        let (name, description, parameters) = resolve_function(tool);
+        let Some(name) = name.filter(|name| !name.is_empty()) else {
+            return Ok(());
+        };
+
+        if let Some(description) = description.filter(|description| !description.is_empty()) {
+            self.output.push_str("// ");
+            self.output.push_str(description);
+            self.output.push('\n');
+        }
+
+        match parameters.filter(|parameters| {
+            parameters
+                .properties
+                .as_ref()
+                .is_some_and(|properties| !properties.is_empty())
+        }) {
+            Some(parameters) => {
+                self.output.push_str("type ");
+                self.output.push_str(name);
+                self.output.push_str(" = (_: {\n");
+                self.write_object_parameters(parameters, 0)?;
+                self.output.push_str("\n}) => any;\n\n");
+            }
+            _ => {
+                self.output.push_str("type ");
+                self.output.push_str(name);
+                self.output.push_str(" = () => any;\n\n");
+            }
+        }
+
+        Ok(())
+    }
+
+    fn write_object_parameters(&mut self, parameters: &Schema, indent: usize) -> Result<(), Error> {
+        let Some(properties) = parameters
+            .properties
+            .as_ref()
+            .filter(|properties| !properties.is_empty())
+        else {
+            return Ok(());
+        };
+        let required = parameters.required.as_deref().unwrap_or_default();
+        for (index, (key, props)) in properties.iter().enumerate() {
+            if index > 0 {
+                self.output.push('\n');
+            }
+            if let Some(description) = props
+                .description
+                .as_deref()
+                .filter(|description| !description.is_empty())
+            {
+                self.write_indent(indent);
+                self.output.push_str("// ");
+                self.output.push_str(description);
+                self.output.push('\n');
+            }
+
+            self.write_indent(indent);
+            self.output.push_str(key);
+            if !required.iter().any(|required| required == key) {
+                self.output.push('?');
+            }
+            self.output.push_str(": ");
+            self.write_type(props, indent)?;
+            self.output.push(',');
+        }
+
+        Ok(())
+    }
+
+    fn write_type(&mut self, props: &Schema, indent: usize) -> Result<(), Error> {
+        let Some(SchemaType::Name(schema_type)) = &props.schema_type else {
+            self.output.push_str("any");
+            return Ok(());
+        };
+
+        match schema_type.as_str() {
+            "string" | "integer" | "number" => match &props.enum_values {
+                Some(values) => self.write_enum(values),
+                None if schema_type == "string" => self.output.push_str("string"),
+                None => self.output.push_str("number"),
+            },
+            "array" => {
+                let items = props.items.as_deref().ok_or(Error::ArrayItems)?;
+                self.write_type(items, indent)?;
+                self.output.push_str("[]");
+            }
+            "object" => {
+                self.output.push_str("{\n");
+                self.write_object_parameters(props, indent + 2)?;
+                self.output.push_str("\n}");
+            }
+            "boolean" => self.output.push_str("boolean"),
+            "null" => self.output.push_str("null"),
+            _ => self.output.push_str("any"),
+        }
+
+        Ok(())
+    }
+
+    fn write_enum(&mut self, values: &[EnumValue]) {
+        for (index, value) in values.iter().enumerate() {
+            if index > 0 {
+                self.output.push_str(" | ");
+            }
+            self.output.push('"');
+            match value {
+                EnumValue::Text(text) => self.output.push_str(text),
+                EnumValue::Integer(number) => {
+                    write!(self.output, "{number}").expect("writing to a String cannot fail");
+                }
+            }
+            self.output.push('"');
+        }
+    }
+
+    fn write_indent(&mut self, indent: usize) {
+        for _ in 0..indent {
+            self.output.push(' ');
+        }
+    }
+}
+
+fn resolve_function(tool: &ToolDefinition) -> (Option<&str>, Option<&str>, Option<&Schema>) {
+    match &tool.function {
+        Some(function) => (
+            function.name.as_deref(),
+            function.description.as_deref(),
+            function.parameters.as_ref(),
+        ),
+        None => (
+            tool.name.as_deref(),
+            tool.description.as_deref(),
+            tool.input_schema.as_ref().or(tool.parameters.as_ref()),
+        ),
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn parse_tools(json: &str) -> Vec<ToolDefinition> {
+        serde_json::from_str(json).expect("tool fixture parses")
+    }
+
+    #[test]
+    fn empty_tool_list_renders_an_empty_namespace() {
+        assert_eq!(
+            format_function_definitions(&[]).expect("empty tool list renders"),
+            "namespace functions {\n\n} // namespace functions"
+        );
+    }
+
+    #[test]
+    fn unnamed_tools_are_skipped_and_function_shape_takes_precedence() {
+        let tools = parse_tools(
+            r#"[
+                {},
+                {"name":""},
+                {"name":"ignored","function":{"description":"missing a function name"}},
+                {"name":"ping","description":""}
+            ]"#,
+        );
+
+        assert_eq!(
+            format_function_definitions(&tools).expect("tools render"),
+            "namespace functions {\n\ntype ping = () => any;\n\n} // namespace functions"
+        );
+    }
+
+    #[test]
+    fn empty_or_missing_properties_render_a_no_argument_function() {
+        let tools = parse_tools(
+            r#"[
+                {"name":"missing","input_schema":{"type":"object"}},
+                {"name":"empty","input_schema":{"type":"object","properties":{}}}
+            ]"#,
+        );
+
+        assert_eq!(
+            format_function_definitions(&tools).expect("tools render"),
+            concat!(
+                "namespace functions {\n\n",
+                "type missing = () => any;\n\n",
+                "type empty = () => any;\n\n",
+                "} // namespace functions"
+            )
+        );
+    }
+
+    #[test]
+    fn anthropic_parameters_render_all_supported_types() {
+        let tools = parse_tools(
+            r#"[{
+                "name":"inspect",
+                "description":"Inspect a value",
+                "input_schema":{
+                    "type":"object",
+                    "properties":{
+                        "text":{"type":"string"},
+                        "count":{"type":"integer","description":"Number of attempts"},
+                        "ratio":{"type":"number"},
+                        "enabled":{"type":"boolean"},
+                        "nothing":{"type":"null"},
+                        "unknown":{"type":"custom"},
+                        "union":{"type":["string","null"]},
+                        "labels":{"type":"array","items":{"type":"string"}},
+                        "config":{"type":"object","properties":{"retries":{"type":"integer"}},"required":["retries"]},
+                        "mode":{"type":"string","enum":["fast",2]}
+                    },
+                    "required":["text"]
+                }
+            }]"#,
+        );
+
+        assert_eq!(
+            format_function_definitions(&tools).expect("tool renders"),
+            concat!(
+                "namespace functions {\n\n",
+                "// Inspect a value\n",
+                "type inspect = (_: {\n",
+                "text: string,\n",
+                "// Number of attempts\n",
+                "count?: number,\n",
+                "ratio?: number,\n",
+                "enabled?: boolean,\n",
+                "nothing?: null,\n",
+                "unknown?: any,\n",
+                "union?: any,\n",
+                "labels?: string[],\n",
+                "config?: {\n",
+                "  retries: number,\n",
+                "},\n",
+                "mode?: \"fast\" | \"2\",\n",
+                "}) => any;\n\n",
+                "} // namespace functions"
+            )
+        );
+    }
+
+    #[test]
+    fn input_schema_takes_precedence_over_legacy_parameters() {
+        let tools = parse_tools(
+            r#"[{
+                "name":"choose",
+                "input_schema":{"type":"object","properties":{"current":{"type":"string"}}},
+                "parameters":{"type":"object","properties":{"legacy":{"type":"string"}}}
+            }]"#,
+        );
+
+        let rendered = format_function_definitions(&tools).expect("tool renders");
+        assert!(rendered.contains("current?: string,"));
+        assert!(!rendered.contains("legacy"));
+    }
+
+    #[test]
+    fn array_without_items_returns_an_error() {
+        let tools = parse_tools(
+            r#"[{"name":"broken","parameters":{"type":"object","properties":{"values":{"type":"array"}}}}]"#,
+        );
+
+        assert!(matches!(
+            format_function_definitions(&tools),
+            Err(Error::ArrayItems)
+        ));
+    }
+}

litellm-rust/crates/token-counter/src/types.rs

@@ -0,0 +1,238 @@
+use std::fmt;
+
+use indexmap::IndexMap;
+use serde::de::{MapAccess, SeqAccess, Visitor};
+use serde::{Deserialize, Deserializer, Serialize};
+use serde_json::Number;
+
+use super::Error;
+
+/// The parts of a request body the host's budget counter reads. Chat and
+/// Anthropic Messages bodies carry `messages`; completions carry `prompt`;
+/// Responses and embeddings carry `input`; rerank carries `query` and
+/// `documents`. The host checks key presence, not nullness, so an explicit
+/// `null` is kept distinct from an absent key. Anything outside this shape is
+/// declined so the host can fall back to its own counter instead of silently
+/// miscounting.
+#[derive(Clone, Debug, Deserialize, PartialEq)]
+pub struct CountableRequest {
+    pub(crate) model: Option<String>,
+    #[serde(default, deserialize_with = "present_messages")]
+    pub(crate) messages: Option<Vec<Message>>,
+    pub(crate) tools: Option<Vec<ToolDefinition>>,
+    pub(crate) tool_choice: Option<ToolChoice>,
+    #[serde(default, deserialize_with = "present_text")]
+    pub(crate) prompt: Option<TextValue>,
+    #[serde(default, deserialize_with = "present_text")]
+    pub(crate) input: Option<TextValue>,
+    #[serde(default, deserialize_with = "present_text")]
+    pub(crate) query: Option<TextValue>,
+    #[serde(default, deserialize_with = "present_text")]
+    pub(crate) documents: Option<TextValue>,
+}
+
+impl CountableRequest {
+    pub fn parse(body: &[u8]) -> Result<Self, Error> {
+        serde_json::from_slice(body).map_err(Error::RequestParse)
+    }
+}
+
+fn present_messages<'de, D: Deserializer<'de>>(
+    deserializer: D,
+) -> Result<Option<Vec<Message>>, D::Error> {
+    Option::<Vec<Message>>::deserialize(deserializer)
+        .map(|messages| Some(messages.unwrap_or_default()))
+}
+
+fn present_text<'de, D: Deserializer<'de>>(deserializer: D) -> Result<Option<TextValue>, D::Error> {
+    TextValue::deserialize(deserializer).map(Some)
+}
+
+/// Free-form JSON the host counts as text: strings and integers via `str()`,
+/// objects via `json.dumps()`, lists flattened. Objects keep document order so
+/// the dumped text matches Python byte for byte.
+#[derive(Clone, Debug, PartialEq, Serialize)]
+#[serde(untagged)]
+pub(crate) enum TextValue {
+    Null,
+    Bool(bool),
+    Number(Number),
+    Text(String),
+    List(Vec<TextValue>),
+    Object(IndexMap<String, TextValue>),
+}
+
+impl<'de> Deserialize<'de> for TextValue {
+    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
+        deserializer.deserialize_any(TextValueVisitor)
+    }
+}
+
+struct TextValueVisitor;
+
+impl<'de> Visitor<'de> for TextValueVisitor {
+    type Value = TextValue;
+
+    fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
+        formatter.write_str("a JSON value")
+    }
+
+    fn visit_unit<E>(self) -> Result<Self::Value, E> {
+        Ok(TextValue::Null)
+    }
+
+    fn visit_none<E>(self) -> Result<Self::Value, E> {
+        Ok(TextValue::Null)
+    }
+
+    fn visit_bool<E>(self, value: bool) -> Result<Self::Value, E> {
+        Ok(TextValue::Bool(value))
+    }
+
+    fn visit_i64<E>(self, value: i64) -> Result<Self::Value, E> {
+        Ok(TextValue::Number(value.into()))
+    }
+
+    fn visit_u64<E>(self, value: u64) -> Result<Self::Value, E> {
+        Ok(TextValue::Number(value.into()))
+    }
+
+    fn visit_f64<E>(self, value: f64) -> Result<Self::Value, E>
+    where
+        E: serde::de::Error,
+    {
+        Number::from_f64(value)
+            .map(TextValue::Number)
+            .ok_or_else(|| E::custom("non-finite JSON number"))
+    }
+
+    fn visit_str<E>(self, value: &str) -> Result<Self::Value, E> {
+        Ok(TextValue::Text(value.to_owned()))
+    }
+
+    fn visit_string<E>(self, value: String) -> Result<Self::Value, E> {
+        Ok(TextValue::Text(value))
+    }
+
+    fn visit_seq<A>(self, mut sequence: A) -> Result<Self::Value, A::Error>
+    where
+        A: SeqAccess<'de>,
+    {
+        let mut items = Vec::with_capacity(sequence.size_hint().unwrap_or(0));
+        while let Some(item) = sequence.next_element()? {
+            items.push(item);
+        }
+        Ok(TextValue::List(items))
+    }
+
+    fn visit_map<A>(self, mut map: A) -> Result<Self::Value, A::Error>
+    where
+        A: MapAccess<'de>,
+    {
+        let mut entries = IndexMap::with_capacity(map.size_hint().unwrap_or(0));
+        while let Some((key, value)) = map.next_entry()? {
+            entries.insert(key, value);
+        }
+        Ok(TextValue::Object(entries))
+    }
+}
+
+/// Python counts every string-valued key of a message, so any key beyond these
+/// makes the shape unsupported rather than silently uncounted.
+#[derive(Clone, Debug, Deserialize, PartialEq)]
+#[serde(deny_unknown_fields)]
+pub(crate) struct Message {
+    pub(crate) role: Option<String>,
+    pub(crate) name: Option<String>,
+    pub(crate) content: Option<MessageContent>,
+}
+
+#[derive(Clone, Debug, Deserialize, PartialEq)]
+#[serde(untagged)]
+pub(crate) enum MessageContent {
+    Text(String),
+    Blocks(Vec<ContentItem>),
+}
+
+#[derive(Clone, Debug, Deserialize, PartialEq)]
+#[serde(untagged)]
+pub(crate) enum ContentItem {
+    Text(String),
+    Block(ContentBlock),
+}
+
+#[derive(Clone, Debug, Deserialize, PartialEq)]
+#[serde(tag = "type")]
+pub(crate) enum ContentBlock {
+    #[serde(rename = "text")]
+    Text { text: String },
+    #[serde(rename = "thinking")]
+    Thinking { thinking: String },
+    #[serde(rename = "tool_reference")]
+    ToolReference { tool_name: Option<String> },
+    /// Images, documents, files and tool use/result blocks price through
+    /// Python-only helpers, so they stay on the Python counter.
+    #[serde(other)]
+    Unsupported,
+}
+
+/// Either the OpenAI `{"type": "function", "function": {...}}` shape or the
+/// Anthropic `{"name", "description", "input_schema"}` shape.
+#[derive(Clone, Debug, Deserialize, PartialEq)]
+pub(crate) struct ToolDefinition {
+    pub(crate) function: Option<FunctionDefinition>,
+    pub(crate) name: Option<String>,
+    pub(crate) description: Option<String>,
+    pub(crate) input_schema: Option<Schema>,
+    pub(crate) parameters: Option<Schema>,
+}
+
+#[derive(Clone, Debug, Deserialize, PartialEq)]
+pub(crate) struct FunctionDefinition {
+    pub(crate) name: Option<String>,
+    pub(crate) description: Option<String>,
+    pub(crate) parameters: Option<Schema>,
+}
+
+#[derive(Clone, Debug, Default, Deserialize, PartialEq)]
+pub(crate) struct Schema {
+    #[serde(rename = "type")]
+    pub(crate) schema_type: Option<SchemaType>,
+    pub(crate) description: Option<String>,
+    #[serde(rename = "enum")]
+    pub(crate) enum_values: Option<Vec<EnumValue>>,
+    pub(crate) items: Option<Box<Schema>>,
+    pub(crate) properties: Option<IndexMap<String, Schema>>,
+    pub(crate) required: Option<Vec<String>>,
+}
+
+#[derive(Clone, Debug, Deserialize, PartialEq)]
+#[serde(untagged)]
+pub(crate) enum SchemaType {
+    Name(String),
+    Union(Vec<String>),
+}
+
+#[derive(Clone, Debug, Deserialize, PartialEq)]
+#[serde(untagged)]
+pub(crate) enum EnumValue {
+    Text(String),
+    Integer(i64),
+}
+
+#[derive(Clone, Debug, Deserialize, PartialEq)]
+#[serde(untagged)]
+pub(crate) enum ToolChoice {
+    Mode(String),
+    Named(NamedToolChoice),
+}
+
+#[derive(Clone, Debug, Deserialize, PartialEq)]
+pub(crate) struct NamedToolChoice {
+    pub(crate) function: NamedFunction,
+}
+
+#[derive(Clone, Debug, Deserialize, PartialEq)]
+pub(crate) struct NamedFunction {
+    pub(crate) name: String,
+}

litellm-rust/crates/token-counter/src/unicode_classes.rs

@@ -0,0 +1,73 @@
+use std::cmp::Ordering;
+use std::sync::LazyLock;
+
+use tokenizers::utils::SysRegex;
+
+struct Ranges(Box<[(u32, u32)]>);
+
+pub(super) struct UnicodeClasses {
+    letters: Ranges,
+    numbers: Ranges,
+    spaces: Ranges,
+}
+
+static CLASSES: LazyLock<Option<UnicodeClasses>> = LazyLock::new(|| {
+    let scalars: String = (0..=u32::from(char::MAX))
+        .filter_map(char::from_u32)
+        .collect();
+    Some(UnicodeClasses {
+        letters: Ranges::load(r"\p{L}+", &scalars)?,
+        numbers: Ranges::load(r"\p{N}+", &scalars)?,
+        spaces: Ranges::load(r"\s+", &scalars)?,
+    })
+});
+
+impl Ranges {
+    fn load(pattern: &str, scalars: &str) -> Option<Self> {
+        let regex = SysRegex::new(pattern).ok()?;
+        let ranges = regex
+            .find_iter(scalars)
+            .map(|(start, end)| {
+                let matched = scalars.get(start..end)?;
+                Some((
+                    u32::from(matched.chars().next()?),
+                    u32::from(matched.chars().next_back()?),
+                ))
+            })
+            .collect::<Option<Box<[_]>>>()?;
+        Some(Self(ranges))
+    }
+
+    fn contains(&self, character: char) -> bool {
+        let code = u32::from(character);
+        self.0
+            .binary_search_by(|(low, high)| {
+                if *high < code {
+                    Ordering::Less
+                } else if *low > code {
+                    Ordering::Greater
+                } else {
+                    Ordering::Equal
+                }
+            })
+            .is_ok()
+    }
+}
+
+impl UnicodeClasses {
+    pub(super) fn get() -> Option<&'static Self> {
+        CLASSES.as_ref()
+    }
+
+    pub(super) fn is_letter(&self, character: char) -> bool {
+        self.letters.contains(character)
+    }
+
+    pub(super) fn is_number(&self, character: char) -> bool {
+        self.numbers.contains(character)
+    }
+
+    pub(super) fn is_space(&self, character: char) -> bool {
+        self.spaces.contains(character)
+    }
+}

litellm-rust/crates/token-counter/tests/token_counter.rs

@@ -0,0 +1,176 @@
+use rstest::rstest;
+
+use litellm_token_counter::{CountableRequest, Error, InputTokenCount, TokenCounter};
+
+/// Expected counts are pinned from `litellm.token_counter(model="claude-sonnet-4-5", ...)`
+/// so this test also guards Python parity.
+fn counter() -> TokenCounter {
+    let path = concat!(
+        env!("CARGO_MANIFEST_DIR"),
+        "/../../../litellm/litellm_core_utils/tokenizers/anthropic_tokenizer.json"
+    );
+    let json = std::fs::read_to_string(path).expect("anthropic tokenizer json is in the repo");
+    TokenCounter::from_json(&json).expect("anthropic tokenizer loads")
+}
+
+const SIMPLE: &str = r#"{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"Hello, how are you today?"}]}"#;
+
+const BLOCKS_AND_SYSTEM: &str = r#"{"model":"claude-sonnet-4-5","messages":[
+  {"role":"system","content":"You are a terse assistant."},
+  {"role":"user","name":"alice","content":[
+    {"type":"text","text":"Summarise this paragraph about ships and harbours."},
+    "plain string item",
+    {"type":"thinking","thinking":"pondering"},
+    {"type":"tool_reference","tool_name":"get_weather"}]},
+  {"role":"assistant","content":[{"type":"text","text":"Sure.","cache_control":{"type":"ephemeral"}}]}]}"#;
+
+const TOOLS_OPENAI: &str = r#"{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"weather?"}],
+  "tools":[
+    {"type":"function","function":{"name":"get_weather","description":"Get weather","parameters":{
+      "type":"object",
+      "properties":{
+        "location":{"type":"string","description":"City name"},
+        "unit":{"type":"string","enum":["celsius","fahrenheit"]},
+        "days":{"type":"integer"},
+        "tags":{"type":"array","items":{"type":"string"}},
+        "opts":{"type":"object","properties":{"verbose":{"type":"boolean"},"level":{"type":"integer","enum":[1,2]}},"required":["verbose"]},
+        "anything":{}},
+      "required":["location"]}}},
+    {"type":"function","function":{"name":"noop"}}],
+  "tool_choice":{"type":"function","function":{"name":"get_weather"}}}"#;
+
+const TOOLS_ANTHROPIC_SYSTEM: &str = r#"{"model":"claude-sonnet-4-5",
+  "messages":[{"role":"system","content":"sys"},{"role":"user","content":"weather?"}],
+  "tools":[{"name":"get_weather","description":"Get weather","input_schema":{
+    "type":"object","properties":{"location":{"type":["string","null"]}},"required":["location"]}}],
+  "tool_choice":"none"}"#;
+
+const COMPLETIONS_PROMPT: &str =
+    r#"{"model":"claude-sonnet-4-5","prompt":"Write a haiku about ships."}"#;
+
+const COMPLETIONS_PROMPT_LIST: &str =
+    r#"{"model":"claude-sonnet-4-5","prompt":["first prompt","second prompt"]}"#;
+
+const RESPONSES_INPUT: &str = r#"{"model":"claude-sonnet-4-5","input":[
+  {"role":"user","content":[{"type":"input_text","text":"Summarise caf\u00e9 menus, na\u00efve \u2014 ok? \"quoted\"\n"}]},
+  {"role":"assistant","content":"Sure."}],"instructions":"be terse"}"#;
+
+const EMBEDDINGS_TOKEN_IDS: &str =
+    r#"{"model":"claude-sonnet-4-5","input":[[101,2023,5],[7]],"encoding_format":"float"}"#;
+
+const RERANK: &str = r#"{"model":"claude-sonnet-4-5","query":"best harbour",
+  "documents":["doc one",{"text":"doc two","title":"T","n":3,"ok":true,"none":null,"tags":["a","b"]}]}"#;
+
+/// Expected counts are pinned from
+/// `litellm.proxy.spend_tracking.budget_reservation._count_input_tokens(body, "claude-sonnet-4-5")`.
+#[rstest]
+#[case::text_only(SIMPLE, 14)]
+#[case::content_blocks_name_and_system(BLOCKS_AND_SYSTEM, 45)]
+#[case::openai_tools_named_choice(TOOLS_OPENAI, 123)]
+#[case::anthropic_tools_system_discount_choice_none(TOOLS_ANTHROPIC_SYSTEM, 53)]
+#[case::completions_prompt(COMPLETIONS_PROMPT, 7)]
+#[case::completions_prompt_list(COMPLETIONS_PROMPT_LIST, 4)]
+#[case::responses_input_items(RESPONSES_INPUT, 62)]
+#[case::embeddings_token_ids(EMBEDDINGS_TOKEN_IDS, 5)]
+#[case::rerank_query_and_documents(RERANK, 41)]
+fn count_request_matches_python_token_counter(#[case] body: &str, #[case] expected: usize) {
+    let request = CountableRequest::parse(body.as_bytes()).expect("fixture parses");
+    let count = counter().count_request(&request).expect("fixture counts");
+    assert_eq!(
+        count,
+        InputTokenCount {
+            model: Some("claude-sonnet-4-5".to_string()),
+            input_tokens: expected,
+        }
+    );
+}
+
+#[rstest]
+#[case::null_messages_win_over_prompt(r#"{"model":"m","messages":null,"prompt":"ignored"}"#, 3)]
+#[case::model_from_route(r#"{"prompt":"hi"}"#, 1)]
+#[case::bools_and_ints_use_python_str(r#"{"model":"m","prompt":[true,false,42]}"#, 3)]
+#[case::null_prompt_counts_zero(r#"{"model":"m","prompt":null}"#, 0)]
+fn key_presence_follows_python(#[case] body: &str, #[case] expected: usize) {
+    let request = CountableRequest::parse(body.as_bytes()).expect("fixture parses");
+    let count = counter().count_request(&request).expect("fixture counts");
+    assert_eq!(count.input_tokens, expected);
+}
+
+#[rstest]
+#[case::not_json(b"not json" as &[u8])]
+#[case::messages_not_a_list(br#"{"model":"m","messages":"hi"}"#)]
+#[case::message_with_tool_calls(
+    br#"{"model":"m","messages":[{"role":"assistant","tool_calls":[{"id":"1","type":"function","function":{"name":"f","arguments":"{}"}}]}]}"#
+)]
+#[case::dict_content(
+    br#"{"model":"m","messages":[{"role":"user","content":{"type":"text","text":"x"}}]}"#
+)]
+#[case::float_enum(
+    br#"{"model":"m","messages":[],"tools":[{"name":"f","input_schema":{"type":"object","properties":{"x":{"type":"number","enum":[1.5]}}}}]}"#
+)]
+#[case::anthropic_tool_choice_without_function(
+    br#"{"model":"m","messages":[],"tool_choice":{"type":"auto"}}"#
+)]
+fn shapes_outside_the_mirror_are_declined_at_parse(#[case] body: &[u8]) {
+    assert!(matches!(
+        CountableRequest::parse(body),
+        Err(Error::RequestParse(_))
+    ));
+}
+
+#[rstest]
+#[case::no_countable_input(br#"{"model":"m","instructions":"hi"}"# as &[u8])]
+#[case::float_prompt(br#"{"model":"m","prompt":1.5}"#)]
+#[case::float_inside_document(br#"{"model":"m","documents":[{"score":0.5}]}"#)]
+#[case::image_block(
+    br#"{"model":"m","messages":[{"role":"user","content":[{"type":"image","source":{"type":"base64","media_type":"image/png","data":"AA=="}}]}]}"#
+)]
+#[case::tool_result_block(
+    br#"{"model":"m","messages":[{"role":"user","content":[{"type":"tool_result","tool_use_id":"1","content":"ok"}]}]}"#
+)]
+#[case::array_without_items(
+    br#"{"model":"m","messages":[],"tools":[{"name":"f","input_schema":{"type":"object","properties":{"x":{"type":"array"}}}}]}"#
+)]
+fn shapes_outside_the_mirror_are_declined_at_count(#[case] body: &[u8]) {
+    let request = CountableRequest::parse(body).expect("shape parses");
+    assert!(matches!(
+        counter().count_request(&request),
+        Err(Error::MissingInput | Error::FloatText | Error::ContentBlock | Error::ArrayItems)
+    ));
+}
+
+#[test]
+fn tool_choice_and_system_discount_change_the_count() {
+    let counter = counter();
+    let count = |body: &str| {
+        counter
+            .count_request(&CountableRequest::parse(body.as_bytes()).expect("parses"))
+            .expect("counts")
+            .input_tokens
+    };
+    let base = count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}]}"#);
+    assert_eq!(
+        count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tool_choice":"none"}"#),
+        base + 1
+    );
+    assert_eq!(
+        count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tool_choice":"auto"}"#),
+        base
+    );
+    let with_tools = count(
+        r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tools":[{"name":"f"}]}"#,
+    );
+    let with_tools_and_system = count(
+        r#"{"model":"m","messages":[{"role":"system","content":"hi"}],"tools":[{"name":"f"}]}"#,
+    );
+    assert_eq!(with_tools - with_tools_and_system, 4);
+    assert_eq!(
+        count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tools":[]}"#),
+        base
+    );
+}
+
+#[test]
+fn loading_a_bad_tokenizer_is_a_load_error() {
+    assert!(matches!(TokenCounter::from_json("{}"), Err(Error::Load(_))));
+}

litellm/_redis.py

@@ -24,6 +24,7 @@ from redis.credentials import CredentialProvider
 from litellm import get_secret, get_secret_str
 from litellm._redis_credential_provider import (
     AzureADCredentialProvider,
+    ElastiCacheIAMCredentialProvider,
     GCPIAMCredentialProvider,
     _generate_gcp_iam_access_token,
 )
@@ -38,6 +39,14 @@ from ._logging import verbose_logger
 
 AZURE_REDIS_SCOPE: Final = "https://redis.azure.com/.default"
 
+_AWS_IAM_KWARG_NAMES: Final = (
+    "aws_iam_auth",
+    "aws_iam_user_name",
+    "aws_iam_cache_name",
+    "aws_iam_region",
+    "aws_iam_serverless",
+)
+
 
 def _unwrapped_init_args(cls: type) -> frozenset[str]:
     """Every parameter on a single class's own ``__init__``, decorator-unwrapped.
@@ -75,6 +84,7 @@ def _get_redis_kwargs():
         "azure_client_id",
         "azure_tenant_id",
         "azure_client_secret",
+        *_AWS_IAM_KWARG_NAMES,
     }
 
     available_args: Final = {x for x in _unwrapped_init_args(redis.Redis) if x not in exclude_args} | include_args
@@ -270,6 +280,42 @@ def _redis_kwargs_from_environment():
     return return_dict
 
 
+def _coerces_to_true(value: object | None) -> bool:
+    return _str_to_bool(value) if isinstance(value, str) else bool(value)
+
+
+def _uses_tls(redis_kwargs: Mapping[str, object]) -> bool:
+    if redis_kwargs.get("startup_nodes") is not None:
+        return _coerces_to_true(redis_kwargs.get("ssl"))
+    url: Final = redis_kwargs.get("url")
+    if isinstance(url, str):
+        return urlsplit(url).scheme.lower() == "rediss"
+    return _coerces_to_true(redis_kwargs.get("ssl"))
+
+
+def _build_elasticache_iam_provider(redis_kwargs: Mapping[str, object]) -> ElastiCacheIAMCredentialProvider:
+    user_name: Final = redis_kwargs.get("aws_iam_user_name")
+    cache_name: Final = redis_kwargs.get("aws_iam_cache_name")
+    region: Final = (
+        redis_kwargs.get("aws_iam_region") or get_secret_str("AWS_REGION") or get_secret_str("AWS_DEFAULT_REGION")
+    )
+    required_settings: Final = (
+        ("aws_iam_user_name", user_name),
+        ("aws_iam_cache_name", cache_name),
+        ("aws_iam_region", region),
+    )
+    missing_settings: Final = tuple(name for name, value in required_settings if not value)
+    if missing_settings:
+        raise ValueError("AWS ElastiCache IAM Redis authentication requires: " + ", ".join(missing_settings))
+
+    return ElastiCacheIAMCredentialProvider(
+        user_name=str(user_name),
+        cache_name=str(cache_name),
+        region=str(region),
+        is_serverless=_coerces_to_true(redis_kwargs.get("aws_iam_serverless")),
+    )
+
+
 def create_gcp_iam_redis_connect_func(
     service_account: str,
     ssl_ca_certs: str | None = None,
@@ -540,6 +586,7 @@ def _get_redis_client_logic(**env_overrides):
         _azure_redis_ad_token: Final = redis_kwargs.get("azure_redis_ad_token") or get_secret("REDIS_AZURE_AD_TOKEN")
 
         _azure_ad_enabled: Final = _azure_redis_ad_token is not None and str(_azure_redis_ad_token).lower() == "true"
+        _aws_iam_enabled: Final = _coerces_to_true(redis_kwargs.get("aws_iam_auth"))
 
         if _azure_ad_enabled and _gcp_service_account is not None:
             verbose_logger.warning(
@@ -567,6 +614,22 @@ def _get_redis_client_logic(**env_overrides):
             # credentials via inspection or logging.
             redis_kwargs["redis_connect_func"]._azure_redis_ad_token = True
 
+        if _aws_iam_enabled and _gcp_service_account is not None:
+            verbose_logger.warning(
+                "Both GCP IAM (gcp_service_account) and AWS ElastiCache IAM (aws_iam_auth) are configured "
+                "for Redis. Using GCP IAM. Remove one to avoid misconfiguration."
+            )
+        elif _aws_iam_enabled and _azure_ad_enabled:
+            verbose_logger.warning(
+                "Both Azure AD (azure_redis_ad_token) and AWS ElastiCache IAM (aws_iam_auth) are configured "
+                "for Redis. Using Azure AD. Remove one to avoid misconfiguration."
+            )
+        elif _aws_iam_enabled:
+            if not _uses_tls(redis_kwargs):
+                raise ValueError("AWS ElastiCache IAM Redis authentication requires TLS")
+            verbose_logger.debug("Setting up AWS ElastiCache IAM authentication for Redis.")
+            redis_kwargs["credential_provider"] = _build_elasticache_iam_provider(redis_kwargs)
+
     redis_kwargs.pop("gcp_service_account", None)
     redis_kwargs.pop("gcp_ssl_ca_certs", None)
 
@@ -575,6 +638,8 @@ def _get_redis_client_logic(**env_overrides):
     redis_kwargs.pop("azure_client_id", None)
     redis_kwargs.pop("azure_tenant_id", None)
     redis_kwargs.pop("azure_client_secret", None)
+    for aws_iam_key in _AWS_IAM_KWARG_NAMES:
+        redis_kwargs.pop(aws_iam_key, None)
 
     if redis_kwargs.get("credential_provider") is not None:
         redis_kwargs.pop("redis_connect_func", None)

litellm/_redis_credential_provider.py

@@ -1,10 +1,17 @@
+from __future__ import annotations
+
 import asyncio
 import threading
 import time
-from typing import Final, Protocol
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Final, Protocol
+from urllib.parse import urlencode
 
 from redis.credentials import CredentialProvider
 
+if TYPE_CHECKING:
+    from botocore.credentials import Credentials
+
 # Azure AD scope for Redis Cache for Azure.
 AZURE_REDIS_SCOPE: Final = "https://redis.azure.com/.default"
 
@@ -117,6 +124,82 @@ class GCPIAMCredentialProvider(CredentialProvider):
         return (token,)
 
 
+_ELASTICACHE_SERVICE_NAME: Final = "elasticache"
+_ELASTICACHE_TOKEN_TTL_SECONDS: Final = 900
+_ELASTICACHE_SERVERLESS_RESOURCE_TYPE: Final = "ServerlessCache"
+
+
+class ElastiCacheIAMCredentialProvider(CredentialProvider):
+    def __init__(
+        self,
+        user_name: str,
+        cache_name: str,
+        region: str,
+        is_serverless: bool = False,
+        credentials_resolver: Callable[[], Credentials | None] | None = None,
+        token_lifetime_seconds: int = _ELASTICACHE_TOKEN_TTL_SECONDS,
+    ) -> None:
+        self._user_name = user_name
+        self._cache_name = cache_name.lower()
+        self._region = region
+        self._is_serverless = is_serverless
+        self._credentials_resolver = credentials_resolver or self._resolve_credentials
+        self._credentials: Credentials | None = None
+        self._token_lifetime_seconds = token_lifetime_seconds
+
+    @staticmethod
+    def _resolve_credentials() -> Credentials | None:
+        try:
+            import botocore.session
+        except ImportError as e:
+            raise ImportError(
+                "botocore is required for ElastiCache IAM Redis authentication. Install it with: pip install boto3"
+            ) from e
+
+        return botocore.session.get_session().get_credentials()
+
+    def _get_credentials(self) -> tuple[str, str]:
+        credentials: Final = self._credentials if self._credentials is not None else self._credentials_resolver()
+        if credentials is None:
+            raise RuntimeError("Unable to resolve AWS credentials for ElastiCache IAM Redis authentication")
+        self._credentials = credentials
+
+        frozen_credentials: Final = credentials.get_frozen_credentials()
+
+        try:
+            from botocore.auth import SigV4QueryAuth
+            from botocore.awsrequest import AWSRequest
+        except ImportError as e:
+            raise ImportError(
+                "botocore is required for ElastiCache IAM Redis authentication. Install it with: pip install boto3"
+            ) from e
+
+        query: Final = urlencode(
+            (
+                ("Action", "connect"),
+                ("User", self._user_name),
+                *((("ResourceType", _ELASTICACHE_SERVERLESS_RESOURCE_TYPE),) if self._is_serverless else ()),
+            )
+        )
+        request: Final = AWSRequest(method="GET", url=f"https://{self._cache_name}/?{query}")
+        SigV4QueryAuth(
+            frozen_credentials,
+            _ELASTICACHE_SERVICE_NAME,
+            self._region,
+            expires=self._token_lifetime_seconds,
+        ).add_auth(request)
+        signed_url: Final = request.url
+        if signed_url is None:
+            raise RuntimeError("Unable to generate AWS ElastiCache IAM credentials")
+        return self._user_name, signed_url.removeprefix("https://")
+
+    def get_credentials(self) -> tuple[str, str]:
+        return self._get_credentials()
+
+    async def get_credentials_async(self) -> tuple[str, str]:
+        return await asyncio.to_thread(self._get_credentials)
+
+
 class AzureADCredentialProvider(CredentialProvider):
     """
     redis.credentials.CredentialProvider implementation that supplies Azure AD

litellm/completion_extras/litellm_responses_transformation/transformation.py

@@ -1201,6 +1201,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
         # Cast to Any to match the expected union type for tools list items
         tools.append(cast(Any, web_search_tool))
 
+    def transform_response_format_to_text_format(self, response_format: object) -> "ResponseText | None":
+        return self._transform_response_format_to_text_format(response_format)
+
     def _transform_response_format_to_text_format(self, response_format: object) -> "ResponseText | None":
         """
         Transform Chat Completion response_format parameter to Responses API text.format parameter.

litellm/constants.py

@@ -1654,6 +1654,7 @@ CLOUDZERO_EXPORT_USAGE_DATA_JOB_NAME: Final = "cloudzero_export_usage_data"
 MAVVRIK_FOCUS_EXPORT_JOB_NAME: Final = "mavvrik_focus_export_usage_data"
 CLOUDZERO_MAX_FETCHED_DATA_RECORDS: Final = int(os.getenv("CLOUDZERO_MAX_FETCHED_DATA_RECORDS", 50000))
 SPEND_LOG_CLEANUP_JOB_NAME: Final = "spend_log_cleanup"
+BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME: Final = "background_health_check_db_save"
 KEY_ROTATION_JOB_NAME: Final = "litellm_key_rotation_job"
 EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME: Final = "litellm_expired_ui_session_key_cleanup_job"
 WEEKLY_SPEND_REPORT_JOB_ID: Final = "weekly_spend_report_job"

litellm/llms/base_llm/responses/transformation.py

@@ -62,6 +62,9 @@ class BaseResponsesAPIConfig(ABC):
         """
         return False
 
+    def supports_encrypted_agent_messages(self) -> bool:
+        return False
+
     def sign_request(
         self,
         headers: dict,

litellm/llms/openai/responses/transformation.py

@@ -110,6 +110,9 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
     def supports_native_file_search(self) -> bool:
         return True
 
+    def supports_encrypted_agent_messages(self) -> bool:
+        return self.custom_llm_provider in (LlmProviders.OPENAI, LlmProviders.AZURE)
+
     @staticmethod
     def _is_gpt_5_model(model: str) -> bool:
         """Return True only for actual OpenAI GPT-5 models.

litellm/llms/voyage/embedding/transformation_contextual.py

@@ -3,6 +3,7 @@ This module is used to transform the request and response for the Voyage context
 This would be used for all the contextualized embeddings models in Voyage.
 """
 
+from collections.abc import Mapping
 from typing import Final
 
 import httpx
@@ -24,7 +25,10 @@ class VoyageError(BaseLLMException):
     ):
         self.status_code = status_code
         self.message = message
-        self.request = httpx.Request(method="POST", url="https://api.voyageai.com/v1/contextualizedembeddings")
+        self.request = httpx.Request(
+            method="POST",
+            url="https://api.voyageai.com/v1/contextualizedembeddings",
+        )
         self.response = httpx.Response(status_code=status_code, request=self.request)
         super().__init__(
             status_code=status_code,
@@ -56,16 +60,16 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig):
             return api_base
         return "https://api.voyageai.com/v1/contextualizedembeddings"
 
-    def get_supported_openai_params(self, model: str) -> list:
+    def get_supported_openai_params(self, model: str) -> list:  # mutable-ok: base class signature
         return ["encoding_format", "dimensions"]
 
     def map_openai_params(
         self,
-        non_default_params: dict,
-        optional_params: dict,
+        non_default_params: dict,  # mutable-ok: base class signature
+        optional_params: dict,  # mutable-ok: base class signature
         model: str,
         drop_params: bool,
-    ) -> dict:
+    ) -> dict:  # mutable-ok: base class signature
         """
         Map OpenAI params to Voyage params
 
@@ -79,7 +83,7 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig):
 
     def validate_environment(
         self,
-        headers: dict,
+        headers: dict,  # mutable-ok: base class signature
         model: str,
         messages: list[AllMessageValues],
         optional_params: dict,
@@ -97,6 +101,8 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig):
             "Authorization": f"Bearer {api_key}",
         }
 
+    AUTO_CHUNK_SIZE: Final = 32000
+
     def transform_embedding_request(
         self,
         model: str,
@@ -105,11 +111,27 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig):
         headers: dict,
     ) -> dict:
         return {
-            "inputs": input,
+            "inputs": [input] if isinstance(input, str) else input,
             "model": model,
+            **self._auto_chunk_params(input, optional_params),
             **optional_params,
         }
 
+    @classmethod
+    def _auto_chunk_params(
+        cls,
+        input: AllEmbeddingInputValues | list[list[str]],
+        optional_params: Mapping[str, object],
+    ) -> Mapping[str, object]:
+        is_flat: Final = isinstance(input, str) or all(isinstance(item, str) for item in input)
+        if not is_flat or optional_params.get("input_type") == "query":
+            return {}
+        return {
+            "enable_auto_chunking": True,
+            "chunk_size": cls.AUTO_CHUNK_SIZE,
+            "input_type": "document",
+        }
+
     def transform_embedding_response(
         self,
         model: str,
@@ -124,9 +146,11 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig):
         try:
             raw_response_json: Final = raw_response.json()
         except Exception:
-            raise VoyageError(message=raw_response.text, status_code=raw_response.status_code)
+            raise VoyageError(
+                message=raw_response.text,
+                status_code=raw_response.status_code,
+            )
 
-        # model_response.usage
         model_response.model = raw_response_json.get("model")
         model_response.data = raw_response_json.get("data")
         model_response.object = raw_response_json.get("object")

litellm/llms/wandb/chat/transformation.py

@@ -6,10 +6,17 @@ This is OpenAI compatible - no translation needed / occurs
 
 from typing import Final
 
+import litellm
 from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig
 
 
 class WandbConfig(OpenAIGPTConfig):
+    def get_supported_openai_params(self, model: str) -> list[str]:  # mutable-ok: inherited contract
+        supported_params: Final = super().get_supported_openai_params(model)
+        if litellm.supports_reasoning(model=model, custom_llm_provider="wandb"):
+            return supported_params + ["reasoning_effort"]  # mutable-ok: inherited contract
+        return supported_params
+
     def map_openai_params(
         self,
         non_default_params: dict,

litellm/model_prices_and_context_window_backup.json

@@ -48655,6 +48655,7 @@
         "output_cost_per_token": 0.0
     },
     "wandb/openai/gpt-oss-120b": {
+        "supports_reasoning": true,
         "max_tokens": 131072,
         "max_input_tokens": 131072,
         "max_output_tokens": 131072,
@@ -48665,6 +48666,7 @@
         "source": "https://wandb.ai/site/pricing/tokens/"
     },
     "wandb/openai/gpt-oss-20b": {
+        "supports_reasoning": true,
         "max_tokens": 131072,
         "max_input_tokens": 131072,
         "max_output_tokens": 131072,
@@ -48675,6 +48677,7 @@
         "source": "https://wandb.ai/site/pricing/tokens/"
     },
     "wandb/zai-org/GLM-4.5": {
+        "supports_reasoning": true,
         "max_tokens": 131072,
         "max_input_tokens": 131072,
         "max_output_tokens": 131072,
@@ -48703,6 +48706,7 @@
         "source": "https://wandb.ai/site/pricing/tokens/"
     },
     "wandb/Qwen/Qwen3-235B-A22B-Thinking-2507": {
+        "supports_reasoning": true,
         "max_tokens": 262144,
         "max_input_tokens": 262144,
         "max_output_tokens": 262144,
@@ -48759,6 +48763,7 @@
         "source": "https://wandb.ai/site/pricing/tokens/"
     },
     "wandb/deepseek-ai/DeepSeek-V3.1": {
+        "supports_reasoning": true,
         "max_tokens": 128000,
         "max_input_tokens": 161000,
         "max_output_tokens": 128000,
@@ -48769,6 +48774,7 @@
         "source": "https://wandb.ai/site/pricing/tokens/"
     },
     "wandb/deepseek-ai/DeepSeek-R1-0528": {
+        "supports_reasoning": true,
         "max_tokens": 161000,
         "max_input_tokens": 161000,
         "max_output_tokens": 161000,
@@ -58630,6 +58636,7 @@
         "supports_vision": false
     },
     "wandb/deepseek-ai/DeepSeek-V4-Flash": {
+        "supports_reasoning": true,
         "max_tokens": 1048576,
         "max_input_tokens": 1048576,
         "input_cost_per_token": 1.4e-07,
@@ -58642,6 +58649,7 @@
         "source": "https://wandb.ai/site/pricing/tokens/"
     },
     "wandb/deepseek-ai/DeepSeek-V4-Flash-0731": {
+        "supports_reasoning": true,
         "max_tokens": 262144,
         "max_input_tokens": 262144,
         "input_cost_per_token": 1.3e-07,
@@ -58654,6 +58662,7 @@
         "source": "https://wandb.ai/site/pricing/tokens/"
     },
     "wandb/deepseek-ai/DeepSeek-V4-Pro": {
+        "supports_reasoning": true,
         "max_tokens": 1048576,
         "max_input_tokens": 1048576,
         "input_cost_per_token": 1.15e-06,
@@ -58666,6 +58675,7 @@
         "source": "https://wandb.ai/site/pricing/tokens/"
     },
     "wandb/google/gemma-4-31B-it": {
+        "supports_reasoning": true,
         "max_tokens": 262144,
         "max_input_tokens": 262144,
         "input_cost_per_token": 1e-07,
@@ -58706,6 +58716,7 @@
         "source": "https://wandb.ai/site/pricing/tokens/"
     },
     "wandb/MiniMaxAI/MiniMax-M3": {
+        "supports_reasoning": true,
         "max_tokens": 262144,
         "max_input_tokens": 262144,
         "input_cost_per_token": 2.3e-07,
@@ -58718,6 +58729,7 @@
         "source": "https://wandb.ai/site/pricing/tokens/"
     },
     "wandb/moonshotai/Kimi-K2.7-Code": {
+        "supports_reasoning": true,
         "max_tokens": 262144,
         "max_input_tokens": 262144,
         "input_cost_per_token": 7.1e-07,
@@ -58730,6 +58742,7 @@
         "source": "https://wandb.ai/site/pricing/tokens/"
     },
     "wandb/moonshotai/Kimi-K2.6": {
+        "supports_reasoning": true,
         "max_tokens": 262144,
         "max_input_tokens": 262144,
         "input_cost_per_token": 6.5e-07,
@@ -58742,6 +58755,7 @@
         "source": "https://wandb.ai/site/pricing/tokens/"
     },
     "wandb/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B": {
+        "supports_reasoning": true,
         "max_tokens": 262144,
         "max_input_tokens": 262144,
         "input_cost_per_token": 1e-07,
@@ -58754,6 +58768,7 @@
         "source": "https://wandb.ai/site/pricing/tokens/"
     },
     "wandb/nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B": {
+        "supports_reasoning": true,
         "max_tokens": 262144,
         "max_input_tokens": 262144,
         "input_cost_per_token": 7.5e-07,
@@ -58776,6 +58791,7 @@
         "source": "https://wandb.ai/site/pricing/tokens/"
     },
     "wandb/Qwen/Qwen3.8-27B": {
+        "supports_reasoning": true,
         "max_tokens": 262144,
         "max_input_tokens": 262144,
         "input_cost_per_token": 4e-07,
@@ -58788,6 +58804,7 @@
         "source": "https://wandb.ai/site/pricing/tokens/"
     },
     "wandb/Qwen/Qwen3.6-35B-A3B": {
+        "supports_reasoning": true,
         "max_tokens": 262144,
         "max_input_tokens": 262144,
         "input_cost_per_token": 2.5e-07,
@@ -58798,6 +58815,7 @@
         "source": "https://wandb.ai/site/pricing/tokens/"
     },
     "wandb/Qwen/Qwen3.6-27B": {
+        "supports_reasoning": true,
         "max_tokens": 262144,
         "max_input_tokens": 262144,
         "input_cost_per_token": 6e-07,
@@ -58810,6 +58828,7 @@
         "source": "https://wandb.ai/site/pricing/tokens/"
     },
     "wandb/Qwen/Qwen3.5-35B-A3B": {
+        "supports_reasoning": true,
         "max_tokens": 262144,
         "max_input_tokens": 262144,
         "input_cost_per_token": 2.5e-07,
@@ -58829,7 +58848,28 @@
         "supports_vision": false,
         "source": "https://wandb.ai/site/pricing/tokens/"
     },
+    "wandb/deepseek-ai/DeepSeek-V4-Pro-0813": {
+        "litellm_provider": "wandb",
+        "mode": "chat",
+        "supports_reasoning": true,
+        "input_cost_per_token": 0.00000131,
+        "output_cost_per_token": 0.00000396,
+        "cache_read_input_token_cost": 0.000000044,
+        "supports_prompt_caching": true,
+        "source": "https://wandb.ai/site/pricing/tokens/"
+    },
+    "wandb/ibm-granite/granite-4.2-8b": {
+        "litellm_provider": "wandb",
+        "mode": "chat",
+        "supports_reasoning": true,
+        "input_cost_per_token": 0.0000001,
+        "output_cost_per_token": 0.00000015,
+        "cache_read_input_token_cost": 0.00000005,
+        "supports_prompt_caching": true,
+        "source": "https://wandb.ai/site/pricing/tokens/"
+    },
     "wandb/zai-org/GLM-5.2": {
+        "supports_reasoning": true,
         "max_tokens": 262144,
         "max_input_tokens": 262144,
         "input_cost_per_token": 7.6e-07,

litellm/models/object_permission.py

@@ -23,3 +23,4 @@ class LiteLLM_ObjectPermissionTable(LiteLLMPydanticObjectBase):
     blocked_tools: list[str] | None = []
     search_tools: list[str] | None = []
     mcp_tool_search_enabled: bool | None = None
+    skills: list[str] | None = None

litellm/proxy/_experimental/mcp_server/db.py

@@ -6,10 +6,18 @@ from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
 from datetime import datetime, timedelta, timezone
 from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypedDict, cast
 
+from fastapi import HTTPException
+from typing_extensions import ReadOnly
+
 from litellm._logging import verbose_proxy_logger
 from litellm._uuid import uuid
 from litellm.constants import MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS
 from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
+from litellm.proxy._experimental.mcp_server.oauth_identity_binding import (
+    RefreshTokenPresented,
+    credential_binding_matches,
+    enforce_oauth_identity_binding,
+)
 from litellm.proxy._experimental.mcp_server.oauth_utils import build_upstream_oauth2_token_request
 from litellm.proxy._types import (
     LiteLLM_MCPServerTable,
@@ -117,6 +125,7 @@ class _OAuthCredentialAccessToken(TypedDict):
 
 
 class OAuthCredentialPayload(_OAuthCredentialAccessToken, total=False):
+    identity_binding_proof: ReadOnly[str]
     type: str
     refresh_token: str
     expires_at: str
@@ -1393,6 +1402,7 @@ async def store_user_oauth_credential(
     expires_in: int | None = None,
     scopes: list[str] | None = None,
     skip_byok_guard: bool = False,
+    identity_binding_proof: str | None = None,
 ) -> None:
     """Persist an OAuth2 access token for a user+server pair.
 
@@ -1409,6 +1419,7 @@ async def store_user_oauth_credential(
         "type": "oauth2",
         "access_token": access_token,
         "connected_at": datetime.now(timezone.utc).isoformat(),
+        **({"identity_binding_proof": identity_binding_proof} if identity_binding_proof else {}),
     }
     if refresh_token:
         payload["refresh_token"] = refresh_token
@@ -1628,6 +1639,11 @@ async def refresh_user_oauth_token(
     warning and returns ``None`` — the caller is responsible for clearing the
     stale credential and triggering re-authentication.
     """
+    binding: Final = server.oauth_identity_binding
+    if binding is not None and binding.mode == "enforce":
+        if not await credential_binding_matches(binding, user_id, server.server_id, cred):
+            return None
+
     refresh_token: Final[str | None] = cred.get("refresh_token")
     token_url: Final[str | None] = getattr(server, "effective_token_url", None) or getattr(server, "token_url", None)
     server_id: Final[str] = getattr(server, "server_id", "")
@@ -1677,6 +1693,19 @@ async def refresh_user_oauth_token(
         )
         return None
 
+    try:
+        binding_proof: Final = await enforce_oauth_identity_binding(
+            server=server,
+            token_response=body,
+            litellm_user_id=user_id,
+            grant_type="refresh_token",
+            refresh_ownership=RefreshTokenPresented(refresh_token),
+        )
+    except HTTPException as exc:
+        if exc.status_code != 403:
+            raise
+        return None
+
     access_token: Final[str | None] = body.get("access_token")
     if not access_token:
         verbose_proxy_logger.warning(
@@ -1709,6 +1738,7 @@ async def refresh_user_oauth_token(
         refresh_token=new_refresh_token,
         expires_in=expires_in,
         scopes=scopes,
+        identity_binding_proof=binding_proof,
         skip_byok_guard=True,  # Row is already OAuth2; skip the extra find_unique check
     )
 
@@ -1742,6 +1772,10 @@ async def resolve_valid_user_oauth_token(
     grant: Final = oauth_grant_state(cred)
     if cred is None or grant == "absent":
         return None
+    binding: Final = server.oauth_identity_binding
+    if binding is not None and binding.mode == "enforce":
+        if not await credential_binding_matches(binding, user_id, server.server_id, cred):
+            return None
     if grant == "valid":
         return cred
     if prisma_client is None:
@@ -1782,7 +1816,17 @@ async def resolve_user_oauth_access_token(
             mcp_per_user_token_cache,
         )
 
-        if prefetched_creds is None:
+        binding: Final = server.oauth_identity_binding
+        enforce_binding: Final = binding is not None and binding.mode == "enforce"
+        if prefetched_creds is None and enforce_binding and binding is not None:
+            bound_token: Final = await mcp_per_user_token_cache.get_token(user_id, server_id)
+            if bound_token is not None:
+                if await credential_binding_matches(
+                    binding, user_id, server_id, {"identity_binding_proof": bound_token.identity_binding_proof}
+                ):
+                    return bound_token.access_token
+                await mcp_per_user_token_cache.delete(user_id, server_id)
+        if prefetched_creds is None and not enforce_binding:
             cached_token: Final = await mcp_per_user_token_cache.get(user_id, server_id)
             if cached_token is not None:
                 return cached_token
@@ -1816,7 +1860,9 @@ async def resolve_user_oauth_access_token(
         access_token: Final[str] = cred["access_token"]
         if prefetched_creds is None:
             ttl: Final = _compute_per_user_token_ttl(server, _remaining_token_seconds(cred.get("expires_at")))
-            await mcp_per_user_token_cache.set(user_id, server_id, access_token, ttl)
+            await mcp_per_user_token_cache.set(
+                user_id, server_id, access_token, ttl, identity_binding_proof=cred.get("identity_binding_proof")
+            )
         return access_token
     except Exception as e:
         verbose_proxy_logger.warning(

litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py

@@ -57,6 +57,11 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import (
     relative_request_url,
     revoke_refresh_token,
 )
+from litellm.proxy._experimental.mcp_server.oauth_identity_binding import (
+    RefreshOwnershipProven,
+    RefreshTokenPresented,
+    enforce_oauth_identity_binding,
+)
 from litellm.proxy._experimental.mcp_server.oauth_utils import (
     TOKEN_NO_CACHE_HEADERS,
     build_upstream_oauth2_token_request,
@@ -139,6 +144,7 @@ def encode_state_with_base_url(
     dcr_client_id: str | None = None,
     dcr_client_secret: str | None = None,
     dcr_token_endpoint_auth_method: MCPTokenEndpointAuthMethod | None = None,
+    oauth_nonce: str | None = None,
 ) -> str:
     """
     Encode the base_url, original state, and PKCE parameters using encryption.
@@ -149,9 +155,8 @@ def encode_state_with_base_url(
         code_challenge: PKCE code challenge from client
         code_challenge_method: PKCE code challenge method from client
         client_redirect_uri: Original redirect_uri from client
-        litellm_user_id: The SSO-authenticated litellm user captured at the bridge authorize
-            (interactive dcr_bridge oauth_delegate only); the callback seals it into the gateway
-            authorization code so the token mint can bind the envelope to this user
+        litellm_user_id: The authenticated user captured for bridge or identity-bound per-user OAuth;
+            the callback seals this credential owner into the authorization code
         mcp_server_id: The server the flow targets, sealed alongside litellm_user_id (bridge) or
             dcr_client_id (ephemeral mint) so the gateway code cannot be replayed against another
             server
@@ -169,6 +174,7 @@ def encode_state_with_base_url(
         An encrypted string that encodes all values
     """
     state_data: Final = {
+        "oauth_nonce": oauth_nonce,
         "base_url": base_url,
         "original_state": original_state,
         "code_challenge": code_challenge,
@@ -210,10 +216,10 @@ _BRIDGE_AUTH_CODE_PREFIX: Final = "llm_bcode_"
 
 
 class _BridgeAuthorizationCode(BaseModel):
-    """The identity and upstream code the gateway seals into the authorization code it hands a DCR
-    client for an interactive dcr_bridge oauth_delegate sign-in, recovered at the token endpoint."""
+    """Authenticated caller and upstream code sealed for bridge or identity-bound per-user OAuth."""
 
     model_config = ConfigDict(frozen=True)
+    oauth_nonce: str | None = None
     upstream_code: str = Field(min_length=1)
     litellm_user_id: str = Field(min_length=1)
     mcp_server_id: str = Field(min_length=1)
@@ -225,7 +231,12 @@ def is_bridge_authorization_code(code: str) -> bool:
     return code.startswith(_BRIDGE_AUTH_CODE_PREFIX)
 
 
-def seal_bridge_authorization_code(upstream_code: str, litellm_user_id: str, mcp_server_id: str) -> str:
+def seal_bridge_authorization_code(
+    upstream_code: str,
+    litellm_user_id: str,
+    mcp_server_id: str,
+    oauth_nonce: str | None = None,
+) -> str:
     """Seal the upstream authorization code and the SSO-captured litellm user into a gateway
     authorization code. The DCR client only echoes this opaque value back at the token endpoint; the
     gateway decrypts it there to recover the user (to bind the envelope) and the upstream code (to
@@ -234,7 +245,12 @@ def seal_bridge_authorization_code(upstream_code: str, litellm_user_id: str, mcp
     authenticated symmetric helper (the same family the OAuth state uses), so the client can neither
     read nor forge it."""
     payload: Final = json.dumps(
-        {"upstream_code": upstream_code, "litellm_user_id": litellm_user_id, "mcp_server_id": mcp_server_id},
+        {
+            "upstream_code": upstream_code,
+            "litellm_user_id": litellm_user_id,
+            "mcp_server_id": mcp_server_id,
+            "oauth_nonce": oauth_nonce,
+        },
         sort_keys=True,
     )
     return _BRIDGE_AUTH_CODE_PREFIX + encrypt_value_helper(payload)
@@ -547,6 +563,7 @@ async def _store_per_user_token_server_side(
     server: MCPServer,
     user_id: str,
     token_response: dict[str, Any],
+    identity_binding_proof: str | None = None,
 ) -> None:
     """Persist the OAuth token server-side and warm the Redis cache.
 
@@ -588,6 +605,7 @@ async def _store_per_user_token_server_side(
             refresh_token=refresh_token,
             expires_in=expires_in,
             scopes=scopes,
+            identity_binding_proof=identity_binding_proof,
         )
         verbose_logger.info(
             "_store_per_user_token_server_side: stored token for user=%s server=%s",
@@ -616,6 +634,7 @@ async def _store_per_user_token_server_side(
         server_id=server.server_id,
         access_token=access_token,
         ttl=ttl,
+        identity_binding_proof=identity_binding_proof,
     )
 
 
@@ -854,6 +873,11 @@ async def authorize_with_server(
             ),
         )
 
+    binding: Final = resolved_server.oauth_identity_binding
+    enforce_binding: Final = binding is not None and binding.mode == "enforce"
+    if enforce_binding:
+        _require_s256_pkce(code_challenge, code_challenge_method)
+
     if resolved_server.is_dcr_bridge:
         # Enforce S256 PKCE on both bridge arms. The relay arm forwards the validated,
         # now-non-optional pair to the upstream authorize; the short-circuit arm keeps
@@ -884,19 +908,16 @@ async def authorize_with_server(
     base_url: Final = urlunparse(parsed._replace(query=""))
     request_base_url: Final = get_request_base_url(request)
 
-    # Interactive dcr_bridge oauth_delegate sign-in: this arm runs the gateway /callback and /token in
-    # the loop, so the gateway can capture the litellm user here (from the browser's UI session) and
-    # carry it to the back-channel token mint. Seal the SSO user and the target server into the state;
-    # the callback reads them back to mint the gateway authorization code. A DCR client cannot present a
-    # litellm key, so the browser session is the only identity source; without one there is nothing to
-    # bind, so send the user through login first. Every other oauth2 server keeps the identity-less state.
+    # Seal the authenticated caller into state so the token exchange cannot select another credential owner.
     litellm_user_id: str | None = None
-    if resolved_server.is_dcr_bridge and resolved_server.is_oauth_delegate:
+    if enforce_binding or (resolved_server.is_dcr_bridge and resolved_server.is_oauth_delegate):
         from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import (  # noqa: PLC0415  # inline import avoids a module-load circular import
             _user_id_from_session_cookie,
         )
 
-        litellm_user_id = _user_id_from_session_cookie(request)
+        litellm_user_id = (
+            await _extract_user_id_from_request(request) if enforce_binding else None
+        ) or _user_id_from_session_cookie(request)
         if litellm_user_id is None:
             return _redirect_to_litellm_login(request)
         denial: Final = await _bridge_authorize_access_denial(
@@ -908,9 +929,11 @@ async def authorize_with_server(
         if denial is not None:
             return denial
 
+    oauth_nonce: Final = secrets.token_urlsafe(32) if enforce_binding else None
     encoded_state: Final = encode_state_with_base_url(
         base_url=base_url,
         original_state=state,
+        oauth_nonce=oauth_nonce,
         code_challenge=code_challenge,
         code_challenge_method=code_challenge_method,
         client_redirect_uri=redirect_uri,
@@ -930,11 +953,16 @@ async def authorize_with_server(
         "state": relay_state,
         "response_type": response_type or "code",
     }
+    if oauth_nonce:
+        params["nonce"] = oauth_nonce
     if scope:
         params["scope"] = scope
     elif resolved_server.scopes:
         params["scope"] = " ".join(resolved_server.scopes)
 
+    if enforce_binding and "openid" not in params.get("scope", "").split():
+        params["scope"] = f"openid {params.get('scope', '')}".strip()
+
     if code_challenge:
         params["code_challenge"] = code_challenge
     if code_challenge_method:
@@ -1015,6 +1043,12 @@ async def exchange_token_with_server(
     except TokenEndpointAuthConfigError as exc:
         raise HTTPException(status_code=400, detail=str(exc)) from exc
 
+    request_user_id: Final = (
+        await _extract_user_id_from_request(request)
+        if resolved_server.needs_user_oauth_token or resolved_server.oauth_identity_binding is not None
+        else None
+    )
+
     bridge_identity: _BridgeAuthorizationCode | None = None
     bridge_mint_ready: _BridgeMintReady | None = None
     bridge_upstream_refresh: SecretStr | None = None
@@ -1051,7 +1085,13 @@ async def exchange_token_with_server(
         refresh_request_scope = scope or bridge_upstream_scope
         if refresh_request_scope:
             token_data["scope"] = refresh_request_scope
+        refresh_ownership = (  # rebind-ok: grant-specific branches assign one ownership value
+            RefreshOwnershipProven()
+            if bridge_upstream_refresh is not None
+            else RefreshTokenPresented(upstream_refresh_token)
+        )
     else:
+        refresh_ownership = None  # rebind-ok: grant-specific branches assign one ownership value
         if not code:
             raise HTTPException(
                 status_code=400,
@@ -1070,6 +1110,14 @@ async def exchange_token_with_server(
                     detail="Authorization code was issued for a different MCP server",
                 )
             code = bridge_identity.upstream_code
+        binding: Final = resolved_server.oauth_identity_binding
+        if binding is not None and binding.mode == "enforce":
+            if bridge_identity is None or not bridge_identity.oauth_nonce:
+                raise HTTPException(status_code=403, detail={"error": "oauth_identity_binding_failed"})
+            if request_user_id is not None and request_user_id != bridge_identity.litellm_user_id:
+                raise HTTPException(status_code=403, detail={"error": "oauth_principal_mismatch"})
+            if not code_verifier:
+                raise HTTPException(status_code=403, detail={"error": "oauth_identity_binding_failed"})
         bridge_token_relay: Final = _dcr_bridge_relays_client_registration(resolved_server)
         if bridge_token_relay and not redirect_uri:
             raise HTTPException(
@@ -1097,6 +1145,16 @@ async def exchange_token_with_server(
                 return _bridge_mint_error_response(prepared)
             bridge_mint_ready = prepared
 
+    refresh_binding: Final = resolved_server.oauth_identity_binding
+    if grant_type == "refresh_token" and refresh_binding is not None and refresh_binding.mode == "enforce":
+        await enforce_oauth_identity_binding(
+            server=resolved_server,
+            token_response={},
+            litellm_user_id=request_user_id,
+            grant_type=grant_type,
+            refresh_ownership=refresh_ownership,
+        )
+
     async_client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check)
     try:
         response: Final = await async_client.post(
@@ -1137,17 +1195,34 @@ async def exchange_token_with_server(
             server_id=resolved_server.server_id,
         )
 
+    # Bind the exchanged token to the LiteLLM caller BEFORE it is returned, stored, or cached, so a
+    # token minted for a different upstream principal never becomes usable under the caller's user_id.
+    resolved_user_id: Final = bridge_identity.litellm_user_id if bridge_identity else request_user_id
+    binding_proof: Final = (
+        await enforce_oauth_identity_binding(
+            server=resolved_server,
+            token_response=token_response,
+            litellm_user_id=resolved_user_id,
+            grant_type=grant_type,
+            refresh_ownership=refresh_ownership,
+            expected_nonce=bridge_identity.oauth_nonce if bridge_identity else None,
+        )
+        if isinstance(token_response, dict)
+        else None
+    )
+
     # Store server-side when the server is configured for per-user OAuth and
     # the calling client has provided a valid LiteLLM identity.
     # Errors are non-fatal: the token is still returned to the client.
     if resolved_server.needs_user_oauth_token:
-        user_id: Final = await _extract_user_id_from_request(request)
+        user_id: Final = resolved_user_id
         if user_id:
             try:
                 await _store_per_user_token_server_side(
                     server=resolved_server,
                     user_id=user_id,
                     token_response=token_response,
+                    identity_binding_proof=binding_proof,
                 )
             except Exception as exc:
                 verbose_logger.warning(
@@ -2134,7 +2209,10 @@ async def callback(
         forwarded_code = code
         if isinstance(litellm_user_id, str) and litellm_user_id and isinstance(mcp_server_id, str) and mcp_server_id:
             forwarded_code = seal_bridge_authorization_code(
-                upstream_code=code, litellm_user_id=litellm_user_id, mcp_server_id=mcp_server_id
+                upstream_code=code,
+                litellm_user_id=litellm_user_id,
+                mcp_server_id=mcp_server_id,
+                oauth_nonce=state_data.get("oauth_nonce"),
             )
         elif isinstance(dcr_client_id, str) and dcr_client_id and isinstance(mcp_server_id, str) and mcp_server_id:
             forwarded_code = seal_passthrough_authorization_code(

litellm/proxy/_experimental/mcp_server/mcp_server_manager.py

@@ -2444,6 +2444,8 @@ class MCPServerManager:
                 allow_elicitation=bool(server_config.get("allow_elicitation", False)),
                 timeout=server_config.get("timeout", None),
                 max_concurrent_requests=server_config.get("max_concurrent_requests", None),
+                token_validation=server_config.get("token_validation", None),
+                oauth_identity_binding=server_config.get("oauth_identity_binding", None),
             )
             self._assign_unique_short_prefix(new_server)
             _warn_internal_delegate_pkce_if_applicable(new_server, source="config")

litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py

@@ -28,6 +28,8 @@ from litellm.proxy._experimental.mcp_server.oauth_utils import (
     build_upstream_oauth2_token_request,
     resolve_upstream_resource,
 )
+from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import OAuthToken
+from litellm.proxy._experimental.mcp_server.outbound_credentials.token_cache_codec import OAuthTokenCacheCodec
 from litellm.proxy.common_utils.encrypt_decrypt_utils import (
     decrypt_value_helper,
     encrypt_value_helper,
@@ -233,8 +235,17 @@ class MCPPerUserTokenCache:
     def _cache_key(self, user_id: str, server_id: str) -> str:
         return f"{MCP_PER_USER_TOKEN_REDIS_KEY_PREFIX}:{user_id}:{server_id}"
 
+    def _codec(self) -> OAuthTokenCacheCodec:
+        return OAuthTokenCacheCodec(
+            encrypt_value_helper,
+            lambda blob: decrypt_value_helper(blob, key="mcp_per_user_token", exception_type="debug"),
+        )
+
     async def get(self, user_id: str, server_id: str) -> str | None:
-        """Return the plaintext access_token, or None on miss/error."""
+        token: Final = await self.get_token(user_id, server_id)
+        return token.access_token if token is not None else None
+
+    async def get_token(self, user_id: str, server_id: str) -> OAuthToken | None:
         try:
             from litellm.proxy.proxy_server import user_api_key_cache  # noqa: PLC0415
 
@@ -242,12 +253,7 @@ class MCPPerUserTokenCache:
             encrypted: Final = await user_api_key_cache.async_get_cache(key)
             if encrypted is None:
                 return None
-            plaintext: Final = decrypt_value_helper(
-                encrypted,
-                key="mcp_per_user_token",
-                exception_type="debug",
-            )
-            return plaintext or None
+            return self._codec().decode(encrypted)
         except Exception as exc:
             verbose_logger.debug(
                 "MCPPerUserTokenCache.get failed for user=%s server=%s: %s",
@@ -263,13 +269,16 @@ class MCPPerUserTokenCache:
         server_id: str,
         access_token: str,
         ttl: int,
+        identity_binding_proof: str | None = None,
     ) -> None:
         """Store NaCl-encrypted access_token in Redis with the given TTL."""
         try:
             from litellm.proxy.proxy_server import user_api_key_cache  # noqa: PLC0415
 
             key: Final = self._cache_key(user_id, server_id)
-            encrypted: Final = encrypt_value_helper(access_token)
+            encrypted: Final = self._codec().encode(
+                OAuthToken(access_token=access_token, identity_binding_proof=identity_binding_proof)
+            )
             await user_api_key_cache.async_set_cache(key, encrypted, ttl=ttl)
             verbose_logger.debug(
                 "MCPPerUserTokenCache.set: cached token for user=%s server=%s ttl=%ds",

litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py

@@ -0,0 +1,423 @@
+"""Per-user OAuth identity binding: verify the upstream OIDC principal matches the LiteLLM caller.
+
+Closes the confused-deputy gap where a browser authenticated upstream as one principal produces a
+token that the relay stores under a different, LiteLLM-authenticated principal: before the token
+endpoint returns, stores, or caches an exchanged token for an identity-bound server, the id_token
+is validated (signature via the pinned issuer's JWKS, issuer, audience, expiry) and its principal
+claim is compared to the caller's trusted LiteLLM identity. Mismatches fail closed in enforce mode
+and are logged in audit mode.
+"""
+
+import hashlib
+import hmac
+import json
+from collections.abc import Awaitable, Callable, Mapping, Sequence
+from dataclasses import dataclass
+from typing import Final, Literal, Protocol, TypeAlias
+
+import jwt
+from fastapi import HTTPException
+from jwt.types import Options
+from typing_extensions import assert_never
+
+from litellm._logging import verbose_logger
+from litellm.caching.in_memory_cache import InMemoryCache
+from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
+from litellm.types.llms.custom_http import httpxSpecialProvider
+from litellm.types.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer
+
+_ALLOWED_ID_TOKEN_ALGORITHMS: Final = (
+    "RS256",
+    "RS384",
+    "RS512",
+    "ES256",
+    "ES384",
+    "ES512",
+    "PS256",
+    "PS384",
+    "PS512",
+)
+_JWKS_CACHE_TTL_SECONDS: Final = 3600
+_jwks_cache: Final = InMemoryCache(default_ttl=_JWKS_CACHE_TTL_SECONDS)
+
+JwksFetcher: TypeAlias = Callable[
+    [MCPOAuthIdentityBinding],  # mutable-ok: Callable parameter syntax requires a list
+    Awaitable[Sequence[Mapping[str, object]]],
+]
+CallerPrincipalLoader: TypeAlias = Callable[
+    [str, MCPOAuthIdentityBinding],  # mutable-ok: Callable parameter syntax requires a list
+    Awaitable[str | None],
+]
+
+
+@dataclass(frozen=True, slots=True)
+class VerifiedRefreshToken:
+    refresh_token: str
+    binding_proof: str
+
+
+StoredRefreshTokenLoader: TypeAlias = Callable[
+    [str, str, MCPOAuthIdentityBinding],  # mutable-ok: Callable parameter syntax requires a list
+    Awaitable[VerifiedRefreshToken | None],
+]
+
+_RejectionCode: TypeAlias = Literal["oauth_principal_mismatch", "oauth_identity_binding_failed"]
+
+
+@dataclass(frozen=True, slots=True)
+class _BindingRejection:
+    code: _RejectionCode
+    description: str
+
+
+@dataclass(frozen=True, slots=True)
+class RefreshOwnershipProven:
+    """The gateway itself unwrapped the upstream refresh token from a sealed per-user envelope."""
+
+
+@dataclass(frozen=True, slots=True)
+class RefreshTokenPresented:
+    refresh_token: str
+
+
+RefreshOwnership: TypeAlias = RefreshOwnershipProven | RefreshTokenPresented | None
+
+
+class BindingValidator(Protocol):
+    async def __call__(
+        self,
+        *,
+        server: MCPServer,
+        token_response: Mapping[str, object],
+        litellm_user_id: str | None,
+        grant_type: str,
+        refresh_ownership: RefreshOwnership,
+    ) -> str | None: ...
+
+
+async def _fetch_issuer_jwks(binding: MCPOAuthIdentityBinding) -> Sequence[Mapping[str, object]]:
+    jwks_url: Final[str] = binding.jwks_url or await _discover_jwks_url(binding.issuer)
+    cached: Final = await _jwks_cache.async_get_cache(jwks_url)
+    if isinstance(cached, list):
+        return cached
+    client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check)
+    response: Final = await client.get(jwks_url)
+    response.raise_for_status()
+    document: Final = response.json()
+    keys: Final = document.get("keys") if isinstance(document, dict) else None
+    if not isinstance(keys, list):
+        raise TypeError(f"JWKS document at {jwks_url} has no 'keys' array")
+    await _jwks_cache.async_set_cache(jwks_url, keys, ttl=_JWKS_CACHE_TTL_SECONDS)
+    return keys
+
+
+async def _discover_jwks_url(issuer: str) -> str:
+    discovery_url: Final = f"{issuer.rstrip('/')}/.well-known/openid-configuration"
+    client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check)
+    response: Final = await client.get(discovery_url)
+    response.raise_for_status()
+    metadata: Final = response.json()
+    jwks_uri: Final = metadata.get("jwks_uri") if isinstance(metadata, dict) else None
+    if not isinstance(jwks_uri, str) or not jwks_uri:
+        raise ValueError(f"OIDC discovery at {discovery_url} returned no jwks_uri")
+    return jwks_uri
+
+
+def _select_signing_key(id_token: str, keys: Sequence[Mapping[str, object]]) -> "jwt.PyJWK | _BindingRejection":
+    header: Final = jwt.get_unverified_header(id_token)
+    kid: Final = header.get("kid")
+    for key in keys:
+        if kid is None or key.get("kid") == kid:
+            return jwt.PyJWK(dict(key))  # mutable-ok: PyJWT requires a concrete JWK dictionary
+    return _BindingRejection(
+        code="oauth_identity_binding_failed",
+        description=f"id_token signing key (kid={kid!r}) not found in the issuer's JWKS",
+    )
+
+
+def _decode_id_token(
+    id_token: str,
+    binding: MCPOAuthIdentityBinding,
+    signing_key: "jwt.PyJWK",
+) -> "Mapping[str, object] | _BindingRejection":
+    try:
+        decode_options: Final[Options] = {"require": ("iss", "exp", "aud", "sub", "iat")}
+        return jwt.decode(
+            id_token,
+            signing_key.key,
+            algorithms=_ALLOWED_ID_TOKEN_ALGORITHMS,
+            issuer=binding.issuer,
+            audience=binding.audiences,
+            options=decode_options,
+        )
+    except jwt.InvalidTokenError as exc:
+        return _BindingRejection(
+            code="oauth_identity_binding_failed",
+            description=f"id_token validation failed: {exc}",
+        )
+
+
+def _upstream_principal(
+    claims: Mapping[str, object],
+    binding: MCPOAuthIdentityBinding,
+) -> "str | _BindingRejection":
+    principal: Final = claims.get(binding.principal_claim)
+    if not isinstance(principal, str) or not principal:
+        return _BindingRejection(
+            code="oauth_identity_binding_failed",
+            description=f"id_token has no usable '{binding.principal_claim}' claim",
+        )
+    if (
+        binding.principal_claim == "email"
+        and binding.require_email_verified
+        and claims.get("email_verified") is not True
+    ):
+        return _BindingRejection(
+            code="oauth_identity_binding_failed",
+            description="id_token email is not verified (email_verified is not true)",
+        )
+    return principal
+
+
+async def _load_caller_principal(litellm_user_id: str, binding: MCPOAuthIdentityBinding) -> str | None:
+    if binding.caller_field == "user_id":
+        return litellm_user_id
+    from litellm.proxy._experimental.mcp_server.bridge_token_flow import (  # noqa: PLC0415  # inline import avoids a module-load circular import
+        load_active_user_by_id,
+    )
+
+    loaded: Final = await load_active_user_by_id(litellm_user_id)
+    if isinstance(loaded, str):
+        return None
+    return loaded.user_email
+
+
+async def _load_stored_refresh_token(
+    litellm_user_id: str, server_id: str, binding: MCPOAuthIdentityBinding
+) -> VerifiedRefreshToken | None:
+    try:
+        from litellm.proxy._experimental.mcp_server.db import (  # noqa: PLC0415  # keep database imports lazy
+            get_user_oauth_credential,
+        )
+        from litellm.proxy.utils import get_prisma_client_or_throw  # noqa: PLC0415  # keep database imports lazy
+
+        prisma_client: Final = get_prisma_client_or_throw(
+            "Database not connected. Cannot verify OAuth refresh token ownership."
+        )
+        cred: Final = await get_user_oauth_credential(
+            prisma_client=prisma_client,
+            user_id=litellm_user_id,
+            server_id=server_id,
+        )
+        if not cred or not await credential_binding_matches(binding, litellm_user_id, server_id, cred):
+            return None
+        refresh_token: Final = cred.get("refresh_token")
+        proof: Final = cred.get("identity_binding_proof")
+        return VerifiedRefreshToken(refresh_token, proof) if refresh_token and proof else None
+    except Exception:  # noqa: BLE001  # a credential lookup failure must fail closed
+        return None
+
+
+async def current_binding_proof(
+    binding: MCPOAuthIdentityBinding,
+    user_id: str,
+    server_id: str,
+    caller_principal_loader: CallerPrincipalLoader = _load_caller_principal,
+) -> str | None:
+    principal: Final = await caller_principal_loader(user_id, binding)
+    if not principal:
+        return None
+    return _binding_proof(binding, user_id, server_id, principal)
+
+
+def _binding_proof(binding: MCPOAuthIdentityBinding, user_id: str, server_id: str, principal: str) -> str:
+    payload: Final = json.dumps(
+        ("oidc-nonce-v1", server_id, user_id, principal, binding.model_dump(mode="json")),
+        sort_keys=True,
+    )
+    return hashlib.sha256(payload.encode()).hexdigest()
+
+
+async def credential_binding_matches(
+    binding: MCPOAuthIdentityBinding,
+    user_id: str,
+    server_id: str,
+    credential: Mapping[str, object],
+    caller_principal_loader: CallerPrincipalLoader = _load_caller_principal,
+) -> bool:
+    stored: Final = credential.get("identity_binding_proof")
+    if not isinstance(stored, str) or not stored:
+        return False
+    expected: Final = await current_binding_proof(binding, user_id, server_id, caller_principal_loader)
+    return expected is not None and hmac.compare_digest(stored, expected)
+
+
+def _principals_match(upstream: str, caller: str, binding: MCPOAuthIdentityBinding) -> bool:
+    if binding.principal_claim == "email" or binding.caller_field == "user_email":
+        return upstream.strip().casefold() == caller.strip().casefold()
+    return upstream == caller
+
+
+async def _evaluate_refresh_ownership(
+    binding: MCPOAuthIdentityBinding,
+    litellm_user_id: str | None,
+    server_id: str,
+    refresh_ownership: RefreshOwnership,
+    stored_refresh_token_loader: StoredRefreshTokenLoader,
+) -> _BindingRejection | str:
+    match refresh_ownership:
+        case RefreshOwnershipProven():
+            return _BindingRejection(
+                code="oauth_identity_binding_failed",
+                description="an identity envelope alone does not prove upstream principal binding",
+            )
+        case None:
+            return _BindingRejection(
+                code="oauth_identity_binding_failed",
+                description="refresh_token grant without an id_token carries no refresh token to prove ownership of",
+            )
+        case RefreshTokenPresented(refresh_token):
+            if not litellm_user_id:
+                return _BindingRejection(
+                    code="oauth_identity_binding_failed",
+                    description="the request carries no resolvable LiteLLM user identity to bind the credential to",
+                )
+            stored: Final = await stored_refresh_token_loader(litellm_user_id, server_id, binding)
+            if stored is None or not hmac.compare_digest(stored.refresh_token, refresh_token):
+                return _BindingRejection(
+                    code="oauth_identity_binding_failed",
+                    description="the presented refresh_token is not the caller's stored credential for this server",
+                )
+            return stored.binding_proof
+    assert_never(refresh_ownership)  # pragma: no cover
+
+
+async def _evaluate_binding(
+    binding: MCPOAuthIdentityBinding,
+    token_response: Mapping[str, object],
+    litellm_user_id: str | None,
+    grant_type: str,
+    server_id: str,
+    refresh_ownership: RefreshOwnership,
+    jwks_fetcher: JwksFetcher,
+    caller_principal_loader: CallerPrincipalLoader,
+    stored_refresh_token_loader: StoredRefreshTokenLoader,
+    expected_nonce: str | None,
+) -> _BindingRejection | str:
+    id_token: Final = token_response.get("id_token")
+    if not isinstance(id_token, str) or not id_token:
+        if grant_type != "refresh_token":
+            return _BindingRejection(
+                code="oauth_identity_binding_failed",
+                description="the upstream token response carries no id_token to bind the credential to a principal",
+            )
+        return await _evaluate_refresh_ownership(
+            binding,
+            litellm_user_id,
+            server_id,
+            refresh_ownership,
+            stored_refresh_token_loader,
+        )
+    if not litellm_user_id:
+        return _BindingRejection(
+            code="oauth_identity_binding_failed",
+            description="the request carries no resolvable LiteLLM user identity to bind the credential to",
+        )
+    try:
+        keys: Final = await jwks_fetcher(binding)
+    except Exception as exc:  # noqa: BLE001  # a JWKS fetch failure must fail closed, not surface as a 500
+        return _BindingRejection(
+            code="oauth_identity_binding_failed",
+            description=f"could not fetch the issuer's JWKS: {exc}",
+        )
+    try:
+        signing_key: Final = _select_signing_key(id_token, keys)
+    except (jwt.PyJWTError, ValueError, TypeError):
+        return _BindingRejection(
+            code="oauth_identity_binding_failed",
+            description="invalid id_token header or issuer signing key",
+        )
+    if isinstance(signing_key, _BindingRejection):
+        return signing_key
+    claims: Final = _decode_id_token(id_token, binding, signing_key)
+    if isinstance(claims, _BindingRejection):
+        return claims
+    if grant_type == "authorization_code" and (binding.mode == "enforce" or expected_nonce is not None):
+        nonce: Final = claims.get("nonce")
+        if not expected_nonce or not isinstance(nonce, str) or not hmac.compare_digest(nonce, expected_nonce):
+            return _BindingRejection(
+                code="oauth_identity_binding_failed",
+                description="id_token nonce does not match the authenticated authorization transaction",
+            )
+    upstream: Final = _upstream_principal(claims, binding)
+    if isinstance(upstream, _BindingRejection):
+        return upstream
+    caller: Final = await caller_principal_loader(litellm_user_id, binding)
+    if not caller:
+        return _BindingRejection(
+            code="oauth_identity_binding_failed",
+            description=f"the LiteLLM user has no '{binding.caller_field}' to compare the upstream principal against",
+        )
+    if not _principals_match(upstream, caller, binding):
+        return _BindingRejection(
+            code="oauth_principal_mismatch",
+            description="The browser account does not match the selected credential owner.",
+        )
+    return _binding_proof(binding, litellm_user_id, server_id, caller)
+
+
+async def enforce_oauth_identity_binding(
+    server: MCPServer,
+    token_response: Mapping[str, object],
+    litellm_user_id: str | None,
+    grant_type: str,
+    refresh_ownership: RefreshOwnership,
+    jwks_fetcher: JwksFetcher = _fetch_issuer_jwks,
+    caller_principal_loader: CallerPrincipalLoader = _load_caller_principal,
+    stored_refresh_token_loader: StoredRefreshTokenLoader = _load_stored_refresh_token,
+    expected_nonce: str | None = None,
+) -> str | None:
+    """Validate the exchanged token's upstream principal against the LiteLLM caller.
+
+    No-op when the server has no binding or it is disabled. In enforce mode a failure raises 403
+    before the caller returns, stores, or caches the token; in audit mode failures are logged only.
+    A refresh_token grant without an id_token is allowed only when the presented refresh token matches
+    the caller's stored credential and that credential was previously identity-validated.
+    """
+    binding: Final = server.oauth_identity_binding
+    if binding is None or binding.mode not in ("audit", "enforce"):
+        return
+    rejection: Final = await _evaluate_binding(
+        binding=binding,
+        token_response=token_response,
+        litellm_user_id=litellm_user_id,
+        grant_type=grant_type,
+        server_id=server.server_id,
+        refresh_ownership=refresh_ownership,
+        jwks_fetcher=jwks_fetcher,
+        caller_principal_loader=caller_principal_loader,
+        stored_refresh_token_loader=stored_refresh_token_loader,
+        expected_nonce=expected_nonce,
+    )
+    if isinstance(rejection, str):
+        return rejection if binding.mode == "enforce" else None
+    if binding.mode == "audit":
+        verbose_logger.warning(
+            "oauth_identity_binding audit: server=%s user=%s grant=%s rejected=%s (%s)",
+            server.server_id,
+            litellm_user_id,
+            grant_type,
+            rejection.code,
+            rejection.description,
+        )
+        return
+    raise HTTPException(
+        status_code=403,
+        detail={
+            "error": rejection.code,
+            "error_description": rejection.description,
+            "server_id": server.server_id,
+            "credential_owner": "caller",
+            "credential_stored": False,
+        },
+    )

litellm/proxy/_experimental/mcp_server/outbound_credentials/authz_code_refresher.py

@@ -14,10 +14,17 @@ import time
 from collections.abc import Awaitable, Callable
 from typing import TYPE_CHECKING, Final, Protocol
 
+from fastapi import HTTPException
+
 from litellm._logging import verbose_logger
 from litellm.proxy._experimental.mcp_server.auth.token_endpoint_auth import (
     TokenEndpointAuthConfigError,
 )
+from litellm.proxy._experimental.mcp_server.oauth_identity_binding import (
+    BindingValidator,
+    RefreshTokenPresented,
+    enforce_oauth_identity_binding,
+)
 from litellm.proxy._experimental.mcp_server.oauth_utils import build_upstream_oauth2_token_request
 from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import (
     OAuthToken,
@@ -39,6 +46,7 @@ class CredentialPersist(Protocol):
         refresh_token: str | None,
         expires_in: int | None,
         scopes: tuple[str, ...] | None,
+        identity_binding_proof: str | None = None,
     ) -> None: ...
 
 
@@ -78,13 +86,23 @@ class AuthorizationCodeRefresher:
         persist: CredentialPersist,
         *,
         clock: Callable[[], float] = time.time,
+        identity_validator: BindingValidator = enforce_oauth_identity_binding,
     ) -> None:
         self._server_lookup = server_lookup
         self._token_endpoint = token_endpoint
         self._persist = persist
         self._clock = clock
+        self._identity_validator = identity_validator
 
     async def refresh(self, user_id: str, server_id: str, token: OAuthToken) -> OAuthToken | None:
+        try:
+            return await self._refresh(user_id, server_id, token)
+        except HTTPException as exc:
+            if exc.status_code != 403:
+                raise
+            return None
+
+    async def _refresh(self, user_id: str, server_id: str, token: OAuthToken) -> OAuthToken | None:
         if token.refresh_token is None:
             return None
         server: Final = self._server_lookup(server_id)
@@ -104,6 +122,15 @@ class AuthorizationCodeRefresher:
         except TokenEndpointAuthConfigError as exc:
             verbose_logger.warning("MCP OAuth refresh misconfigured for server %s: %s", server_id, exc)
             return None
+        binding: Final = server.oauth_identity_binding
+        if binding is not None and binding.mode == "enforce":
+            await self._identity_validator(
+                server=server,
+                token_response={},
+                litellm_user_id=user_id,
+                grant_type="refresh_token",
+                refresh_ownership=RefreshTokenPresented(token.refresh_token),
+            )
         form: Final = {
             "grant_type": "refresh_token",
             "refresh_token": token.refresh_token,
@@ -116,15 +143,34 @@ class AuthorizationCodeRefresher:
         if not isinstance(access_token, str) or not access_token:
             return None
 
+        binding_proof: Final = await self._identity_validator(
+            server=server,
+            token_response=body,
+            litellm_user_id=user_id,
+            grant_type="refresh_token",
+            refresh_ownership=RefreshTokenPresented(token.refresh_token),
+        )
         rotated: Final = body.get("refresh_token")
         new_refresh: Final = rotated if isinstance(rotated, str) and rotated else token.refresh_token
         expires_in: Final = _parse_expires_in(body.get("expires_in"))
         scopes: Final = _parse_scopes(body.get("scope")) or token.scopes
 
-        await self._persist(user_id, server_id, access_token, new_refresh, expires_in, scopes or None)
+        if binding_proof is not None:
+            await self._persist(
+                user_id,
+                server_id,
+                access_token,
+                new_refresh,
+                expires_in,
+                scopes or None,
+                identity_binding_proof=binding_proof,
+            )
+        else:
+            await self._persist(user_id, server_id, access_token, new_refresh, expires_in, scopes or None)
         return OAuthToken(
             access_token=access_token,
             expires_at=self._clock() + expires_in if expires_in is not None else None,
             refresh_token=new_refresh,
             scopes=scopes,
+            identity_binding_proof=binding_proof,
         )

litellm/proxy/_experimental/mcp_server/outbound_credentials/oauth_token_store.py

@@ -41,6 +41,7 @@ class OAuthToken:
     expires_at: float | None = None
     refresh_token: str | None = None
     scopes: tuple[str, ...] = ()
+    identity_binding_proof: str | None = None
 
     def __repr__(self) -> str:
         has_refresh: Final = self.refresh_token is not None

litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py

@@ -12,9 +12,11 @@ from __future__ import annotations
 
 import asyncio
 from collections.abc import Callable, Mapping
+from functools import partial
 from typing import TYPE_CHECKING, Final
 
 from litellm._logging import verbose_logger
+from litellm.proxy._experimental.mcp_server.oauth_identity_binding import credential_binding_matches
 from litellm.proxy._experimental.mcp_server.outbound_credentials.authz_code_refresher import (
     AuthorizationCodeRefresher,
 )
@@ -69,6 +71,7 @@ async def _persist_credential(
     refresh_token: str | None,
     expires_in: int | None,
     scopes: tuple[str, ...] | None,
+    identity_binding_proof: str | None = None,
 ) -> None:
     from litellm.proxy._experimental.mcp_server.db import (  # noqa: PLC0415
         store_user_oauth_credential,
@@ -86,6 +89,7 @@ async def _persist_credential(
         expires_in=expires_in,
         scopes=list(scopes) if scopes else None,
         skip_byok_guard=True,
+        identity_binding_proof=identity_binding_proof,
     )
 
 
@@ -142,12 +146,26 @@ def _runtime_backend_and_coordinator() -> tuple[TokenCacheBackend | None, Refres
     return backend, coordinator, True
 
 
+async def _read_bound_credential(
+    server_lookup: ServerLookup, user_id: str, server_id: str
+) -> Mapping[str, object] | None:
+    credential: Final = await _read_credential(user_id, server_id)
+    server: Final = server_lookup(server_id)
+    binding: Final = server.oauth_identity_binding if server else None
+    if credential is not None and binding is not None and binding.mode == "enforce":
+        if not await credential_binding_matches(binding, user_id, server_id, credential):
+            return None
+    return credential
+
+
 def _build_per_user_oauth_token_store(
     server_lookup: ServerLookup,
 ) -> tuple[CachedOAuthTokenStore, bool]:
     backend, coordinator, uses_redis = _runtime_backend_and_coordinator()
     refresher: Final = AuthorizationCodeRefresher(server_lookup, _post_token_endpoint, _persist_credential)
-    refreshing: Final = RefreshingTokenStore(V2PerUserTokenStore(_read_credential), refresher, coordinator=coordinator)
+    refreshing: Final = RefreshingTokenStore(
+        V2PerUserTokenStore(partial(_read_bound_credential, server_lookup)), refresher, coordinator=coordinator
+    )
     return CachedOAuthTokenStore(refreshing, default_ttl_seconds=_DEFAULT_TTL_SECONDS, backend=backend), uses_redis
 
 
@@ -182,6 +200,18 @@ class LazyPerUserOAuthTokenStore:
         self._local_fetches = 0
 
     async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None:
+        token: Final = await self._fetch_token(user_id, server_id)
+        server: Final = self._server_lookup(server_id)
+        binding: Final = server.oauth_identity_binding if server else None
+        if token is not None and binding is not None and binding.mode == "enforce":
+            if not await credential_binding_matches(
+                binding, user_id, server_id, {"identity_binding_proof": token.identity_binding_proof}
+            ):
+                await self.invalidate(user_id, server_id)
+                return None
+        return token
+
+    async def _fetch_token(self, user_id: str, server_id: str) -> OAuthToken | None:
         if self._uses_redis:
             store = self._store
             if store is not None:

litellm/proxy/_experimental/mcp_server/outbound_credentials/token_cache_codec.py

@@ -1,24 +1,26 @@
 """Serialize + encrypt boundary for caching an OAuth token in a shared (Redis) cache.
 
-A cross-replica cache must serialize the token, and a plaintext bearer in Redis is a leak, so this
-encrypts the value (NaCl in production via the injected ``encrypt``, identity in tests). It caches
-**only** the ``access_token``: the hot path needs just the bearer, expiry is carried by the cache
-entry's TTL (set from the token's ``expires_at`` by the cache), and the long-lived refresh_token stays
-in the DB - the refresh path is always a cache miss that re-reads it - so it never reaches Redis. A
-decoded token therefore carries only the bearer (``expires_at`` and ``refresh_token`` both None); the
-TTL, not the value, bounds its life. An empty/undecryptable blob (e.g. master-key rotation) is a miss.
+Shared cache values contain an encrypted access token and optional identity-binding proof.
+Refresh tokens remain in the database; cache TTL bounds the access token's lifetime.
+Legacy bearer-only entries decode without proof and cannot satisfy identity enforcement.
 """
 
 from __future__ import annotations
 
+import json
 from collections.abc import Callable
 from dataclasses import dataclass
 from typing import Final
 
+from pydantic import TypeAdapter, ValidationError
+
 from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import (
     OAuthToken,
 )
 
+_BOUND_PREFIX: Final = "litellm-bound-oauth-v1:"
+_BOUND_PAYLOAD: Final = TypeAdapter(dict[str, str])
+
 
 @dataclass(frozen=True, slots=True)
 class OAuthTokenCacheCodec:
@@ -26,10 +28,30 @@ class OAuthTokenCacheCodec:
     decrypt: Callable[[str], str | None]
 
     def encode(self, token: OAuthToken) -> str:
+        if token.identity_binding_proof is not None:
+            return self.encrypt(
+                _BOUND_PREFIX
+                + json.dumps(
+                    {
+                        "access_token": token.access_token,
+                        "identity_binding_proof": token.identity_binding_proof,
+                    }
+                )
+            )
         return self.encrypt(token.access_token)
 
     def decode(self, blob: str) -> OAuthToken | None:
         access_token: Final = self.decrypt(blob)
         if not access_token:
             return None
+        if access_token.startswith(_BOUND_PREFIX):
+            try:
+                payload: Final = _BOUND_PAYLOAD.validate_json(access_token[len(_BOUND_PREFIX) :])
+            except ValidationError:
+                return None
+            bearer: Final = payload.get("access_token")
+            proof: Final = payload.get("identity_binding_proof")
+            if not bearer or not proof:
+                return None
+            return OAuthToken(access_token=bearer, identity_binding_proof=proof)
         return OAuthToken(access_token=access_token, refresh_token=None)

litellm/proxy/_experimental/mcp_server/outbound_credentials/v2_token_store.py

@@ -46,11 +46,13 @@ def _to_oauth_token(payload: Mapping[str, object]) -> OAuthToken | None:
         return None
     refresh_token: Final = payload.get("refresh_token")
     expires_at: Final = payload.get("expires_at")
+    binding_proof: Final = payload.get("identity_binding_proof")
     return OAuthToken(
         access_token=access_token,
         expires_at=_iso_to_epoch(expires_at) if isinstance(expires_at, str) else None,
         refresh_token=refresh_token if isinstance(refresh_token, str) else None,
         scopes=_to_scopes(payload.get("scopes")),
+        identity_binding_proof=binding_proof if isinstance(binding_proof, str) else None,
     )
 
 

litellm/proxy/_experimental/mcp_server/server.py

@@ -2048,18 +2048,44 @@ if MCP_AVAILABLE:
             return texts[0][1]
         return "\n\n---\n\n".join(f"[{lbl}]\n{txt}" for lbl, txt in texts)
 
+    async def _raise_if_initialize_grants_no_mcp_servers(
+        allowed: Sequence[MCPServer],
+        user_api_key_auth: UserAPIKeyAuth | None,
+        mcp_servers: Sequence[str] | None,
+        client_ip: str | None,
+    ) -> None:
+        if allowed or user_api_key_auth is None or not user_api_key_auth.api_key:
+            return
+        if mcp_servers:
+            await raise_denied_scoped_mcp_access(
+                requested_names=mcp_servers,
+                user_api_key_auth=user_api_key_auth,
+                client_ip=client_ip,
+            )
+        no_servers_denial: Final[_McpDeniedDetail] = {
+            "error": (
+                "The key has no MCP servers granted, or none of its granted servers is loaded and allowed for "
+                "this client IP. Grant servers or access groups to the key, its team, or its organization "
+                "(object_permission.mcp_servers), check the server's allowed IPs, and reconnect."
+            )
+        }
+        raise HTTPException(status_code=403, detail=no_servers_denial)
+
     @contextlib.asynccontextmanager
     async def _gateway_initialize_instructions_request_scope(
         user_api_key_auth: UserAPIKeyAuth | None,
         mcp_servers: list[str] | None,
         client_ip: str | None,
         scoped_server_endpoint: bool = False,
+        is_initialize: bool = False,
     ) -> AsyncIterator[None]:
         allowed: Final = await _get_allowed_mcp_servers(
             user_api_key_auth=user_api_key_auth,
             mcp_servers=mcp_servers,
             client_ip=client_ip,
         )
+        if is_initialize:
+            await _raise_if_initialize_grants_no_mcp_servers(allowed, user_api_key_auth, mcp_servers, client_ip)
         if allowed:
             # return_exceptions=True: a per-server probe failure (incl. CancelledError
             # bubbled from anyio task group teardown on connection refused) must not
@@ -4683,6 +4709,7 @@ if MCP_AVAILABLE:
                     mcp_servers,
                     _client_ip,
                     scoped_server_endpoint=scoped_server_endpoint,
+                    is_initialize=is_initialize,
                 ):
                     await target_manager.handle_request(scope, receive, local_send)
                     if use_stateful and session_id and scope.get("method") == "DELETE":
@@ -4819,6 +4846,7 @@ if MCP_AVAILABLE:
                 mcp_servers,
                 _sse_client_ip,
                 scoped_server_endpoint=scoped_server_endpoint,
+                is_initialize=scope.get("method") == "GET",
             ):
                 await sse_session_manager.handle_request(scope, receive, send)
         except MCPUpstreamAuthError as e:

litellm/proxy/_lazy_openapi_snapshot.json

@@ -5384,7 +5384,7 @@
               "additionalProperties": {
                 "type": "string"
               },
-              "description": "Git source reference",
+              "description": "Plugin source reference",
               "title": "Source",
               "type": "object"
             },
@@ -5411,7 +5411,7 @@
           "type": "object"
         },
         "RegisterPluginRequest": {
-          "description": "Request body for registering a plugin in the marketplace.\n\nLiteLLM acts as a registry/discovery layer. Plugins are hosted on\nGitHub/GitLab/Bitbucket and referenced by their git source.",
+          "description": "Request body for registering a plugin in the marketplace.\n\nLiteLLM acts as a registry/discovery layer. Plugins are hosted on\nGitHub/GitLab/Bitbucket or as a zip archive on any https host and referenced by their source.",
           "properties": {
             "author": {
               "anyOf": [
@@ -5509,7 +5509,7 @@
               "additionalProperties": {
                 "type": "string"
               },
-              "description": "Git source reference. Supported formats:\n- GitHub: {'source': 'github', 'repo': 'org/repo'}\n- Git URL: {'source': 'url', 'url': 'https://github.com/org/repo.git'}\n- Git Subdir: {'source': 'git-subdir', 'url': 'https://github.com/org/repo.git', 'path': 'plugins/plugin-name'}",
+              "description": "Plugin source reference. Supported formats:\n- GitHub: {'source': 'github', 'repo': 'org/repo'}\n- Git URL: {'source': 'url', 'url': 'https://github.com/org/repo.git'}\n- Git Subdir: {'source': 'git-subdir', 'url': 'https://github.com/org/repo.git', 'path': 'plugins/plugin-name'}\n- Zip archive on any https host (e.g. S3): {'source': 'archive', 'url': 'https://bucket.s3.amazonaws.com/plugin.zip', 'sha256': '<optional hex digest>'}",
               "title": "Source",
               "type": "object"
             },
@@ -5653,7 +5653,7 @@
               "additionalProperties": {
                 "type": "string"
               },
-              "description": "Git source reference. Supported formats:\n- GitHub: {'source': 'github', 'repo': 'org/repo'}\n- Git URL: {'source': 'url', 'url': 'https://github.com/org/repo.git'}\n- Git Subdir: {'source': 'git-subdir', 'url': 'https://github.com/org/repo.git', 'path': 'plugins/plugin-name'}",
+              "description": "Plugin source reference. Supported formats:\n- GitHub: {'source': 'github', 'repo': 'org/repo'}\n- Git URL: {'source': 'url', 'url': 'https://github.com/org/repo.git'}\n- Git Subdir: {'source': 'git-subdir', 'url': 'https://github.com/org/repo.git', 'path': 'plugins/plugin-name'}\n- Zip archive on any https host (e.g. S3): {'source': 'archive', 'url': 'https://bucket.s3.amazonaws.com/plugin.zip', 'sha256': '<optional hex digest>'}",
               "title": "Source",
               "type": "object"
             },
@@ -5721,8 +5721,26 @@
     "paths": {
       "/claude-code/marketplace.json": {
         "get": {
-          "description": "Serve marketplace.json for Claude Code plugin discovery.\n\nThis endpoint is accessed by Claude Code CLI when users run:\n- claude plugin marketplace add <url>\n- claude plugin install <name>@<marketplace>\n\nReturns:\n    Marketplace catalog with list of available plugins and their git sources.\n\nExample:\n    ```bash\n    claude plugin marketplace add http://localhost:4000/claude-code/marketplace.json\n    claude plugin install my-plugin@litellm\n    ```",
+          "description": "Serve marketplace.json for Claude Code plugin discovery.\n\nThis endpoint is accessed by Claude Code CLI when users run:\n- claude plugin marketplace add <url>\n- claude plugin install <name>@<marketplace>\n\nWithout `key` the catalog holds the enabled (public) plugins. With `?key=sk-...`\nthe key is authenticated and the catalog also holds the disabled plugins granted\nto it through `object_permission.skills` on the key or its team.\n\nReturns:\n    Marketplace catalog with list of available plugins and their git sources.\n\nExample:\n    ```bash\n    claude plugin marketplace add http://localhost:4000/claude-code/marketplace.json\n    claude plugin marketplace add \"http://localhost:4000/claude-code/marketplace.json?key=sk-...\"\n    claude plugin install my-plugin@litellm\n    ```",
           "operationId": "get_marketplace_claude_code_marketplace_json_get",
+          "parameters": [
+            {
+              "in": "query",
+              "name": "key",
+              "required": false,
+              "schema": {
+                "anyOf": [
+                  {
+                    "type": "string"
+                  },
+                  {
+                    "type": "null"
+                  }
+                ],
+                "title": "Key"
+              }
+            }
+          ],
           "responses": {
             "200": {
               "content": {
@@ -5731,6 +5749,16 @@
                 }
               },
               "description": "Successful Response"
+            },
+            "422": {
+              "content": {
+                "application/json": {
+                  "schema": {
+                    "$ref": "#/components/schemas/HTTPValidationError"
+                  }
+                }
+              },
+              "description": "Validation Error"
             }
           },
           "summary": "Get Marketplace",
@@ -5788,7 +5816,7 @@
           ]
         },
         "post": {
-          "description": "Register a new plugin in the LiteLLM marketplace.\n\nLiteLLM acts as a registry/discovery layer. Plugins are hosted on\nGitHub/GitLab/Bitbucket. Claude Code will clone from the git source\nwhen users install.\n\nThis endpoint is create-only and never overwrites. If a plugin with\nthe same name already exists it returns 409 Conflict; use\nPUT /claude-code/plugins/{plugin_name} to update an existing plugin.\n\nRequires a proxy admin API key.\n\nParameters:\n    - name: Plugin name (kebab-case)\n    - source: Git source reference (github, url, or git-subdir format)\n    - version: Semantic version (optional)\n    - description: Plugin description (optional)\n    - author: Author information (optional)\n    - homepage: Plugin homepage URL (optional)\n    - keywords: Search keywords (optional)\n    - category: Plugin category (optional)\n\nReturns:\n    Registration status (action is always \"created\") and plugin information.\n\nExample:\n    ```bash\n    curl -X POST http://localhost:4000/claude-code/plugins \\\n      -H \"Authorization: Bearer sk-...\" \\\n      -H \"Content-Type: application/json\" \\\n      -d '{\n        \"name\": \"my-plugin\",\n        \"source\": {\"source\": \"github\", \"repo\": \"org/my-plugin\"},\n        \"version\": \"1.0.0\",\n        \"description\": \"My awesome plugin\"\n      }'\n    ```",
+          "description": "Register a new plugin in the LiteLLM marketplace.\n\nLiteLLM acts as a registry/discovery layer. Plugins are hosted on\nGitHub/GitLab/Bitbucket or as a zip archive on any https host (e.g. S3).\nClaude Code clones the git source or downloads the archive when users install.\n\nThis endpoint is create-only and never overwrites. If a plugin with\nthe same name already exists it returns 409 Conflict; use\nPUT /claude-code/plugins/{plugin_name} to update an existing plugin.\n\nRequires a proxy admin API key.\n\nParameters:\n    - name: Plugin name (kebab-case)\n    - source: Plugin source reference (github, url, git-subdir, or archive format)\n    - version: Semantic version (optional)\n    - description: Plugin description (optional)\n    - author: Author information (optional)\n    - homepage: Plugin homepage URL (optional)\n    - keywords: Search keywords (optional)\n    - category: Plugin category (optional)\n\nReturns:\n    Registration status (action is always \"created\") and plugin information.\n\nExample:\n    ```bash\n    curl -X POST http://localhost:4000/claude-code/plugins \\\n      -H \"Authorization: Bearer sk-...\" \\\n      -H \"Content-Type: application/json\" \\\n      -d '{\n        \"name\": \"my-plugin\",\n        \"source\": {\"source\": \"github\", \"repo\": \"org/my-plugin\"},\n        \"version\": \"1.0.0\",\n        \"description\": \"My awesome plugin\"\n      }'\n    ```",
           "operationId": "register_plugin_claude_code_plugins_post",
           "requestBody": {
             "content": {
@@ -5923,7 +5951,7 @@
           ]
         },
         "put": {
-          "description": "Update an existing plugin in the LiteLLM marketplace.\n\nThe plugin is identified by its name in the path, which is the resource\nidentity and cannot be changed here. This is a full replace, not a merge:\nthe manifest is rebuilt from the request body, so any optional field left\nout is reset to its default (e.g. an omitted version is cleared, not kept).\nSend the full desired state.\n\nReturns 404 if no plugin with the given name exists; use\nPOST /claude-code/plugins to create a new plugin.\n\nRequires a proxy admin API key.\n\nParameters:\n    - plugin_name: Name of the plugin to update (path parameter)\n    - source: Git source reference (github, url, or git-subdir format)\n    - version: Semantic version (optional)\n    - description: Plugin description (optional)\n    - author: Author information (optional)\n    - homepage: Plugin homepage URL (optional)\n    - keywords: Search keywords (optional)\n    - category: Plugin category (optional)\n\nReturns:\n    Update status (action is always \"updated\") and plugin information.\n\nExample:\n    ```bash\n    curl -X PUT http://localhost:4000/claude-code/plugins/my-plugin \\\n      -H \"Authorization: Bearer sk-...\" \\\n      -H \"Content-Type: application/json\" \\\n      -d '{\n        \"source\": {\"source\": \"github\", \"repo\": \"org/my-plugin\"},\n        \"version\": \"2.0.0\",\n        \"description\": \"My awesome plugin\"\n      }'\n    ```",
+          "description": "Update an existing plugin in the LiteLLM marketplace.\n\nThe plugin is identified by its name in the path, which is the resource\nidentity and cannot be changed here. This is a full replace, not a merge:\nthe manifest is rebuilt from the request body, so any optional field left\nout is reset to its default (e.g. an omitted version is cleared, not kept).\nSend the full desired state.\n\nReturns 404 if no plugin with the given name exists; use\nPOST /claude-code/plugins to create a new plugin.\n\nRequires a proxy admin API key.\n\nParameters:\n    - plugin_name: Name of the plugin to update (path parameter)\n    - source: Plugin source reference (github, url, git-subdir, or archive format)\n    - version: Semantic version (optional)\n    - description: Plugin description (optional)\n    - author: Author information (optional)\n    - homepage: Plugin homepage URL (optional)\n    - keywords: Search keywords (optional)\n    - category: Plugin category (optional)\n\nReturns:\n    Update status (action is always \"updated\") and plugin information.\n\nExample:\n    ```bash\n    curl -X PUT http://localhost:4000/claude-code/plugins/my-plugin \\\n      -H \"Authorization: Bearer sk-...\" \\\n      -H \"Content-Type: application/json\" \\\n      -d '{\n        \"source\": {\"source\": \"github\", \"repo\": \"org/my-plugin\"},\n        \"version\": \"2.0.0\",\n        \"description\": \"My awesome plugin\"\n      }'\n    ```",
           "operationId": "update_plugin_claude_code_plugins__plugin_name__put",
           "parameters": [
             {