scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 124

Merge branch 'litellm_internal_staging' into litellm_redis_breaker_open_silent_miss

Browse source 
Take the staging tree outright. #40620 already landed the typed
RedisCircuitBreakerOpenError, DEBUG-level refusal logging at every
DualCache, Cache, limiter, and router catch site, and sync timeout
classification, so the overlapping parts of this branch are superseded.
The pieces that are still missing on staging are rebuilt on top in the
next commit
This commit is contained in:
mateo-berri 2026-09-10 17:34:17 -07:00
commit 68d3016a11
261 changed files with 20900 additions and 1805 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
.github/ISSUE_TEMPLATE/config.yml vendored
View file

@ -1,4 +1,4 @@
blank_issues_enabled: true
blank_issues_enabled: false
contact_links:
- name: Schedule Demo
url: https://enterprise.litellm.ai/demo

28
.github/workflows/test-rust.yml vendored
View file

@ -4,8 +4,22 @@ on:
push:
paths:
- "litellm-rust/**"
- "litellm/rust_bridge/**"
- "tests/test_litellm_rust/**"
- "litellm/integrations/custom_logger.py"
- "litellm/litellm_core_utils/litellm_logging.py"
- "litellm/litellm_core_utils/logging_worker.py"
- "litellm/proxy/guardrails/**"
- "litellm/utils.py"
- "litellm/ocr/**"
- "litellm/llms/base_llm/ocr/**"
- "litellm/llms/custom_httpx/llm_http_handler.py"
- "tests/test_litellm/ocr/**"
- "tests/test_litellm/conftest.py"
- "Makefile"
- ".cargo/**"
- "pyproject.toml"
- "uv.lock"
- "rust-toolchain.toml"
- ".github/actions/setup-uv-with-retries/**"
- ".github/scripts/smoke_test_native_wheel.py"
@ -20,8 +34,22 @@ on:
- "litellm_**"
paths:
- "litellm-rust/**"
- "litellm/rust_bridge/**"
- "tests/test_litellm_rust/**"
- "litellm/integrations/custom_logger.py"
- "litellm/litellm_core_utils/litellm_logging.py"
- "litellm/litellm_core_utils/logging_worker.py"
- "litellm/proxy/guardrails/**"
- "litellm/utils.py"
- "litellm/ocr/**"
- "litellm/llms/base_llm/ocr/**"
- "litellm/llms/custom_httpx/llm_http_handler.py"
- "tests/test_litellm/ocr/**"
- "tests/test_litellm/conftest.py"
- "Makefile"
- ".cargo/**"
- "pyproject.toml"
- "uv.lock"
- "rust-toolchain.toml"
- ".github/actions/setup-uv-with-retries/**"
- ".github/scripts/smoke_test_native_wheel.py"

40
.github/workflows/test-terraform-modules.yml vendored
View file

@ -25,13 +25,17 @@ concurrency:
cancel-in-progress: true
jobs:
aws-module:
name: fmt, validate, test (aws)
module:
name: fmt, validate, test (${{ matrix.module }})
runs-on: ubuntu-latest
timeout-minutes: 15
strategy:
fail-fast: false
matrix:
module: [aws, gcp]
defaults:
run:
working-directory: terraform/litellm/aws
working-directory: terraform/litellm/${{ matrix.module }}
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
@ -51,35 +55,7 @@ jobs:
- name: validate
run: terraform validate
# Plan-only, mock_provider-backed: no AWS credentials, no API calls.
# Plan-only, mock_provider-backed: no cloud credentials, no API calls.
- name: test
run: terraform test
gcp-module:
name: fmt, validate, test (gcp)
runs-on: ubuntu-latest
timeout-minutes: 15
defaults:
run:
working-directory: terraform/litellm/gcp
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- uses: hashicorp/setup-terraform@b9cd54a3c349d3f38e8881555d616ced269862dd # v3.1.2
with:
terraform_version: 1.13.3
terraform_wrapper: false
- name: fmt
run: terraform fmt -recursive -check -diff
- name: init
run: terraform init -backend=false -input=false
- name: validate
run: terraform validate
- name: test
run: terraform test

1
.github/workflows/test-unit.yml vendored
View file

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

2
AGENTS.md
View file

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

19
Dockerfile
View file

@ -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}" \

19
docker/Dockerfile.database
View file

@ -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}" \

19
docker/Dockerfile.non_root
View file

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

28
gateway/Dockerfile
View file

@ -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"]

77
gateway/launch.py Normal file
View file

@ -0,0 +1,77 @@
"""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, and one marked pooled
wins under token auth too, so exporting 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,
export_pooled_database_url,
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``). Under token
auth the pooler mints and renews the upstream token itself.
"""
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=settings.token_auth())
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:
export_pooled_database_url(pooled_url)
serve(uvicorn_argv(argv, os.environ))
if __name__ == "__main__":
main(sys.argv[1:])

160
helm/litellm-helm/templates/_helpers.tpl
View file

@ -161,3 +161,163 @@ taken before the change, which by that point no longer exists.
{{- fail (printf "postgresql.image.tag must be pinned to an explicit version when db.deployStandalone is true (got %q). An unpinned tag can start a different PostgreSQL major against the existing data directory, which makes the database unreadable and is not recoverable in place. Crossing a major version requires a dump and restore." $tag) -}}
{{- end -}}
{{- end -}}
{{/*
Environment shared by the proxy container and the opt-in collector sidecar:
database, pgbouncer, master key, redis, user envVars. Both containers must see
the same DATABASE_URL and REDIS_* so the sidecar reaches the pod's pgbouncer
and the same spend transaction buffer.
*/}}
{{- define "litellm.proxyEnv" -}}
- name: HOST
value: "{{ .Values.listen | default "0.0.0.0" }}"
- name: PORT
value: {{ .Values.service.port | quote}}
{{- if .Values.db.deployStandalone }}
- name: DATABASE_USERNAME
valueFrom:
secretKeyRef:
name: {{ include "litellm.fullname" . }}-dbcredentials
key: username
- name: DATABASE_PASSWORD
valueFrom:
secretKeyRef:
name: {{ include "litellm.fullname" . }}-dbcredentials
key: password
- name: DATABASE_HOST
value: {{ .Release.Name }}-postgresql
- name: DATABASE_NAME
value: litellm
{{- else if .Values.db.useExisting }}
- name: DATABASE_USERNAME
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.usernameKey }}
- name: DATABASE_PASSWORD
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.passwordKey }}
- name: DATABASE_HOST
{{- if .Values.db.secret.endpointKey }}
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.endpointKey }}
{{- else }}
value: {{ .Values.db.endpoint }}
{{- end }}
- name: DATABASE_NAME
value: {{ .Values.db.database }}
- name: DATABASE_URL
value: {{ .Values.db.url | quote }}
{{- end }}
{{- if and .Values.db.useExisting .Values.db.readReplicaUrl .Values.db.secret.readReplicaEndpointKey (not .Values.db.secret.readReplicaUrlKey) }}
- name: DATABASE_READER_HOST
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.readReplicaEndpointKey }}
{{- end }}
{{- if and .Values.db.useExisting .Values.db.secret.readReplicaUrlKey }}
- name: DATABASE_URL_READ_REPLICA
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.readReplicaUrlKey }}
{{- else if .Values.db.readReplicaUrl }}
- 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:
name: {{ .Values.masterkeySecretName | default (printf "%s-masterkey" (include "litellm.fullname" .)) }}
key: {{ .Values.masterkeySecretKey | default "masterkey" }}
{{- if .Values.redis.enabled }}
- name: REDIS_HOST
value: {{ include "litellm.redis.serviceName" . }}
- name: REDIS_PORT
value: {{ include "litellm.redis.port" . | quote }}
- name: REDIS_PASSWORD
valueFrom:
secretKeyRef:
name: {{ include "redis.secretName" .Subcharts.redis }}
key: {{include "redis.secretPasswordKey" .Subcharts.redis }}
{{- end }}
{{- /*
Inject LITELLM_LOG only when envVars does not already define it.
*/}}
{{- if and .Values.logLevel (not (hasKey (default dict .Values.envVars) "LITELLM_LOG")) }}
- name: LITELLM_LOG
value: {{ .Values.logLevel | quote }}
{{- end }}
{{- if .Values.envVars }}
{{- range $key, $val := .Values.envVars }}
- name: {{ $key }}
value: {{ $val | quote }}
{{- end }}
{{- end }}
{{- with .Values.extraEnvVars }}
{{ toYaml . }}
{{- end }}
{{- if .Values.migrationJob.enabled }}
# Schema updates are owned by the dedicated migrations Job; skip
# the proxy's startup `prisma db push` so N replicas don't race
# one DB on every rollout. Placed last (after envVars and
# extraEnvVars) so this override can't be silently shadowed by a
# user-supplied DISABLE_SCHEMA_UPDATE under last-wins duplicate-env
# semantics — same pattern the migrations Job uses.
- name: DISABLE_SCHEMA_UPDATE
value: "true"
{{- end }}
{{- end -}}
{{/*
Proxy-only metering and metrics env. The collector sidecar serves no HTTP
traffic, so it gets neither.
*/}}
{{- define "litellm.proxyMetricsEnv" -}}
{{- if .Values.billingMetrics.enabled }}
{{ include "litellm.billingMetricsEnv" . }}
{{- end }}
{{- if .Values.metricsServer.enabled }}
{{- if eq (int .Values.metricsServer.port) (int .Values.service.port) }}
{{- fail "metricsServer.port must differ from service.port" }}
{{- end }}
- name: PROMETHEUS_METRICS_PORT
value: {{ .Values.metricsServer.port | quote }}
{{- end }}
{{- end -}}
{{/*
Directory of the collector's unix socket, shared between the two containers
through an emptyDir. Empty when the sidecar is off or uses 127.0.0.1 TCP.
*/}}
{{- define "litellm.collector.socketDir" -}}
{{- if and .Values.collector.enabled (hasPrefix "unix://" .Values.collector.address) -}}
{{- dir (trimPrefix "unix://" .Values.collector.address) -}}
{{- end -}}
{{- end -}}
{{- define "litellm.collectorEnv" -}}
- name: LITELLM_COLLECTOR_ENABLED
value: "true"
- name: LITELLM_COLLECTOR_ADDRESS
value: {{ .Values.collector.address | quote }}
- name: LITELLM_COLLECTOR_BUFFER_SIZE
value: {{ .Values.collector.bufferSize | quote }}
- name: LITELLM_COLLECTOR_ON_UNAVAILABLE
value: {{ .Values.collector.onUnavailable | quote }}
- name: LITELLM_COLLECTOR_DRAIN_TIMEOUT_SECONDS
value: {{ .Values.collector.drainTimeoutSeconds | quote }}
{{- end -}}

172
helm/litellm-helm/templates/deployment.yaml
View file

@ -56,118 +56,10 @@ spec:
image: "{{ .Values.image.repository }}:{{ .Values.image.tag | default .Chart.AppVersion }}"
imagePullPolicy: {{ .Values.image.pullPolicy }}
env:
- name: HOST
value: "{{ .Values.listen | default "0.0.0.0" }}"
- name: PORT
value: {{ .Values.service.port | quote}}
{{- if .Values.db.deployStandalone }}
- name: DATABASE_USERNAME
valueFrom:
secretKeyRef:
name: {{ include "litellm.fullname" . }}-dbcredentials
key: username
- name: DATABASE_PASSWORD
valueFrom:
secretKeyRef:
name: {{ include "litellm.fullname" . }}-dbcredentials
key: password
- name: DATABASE_HOST
value: {{ .Release.Name }}-postgresql
- name: DATABASE_NAME
value: litellm
{{- else if .Values.db.useExisting }}
- name: DATABASE_USERNAME
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.usernameKey }}
- name: DATABASE_PASSWORD
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.passwordKey }}
- name: DATABASE_HOST
{{- if .Values.db.secret.endpointKey }}
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.endpointKey }}
{{- else }}
value: {{ .Values.db.endpoint }}
{{- end }}
- name: DATABASE_NAME
value: {{ .Values.db.database }}
- name: DATABASE_URL
value: {{ .Values.db.url | quote }}
{{- end }}
{{- if and .Values.db.useExisting .Values.db.readReplicaUrl .Values.db.secret.readReplicaEndpointKey (not .Values.db.secret.readReplicaUrlKey) }}
- name: DATABASE_READER_HOST
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.readReplicaEndpointKey }}
{{- end }}
{{- if and .Values.db.useExisting .Values.db.secret.readReplicaUrlKey }}
- name: DATABASE_URL_READ_REPLICA
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.readReplicaUrlKey }}
{{- else if .Values.db.readReplicaUrl }}
- name: DATABASE_URL_READ_REPLICA
value: {{ .Values.db.readReplicaUrl | quote }}
{{- end }}
- name: PROXY_MASTER_KEY
valueFrom:
secretKeyRef:
name: {{ .Values.masterkeySecretName | default (printf "%s-masterkey" (include "litellm.fullname" .)) }}
key: {{ .Values.masterkeySecretKey | default "masterkey" }}
{{- if .Values.redis.enabled }}
- name: REDIS_HOST
value: {{ include "litellm.redis.serviceName" . }}
- name: REDIS_PORT
value: {{ include "litellm.redis.port" . | quote }}
- name: REDIS_PASSWORD
valueFrom:
secretKeyRef:
name: {{ include "redis.secretName" .Subcharts.redis }}
key: {{include "redis.secretPasswordKey" .Subcharts.redis }}
{{- end }}
{{- /*
Inject LITELLM_LOG only when envVars does not already define it.
*/}}
{{- if and .Values.logLevel (not (hasKey (default dict .Values.envVars) "LITELLM_LOG")) }}
- name: LITELLM_LOG
value: {{ .Values.logLevel | quote }}
{{- end }}
{{- if .Values.envVars }}
{{- range $key, $val := .Values.envVars }}
- name: {{ $key }}
value: {{ $val | quote }}
{{- end }}
{{- end }}
{{- with .Values.extraEnvVars }}
{{- toYaml . | nindent 12 }}
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsEnv" . | nindent 12 }}
{{- end }}
{{- if .Values.metricsServer.enabled }}
{{- if eq (int .Values.metricsServer.port) (int .Values.service.port) }}
{{- fail "metricsServer.port must differ from service.port" }}
{{- end }}
- name: PROMETHEUS_METRICS_PORT
value: {{ .Values.metricsServer.port | quote }}
{{- end }}
{{- if .Values.migrationJob.enabled }}
# Schema updates are owned by the dedicated migrations Job; skip
# the proxy's startup `prisma db push` so N replicas don't race
# one DB on every rollout. Placed last (after envVars and
# extraEnvVars) so this override can't be silently shadowed by a
# user-supplied DISABLE_SCHEMA_UPDATE under last-wins duplicate-env
# semantics — same pattern the migrations Job uses.
- name: DISABLE_SCHEMA_UPDATE
value: "true"
{{- include "litellm.proxyEnv" . | nindent 12 }}
{{- include "litellm.proxyMetricsEnv" . | nindent 12 }}
{{- if .Values.collector.enabled }}
{{- include "litellm.collectorEnv" . | nindent 12 }}
{{- end }}
envFrom:
{{- range .Values.environmentSecrets }}
@ -245,6 +137,10 @@ spec:
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }}
{{- end }}
{{- if include "litellm.collector.socketDir" . }}
- name: collector-socket
mountPath: {{ include "litellm.collector.socketDir" . }}
{{- end }}
{{- with .Values.volumeMounts }}
{{- toYaml . | nindent 12 }}
{{- end }}
@ -252,6 +148,53 @@ spec:
lifecycle:
{{- toYaml . | nindent 12 }}
{{- end }}
{{- if .Values.collector.enabled }}
- name: {{ include "litellm.name" . }}-collector
securityContext:
{{- toYaml .Values.securityContext | nindent 12 }}
image: "{{ .Values.image.repository }}:{{ .Values.image.tag | default .Chart.AppVersion }}"
imagePullPolicy: {{ .Values.image.pullPolicy }}
command: {{ toYaml .Values.collector.command | nindent 12 }}
env:
{{- include "litellm.proxyEnv" . | nindent 12 }}
{{- include "litellm.collectorEnv" . | nindent 12 }}
- name: LITELLM_JOB_ROLE
value: collector
{{- if not (hasKey (default dict .Values.envVars) "CONFIG_FILE_PATH") }}
- name: CONFIG_FILE_PATH
value: /etc/litellm/config.yaml
{{- end }}
envFrom:
{{- range .Values.environmentSecrets }}
- secretRef:
name: {{ . }}
{{- end }}
{{- range .Values.environmentConfigMaps }}
- configMapRef:
name: {{ . }}
{{- end }}
resources:
{{- toYaml .Values.collector.resources | nindent 12 }}
volumeMounts:
- name: litellm-config
mountPath: /etc/litellm/config.yaml
subPath: config.yaml
{{- if include "litellm.collector.socketDir" . }}
- name: collector-socket
mountPath: {{ include "litellm.collector.socketDir" . }}
{{- end }}
{{ if .Values.securityContext.readOnlyRootFilesystem }}
- name: tmp
mountPath: /tmp
- name: cache
mountPath: /.cache
- name: npm
mountPath: /.npm
{{- end }}
{{- with .Values.volumeMounts }}
{{- toYaml . | nindent 12 }}
{{- end }}
{{- end }}
{{- with .Values.extraContainers }}
{{- tpl (toYaml .) $ | nindent 8 }}
{{- end }}
@ -280,6 +223,11 @@ spec:
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsVolumes" . | nindent 8 }}
{{- end }}
{{- if include "litellm.collector.socketDir" . }}
- name: collector-socket
emptyDir:
sizeLimit: 1Mi
{{- end }}
{{- with .Values.volumes }}
{{- toYaml . | nindent 8 }}
{{- end }}

10
helm/litellm-helm/templates/hpa.yaml
View file

@ -18,6 +18,15 @@ spec:
{{- end }}
metrics:
{{- if .Values.autoscaling.targetCPUUtilizationPercentage }}
{{- if and .Values.collector.enabled .Values.collector.scaleOnProxyContainerCpu }}
- type: ContainerResource
containerResource:
name: cpu
container: {{ include "litellm.name" . }}
target:
type: Utilization
averageUtilization: {{ .Values.autoscaling.targetCPUUtilizationPercentage }}
{{- else }}
- type: Resource
resource:
name: cpu
@ -25,6 +34,7 @@ spec:
type: Utilization
averageUtilization: {{ .Values.autoscaling.targetCPUUtilizationPercentage }}
{{- end }}
{{- end }}
{{- if .Values.autoscaling.targetMemoryUtilizationPercentage }}
- type: Resource
resource:

272
helm/litellm-helm/tests/collector_tests.yaml Normal file
View file

@ -0,0 +1,272 @@
suite: test collector sidecar
templates:
- deployment.yaml
- hpa.yaml
- configmap-litellm.yaml
tests:
- it: should run the proxy alone with no collector env by default
template: deployment.yaml
asserts:
- lengthEqual:
path: spec.template.spec.containers
count: 1
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_COLLECTOR_ENABLED
value: "true"
- notContains:
path: spec.template.spec.volumes
content:
name: collector-socket
any: true
- it: should add the sidecar on the same image and point both containers at the unix socket
template: deployment.yaml
set:
image.tag: test
db.connectionPool.enabled: true
collector.enabled: true
collector.resources:
requests:
cpu: 500m
memory: 1Gi
limits:
cpu: "1"
memory: 2Gi
asserts:
- lengthEqual:
path: spec.template.spec.containers
count: 2
- equal:
path: spec.template.spec.containers[1].name
value: litellm-collector
- equal:
path: spec.template.spec.containers[1].image
value: ghcr.io/berriai/litellm:test
- equal:
path: spec.template.spec.containers[1].command
value: [python, -m, litellm.proxy.collector]
- equal:
path: spec.template.spec.containers[1].resources.requests.cpu
value: 500m
- equal:
path: spec.template.spec.containers[1].resources.limits.memory
value: 2Gi
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_COLLECTOR_ENABLED
value: "true"
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_COLLECTOR_ADDRESS
value: unix:///var/run/litellm/collector.sock
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_COLLECTOR_BUFFER_SIZE
value: "1000"
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_COLLECTOR_ON_UNAVAILABLE
value: fallback
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_JOB_ROLE
value: collector
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_JOB_ROLE
value: collector
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_COLLECTOR_ADDRESS
value: unix:///var/run/litellm/collector.sock
- contains:
path: spec.template.spec.containers[1].env
content:
name: CONFIG_FILE_PATH
value: /etc/litellm/config.yaml
- contains:
path: spec.template.spec.containers[1].env
content:
name: DATABASE_HOST
value: RELEASE-NAME-postgresql
- contains:
path: spec.template.spec.containers[1].env
content:
name: DATABASE_PASSWORD
valueFrom:
secretKeyRef:
name: RELEASE-NAME-litellm-dbcredentials
key: password
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_ENABLED
value: "true"
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
value: "20"
- contains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: collector-socket
mountPath: /var/run/litellm
- contains:
path: spec.template.spec.containers[1].volumeMounts
content:
name: collector-socket
mountPath: /var/run/litellm
- contains:
path: spec.template.spec.containers[1].volumeMounts
content:
name: litellm-config
mountPath: /etc/litellm/config.yaml
subPath: config.yaml
- contains:
path: spec.template.spec.volumes
content:
name: collector-socket
emptyDir:
sizeLimit: 1Mi
- it: should skip the socket volume and pass the policy through on tcp transport
template: deployment.yaml
set:
collector.enabled: true
collector.address: tcp://127.0.0.1:4100
collector.onUnavailable: drop
collector.bufferSize: 50
envVars:
CONFIG_FILE_PATH: /custom/config.yaml
asserts:
- lengthEqual:
path: spec.template.spec.containers
count: 2
- notContains:
path: spec.template.spec.containers[1].env
content:
name: CONFIG_FILE_PATH
value: /etc/litellm/config.yaml
- contains:
path: spec.template.spec.containers[1].env
content:
name: CONFIG_FILE_PATH
value: /custom/config.yaml
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_COLLECTOR_ADDRESS
value: tcp://127.0.0.1:4100
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_COLLECTOR_ON_UNAVAILABLE
value: drop
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_COLLECTOR_BUFFER_SIZE
value: "50"
- notContains:
path: spec.template.spec.volumes
content:
name: collector-socket
any: true
- it: should keep metrics and billing env on the proxy container only
template: deployment.yaml
set:
collector.enabled: true
metricsServer.enabled: true
metricsServer.port: 9090
billingMetrics.enabled: true
billingMetrics.endpoint: https://metering.example.com
billingMetrics.secretName: billing-mtls
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: PROMETHEUS_METRICS_PORT
value: "9090"
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_ENDPOINT
value: https://metering.example.com
- notContains:
path: spec.template.spec.containers[1].env
content:
name: PROMETHEUS_METRICS_PORT
any: true
- notContains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_BILLING_METRICS_ENDPOINT
any: true
- notContains:
path: spec.template.spec.containers[1].volumeMounts
content:
name: billing-metrics-mtls
any: true
- it: should give the sidecar the same scratch mounts as the proxy on a read-only root
template: deployment.yaml
set:
collector.enabled: true
securityContext.readOnlyRootFilesystem: true
asserts:
- contains:
path: spec.template.spec.containers[1].volumeMounts
content:
name: npm
mountPath: /.npm
- contains:
path: spec.template.spec.containers[1].volumeMounts
content:
name: cache
mountPath: /.cache
- contains:
path: spec.template.spec.containers[1].volumeMounts
content:
name: tmp
mountPath: /tmp
- it: should keep the pod-wide cpu metric unless asked to scale on the proxy container
template: hpa.yaml
set:
autoscaling.enabled: true
collector.enabled: true
asserts:
- equal: { path: "spec.metrics[0].type", value: Resource }
- equal: { path: "spec.metrics[0].resource.name", value: cpu }
- it: should scale on the proxy container's cpu only when opted in
template: hpa.yaml
set:
autoscaling.enabled: true
collector.enabled: true
collector.scaleOnProxyContainerCpu: true
asserts:
- equal: { path: "spec.metrics[0].type", value: ContainerResource }
- equal: { path: "spec.metrics[0].containerResource.name", value: cpu }
- equal: { path: "spec.metrics[0].containerResource.container", value: litellm }
- equal: { path: "spec.metrics[0].containerResource.target.averageUtilization", value: 60 }
- isNull: { path: "spec.metrics[0].resource" }
- it: should not switch to the container metric while the sidecar is off
template: hpa.yaml
set:
autoscaling.enabled: true
collector.scaleOnProxyContainerCpu: true
asserts:
- equal: { path: "spec.metrics[0].type", value: Resource }

61
helm/litellm-helm/tests/connection_pool_tests.yaml Normal file
View file

@ -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"

56
helm/litellm-helm/values.yaml
View file

@ -190,6 +190,48 @@ metricsServer:
enabled: false
port: 4001
# Opt-in sidecar that runs the post-response spend pipeline (cost calculation,
# spend logs, spend counters, budget reservation reconciliation) so the proxy's
# uvicorn workers only serialise a compact typed event and go back to serving
# inference. Same image and tag as the proxy, second container in the same pod,
# fed over loopback (a unix socket on a shared emptyDir, or 127.0.0.1 TCP). It
# reuses the pod's in-container pgbouncer (db.connectionPool) and the same Redis
# spend transaction buffer, so the per-pod DB connection budget is unchanged.
# Delivery is at-most-once inside the pod: events already handed to the sidecar
# are lost if it crashes before writing them; events the workers could not hand
# over follow onUnavailable. Both containers drain on SIGTERM within
# terminationGracePeriodSeconds
collector:
enabled: false
# unix:///<dir>/<file>.sock (the <dir> becomes a shared emptyDir) or tcp://127.0.0.1:<port>
address: unix:///var/run/litellm/collector.sock
# Events each uvicorn worker holds in memory while the sidecar is slow or restarting
bufferSize: 1000
# fallback: run the pipeline in the worker when the sidecar is unreachable or the
# buffer is full (spend stays exact, that request costs proxy CPU again)
# drop: count and discard the event instead (spend under-reports)
onUnavailable: fallback
# How long the workers keep pushing buffered events on shutdown, and how long the
# sidecar keeps serving its open connections after SIGTERM
drainTimeoutSeconds: 10
command:
- python
- -m
- litellm.proxy.collector
# Sized independently of the proxy container; the pipeline is CPU bound
resources: {}
# requests:
# cpu: 500m
# memory: 1Gi
# limits:
# cpu: "1"
# memory: 2Gi
# When autoscaling.enabled, swap the pod-wide cpu Resource metric for an
# autoscaling/v2 ContainerResource metric on the proxy container only, so the
# sidecar's CPU never scales inference replicas. Needs Kubernetes 1.30+ (or the
# HPAContainerMetrics feature gate on 1.27 to 1.29)
scaleOnProxyContainerCpu: false
resources:
{}
# Unset by default so the chart installs on small clusters such as Minikube, and so an
@ -355,6 +397,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.

45
helm/litellm/templates/_helpers.tpl
View file

@ -360,6 +360,20 @@ harmless no-op for the Job and authoritative for the app pods.
{{- end }}
{{- end -}}
{{/*
In-container PgBouncer env for the gateway container. Under IAM or Entra auth the pooler mints and renews the database token itself.
*/}}
{{- define "litellm.connectionPoolEnv" -}}
{{- 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.
@ -443,3 +457,34 @@ ImplementationSpecific
{{- end -}}
{{- define "litellm.gateway.prometheusMultiprocDir" -}}/tmp/litellm_prometheus_multiproc{{- end -}}
{{/*
Directory of the collector's unix socket, shared by the gateway and
collector containers through an emptyDir. Empty when the sidecar is off
or gateway.collector.address is a tcp://127.0.0.1:<port> address.
*/}}
{{- define "litellm.gateway.collectorSocketDir" -}}
{{- if and .Values.gateway.collector.enabled (hasPrefix "unix://" .Values.gateway.collector.address) -}}
{{- dir (trimPrefix "unix://" .Values.gateway.collector.address) -}}
{{- end -}}
{{- end -}}
{{/*
LITELLM_COLLECTOR_* env shared by the producer (gateway container) and the
consumer (collector container), so both agree on the transport and the
shutdown drain window.
*/}}
{{- define "litellm.gateway.collectorEnv" -}}
{{- with .Values.gateway.collector }}
- name: LITELLM_COLLECTOR_ENABLED
value: "true"
- name: LITELLM_COLLECTOR_ADDRESS
value: {{ .address | quote }}
- name: LITELLM_COLLECTOR_BUFFER_SIZE
value: {{ .bufferSize | quote }}
- name: LITELLM_COLLECTOR_ON_UNAVAILABLE
value: {{ .onUnavailable | quote }}
- name: LITELLM_COLLECTOR_DRAIN_TIMEOUT_SECONDS
value: {{ .drainTimeoutSeconds | quote }}
{{- end }}
{{- end -}}

59
helm/litellm/templates/gateway/deployment.yaml
View file

@ -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 }}
@ -71,8 +74,11 @@ spec:
- name: PROMETHEUS_MULTIPROC_DIR
value: {{ include "litellm.gateway.prometheusMultiprocDir" . }}
{{- end }}
{{- if .Values.gateway.collector.enabled }}
{{- include "litellm.gateway.collectorEnv" . | nindent 12 }}
{{- end }}
{{- include "litellm.envFrom" .Values.gateway | nindent 10 }}
{{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled }}
{{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled (include "litellm.gateway.collectorSocketDir" .) }}
volumeMounts:
{{- if .Values.gateway.config.create }}
- name: gateway-config
@ -83,6 +89,10 @@ spec:
- name: prometheus-multiproc
mountPath: {{ include "litellm.gateway.prometheusMultiprocDir" . }}
{{- end }}
{{- if include "litellm.gateway.collectorSocketDir" . }}
- name: collector-socket
mountPath: {{ include "litellm.gateway.collectorSocketDir" . }}
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }}
{{- end }}
@ -142,10 +152,50 @@ spec:
resources:
{{- toYaml .Values.gateway.metricsServer.resources | nindent 12 }}
{{- end }}
{{- if .Values.gateway.collector.enabled }}
- name: collector
image: "{{ .Values.gateway.image.repository }}:{{ .Values.gateway.image.tag | default .Chart.AppVersion }}"
imagePullPolicy: {{ .Values.gateway.image.pullPolicy }}
{{- with .Values.gateway.securityContext }}
securityContext:
{{- toYaml . | nindent 12 }}
{{- end }}
command:
- python
- -m
- litellm.proxy.collector
env:
{{- include "litellm.serverEnv" (dict "root" $ "component" .Values.gateway) | nindent 12 }}
{{- if .Values.gateway.config.create }}
- name: CONFIG_FILE_PATH
value: /app/config/config.yaml
{{- end }}
{{- include "litellm.gateway.collectorEnv" . | nindent 12 }}
- name: LITELLM_JOB_ROLE
value: collector
{{- include "litellm.envFrom" .Values.gateway | nindent 10 }}
{{- if or .Values.gateway.config.create .Values.gateway.volumeMounts (include "litellm.gateway.collectorSocketDir" .) }}
volumeMounts:
{{- if .Values.gateway.config.create }}
- name: gateway-config
mountPath: /app/config/config.yaml
subPath: config.yaml
{{- end }}
{{- if include "litellm.gateway.collectorSocketDir" . }}
- name: collector-socket
mountPath: {{ include "litellm.gateway.collectorSocketDir" . }}
{{- end }}
{{- with .Values.gateway.volumeMounts }}
{{- toYaml . | nindent 12 }}
{{- end }}
{{- end }}
resources:
{{- toYaml .Values.gateway.collector.resources | nindent 12 }}
{{- end }}
{{- with .Values.gateway.extraContainers }}
{{- tpl (toYaml .) $ | nindent 8 }}
{{- end }}
{{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled }}
{{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled (include "litellm.gateway.collectorSocketDir" .) }}
volumes:
{{- if .Values.gateway.config.create }}
- name: gateway-config
@ -156,6 +206,11 @@ spec:
- name: prometheus-multiproc
emptyDir: {}
{{- end }}
{{- if include "litellm.gateway.collectorSocketDir" . }}
- name: collector-socket
emptyDir:
sizeLimit: 1Mi
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsVolumes" . | nindent 8 }}
{{- end }}

10
helm/litellm/templates/gateway/hpa.yaml
View file

@ -15,6 +15,15 @@ spec:
maxReplicas: {{ .Values.gateway.hpa.maxReplicas }}
metrics:
{{- if .Values.gateway.hpa.targetCPUUtilizationPercentage }}
{{- if and .Values.gateway.collector.enabled .Values.gateway.collector.scaleOnGatewayContainerCpu }}
- type: ContainerResource
containerResource:
name: cpu
container: gateway
target:
type: Utilization
averageUtilization: {{ .Values.gateway.hpa.targetCPUUtilizationPercentage }}
{{- else }}
- type: Resource
resource:
name: cpu
@ -22,6 +31,7 @@ spec:
type: Utilization
averageUtilization: {{ .Values.gateway.hpa.targetCPUUtilizationPercentage }}
{{- end }}
{{- end }}
{{- if .Values.gateway.hpa.targetMemoryUtilizationPercentage }}
- type: Resource
resource:

204
helm/litellm/tests/collector_tests.yaml Normal file
View file

@ -0,0 +1,204 @@
suite: test gateway collector sidecar
templates:
- gateway/configmap.yaml
- gateway/deployment.yaml
- gateway/hpa.yaml
values:
- ./values/required.yaml
tests:
- it: adds no sidecar, env, volume or container metric when the collector is off
asserts:
- lengthEqual:
path: spec.template.spec.containers
count: 1
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_COLLECTOR_ENABLED
value: "true"
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.volumes
content:
name: collector-socket
any: true
template: gateway/deployment.yaml
- equal:
path: spec.metrics[0].type
value: Resource
template: gateway/hpa.yaml
- it: runs the collector as a sidecar sharing env, config and a unix socket emptyDir, and scales on the gateway container only
set:
gateway.collector.enabled: true
gateway.collector.bufferSize: 250
gateway.collector.onUnavailable: drop
gateway.image.tag: v1.102.0
gateway.numWorkers: 4
gateway.extraEnv:
- name: LITELLM_PGBOUNCER_ENABLED
value: "true"
gateway.envSecrets:
- litellm-license
gateway.volumes:
- name: redis-ca
secret:
secretName: redis-ca
gateway.volumeMounts:
- name: redis-ca
mountPath: /etc/litellm/redis-ca
readOnly: true
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_COLLECTOR_ADDRESS
value: unix:///var/run/litellm/collector.sock
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_COLLECTOR_BUFFER_SIZE
value: "250"
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_COLLECTOR_ON_UNAVAILABLE
value: drop
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: collector-socket
mountPath: /var/run/litellm
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].name
value: collector
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].image
value: ghcr.io/berriai/litellm-gateway:v1.102.0
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].command
value:
- python
- -m
- litellm.proxy.collector
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_JOB_ROLE
value: collector
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[1].env
content:
name: CONFIG_FILE_PATH
value: /app/config/config.yaml
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_ENABLED
value: "true"
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[1].env
content:
name: DATABASE_HOST
value: postgres.example.com
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_COLLECTOR_ADDRESS
value: unix:///var/run/litellm/collector.sock
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[1].env
content:
name: NUM_WORKERS
any: true
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].envFrom
value:
- secretRef:
name: litellm-license
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[1].volumeMounts
content:
name: gateway-config
mountPath: /app/config/config.yaml
subPath: config.yaml
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[1].volumeMounts
content:
name: collector-socket
mountPath: /var/run/litellm
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[1].volumeMounts
content:
name: redis-ca
mountPath: /etc/litellm/redis-ca
readOnly: true
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].resources.limits.cpu
value: "1"
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.volumes
content:
name: collector-socket
emptyDir:
sizeLimit: 1Mi
template: gateway/deployment.yaml
- equal:
path: spec.metrics[0]
value:
type: ContainerResource
containerResource:
name: cpu
container: gateway
target:
type: Utilization
averageUtilization: 70
template: gateway/hpa.yaml
- it: uses loopback tcp without a socket volume and keeps the pod-wide cpu metric when asked
set:
gateway.collector.enabled: true
gateway.collector.address: tcp://127.0.0.1:4010
gateway.collector.scaleOnGatewayContainerCpu: false
asserts:
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_COLLECTOR_ADDRESS
value: tcp://127.0.0.1:4010
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.volumes
content:
name: collector-socket
any: true
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[1].volumeMounts
content:
name: collector-socket
any: true
template: gateway/deployment.yaml
- equal:
path: spec.metrics[0].type
value: Resource
template: gateway/hpa.yaml

133
helm/litellm/tests/connection_pool_tests.yaml Normal file
View file

@ -0,0 +1,133 @@
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 renders both the pool and the token auth flag
template: gateway/deployment.yaml
set:
database.connectionPool.enabled: true
database.writer.useIAMAuth: 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: IAM_TOKEN_DB_AUTH
value: "true"
- it: pool with Entra auth renders both the pool and the token auth flag
template: gateway/deployment.yaml
set:
database.connectionPool.enabled: true
database.writer.useAzureEntraAuth: 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: AZURE_POSTGRESQL_AUTH
value: "true"
- 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

56
helm/litellm/values.yaml
View file

@ -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. With
# `database.writer.useIAMAuth` or `useAzureEntraAuth` the pool mints and
# renews the database token itself, so the workers never see it. 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
@ -294,6 +314,42 @@ gateway:
labels: {}
interval: 15s
scrapeTimeout: 10s
# Opt-in `collector` sidecar (same image, `python -m litellm.proxy.collector`)
# that runs the post-response spend pipeline (cost calculation, spend logs,
# spend counters, budget reservation reconciliation) so the uvicorn workers
# only serialise a compact event over loopback and go back to serving
# requests. It shares the pod's env, proxy config, in-container pgbouncer and
# Redis spend buffer, so the per-pod DB connection budget is unchanged.
# Delivery is at-most-once inside the pod: events already handed over are
# lost if the sidecar dies before writing them; events the workers cannot
# hand over follow `onUnavailable`.
collector:
enabled: false
# unix:///<dir>/<file>.sock (the <dir> becomes a shared emptyDir) or
# tcp://127.0.0.1:<port>
address: unix:///var/run/litellm/collector.sock
# Events each uvicorn worker holds in memory while the sidecar is slow or
# restarting.
bufferSize: 1000
# fallback: run the pipeline in the worker when the sidecar is unreachable
# or the buffer is full (spend stays exact, that request costs gateway CPU
# again). drop: count and discard the event instead (spend under-reports).
onUnavailable: fallback
# How long the workers keep pushing buffered events on shutdown, and how
# long the sidecar keeps serving open connections after SIGTERM.
drainTimeoutSeconds: 10
# Sized independently of the gateway container; the pipeline is CPU bound.
resources:
requests:
cpu: 500m
memory: 1Gi
limits:
cpu: "1"
memory: 2Gi
# With hpa.targetCPUUtilizationPercentage set, scale on an autoscaling/v2
# ContainerResource metric of the `gateway` container only, so the
# sidecar's CPU never drives inference replicas. Needs Kubernetes 1.30+.
scaleOnGatewayContainerCpu: true
image:
repository: ghcr.io/berriai/litellm-gateway
tag: "" # defaults to .Chart.AppVersion

1
litellm-proxy-extras/litellm_proxy_extras/migrations/20260910000000_skills_on_object_permission/migration.sql Normal file
View file

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

1
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -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[]

5
litellm-rust/AGENTS.md
View file

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

422
litellm-rust/Cargo.lock generated
View file

@ -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,7 +1607,7 @@ dependencies = [
"aws-sigv4",
"aws-smithy-runtime-api",
"aws-types",
"base64",
"base64 0.22.1",
"rand 0.8.7",
"reqwest",
"rstest",
@ -1471,6 +1631,7 @@ dependencies = [
"litellm-ai-gateway",
"litellm-core",
"litellm-python-interop",
"litellm-token-counter",
"pyo3",
"pyo3-async-runtimes",
"serde",
@ -1491,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"
@ -1509,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"
@ -1537,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"
@ -1548,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"
@ -1578,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"
@ -1606,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"
@ -1852,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"
@ -1865,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"
@ -1890,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"
@ -1899,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"
@ -1924,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"
@ -1981,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",
@ -2363,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"
@ -2537,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"
@ -2775,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"
@ -2858,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"
@ -3083,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"

3
litellm-rust/Cargo.toml
View file

@ -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" }
@ -40,6 +42,7 @@ tokio-tungstenite = { version = "0.24", default-features = false, features = ["c
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

5
litellm-rust/crates/ai-gateway/README.md
View file

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

7
litellm-rust/crates/core/tests/workspace_crate_allowlist.rs
View file

@ -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",

5
litellm-rust/crates/python-bridge/Cargo.toml
View file

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

2
litellm-rust/crates/python-bridge/src/constants.rs Normal file
View file

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

39
litellm-rust/crates/python-bridge/src/execution.rs
View file

@ -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,

4
litellm-rust/crates/python-bridge/src/lib.rs
View file

@ -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",
];

87
litellm-rust/crates/python-bridge/src/token_counter.rs Normal file
View file

@ -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>()
}

28
litellm-rust/crates/token-counter/Cargo.toml Normal file
View file

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

128
litellm-rust/crates/token-counter/benches/allocations.rs Normal file
View file

@ -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,
},
);
}

100
litellm-rust/crates/token-counter/benches/token_counter.rs Normal file
View file

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

650
litellm-rust/crates/token-counter/src/byte_level.rs Normal file
View file

@ -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);
}
}

194
litellm-rust/crates/token-counter/src/counter.rs Normal file
View file

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

31
litellm-rust/crates/token-counter/src/error.rs Normal file
View file

@ -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),
}

17
litellm-rust/crates/token-counter/src/lib.rs Normal file
View file

@ -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;

155
litellm-rust/crates/token-counter/src/python_json.rs Normal file
View file

@ -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"
);
}
}

304
litellm-rust/crates/token-counter/src/tools.rs Normal file
View file

@ -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)
));
}
}

238
litellm-rust/crates/token-counter/src/types.rs Normal file
View file

@ -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,
}

73
litellm-rust/crates/token-counter/src/unicode_classes.rs Normal file
View file

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

176
litellm-rust/crates/token-counter/tests/token_counter.rs Normal file
View file

@ -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(_))));
}

65
litellm/_redis.py
View file

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

85
litellm/_redis_credential_provider.py
View file

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

15
litellm/caching/caching.py
View file

@ -10,6 +10,7 @@
import ast
import hashlib
import json
import logging
import time
import traceback
from collections.abc import Mapping
@ -32,7 +33,7 @@ from .dual_cache import DualCache # noqa: F401
from .gcs_cache import GCSCache
from .in_memory_cache import InMemoryCache
from .qdrant_semantic_cache import QdrantSemanticCache
from .redis_cache import RedisCache, RedisCircuitBreakerOpenError
from .redis_cache import RedisCache, log_redis_failure
from .redis_cluster_cache import RedisClusterCache
from .redis_semantic_cache import RedisSemanticCache
from .s3_cache import S3Cache
@ -677,10 +678,8 @@ class Cache:
return
cache_key, cached_data, kwargs = self._add_cache_logic(result=result, **kwargs)
self.cache.set_cache(cache_key, cached_data, **kwargs)
except RedisCircuitBreakerOpenError as e:
verbose_logger.debug("LiteLLM Cache: skipped add_cache: %s", e)
except Exception as e:
verbose_logger.exception("LiteLLM Cache: Excepton add_cache: %s", e)
log_redis_failure(verbose_logger, logging.ERROR, "LiteLLM Cache: exception in add_cache", e)
async def async_add_cache(self, result, dynamic_cache_object: BaseCache | None = None, **kwargs):
"""
@ -698,10 +697,8 @@ class Cache:
await dynamic_cache_object.async_set_cache(cache_key, cached_data, **kwargs)
else:
await self.cache.async_set_cache(cache_key, cached_data, **kwargs)
except RedisCircuitBreakerOpenError as e:
verbose_logger.debug("LiteLLM Cache: skipped add_cache: %s", e)
except Exception as e:
verbose_logger.exception("LiteLLM Cache: Excepton add_cache: %s", e)
log_redis_failure(verbose_logger, logging.ERROR, "LiteLLM Cache: exception in add_cache", e)
def _convert_to_cached_embedding(
self,
@ -879,10 +876,8 @@ class Cache:
await dynamic_cache_object.async_set_cache_pipeline(cache_list=cache_list, **kwargs)
else:
await self.cache.async_set_cache_pipeline(cache_list=cache_list, **kwargs)
except RedisCircuitBreakerOpenError as e:
verbose_logger.debug("LiteLLM Cache: skipped add_cache: %s", e)
except Exception as e:
verbose_logger.exception("LiteLLM Cache: Excepton add_cache: %s", e)
log_redis_failure(verbose_logger, logging.ERROR, "LiteLLM Cache: exception in add_cache", e)
def should_use_cache(self, **kwargs):
"""

4
litellm/caching/caching_handler.py
View file

@ -1217,7 +1217,9 @@ class LLMCachingHandler:
}
if litellm.cache is not None:
litellm_params["preset_cache_key"] = litellm.cache._get_preset_cache_key_from_kwargs(**kwargs)
litellm_params["preset_cache_key"] = (
self.preset_cache_key or litellm.cache._get_preset_cache_key_from_kwargs(**kwargs)
)
else:
litellm_params["preset_cache_key"] = None

70
litellm/caching/dual_cache.py
View file

@ -8,8 +8,8 @@ Has 4 primary methods:
- async_get_cache
"""
import logging
import time
import traceback
from collections.abc import Sequence
from threading import Lock
from typing import TYPE_CHECKING, Any, Final
@ -23,7 +23,7 @@ from litellm.constants import DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE
from .base_cache import BaseCache
from .in_memory_cache import InMemoryCache
from .redis_cache import RedisCache, RedisCircuitBreakerOpenError
from .redis_cache import RedisCache, log_redis_failure
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
@ -177,10 +177,10 @@ class DualCache(BaseCache):
print_verbose(f"get cache: cache result: {result}")
return result
except RedisCircuitBreakerOpenError:
return None
except Exception:
verbose_logger.error(traceback.format_exc())
except Exception as e:
log_redis_failure(
verbose_logger, logging.ERROR, "LiteLLM Cache: exception in get_cache", e, with_traceback=True
)
def batch_get_cache(
self,
@ -206,9 +206,6 @@ class DualCache(BaseCache):
redis_result: Final = self.redis_cache.batch_get_cache(
key_list=sublist_keys, parent_otel_span=parent_otel_span
)
except RedisCircuitBreakerOpenError:
self._rollback_redis_batch_key_reservations(previous_access_times)
return result
except Exception:
# Do not throttle subsequent callers if the Redis read fails.
self._rollback_redis_batch_key_reservations(previous_access_times)
@ -222,8 +219,10 @@ class DualCache(BaseCache):
return list( # mutable-ok: public list contract
redis_result.get(key) if value is None else value for key, value in zip(keys, result)
)
except Exception:
verbose_logger.error(traceback.format_exc())
except Exception as e:
log_redis_failure(
verbose_logger, logging.ERROR, "LiteLLM Cache: exception in batch_get_cache", e, with_traceback=True
)
async def async_get_cache(
self,
@ -255,10 +254,10 @@ class DualCache(BaseCache):
print_verbose(f"get cache: cache result: {result}")
return result
except RedisCircuitBreakerOpenError:
return None
except Exception:
verbose_logger.error(traceback.format_exc())
except Exception as e:
log_redis_failure(
verbose_logger, logging.ERROR, "LiteLLM Cache: exception in async_get_cache", e, with_traceback=True
)
def _reserve_redis_batch_keys(
self,
@ -326,9 +325,6 @@ class DualCache(BaseCache):
redis_result: Final = await self.redis_cache.async_batch_get_cache(
sublist_keys, parent_otel_span=parent_otel_span
)
except RedisCircuitBreakerOpenError:
self._rollback_redis_batch_key_reservations(previous_access_times)
return result
except Exception:
# Do not throttle subsequent callers if the Redis read fails.
self._rollback_redis_batch_key_reservations(previous_access_times)
@ -349,8 +345,14 @@ class DualCache(BaseCache):
await self.in_memory_cache.async_set_cache(key, value, **self._backfill_kwargs(kwargs))
return result
except Exception:
verbose_logger.error(traceback.format_exc())
except Exception as e:
log_redis_failure(
verbose_logger,
logging.ERROR,
"LiteLLM Cache: exception in async_batch_get_cache",
e,
with_traceback=True,
)
async def async_set_cache(self, key, value, local_only: bool = False, **kwargs):
print_verbose(f"async set cache: cache key: {key}; local_only: {local_only}; value: {value}")
@ -362,10 +364,10 @@ class DualCache(BaseCache):
if self.redis_cache is not None and local_only is False:
await self.redis_cache.async_set_cache(key, value, **kwargs)
except RedisCircuitBreakerOpenError:
return
except Exception as e:
verbose_logger.exception("LiteLLM Cache: Excepton async add_cache: %s", e)
log_redis_failure(
verbose_logger, logging.ERROR, "LiteLLM Cache: exception in async add_cache", e, with_traceback=True
)
# async_batch_set_cache
async def async_set_cache_pipeline(self, cache_list: list, local_only: bool = False, **kwargs):
@ -383,10 +385,10 @@ class DualCache(BaseCache):
await self.redis_cache.async_set_cache_pipeline(
cache_list=cache_list, ttl=kwargs.pop("ttl", None), **kwargs
)
except RedisCircuitBreakerOpenError:
return
except Exception as e:
verbose_logger.exception("LiteLLM Cache: Excepton async add_cache: %s", e)
log_redis_failure(
verbose_logger, logging.ERROR, "LiteLLM Cache: exception in async add_cache", e, with_traceback=True
)
async def async_increment_cache(
self,
@ -422,12 +424,12 @@ class DualCache(BaseCache):
refresh_ttl=refresh_ttl,
)
return result
except RedisCircuitBreakerOpenError:
return result
except Exception as e:
verbose_logger.warning(
"Redis async_increment_cache failed, falling back to in-memory result: %s",
log_redis_failure(
verbose_logger,
logging.WARNING,
"Redis async_increment_cache failed, falling back to in-memory result",
e,
)
return result
@ -453,12 +455,12 @@ class DualCache(BaseCache):
parent_otel_span=parent_otel_span,
)
return result
except RedisCircuitBreakerOpenError:
return result
except Exception as e:
verbose_logger.warning(
"Redis async_increment_cache_pipeline failed, falling back to in-memory result: %s",
log_redis_failure(
verbose_logger,
logging.WARNING,
"Redis async_increment_cache_pipeline failed, falling back to in-memory result",
e,
)
return result

173
litellm/caching/redis_cache.py
View file

@ -14,10 +14,10 @@ import functools
import hashlib
import inspect
import json
import logging
import time
from collections.abc import Awaitable, Callable, Iterator, Sequence
from collections.abc import Awaitable, Callable, Sequence
from contextvars import ContextVar
from dataclasses import dataclass
from datetime import timedelta
from typing import TYPE_CHECKING, Any, Final, Protocol, TypeVar, cast
@ -149,10 +149,6 @@ def _get_call_stack_info(num_frames: int = 2) -> str:
return "unknown"
class RedisCircuitBreakerOpenError(Exception):
"""Expected fast-fail while the breaker is open; optional-cache callers treat it as a miss."""
class RedisCircuitBreaker:
"""
Tracks Redis health for a RedisCache instance.
@ -168,12 +164,8 @@ class RedisCircuitBreaker:
(no success or hard failure in between) that reaches
failure_threshold and spans timeout_min_duration seconds
OPEN -> HALF_OPEN after recovery_timeout seconds
HALF_OPEN -> CLOSED on the recovery probe's success
HALF_OPEN -> OPEN on the recovery probe's failure (resets timer)
Every OPEN transition starts a new generation. A call reports its outcome only for the
generation it was admitted under, so a success from a call that was already in flight
when the breaker opened cannot close it and the HALF_OPEN probe is the only call that can.
HALF_OPEN -> CLOSED on success
HALF_OPEN -> OPEN on failure (resets timer)
Timeouts are accounted separately from hard connectivity failures because the async
Redis timeout includes time waiting for the worker event loop to resume: one loop
@ -203,14 +195,9 @@ class RedisCircuitBreaker:
self._timeout_count = 0
self._timeout_streak_started_at: float | None = None
self._opened_at: float | None = None
self._generation = 0
self._state = self.CLOSED
_breaker_metrics().record_state_change(None, self._state)
@property
def generation(self) -> int:
return self._generation
def is_open(self) -> bool:
"""Returns True if Redis calls should be skipped."""
if not self.enabled:
@ -255,19 +242,18 @@ class RedisCircuitBreaker:
if self._state != self.OPEN:
verbose_logger.warning(
"Redis circuit breaker OPENED after %d consecutive failures"
" (%d hard connectivity), fast-failing Redis calls for %ds",
" (%d hard connectivity) — fast-failing Redis calls for %ds",
self._failure_count,
self._hard_failure_count,
self.recovery_timeout,
)
self._generation += 1
self._set_state(self.OPEN)
def record_success(self) -> None:
if not self.enabled:
return
if self._state == self.HALF_OPEN:
verbose_logger.info("Redis circuit breaker CLOSED, Redis recovered")
verbose_logger.info("Redis circuit breaker CLOSED — Redis recovered")
self._failure_count = 0
self._hard_failure_count = 0
self._timeout_count = 0
@ -285,10 +271,7 @@ class RedisCircuitBreaker:
_RedisCallResult = TypeVar("_RedisCallResult")
_swallowed_redis_failures: Final[ContextVar[tuple[bool, ...]]] = ContextVar(
"litellm_swallowed_redis_failures", default=()
)
_breaker_depth: Final[ContextVar[int]] = ContextVar("litellm_redis_breaker_depth", default=0)
_swallowed_redis_failures: Final[ContextVar[int]] = ContextVar("litellm_swallowed_redis_failures", default=0)
def _opaque_kwarg_key(value: object) -> str:
@ -337,22 +320,8 @@ def _redis_timeout_error_types() -> tuple[type, ...]:
return (RedisTimeoutError, TimeoutError)
_MAX_EXCEPTION_CAUSE_DEPTH: Final = 20
def _exception_cause_chain(exc: BaseException) -> Iterator[BaseException]:
current = exc # rebind-ok: advances one link per iteration of the bounded walk
for _ in range(_MAX_EXCEPTION_CAUSE_DEPTH):
yield current
if current.__cause__ is None:
return
current = current.__cause__
def _is_redis_timeout_failure(exc: BaseException) -> bool:
"""Walks __cause__: a blocking pool wait raises ConnectionError from asyncio.TimeoutError."""
timeout_types: Final = _redis_timeout_error_types()
return any(isinstance(cause, timeout_types) for cause in _exception_cause_chain(exc))
return isinstance(exc, _redis_timeout_error_types())
class _BreakerMetrics:
@ -408,8 +377,8 @@ def _breaker_metrics() -> _BreakerMetrics:
return _BreakerMetrics()
def _record_swallowed_redis_failure(exc: BaseException) -> None:
"""Note a Redis failure that the calling method is about to swallow, for the breaker exit to judge.
def _record_swallowed_redis_failure(breaker: RedisCircuitBreaker, exc: BaseException) -> None:
"""Record a Redis failure that the calling method is about to swallow.
The marker is a ContextVar rather than a counter on the breaker because breakers are
shared by every concurrent caller. A plain shared counter cannot tell "my call failed"
@ -419,75 +388,38 @@ def _record_swallowed_redis_failure(exc: BaseException) -> None:
"""
if not _is_redis_health_failure(exc):
return
_swallowed_redis_failures.set((*_swallowed_redis_failures.get(), _is_redis_timeout_failure(exc)))
breaker.record_failure(is_timeout=_is_redis_timeout_failure(exc))
_swallowed_redis_failures.set(_swallowed_redis_failures.get() + 1)
@dataclass(frozen=True, slots=True)
class _BreakerAdmission:
swallowed_before: int
generation: int
nested: bool
class RedisCircuitBreakerOpenError(Exception):
pass
def _enter_circuit_breaker(breaker: RedisCircuitBreaker, name: str) -> _BreakerAdmission:
"""Reject the call if the breaker is open, else snapshot what its outcome will be judged against."""
def log_redis_failure(
logger: logging.Logger, level: int, message: str, exc: BaseException, with_traceback: bool = False
) -> None:
if isinstance(exc, RedisCircuitBreakerOpenError):
logger.debug("%s: %s", message, exc)
return
logger.log(level, "%s: %s", message, exc, exc_info=exc if with_traceback else None)
def _enter_circuit_breaker(breaker: RedisCircuitBreaker, name: str) -> int:
"""Reject the call if the breaker is open, else return the swallowed-failure count to compare against."""
if breaker.is_open():
raise RedisCircuitBreakerOpenError(f"Redis circuit breaker is open, skipping {name}")
depth: Final = _breaker_depth.get()
_breaker_depth.set(depth + 1)
return _BreakerAdmission(
swallowed_before=len(_swallowed_redis_failures.get()),
generation=breaker.generation,
nested=depth > 0,
)
raise RedisCircuitBreakerOpenError(f"Redis circuit breaker is open — skipping {name}")
return _swallowed_redis_failures.get()
def _take_swallowed_failures(admission: _BreakerAdmission) -> tuple[bool, ...]:
"""Return the is_timeout flag of every failure this call swallowed, and drop them from the context."""
all_swallowed: Final = _swallowed_redis_failures.get()
_swallowed_redis_failures.set(all_swallowed[: admission.swallowed_before])
return all_swallowed[admission.swallowed_before :]
def _exit_circuit_breaker(breaker: RedisCircuitBreaker, admission: _BreakerAdmission) -> None:
"""Report the call's outcome to the breaker generation it was admitted under.
def _exit_circuit_breaker(breaker: RedisCircuitBreaker, swallowed_before: int) -> None:
"""Record success only when nothing failed while the call ran.
Several Redis methods catch their own connection errors and return a default, so a
method that returned is not on its own proof of a healthy Redis. A call admitted
before the breaker opened reports nothing: its failures would refresh the open
timer or knock out the recovery probe, and its success would close it early.
A guarded method calling another guarded method is one Redis interaction, so only
the outermost admission reports. The inner one leaves its swallowed failures in the
context for the outer to judge, otherwise the outer would read a clean context and
reset the streak the inner just fed.
method that returned is not on its own proof of a healthy Redis.
"""
if admission.nested:
return
swallowed: Final = _take_swallowed_failures(admission)
if admission.generation != breaker.generation:
return
if not swallowed:
if _swallowed_redis_failures.get() == swallowed_before:
breaker.record_success()
return
for is_timeout in swallowed:
breaker.record_failure(is_timeout=is_timeout)
def _leave_circuit_breaker() -> None:
_breaker_depth.set(_breaker_depth.get() - 1)
def _fail_circuit_breaker(breaker: RedisCircuitBreaker, admission: _BreakerAdmission, exc: BaseException) -> None:
if admission.nested:
return
swallowed: Final = _take_swallowed_failures(admission)
if admission.generation != breaker.generation:
return
for is_timeout in swallowed:
breaker.record_failure(is_timeout=is_timeout)
if _is_redis_health_failure(exc):
breaker.record_failure(is_timeout=_is_redis_timeout_failure(exc))
async def _run_under_circuit_breaker(
@ -500,15 +432,14 @@ async def _run_under_circuit_breaker(
Shared by the method decorator and the Lua script executor so both feed the same
health signal.
"""
admission: Final = _enter_circuit_breaker(breaker, name)
swallowed_before: Final = _enter_circuit_breaker(breaker, name)
try:
result: Final = await call()
except Exception as e:
_fail_circuit_breaker(breaker, admission, e)
if _is_redis_health_failure(e):
breaker.record_failure(is_timeout=_is_redis_timeout_failure(e))
raise
finally:
_leave_circuit_breaker()
_exit_circuit_breaker(breaker, admission)
_exit_circuit_breaker(breaker, swallowed_before)
return result
@ -518,15 +449,14 @@ def _run_under_circuit_breaker_sync(
call: Callable[[], _RedisCallResult],
) -> _RedisCallResult:
"""Run one blocking Redis call under a circuit breaker, feeding the same health signal as the async path."""
admission: Final = _enter_circuit_breaker(breaker, name)
swallowed_before: Final = _enter_circuit_breaker(breaker, name)
try:
result: Final = call()
except Exception as e:
_fail_circuit_breaker(breaker, admission, e)
if _is_redis_health_failure(e):
breaker.record_failure(is_timeout=_is_redis_timeout_failure(e))
raise
finally:
_leave_circuit_breaker()
_exit_circuit_breaker(breaker, admission)
_exit_circuit_breaker(breaker, swallowed_before)
return result
@ -1099,7 +1029,7 @@ class RedisCache(BaseCache):
str(e),
value,
)
_record_swallowed_redis_failure(e)
_record_swallowed_redis_failure(self._circuit_breaker, e)
async def _pipeline_helper(
self,
@ -1186,7 +1116,7 @@ class RedisCache(BaseCache):
str(e),
cache_value,
)
_record_swallowed_redis_failure(e)
_record_swallowed_redis_failure(self._circuit_breaker, e)
async def _set_cache_sadd_helper(
self,
@ -1271,7 +1201,7 @@ class RedisCache(BaseCache):
str(e),
value,
)
_record_swallowed_redis_failure(e)
_record_swallowed_redis_failure(self._circuit_breaker, e)
@_redis_circuit_breaker_guard
async def batch_cache_write(self, key, value, **kwargs):
@ -1407,7 +1337,6 @@ class RedisCache(BaseCache):
except Exception:
return ast.literal_eval(decoded)
@_redis_circuit_breaker_guard_sync
def get_cache(self, key, parent_otel_span: Span | None = None, **kwargs):
try:
key = self.check_and_fix_namespace(key=key)
@ -1428,8 +1357,7 @@ class RedisCache(BaseCache):
return self._get_cache_logic(cached_response=cached_response)
except Exception as e:
# NON blocking - notify users Redis is throwing an exception
verbose_logger.error("litellm.caching.caching: get() - Got exception from REDIS: %s", e)
_record_swallowed_redis_failure(e)
verbose_logger.error("litellm.caching.caching: get() - Got exception from REDIS: ", e)
def _run_redis_mget_operation(self, keys: list[str]) -> Sequence[bytes | str | None]:
"""
@ -1468,10 +1396,10 @@ class RedisCache(BaseCache):
start_time: Final = time.time()
try:
swallowed_before: Final = _enter_circuit_breaker(self._circuit_breaker, "batch_get_cache")
_keys: Final = [self.check_and_fix_namespace(key=cache_key or "") for cache_key in _key_list]
results: Final = _run_under_circuit_breaker_sync(
self._circuit_breaker, "batch_get_cache", lambda: self._run_redis_mget_operation(keys=_keys)
)
results: Final = self._run_redis_mget_operation(keys=_keys)
_exit_circuit_breaker(self._circuit_breaker, swallowed_before)
end_time: Final = time.time()
_duration: Final = end_time - start_time
self.service_logger_obj.service_success_hook(
@ -1495,8 +1423,6 @@ class RedisCache(BaseCache):
decoded_results[k] = v
return decoded_results
except RedisCircuitBreakerOpenError:
raise
except Exception as e:
failed_at: Final = time.time()
self.service_logger_obj.service_failure_hook(
@ -1509,6 +1435,7 @@ class RedisCache(BaseCache):
parent_otel_span=parent_otel_span,
)
verbose_logger.error("Error occurred in batch get cache - %s", e)
_record_swallowed_redis_failure(self._circuit_breaker, e)
return key_value_dict
@_redis_circuit_breaker_guard
@ -1555,7 +1482,7 @@ class RedisCache(BaseCache):
)
)
print_verbose(f"litellm.caching.caching: async get() - Got exception from REDIS: {e}")
_record_swallowed_redis_failure(e)
_record_swallowed_redis_failure(self._circuit_breaker, e)
@_redis_circuit_breaker_guard
async def async_batch_get_cache(
@ -1627,7 +1554,7 @@ class RedisCache(BaseCache):
)
)
verbose_logger.error("Error occurred in async batch get cache - %s", e)
_record_swallowed_redis_failure(e)
_record_swallowed_redis_failure(self._circuit_breaker, e)
return key_value_dict
def sync_ping(self) -> bool:
@ -1879,7 +1806,7 @@ class RedisCache(BaseCache):
return ttl
except Exception as e:
verbose_logger.debug("Redis TTL Error: %s", e)
_record_swallowed_redis_failure(e)
_record_swallowed_redis_failure(self._circuit_breaker, e)
return None
@_redis_circuit_breaker_guard

3
litellm/completion_extras/litellm_responses_transformation/transformation.py
View file

@ -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.

5
litellm/constants.py
View file

@ -571,7 +571,6 @@ ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS: Final = int(
LOGGING_WORKER_CONCURRENCY: Final = int(os.getenv("LOGGING_WORKER_CONCURRENCY", 100)) # Must be above 0
LOGGING_WORKER_MAX_QUEUE_SIZE: Final = int(os.getenv("LOGGING_WORKER_MAX_QUEUE_SIZE", 50_000))
LOGGING_WORKER_MAX_TIME_PER_COROUTINE: Final = float(os.getenv("LOGGING_WORKER_MAX_TIME_PER_COROUTINE", 20.0))
LOGGING_WORKER_ERROR_TRACEBACK_INTERVAL_SECONDS: Final = 60.0
LOGGING_WORKER_CLEAR_PERCENTAGE: Final = int(
os.getenv("LOGGING_WORKER_CLEAR_PERCENTAGE", 50)
) # Percentage of queue to clear (default: 50%)
@ -1655,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"
@ -2000,6 +2000,9 @@ NON_INFERENCE_CALL_TYPES: Final[frozenset[str]] = frozenset(
}
)
UNKNOWN_MODEL_SPEND_LOG_MODEL: Final[str] = "unknown-model"
MAX_SPEND_LOG_MODEL_NAME_LENGTH: Final[int] = 256
# PTU reservation rollup writes rows to LiteLLM_DailyTeamSpend with this
# sentinel api_key so PTU flat cost stays distinguishable from real per-request
# spend under the table's composite unique constraint.

41
litellm/integrations/custom_logger.py
View file

@ -3,12 +3,13 @@
import re
import traceback
from collections.abc import AsyncGenerator, Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional
from pydantic import BaseModel
from litellm._logging import verbose_logger
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER, EMPTY_MAPPING
from litellm.types.integrations.argilla import ArgillaItem
from litellm.types.integrations.custom_logger import AgenticLoopPlan
from litellm.types.llms.openai import AllMessageValues, ChatCompletionRequest
@ -897,10 +898,9 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
This is useful for logging payloads that contain sensitive information.
"""
from copy import copy
import litellm
from litellm import Choices, Message, ModelResponse
from litellm.litellm_core_utils.classifier_logging import CLASSIFIER_AUDIT_FIELDS, without_classifier_audit
turn_off_message_logging: Final[bool] = getattr(self, "turn_off_message_logging", False)
excluded_fields: Final[list[str] | None] = getattr(litellm, "standard_logging_payload_excluded_fields", None)
@ -909,30 +909,25 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
if turn_off_message_logging is False and not excluded_fields:
return model_call_details
# Only make a shallow copy of the top-level dict to avoid deepcopy issues
# with complex objects like AuthenticationError that may be present
model_call_details_copy: Final = copy(model_call_details)
standard_logging_object: Final = model_call_details.get("standard_logging_object")
if standard_logging_object is None:
return model_call_details_copy
return model_call_details.copy()
# Make a copy of just the standard_logging_object to avoid modifying the original
standard_logging_object_copy: Final = copy(standard_logging_object)
# Handle excluded fields - remove them entirely from the payload
if excluded_fields:
for field in excluded_fields:
if field in standard_logging_object_copy:
del standard_logging_object_copy[field]
standard_logging_object_copy: Final = {
key: value
for key, value in standard_logging_object.items()
if key not in (excluded_fields or ()) and not (turn_off_message_logging and key in CLASSIFIER_AUDIT_FIELDS)
}
# Handle turn_off_message_logging - redact messages and responses (if not already excluded)
if turn_off_message_logging:
redacted_str: Final = "redacted-by-litellm"
if "messages" not in (excluded_fields or []) and standard_logging_object_copy.get("messages") is not None:
if "messages" not in (excluded_fields or ()) and standard_logging_object_copy.get("messages") is not None:
standard_logging_object_copy["messages"] = [Message(content=redacted_str).model_dump()]
if "response" not in (excluded_fields or []) and standard_logging_object_copy.get("response") is not None:
if "response" not in (excluded_fields or ()) and standard_logging_object_copy.get("response") is not None:
response: Final = standard_logging_object_copy["response"]
# Check if this is a ResponsesAPIResponse (has "output" field)
if isinstance(response, dict) and "output" in response:
@ -956,8 +951,18 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
model_response_dict: Final = model_response.model_dump()
standard_logging_object_copy["response"] = model_response_dict
model_call_details_copy["standard_logging_object"] = standard_logging_object_copy
return model_call_details_copy
params: Final = model_call_details.get("litellm_params")
request: Final = params.get("proxy_server_request") if isinstance(params, dict) else None
redacted_params: Final = (
MappingProxyType({"litellm_params": {**params, "proxy_server_request": without_classifier_audit(request)}})
if turn_off_message_logging and isinstance(params, dict) and isinstance(request, dict)
else EMPTY_MAPPING
)
return {
**model_call_details,
**redacted_params,
"standard_logging_object": standard_logging_object_copy,
}
async def get_proxy_server_request_from_cold_storage_with_object_key(
self,

39
litellm/integrations/newrelic/newrelic_metrics.py
View file

@ -5,8 +5,9 @@ NR Reference API: https://docs.newrelic.com/docs/data-apis/ingest-apis/metric-ap
`async_log_success_event` / `async_log_failure_event` queue one record per request;
at flush the queue is aggregated by (team, model group, model, provider, status)
into count/summary metrics. `interval.ms` is the real window between flushes,
computed at flush time.
into count/summary metrics, plus one max/remaining budget gauge pair per team
taken from the team's latest record. `interval.ms` is the real window between
flushes, computed at flush time.
Team-scoped by construction: the ingest key is injected explicitly and there is
deliberately no environment-variable fallback, so a team's metrics are never sent
@ -47,11 +48,14 @@ from litellm.types.integrations.newrelic import (
NEWRELIC_METRIC_PROMPT_TOKENS,
NEWRELIC_METRIC_REQUEST_DURATION_MS,
NEWRELIC_METRIC_REQUESTS,
NEWRELIC_METRIC_TEAM_MAX_BUDGET,
NEWRELIC_METRIC_TEAM_REMAINING_BUDGET,
NEWRELIC_METRIC_TOTAL_TOKENS,
NEWRELIC_METRICS_MAX_BATCH_SIZE,
NEWRELIC_METRICS_MAX_DRAIN_PASSES,
NEWRELIC_METRICS_MAX_RETRY_QUEUE_SIZE,
NewRelicCountMetric,
NewRelicGaugeMetric,
NewRelicMetric,
NewRelicMetricCommon,
NewRelicMetricEnvelope,
@ -98,6 +102,8 @@ def _metric_record_from_payload(standard_logging_object: StandardLoggingPayload)
completion_tokens=int(standard_logging_object.get("completion_tokens") or 0),
total_tokens=int(standard_logging_object.get("total_tokens") or 0),
duration_ms=float(standard_logging_object.get("response_time") or 0.0) * 1000.0,
team_max_budget=metadata.get("user_api_key_team_max_budget") if metadata else None,
team_spend=metadata.get("user_api_key_team_spend") if metadata else None,
)
@ -140,6 +146,33 @@ def _bucket_metrics(bucket_records: tuple[NewRelicMetricRecord, ...]) -> tuple[N
return (*count_metrics, summary_metric)
def _team_budget_gauges(record: NewRelicMetricRecord) -> tuple[NewRelicMetric, ...]:
team_max_budget: Final = record.team_max_budget
if team_max_budget is None:
return ()
attributes: Final[Mapping[str, str]] = { # mutable-ok: JSON leaf; safe_dumps stringifies MappingProxyType
key: value[:NEWRELIC_METRIC_ATTRIBUTE_MAX_LEN]
for key, value in (("team_id", record.team_id), ("team_alias", record.team_alias))
if value
}
remaining_budget: Final = team_max_budget - (record.team_spend or 0.0) - record.response_cost
return (
NewRelicGaugeMetric(
name=NEWRELIC_METRIC_TEAM_MAX_BUDGET, type="gauge", value=team_max_budget, attributes=attributes
),
NewRelicGaugeMetric(
name=NEWRELIC_METRIC_TEAM_REMAINING_BUDGET, type="gauge", value=remaining_budget, attributes=attributes
),
)
def _team_budget_metrics(records: tuple[NewRelicMetricRecord, ...]) -> tuple[NewRelicMetric, ...]:
latest_by_team: Final[Mapping[str, NewRelicMetricRecord]] = MappingProxyType(
{record.team_id: record for record in records if record.team_id}
)
return tuple(gauge for record in latest_by_team.values() for gauge in _team_budget_gauges(record))
def build_metric_payload(
records: tuple[NewRelicMetricRecord, ...],
*,
@ -158,7 +191,7 @@ def build_metric_payload(
"timestamp": int(window_start * 1000),
"interval.ms": interval_ms,
}
return (NewRelicMetricEnvelope(common=common, metrics=metrics),)
return (NewRelicMetricEnvelope(common=common, metrics=(*metrics, *_team_budget_metrics(records))),)
class NewRelicMetricsLogger(CustomBatchLogger):

67
litellm/litellm_core_utils/classifier_logging.py Normal file
View file

@ -0,0 +1,67 @@
from collections.abc import Mapping
from types import MappingProxyType
from typing import Final
from pydantic import JsonValue, TypeAdapter, ValidationError
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload
from litellm.types.utils import AUTOROUTER_CLASSIFIER_CALL_ORIGIN, ClassifierAudit
CLASSIFIER_AUDIT_FIELDS: Final = ("classifier_input", "originating_request_masked")
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
def classifier_input_snapshot(value: object, *, openai_sdk: bool = False) -> Mapping[str, JsonValue] | None:
try:
if openai_sdk and isinstance(value, Mapping):
body: Final = MappingProxyType(
{key: item for key, item in value.items() if key not in ("extra_headers", "extra_query", "extra_body")}
)
extra_body: Final = value.get("extra_body")
return _JSON_OBJECT.validate_python(
MappingProxyType({**body, **extra_body}) if isinstance(extra_body, Mapping) else body
)
return (
_JSON_OBJECT.validate_json(value)
if isinstance(value, (str, bytes))
else _JSON_OBJECT.validate_python(value)
)
except ValidationError:
return None
def is_classifier_call(call_type: str, params: Mapping[str, object]) -> bool:
return call_type in ("completion", "acompletion", "responses", "aresponses") and any(
isinstance(metadata := params.get(key), Mapping)
and metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == AUTOROUTER_CLASSIFIER_CALL_ORIGIN
for key in ("metadata", "litellm_metadata")
)
def masked_originating_request(request_kwargs: Mapping[str, object] | None) -> Mapping[str, JsonValue] | None:
request: Final = request_kwargs.get("proxy_server_request") if request_kwargs is not None else None
body: Final = request.get("body") if isinstance(request, Mapping) else None
if not isinstance(body, Mapping):
return None
serializable: Final = classifier_input_snapshot(safe_dumps(body))
return classifier_input_snapshot(redact_credentials_in_payload(serializable)) if serializable is not None else None
def classifier_audit_fields(payload: Mapping[str, object]) -> ClassifierAudit:
classifier_input: Final = classifier_input_snapshot(payload.get("classifier_input"))
originating_request: Final = classifier_input_snapshot(payload.get("originating_request_masked"))
if classifier_input is None:
return (
ClassifierAudit(originating_request_masked=originating_request)
if originating_request is not None
else ClassifierAudit()
)
if originating_request is None:
return ClassifierAudit(classifier_input=classifier_input)
return ClassifierAudit(classifier_input=classifier_input, originating_request_masked=originating_request)
def without_classifier_audit(payload: Mapping[str, object]) -> dict[str, object]:
return {key: value for key, value in payload.items() if key not in CLASSIFIER_AUDIT_FIELDS}

30
litellm/litellm_core_utils/litellm_logging.py
View file

@ -17,7 +17,7 @@ from types import MappingProxyType, TracebackType
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, cast
from httpx import Response
from pydantic import BaseModel
from pydantic import BaseModel, JsonValue
import litellm
from litellm import (
@ -42,6 +42,7 @@ from litellm.caching.caching_handler import LLMCachingHandler
from litellm.constants import (
DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT,
DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT,
EMPTY_MAPPING,
PROVIDER_REQUEST_ID_HEADERS,
SENTRY_DENYLIST,
SENTRY_PII_DENYLIST,
@ -64,6 +65,11 @@ from litellm.integrations.custom_logger import CustomLogger
from litellm.integrations.deepeval.deepeval import DeepEvalLogger
from litellm.integrations.mlflow import MlflowLogger
from litellm.integrations.sqs import SQSLogger
from litellm.litellm_core_utils.classifier_logging import (
classifier_audit_fields,
classifier_input_snapshot,
is_classifier_call,
)
from litellm.litellm_core_utils.core_helpers import is_expected_client_error, reconstruct_model_name
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
from litellm.litellm_core_utils.internal_call_metadata import (
@ -89,6 +95,7 @@ from litellm.litellm_core_utils.redact_messages import (
redact_message_input_output_from_custom_logger,
redact_message_input_output_from_logging,
redact_streaming_responses_for_custom_logger,
should_redact_message_logging,
)
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from litellm.llms.base_llm.search.transformation import SearchResponse
@ -473,6 +480,7 @@ class Logging(LiteLLMLoggingBaseClass):
stream_options = None
litellm_request_debug: bool = False
streamed_anthropic_message_id: str | None = None
classifier_input: Mapping[str, JsonValue] | None = None
def __init__(
self,
@ -1211,6 +1219,14 @@ class Logging(LiteLLMLoggingBaseClass):
self.model_call_details["api_key"] = api_key
self.model_call_details["additional_args"] = additional_args
self.model_call_details["log_event_type"] = "pre_api_call"
if is_classifier_call(self.call_type, self.model_call_details.get("litellm_params") or EMPTY_MAPPING):
self.classifier_input = (
None
if should_redact_message_logging(self.model_call_details)
else classifier_input_snapshot(
additional_args.get("complete_input_dict"), openai_sdk=additional_args.get("openai_sdk") is True
)
)
if model: # if model name was changes pre-call, overwrite the initial model call name with the new one
self.model_call_details["model"] = model
self.model_call_details["litellm_params"]["api_base"] = self._get_masked_api_base(
@ -6293,6 +6309,18 @@ def get_standard_logging_object_payload(
)
payload: Final[StandardLoggingPayload] = StandardLoggingPayload(
**(
classifier_audit_fields(
MappingProxyType(
{
"classifier_input": logging_obj.classifier_input,
"originating_request_masked": proxy_server_request.get("originating_request_masked"),
}
)
)
if is_classifier_call(call_type or "", litellm_params) and not should_redact_message_logging(kwargs)
else EMPTY_MAPPING
),
id=str(id),
litellm_call_id=kwargs.get("litellm_call_id") or litellm_params.get("litellm_call_id"),
trace_id=StandardLoggingPayloadSetup.get_standard_logging_payload_trace_id(

25
litellm/litellm_core_utils/logging_worker.py
View file

@ -6,7 +6,6 @@ import atexit
import contextvars
import inspect
import logging
import time
from collections.abc import Coroutine, Iterator
from typing import Final
@ -17,7 +16,6 @@ from litellm.constants import (
LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS,
LOGGING_WORKER_CLEAR_PERCENTAGE,
LOGGING_WORKER_CONCURRENCY,
LOGGING_WORKER_ERROR_TRACEBACK_INTERVAL_SECONDS,
LOGGING_WORKER_MAX_QUEUE_SIZE,
LOGGING_WORKER_MAX_TIME_PER_COROUTINE,
MAX_ITERATIONS_TO_CLEAR_QUEUE,
@ -49,14 +47,10 @@ class LoggingWorker:
timeout: float = LOGGING_WORKER_MAX_TIME_PER_COROUTINE,
max_queue_size: int = LOGGING_WORKER_MAX_QUEUE_SIZE,
concurrency: int = LOGGING_WORKER_CONCURRENCY,
error_traceback_interval: float = LOGGING_WORKER_ERROR_TRACEBACK_INTERVAL_SECONDS,
):
self.timeout = timeout
self.max_queue_size = max_queue_size
self.concurrency = concurrency
self.error_traceback_interval = error_traceback_interval
self._last_error_traceback_at: dict[type[BaseException], float] = {}
self._errors_since_traceback: dict[type[BaseException], int] = {}
self._queue: asyncio.Queue[LoggingTask] | None = None
self._worker_task: asyncio.Task | None = None
self._running_tasks: set[asyncio.Task] = set()
@ -169,7 +163,7 @@ class LoggingWorker:
timeout=self.timeout,
)
except Exception as e:
self._log_task_error(e)
verbose_logger.exception("LoggingWorker error: %s", e)
finally:
self._untrack_dequeued(task)
self._queue.task_done()
@ -177,23 +171,6 @@ class LoggingWorker:
# Always release semaphore, even if queue is None
sem.release()
def _log_task_error(self, error: Exception) -> None:
"""One traceback per error type per interval: a stalled backend fails every in-flight task at once."""
now: Final = time.monotonic()
error_type: Final = type(error)
last_traceback_at: Final = self._last_error_traceback_at.get(error_type)
if last_traceback_at is not None and now - last_traceback_at < self.error_traceback_interval:
self._errors_since_traceback[error_type] = self._errors_since_traceback.get(error_type, 0) + 1
return
verbose_logger.exception(
"LoggingWorker error (%d more %s suppressed since the last traceback): %r",
self._errors_since_traceback.get(error_type, 0),
error_type.__name__,
error,
)
self._last_error_traceback_at[error_type] = now
self._errors_since_traceback[error_type] = 0
async def _worker_loop(self) -> None:
"""Main worker loop that gets tasks and schedules them to run concurrently."""
try:

22
litellm/litellm_core_utils/redact_messages.py
View file

@ -16,6 +16,7 @@ from typing import TYPE_CHECKING, Any, Final
import litellm
from litellm.constants import REDACTED_BY_LITELLM
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.classifier_logging import without_classifier_audit
from litellm.litellm_core_utils.core_helpers import (
get_metadata_variable_name_from_kwargs,
)
@ -170,17 +171,11 @@ def redacted_standard_logging_payload(payload: Mapping[str, object]) -> Mapping[
payload, but the failure path does not, so a callback that batches both has to redact
the ones it is handed.
"""
redacted: Final = copy.deepcopy(dict(payload)) # mutable-ok: redacted in place below
_redact_standard_logging_object({"standard_logging_object": redacted}) # mutable-ok: the callee's shape
return redacted
return _redact_standard_logging_object(payload)
def _redact_standard_logging_object(model_call_details: dict):
"""Redact messages and response inside standard_logging_object if present."""
standard_logging_object: Final = model_call_details.get("standard_logging_object")
if standard_logging_object is None:
return
def _redact_standard_logging_object(payload: Mapping[str, object]) -> dict[str, object]:
standard_logging_object: Final = copy.deepcopy(without_classifier_audit(payload))
redacted_str: Final = REDACTED_BY_LITELLM
if standard_logging_object.get("messages") is not None:
@ -203,6 +198,7 @@ def _redact_standard_logging_object(model_call_details: dict):
else:
# For other formats (empty dict, None, etc.), use simple text format
standard_logging_object["response"] = {"text": redacted_str}
return standard_logging_object
def _redact_tool_calls_dict(message: Mapping[str, object]) -> None:
@ -254,10 +250,16 @@ def perform_redaction(model_call_details: dict, result, redact_streaming_respons
copy via redact_streaming_responses_for_custom_logger instead.
"""
# Redact model_call_details
params: Final = model_call_details.get("litellm_params")
request: Final = params.get("proxy_server_request") if isinstance(params, dict) else None
if isinstance(params, dict) and isinstance(request, Mapping):
model_call_details["litellm_params"] = {**params, "proxy_server_request": without_classifier_audit(request)}
model_call_details["messages"] = [{"role": "user", "content": REDACTED_BY_LITELLM}]
model_call_details["prompt"] = ""
model_call_details["input"] = ""
_redact_standard_logging_object(model_call_details)
standard_logging_object: Final = model_call_details.get("standard_logging_object")
if isinstance(standard_logging_object, Mapping):
model_call_details["standard_logging_object"] = _redact_standard_logging_object(standard_logging_object)
redact_vertex_ai_metadata_from_litellm_params(model_call_details)
# Redact streaming response

1
litellm/litellm_core_utils/sensitive_data_masker.py
View file

@ -15,6 +15,7 @@ _DEFAULT_SENSITIVE_PATTERNS: Final = frozenset(
"token",
"auth",
"authorization",
"cookie",
"credential",
# Plural form: Vertex uses ``vertex_credentials``; segment-exact
# matching otherwise misses it because "credential" != "credentials".

2
litellm/llms/azure/azure.py
View file

@ -323,6 +323,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
"api_version": api_version,
"api_base": api_base,
"complete_input_dict": data,
"openai_sdk": True,
},
)
if not isinstance(max_retries, int):
@ -429,6 +430,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
"api_base": api_base,
"acompletion": True,
"complete_input_dict": data,
"openai_sdk": True,
},
)

3
litellm/llms/base_llm/responses/transformation.py
View file

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

2
litellm/llms/openai/openai.py
View file

@ -797,6 +797,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
"api_base": openai_client._base_url._uri_reference,
"acompletion": acompletion,
"complete_input_dict": data,
"openai_sdk": True,
},
)
@ -938,6 +939,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
"api_base": openai_aclient._base_url._uri_reference,
"acompletion": True,
"complete_input_dict": data,
"openai_sdk": True,
},
)

3
litellm/llms/openai/responses/transformation.py
View file

@ -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.

42
litellm/llms/voyage/embedding/transformation_contextual.py
View file

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

7
litellm/llms/wandb/chat/transformation.py
View file

@ -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,

48
litellm/model_prices_and_context_window_backup.json
View file

@ -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,
@ -57252,6 +57258,14 @@
"model_info": {
"supports_mid_conversation_system": true
}
},
{
"name": "wandb-reasoning-baseline",
"pattern": "^wandb/",
"description": "Any Weights & Biases Inference model id, anchored to the wandb/ namespace so only that provider's ids match. W&B's serverless catalog is reasoning-first and grows faster than this registry names it, so an id the map has not described yet is treated as reasoning-capable and keeps the caller's reasoning_effort instead of dropping it or raising UnsupportedParamsError. Rules lose to exact entries, so a mapped non-reasoning model such as wandb/meta-llama/Llama-3.1-8B-Instruct is unaffected. Carries no mode and no pricing, so cost stays on the standard unpriced behavior and the deployment does not read as catalog-mapped to the router's reasoning-effort resolver.",
"model_info": {
"supports_reasoning": true
}
}
]
},
@ -58630,6 +58644,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 +58657,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 +58670,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 +58683,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 +58724,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 +58737,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 +58750,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 +58763,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 +58776,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 +58799,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 +58812,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 +58823,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 +58836,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 +58856,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,

1
litellm/models/object_permission.py
View file

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

50
litellm/proxy/_experimental/mcp_server/db.py
View file

@ -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(

112
litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py
View file

@ -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(

2
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

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

25
litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py
View file

@ -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",

423
litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py Normal file
View file

@ -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,
},
)

48
litellm/proxy/_experimental/mcp_server/outbound_credentials/authz_code_refresher.py
View file

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

1
litellm/proxy/_experimental/mcp_server/outbound_credentials/oauth_token_store.py
View file

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

32
litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py
View file

@ -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:

36
litellm/proxy/_experimental/mcp_server/outbound_credentials/token_cache_codec.py
View file

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

2
litellm/proxy/_experimental/mcp_server/outbound_credentials/v2_token_store.py
View file

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

28
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -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:

42
litellm/proxy/_lazy_openapi_snapshot.json
View file

@ -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": [
{

20
litellm/proxy/_types.py
View file

@ -497,6 +497,9 @@ class LiteLLMRoutes(enum.Enum):
"/v1/messages/count_tokens",
"/v1/skills",
"/v1/skills/{skill_id}",
"/claude-code/marketplace.json",
"/claude-code/plugins",
"/claude-code/plugins/{plugin_name}",
]
# MCP tool-call / passthrough routes — data-plane. Gated by DISABLE_LLM_API_ENDPOINTS.
@ -700,6 +703,7 @@ class LiteLLMRoutes(enum.Enum):
"/spend/logs",
"/spend/logs/v2",
"/spend/logs/ui",
"/spend/logs/ui/{request_id}",
"/spend/logs/session/ui",
"/key/spend/report",
"/user/spend/report",
@ -929,10 +933,10 @@ class LiteLLMRoutes(enum.Enum):
# PROXY_ADMIN_VIEW_ONLY — the route gate must match).
"/customer/list",
"/customer/info",
# UI Logs page detail drawer (single + session) and the filter facets.
# The list endpoint `/spend/logs/ui` is covered via
# spend_tracking_routes below.
"/spend/logs/ui/{logId}",
# UI Logs page session detail drawer and the end-user filter facet.
# The list endpoint `/spend/logs/ui` and the single-log detail route
# `/spend/logs/ui/{request_id}` are covered via spend_tracking_routes
# below.
"/spend/logs/session/ui",
"/management/v1/spend_logs/end_users",
"/management/v1/spend_logs/users",
@ -1124,6 +1128,7 @@ class LiteLLM_ObjectPermissionBase(LiteLLMPydanticObjectBase):
models: list[str] | None = None
search_tools: list[str] | None = None
mcp_tool_search_enabled: bool | None = None
skills: list[str] | None = None
from litellm.models.team import BudgetLimitEntry as BudgetLimitEntry # noqa: E402
@ -2427,6 +2432,13 @@ class CoordinationRedisParams(LiteLLMPydanticObjectBase):
)
sentinel_password: str | None = Field(None, description="password for the sentinel nodes")
service_name: str | None = Field(None, description="sentinel service name")
aws_iam_auth: bool | str | None = Field(None, description="enable AWS ElastiCache IAM authentication")
aws_iam_user_name: str | None = Field(None, description="AWS ElastiCache IAM user name")
aws_iam_cache_name: str | None = Field(None, description="AWS ElastiCache cache name")
aws_iam_region: str | None = Field(None, description="AWS region for ElastiCache IAM authentication")
aws_iam_serverless: bool | str | None = Field(
None, description="the ElastiCache cache is serverless rather than a self-designed cluster"
)
def has_connection_target(self) -> bool:
return any(value is not None for value in (self.host, self.url, self.startup_nodes, self.sentinel_nodes))

82
litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py
View file

@ -2,14 +2,15 @@
CLAUDE CODE MARKETPLACE
Provides a registry/discovery layer for Claude Code plugins.
Plugins are stored as metadata + git source references in LiteLLM database.
Actual plugin files are hosted on GitHub/GitLab/Bitbucket.
Plugins are stored as metadata + source references in LiteLLM database.
Actual plugin files are hosted on GitHub/GitLab/Bitbucket or as a zip archive on
any HTTPS host (S3, Artifactory, a static file server).
Endpoints:
/claude-code/marketplace.json - GET - List plugins for Claude Code discovery (unauthenticated)
/claude-code/marketplace.json - GET - List plugins for Claude Code discovery (unauthenticated; `?key=` adds the key's granted skills)
/claude-code/plugins - POST - Register a new plugin (create-only, proxy admin only)
/claude-code/plugins - GET - List plugins (any authenticated key)
/claude-code/plugins/{name} - GET - Get plugin details (any authenticated key)
/claude-code/plugins - GET - List plugins visible to the key (enabled, plus granted disabled ones)
/claude-code/plugins/{name} - GET - Get plugin details (403 on a disabled plugin the key is not granted)
/claude-code/plugins/{name} - PUT - Update an existing plugin (proxy admin only)
/claude-code/plugins/{name}/enable - POST - Enable a plugin (proxy admin only)
/claude-code/plugins/{name}/disable - POST - Disable a plugin (proxy admin only)
@ -21,12 +22,17 @@ import re
from collections.abc import Mapping, Sequence
from datetime import datetime, timezone
from typing import Annotated, Final, Protocol, TypedDict
from urllib.parse import urlsplit
from fastapi import APIRouter, Depends, HTTPException
from fastapi import APIRouter, Depends, HTTPException, Request
from fastapi.responses import JSONResponse
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth
from litellm.proxy._types import CommonProxyErrors, ProxyException, UserAPIKeyAuth
from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_skill_access import (
SkillVisibility,
skill_visibility,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.resource_ownership import is_proxy_admin
from litellm.repositories.table_repositories import ClaudeCodePluginRepository
@ -82,7 +88,7 @@ async def _get_prisma_client() -> object:
"/claude-code/marketplace.json",
tags=["Claude Code Marketplace"],
)
async def get_marketplace():
async def get_marketplace(request: Request, key: str | None = None):
"""
Serve marketplace.json for Claude Code plugin discovery.
@ -90,24 +96,35 @@ async def get_marketplace():
- claude plugin marketplace add <url>
- claude plugin install <name>@<marketplace>
Without `key` the catalog holds the enabled (public) plugins. With `?key=sk-...`
the key is authenticated and the catalog also holds the disabled plugins granted
to it through `object_permission.skills` on the key or its team.
Returns:
Marketplace catalog with list of available plugins and their git sources.
Example:
```bash
claude plugin marketplace add http://localhost:4000/claude-code/marketplace.json
claude plugin marketplace add "http://localhost:4000/claude-code/marketplace.json?key=sk-..."
claude plugin install my-plugin@litellm
```
"""
try:
prisma_client: Final = await _get_prisma_client()
caller: Final[UserAPIKeyAuth | None] = (
await user_api_key_auth(request=request, api_key=f"Bearer {key}") if key else None
)
visibility: Final[SkillVisibility] = skill_visibility(caller)
plugins: Final[Sequence[_PluginRecord]] = await ClaudeCodePluginRepository(prisma_client).table.find_many(
where={"enabled": True}
where=visibility.where()
)
plugin_list: Final = []
for plugin in plugins:
if not visibility.allows(plugin):
continue
try:
manifest: Mapping[str, object] = json.loads(plugin.manifest_json or "{}")
except json.JSONDecodeError:
@ -147,7 +164,7 @@ async def get_marketplace():
return JSONResponse(content=marketplace)
except HTTPException:
except (HTTPException, ProxyException):
raise
except Exception as e:
verbose_proxy_logger.exception("Error generating marketplace: %s", e)
@ -162,6 +179,15 @@ async def get_marketplace():
# alphanumeric characters, dots, hyphens, and underscores.
# This implicitly blocks '..', leading '/', backslashes, and percent-encoded sequences.
_VALID_GIT_SUBDIR_PATH_RE: Final = re.compile(r"^[a-zA-Z0-9][a-zA-Z0-9._-]*(/[a-zA-Z0-9][a-zA-Z0-9._-]*)*$")
_VALID_SHA256_RE: Final = re.compile(r"^[0-9a-fA-F]{64}$")
def _is_https_url_with_host(url: str) -> bool:
try:
parts: Final = urlsplit(url)
except ValueError:
return False
return parts.scheme == "https" and bool(parts.hostname)
def _validate_plugin_source(source: Mapping[str, str]) -> None:
@ -199,10 +225,24 @@ def _validate_plugin_source(source: Mapping[str, str]) -> None:
"error": "git-subdir 'path' must be a relative path of the form 'segment/segment' (alphanumeric, dots, hyphens, underscores only)"
},
)
elif source_type == "archive":
if not _is_https_url_with_host(source.get("url", "")):
raise HTTPException(
status_code=400,
detail={
"error": "archive source must include an https 'url' field "
"(e.g., 'https://bucket.s3.amazonaws.com/plugins/plugin-name.zip')"
},
)
if "sha256" in source and not _VALID_SHA256_RE.match(source["sha256"]):
raise HTTPException(
status_code=400,
detail={"error": "archive 'sha256' must be a 64-character hex digest"},
)
else:
raise HTTPException(
status_code=400,
detail={"error": "source.source must be 'github', 'url', or 'git-subdir'"},
detail={"error": "source.source must be 'github', 'url', 'git-subdir', or 'archive'"},
)
@ -248,8 +288,8 @@ async def register_plugin(
Register a new plugin in the LiteLLM marketplace.
LiteLLM acts as a registry/discovery layer. Plugins are hosted on
GitHub/GitLab/Bitbucket. Claude Code will clone from the git source
when users install.
GitHub/GitLab/Bitbucket or as a zip archive on any https host (e.g. S3).
Claude Code clones the git source or downloads the archive when users install.
This endpoint is create-only and never overwrites. If a plugin with
the same name already exists it returns 409 Conflict; use
@ -259,7 +299,7 @@ async def register_plugin(
Parameters:
- name: Plugin name (kebab-case)
- source: Git source reference (github, url, or git-subdir format)
- source: Plugin source reference (github, url, git-subdir, or archive format)
- version: Semantic version (optional)
- description: Plugin description (optional)
- author: Author information (optional)
@ -370,13 +410,15 @@ async def list_plugins(
try:
prisma_client: Final = await _get_prisma_client()
where: Final = {"enabled": True} if enabled_only else {}
visibility: Final[SkillVisibility] = skill_visibility(user_api_key_dict)
plugins: Final[Sequence[_PluginRecord]] = await ClaudeCodePluginRepository(prisma_client).table.find_many(
where=where
where={"enabled": True} if enabled_only else visibility.where()
)
plugin_list: Final = []
for p in plugins:
if not visibility.allows(p):
continue
# Parse manifest to get additional fields
manifest = json.loads(p.manifest_json) if p.manifest_json else {}
@ -448,6 +490,12 @@ async def get_plugin(
detail={"error": f"Plugin '{plugin_name}' not found"},
)
if not skill_visibility(user_api_key_dict).allows(plugin):
raise HTTPException(
status_code=403,
detail={"error": f"Plugin '{plugin_name}' is not granted to this key"},
)
manifest: Final[Mapping[str, object]] = json.loads(plugin.manifest_json or "{}") if plugin.manifest_json else {}
return {
@ -503,7 +551,7 @@ async def update_plugin(
Parameters:
- plugin_name: Name of the plugin to update (path parameter)
- source: Git source reference (github, url, or git-subdir format)
- source: Plugin source reference (github, url, git-subdir, or archive format)
- version: Semantic version (optional)
- description: Plugin description (optional)
- author: Author information (optional)

67
litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_skill_access.py Normal file
View file

@ -0,0 +1,67 @@
"""
Claude Code marketplace visibility: enabled plugins are public, disabled plugins
are private and resolve only for proxy admins or keys granted them via
``object_permission.skills``.
"""
from dataclasses import dataclass
from typing import TYPE_CHECKING, Final, Protocol
from litellm.proxy._types import LiteLLM_ObjectPermissionTable, UserAPIKeyAuth
from litellm.proxy.common_utils.resource_ownership import is_proxy_admin
if TYPE_CHECKING:
from prisma.types import LiteLLM_ClaudeCodePluginTableWhereInput
class _SkillRecord(Protocol):
name: str
enabled: bool
def _skills_of(permission: LiteLLM_ObjectPermissionTable | None) -> frozenset[str]:
return frozenset(permission.skills or ()) if permission is not None else frozenset()
def granted_skills(user_api_key_dict: UserAPIKeyAuth) -> frozenset[str]:
"""Key grant intersected with the team grant when both are non-empty; either alone applies as is.
An empty list is the Prisma column default for every object-permission row, so it means
"no private grants configured here" and defers to the other scope, same as the agents check.
"""
key_skills: Final = _skills_of(user_api_key_dict.object_permission)
team_skills: Final = _skills_of(user_api_key_dict.team_object_permission)
match (bool(key_skills), bool(team_skills)):
case (True, True):
return key_skills & team_skills
case (True, False):
return key_skills
case _:
return team_skills
@dataclass(frozen=True, slots=True)
class SkillVisibility:
granted: frozenset[str]
sees_private: bool
def allows(self, skill: _SkillRecord) -> bool:
return skill.enabled or self.sees_private or skill.name in self.granted
def where(self) -> "LiteLLM_ClaudeCodePluginTableWhereInput":
if self.sees_private:
return {}
if not self.granted:
return {"enabled": True}
return {"OR": [{"enabled": True}, {"name": {"in": sorted(self.granted)}}]}
PUBLIC_ONLY: Final = SkillVisibility(granted=frozenset(), sees_private=False)
def skill_visibility(user_api_key_dict: UserAPIKeyAuth | None) -> SkillVisibility:
if user_api_key_dict is None:
return PUBLIC_ONLY
if is_proxy_admin(user_api_key_dict):
return SkillVisibility(granted=frozenset(), sees_private=True)
return SkillVisibility(granted=granted_skills(user_api_key_dict), sees_private=False)

4
litellm/proxy/auth/user_api_key_auth.py
View file

@ -91,6 +91,7 @@ from litellm.proxy.common_utils.http_parsing_utils import (
_safe_get_request_query_params,
_safe_set_request_parsed_body,
populate_request_with_path_params,
read_raw_json_body,
)
from litellm.proxy.common_utils.model_listing_utils import claude_code_requested_group
from litellm.proxy.common_utils.realtime_utils import _realtime_request_body
@ -2689,6 +2690,7 @@ async def _run_centralized_common_checks(
await _reserve_budget_after_common_checks(
user_api_key_auth_obj=user_api_key_auth_obj,
request=request,
request_data=request_data,
route=route,
llm_router=llm_router,
@ -2724,6 +2726,7 @@ async def _reserve_budget_after_common_checks(
general_settings: dict,
end_user_id: str | None = None,
end_user_object: LiteLLM_EndUserTable | None = None,
request: Request | None = None,
) -> None:
user_api_key_auth_obj.budget_reservation = None
if skip_budget_checks:
@ -2749,6 +2752,7 @@ async def _reserve_budget_after_common_checks(
end_user_object=end_user_object,
apply_user_budget_to_team_keys=general_settings.get("apply_user_budget_to_team_keys") is True,
fail_closed_budget_enforcement=general_settings.get("fail_closed_budget_enforcement") is True,
raw_body=await read_raw_json_body(request=request),
)

220
litellm/proxy/collector.py Normal file
View file

@ -0,0 +1,220 @@
"""Collector sidecar: consume spend events from the pod's inference workers and run the cost pipeline.
Runs the proxy startup lifespan (config, Prisma, Redis transaction buffer, scheduled spend flushes)
without serving HTTP, then listens on ``LITELLM_COLLECTOR_ADDRESS`` for newline-delimited spend
events. Each event goes through the unchanged ``_ProxyDBLogger._PROXY_track_cost_callback``, so
spend logs, spend counters, budget reservation reconciliation and cache updates happen exactly as
they would in-process, just in this container. Events are handled in order per producer connection
(one per uvicorn worker); a slow pipeline fills the socket buffer and the producer's bounded queue,
which is the backpressure that triggers its fallback or drop policy. ``SIGTERM`` stops accepting
connections, half-closes every producer connection so the producers switch to their unavailable
policy, finishes the events already sent, then runs the proxy shutdown (which flushes the buffered
spend transactions).
``DATABASE_URL`` is assembled from the same ``DATABASE_*`` inputs as the proxy container, and when
``LITELLM_PGBOUNCER_ENABLED`` is set it points at the PgBouncer that container already runs on the
pod's loopback, so the sidecar must see the same env as the proxy. Under ``IAM_TOKEN_DB_AUTH`` or
``AZURE_POSTGRESQL_AUTH`` that PgBouncer only accepts the token the proxy container minted, so the
sidecar goes to Postgres directly and mints its own. Works from any image that has ``litellm``
installed:
python -m litellm.proxy.collector [--address unix:///path.sock]
"""
import asyncio
import logging
import os
import signal
import sys
from collections.abc import Awaitable, Callable, Mapping, Sequence
from pathlib import Path
from typing import Final
from litellm._logging import verbose_logger, verbose_proxy_logger, verbose_router_logger
from litellm.proxy.db.db_url_settings import DatabaseURLSettings
from litellm.proxy.db.pgbouncer import (
PgBouncerError,
PgBouncerSettings,
export_pooled_database_url,
pooled_database_url,
)
from litellm.proxy.spend_tracking.spend_event_producer import (
COLLECTOR_JOB_ROLE,
AddressError,
CollectorAddress,
CollectorSettings,
TcpAddress,
UnixAddress,
parse_collector_address,
)
MAX_EVENT_BYTES: Final = 64 * 1024 * 1024
class SpendEventConsumer:
"""Accepts producer connections and runs ``handler`` on every line each one sends, in order."""
def __init__(self, handler: Callable[[bytes], Awaitable[None]]) -> None:
self._handler = handler
self._open_connections: set[asyncio.StreamWriter] = set() # mutable-ok: live producer connections
self._idle = asyncio.Event()
self._idle.set()
self._received = 0
self._handled = 0
self._failed = 0
@property
def received(self) -> int:
return self._received
@property
def handled(self) -> int:
return self._handled
@property
def failed(self) -> int:
return self._failed
async def serve(self, address: CollectorAddress) -> asyncio.Server:
match address:
case UnixAddress(path=path):
socket_path: Final = Path(path)
socket_path.parent.mkdir(parents=True, exist_ok=True)
socket_path.unlink(missing_ok=True)
return await asyncio.start_unix_server(self._on_connection, path=path, limit=MAX_EVENT_BYTES)
case TcpAddress(host=host, port=port):
return await asyncio.start_server(self._on_connection, host=host, port=port, limit=MAX_EVENT_BYTES)
async def _on_connection(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
self._open_connections.add(writer)
self._idle.clear()
try:
while line := await reader.readline():
if not line.endswith(b"\n"):
verbose_proxy_logger.error("collector: discarding truncated spend event (%d bytes)", len(line))
break
self._received += 1
await self._handle(line)
except (ConnectionError, asyncio.IncompleteReadError, asyncio.LimitOverrunError) as error:
verbose_proxy_logger.warning("collector: producer connection ended abnormally: %s", error)
finally:
writer.close()
self._open_connections.discard(writer)
if not self._open_connections:
self._idle.set()
async def _handle(self, line: bytes) -> None:
try:
await self._handler(line)
self._handled += 1
except Exception: # noqa: BLE001 # the cost pipeline raises anything; one bad event must not stop the sidecar
self._failed += 1
verbose_proxy_logger.exception("collector: spend event failed")
async def drain(self, timeout: float) -> int:
"""Half-close every producer connection, then keep reading until each producer hangs up or ``timeout``.
Returns how many producer connections were still open when the timeout hit.
"""
for writer in tuple(self._open_connections):
if writer.is_closing() or not writer.can_write_eof():
continue
try:
writer.write_eof()
except (OSError, RuntimeError) as error:
verbose_proxy_logger.debug("collector: producer already gone before half-close: %s", error)
try:
await asyncio.wait_for(self._idle.wait(), timeout)
except TimeoutError:
pass
return len(self._open_connections)
def _install_stop_signals(loop: asyncio.AbstractEventLoop, stop: asyncio.Event) -> None:
for signum in (signal.SIGTERM, signal.SIGINT):
loop.add_signal_handler(signum, stop.set)
async def run_collector(address: CollectorAddress, drain_timeout: float) -> None:
from fastapi import FastAPI
from litellm.proxy.hooks.proxy_track_cost_callback import run_spend_event
from litellm.proxy.proxy_server import proxy_startup_event
stop: Final = asyncio.Event()
_install_stop_signals(asyncio.get_running_loop(), stop)
consumer: Final = SpendEventConsumer(handler=run_spend_event)
async with proxy_startup_event(FastAPI()):
server: Final = await consumer.serve(address)
verbose_proxy_logger.info("collector: listening on %s", address)
await stop.wait()
server.close()
still_open: Final = await consumer.drain(drain_timeout)
verbose_proxy_logger.info(
"collector: stopping. received=%d handled=%d failed=%d connections_cut=%d",
consumer.received,
consumer.handled,
consumer.failed,
still_open,
)
def address_argument(argv: Sequence[str], default: str) -> str | AddressError:
match tuple(argv):
case ():
return default
case ("--address", value):
return value
case _:
return AddressError(f"usage: python -m litellm.proxy.collector [--address ADDRESS], got {tuple(argv)}")
def apply_log_level(litellm_log: str | None) -> None:
"""Mirror the proxy's ``LITELLM_LOG`` handling: the sidecar has no CLI flags to turn logging on."""
level: Final = logging.getLevelNamesMapping().get((litellm_log or "").upper())
if level is None:
return
for logger in (verbose_logger, verbose_router_logger, verbose_proxy_logger):
logger.setLevel(level)
def pod_pgbouncer_database_url(
pgbouncer: PgBouncerSettings, environ: Mapping[str, str], *, token_auth: bool
) -> str | PgBouncerError | None:
"""The proxy container's PgBouncer URL for ``environ["DATABASE_URL"]``, or None to connect to Postgres directly.
Direct is the answer when PgBouncer is off, and also under token auth: that PgBouncer's auth file
only holds the token its own container minted, which this container cannot present.
"""
if not pgbouncer.enabled or token_auth:
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 pooled_database_url(upstream_url, pgbouncer)
def main(argv: Sequence[str]) -> None:
os.environ.setdefault("LITELLM_JOB_ROLE", COLLECTOR_JOB_ROLE)
apply_log_level(os.environ.get("LITELLM_LOG"))
database: Final = DatabaseURLSettings.from_env()
database.apply_to_env()
pooled: Final = pod_pgbouncer_database_url(
PgBouncerSettings(),
os.environ,
token_auth=database.iam_token_db_auth or database.azure_postgresql_auth,
)
if isinstance(pooled, PgBouncerError):
sys.exit(f"LiteLLM collector: cannot use the pod's pgbouncer: {pooled.reason}")
if pooled is not None:
export_pooled_database_url(pooled)
settings: Final = CollectorSettings()
raw_address: Final = address_argument(argv, default=settings.address)
address: Final = raw_address if isinstance(raw_address, AddressError) else parse_collector_address(raw_address)
if isinstance(address, AddressError):
sys.exit(f"LiteLLM collector: {address.reason}")
asyncio.run(run_collector(address, drain_timeout=settings.drain_timeout_seconds))
if __name__ == "__main__":
main(sys.argv[1:])

12
litellm/proxy/common_utils/http_parsing_utils.py
View file

@ -213,6 +213,18 @@ async def _read_request_body(request: Request | None) -> dict:
return {}
async def read_raw_json_body(request: Request | None) -> bytes | None:
if request is None or _safe_get_request_parsed_body(request=request) is None:
return None
content_type: Final = _safe_get_request_headers(request=request).get("content-type", "")
if _is_form_content_type(content_type):
return None
try:
return await request.body()
except RuntimeError:
return None
def _safe_get_request_parsed_body(request: Request | None) -> dict | None:
if request is None:
return None

46
litellm/proxy/common_utils/reset_budget_job.py
View file

@ -224,33 +224,31 @@ def _queue_budget_linked_resets(
def _queue_enduser_resets(writes: LinkedSpendResetWrites, cascade: "_BudgetCascade") -> None:
"""End users are matched by id rather than budget link: rows with no
budget_id ride the default budget tier (litellm.max_end_user_budget_id).
Zero-before-decrement ordering matters here too (see
_queue_budget_linked_resets)."""
if not cascade.rollover_caps:
if cascade.endusers:
writes.queue_spend_zero(
where={"user_id": {"in": [row.user_id for row in cascade.endusers]}}
) # mutable-ok: prisma where filter must be a dict
"""End users reset on the budget link like every other gated table, plus a
NULL-budget_id branch: rows created implicitly persist no link and ride the
default tier (litellm.max_end_user_budget_id).
Matching on the link rather than enumerating user ids keeps a statement's
bind count proportional to the expiring tiers instead of the customer
population, which past ~32,700 dependents exceeds PostgreSQL's per-statement
bind ceiling and wedges the cascade permanently (#40564).
"""
_queue_budget_linked_resets(writes, cascade, extra=_SPENT_ROWS_WHERE)
default_budget_id: Final = litellm.max_end_user_budget_id
if default_budget_id is None or default_budget_id not in cascade.budget_ids:
return
tiered: Final = tuple((row.budget_id or litellm.max_end_user_budget_id, row.user_id) for row in cascade.endusers)
for budget_id, cap in cascade.rollover_caps.items():
if not (
user_ids := [uid for bid, uid in tiered if bid == budget_id]
): # mutable-ok: prisma "in" filter takes a list
continue
cap: Final = cascade.rollover_caps.get(default_budget_id)
if cap is None:
writes.queue_spend_zero(
where={"user_id": {"in": user_ids}, "spend": {"lte": cap}}
where={"budget_id": None, **_SPENT_ROWS_WHERE}
) # mutable-ok: prisma where filter must be a dict
writes.queue_spend_decrement(
where={"user_id": {"in": user_ids}, "spend": {"gt": cap}}, amount=cap
) # mutable-ok: prisma where filter must be a dict
plain: Final = [
uid for bid, uid in tiered if bid is None or bid not in cascade.rollover_caps
] # mutable-ok: prisma "in" filter takes a list
if plain:
writes.queue_spend_zero(where={"user_id": {"in": plain}}) # mutable-ok: prisma where filter must be a dict
return
writes.queue_spend_zero(
where={"budget_id": None, "spend": {"gt": 0, "lte": cap}}
) # mutable-ok: prisma where filter must be a dict
writes.queue_spend_decrement(
where={"budget_id": None, "spend": {"gt": cap}}, amount=cap
) # mutable-ok: prisma where filter must be a dict
@dataclass(frozen=True, slots=True)

5
litellm/proxy/db/db_url_settings.py
View file

@ -52,6 +52,7 @@ from typing import Annotated, Final, Protocol, TypeAlias, cast
from pydantic import AliasChoices, BeforeValidator, Field
from pydantic_settings import BaseSettings, SettingsConfigDict
from litellm.proxy.db.pgbouncer import database_url_is_pooled
from litellm.proxy.db.token_auth import (
AZURE_POSTGRESQL_AUTH_ENV_VAR,
DEFAULT_POSTGRES_PORT,
@ -358,8 +359,12 @@ class DatabaseURLSettings(BaseSettings):
Raises ``RuntimeError`` (naming the offending vars) when token auth is
enabled but a required field is missing — the proxy cannot recover
from this and a clear startup error beats a Prisma connect failure.
A ``DATABASE_URL`` the supervisor pointed at the in-container PgBouncer
is kept even under token auth: the pooler renews the token upstream.
"""
auth: Final = self.token_auth()
if auth is not None and database_url_is_pooled():
return None
if auth is not None:
missing: Final = tuple(
env

12
litellm/proxy/db/exception_handler.py
View file

@ -221,6 +221,18 @@ class PrismaDBExceptionHandler:
or "write conflict or a deadlock" in error_message
)
@staticmethod
def is_read_only_transaction_error(e: Exception) -> bool:
"""True iff ``e`` is Postgres SQLSTATE 25006 surfaced through prisma: the
pooled session answers reads but rejects writes, so the connection is
poisoned until the client is recreated."""
import prisma
if not isinstance(e, _exception_types(prisma.errors.PrismaError)):
return False
error_message: Final = str(e).lower()
return '"25006"' in error_message or "read-only transaction" in error_message
@staticmethod
def is_prisma_engine_internal_error(e: Exception) -> bool:
"""True iff ``e`` is a non-``PrismaError`` exception raised from inside

96
litellm/proxy/db/health_check_latest.py Normal file
View file

@ -0,0 +1,96 @@
"""
Latest health-check row per model, deduplicated by Postgres.
prisma-client-py's ``find_many(distinct=...)`` dedups client-side: the emitted
SQL carries no DISTINCT, so the whole append-only history table streams to the
worker on every call. ``SELECT DISTINCT ON`` keeps the transfer at one row per
(model_id, model_name) and is served by the matching descending index.
"""
from __future__ import annotations
import json
from collections.abc import Sequence
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Final
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, field_validator
from litellm._logging import verbose_proxy_logger
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
LATEST_HEALTH_CHECKS_SQL: Final = """
SELECT DISTINCT ON ("model_id", "model_name")
"health_check_id", "model_name", "model_id", "status",
"healthy_count", "unhealthy_count", "error_message",
"response_time_ms", "details", "checked_by",
"checked_at", "created_at", "updated_at"
FROM "LiteLLM_HealthCheckTable"
ORDER BY "model_id" ASC, "model_name" ASC, "checked_at" DESC
"""
LATEST_HEALTH_CHECKS_FOR_MODELS_SQL: Final = """
SELECT DISTINCT ON ("model_id", "model_name")
"health_check_id", "model_name", "model_id", "status",
"healthy_count", "unhealthy_count", "error_message",
"response_time_ms", "details", "checked_by",
"checked_at", "created_at", "updated_at"
FROM "LiteLLM_HealthCheckTable"
WHERE "model_name" = ANY($1)
ORDER BY "model_id" ASC, "model_name" ASC, "checked_at" DESC
"""
class LatestHealthCheckRow(BaseModel):
model_config = ConfigDict(frozen=True, protected_namespaces=())
health_check_id: str
model_name: str
model_id: str | None = None
status: str
healthy_count: int = 0
unhealthy_count: int = 0
error_message: str | None = None
response_time_ms: float | None = None
details: JsonValue | None = None
checked_by: str | None = None
checked_at: datetime
created_at: datetime
updated_at: datetime
@field_validator("details", mode="before")
@classmethod
def _decode_json_text(cls, value: object) -> object:
return json.loads(value) if isinstance(value, str) else value
@field_validator("checked_at", "created_at", "updated_at")
@classmethod
def _assume_utc(cls, value: datetime) -> datetime:
return value.replace(tzinfo=timezone.utc) if value.tzinfo is None else value
_ROWS_ADAPTER: Final = TypeAdapter(tuple[LatestHealthCheckRow, ...])
async def fetch_latest_health_checks(prisma_client: PrismaClient) -> tuple[LatestHealthCheckRow, ...]:
try:
rows: Final = await prisma_client.db.query_raw(LATEST_HEALTH_CHECKS_SQL)
return _ROWS_ADAPTER.validate_python(rows)
except Exception as query_err: # noqa: BLE001 # health decorates other reads; a driver error must not fail them
verbose_proxy_logger.error("Error getting all latest health checks: %s", query_err)
return ()
async def fetch_latest_health_checks_for_models(
prisma_client: PrismaClient, model_names: Sequence[str]
) -> tuple[LatestHealthCheckRow, ...]:
if not model_names:
return ()
try:
rows: Final = await prisma_client.db.query_raw(LATEST_HEALTH_CHECKS_FOR_MODELS_SQL, list(model_names))
return _ROWS_ADAPTER.validate_python(rows)
except Exception as query_err: # noqa: BLE001 # a paged model list must not fail on its health decoration
verbose_proxy_logger.error("Error getting latest health checks for models: %s", query_err)
return ()

743
litellm/proxy/db/pgbouncer.py Normal file
View file

@ -0,0 +1,743 @@
"""In-container PgBouncer shared by every proxy worker.
Each uvicorn worker owns a Prisma query engine with its own pool of
``connection_limit`` server connections, so the connections a pod holds open
against Postgres scale as ``workers * connection_limit`` and a database with a
fixed connection ceiling runs out of room as pods and workers are added.
When ``LITELLM_PGBOUNCER_ENABLED`` is set, the supervisor process starts one
PgBouncer next to the workers (no extra network hop: it listens on loopback
inside the pod) in transaction pooling mode, points ``DATABASE_URL`` at it
with ``pgbouncer=true`` so Prisma stops using server-side prepared statements,
and keeps it running for the life of the proxy. Every worker's pool then
becomes cheap client connections to PgBouncer while the upstream connection
count is capped at ``LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS`` per pod, no matter
how many workers run.
Migrations and the schema diff run in the supervisor before the pooler is
started, so they always go straight to Postgres. ``DATABASE_URL_READ_REPLICA``
is left untouched.
The workers never hold the upstream credential: they log in to PgBouncer as
``litellm_pgbouncer`` with a random password made at startup, and PgBouncer
takes the database user's password from its auth file. Under
``IAM_TOKEN_DB_AUTH`` or ``AZURE_POSTGRESQL_AUTH`` that password is a
short-lived token, so the supervisor mints a new one before it expires,
rewrites the auth file and asks PgBouncer to reload; only new upstream
connections authenticate, so live ones are unaffected. The pooled
``DATABASE_URL`` then carries a static password, and the workers must not run
their own token refresh against it: ``LITELLM_PGBOUNCER_POOLED_DATABASE_URL``
tells them so, while a read replica keeps refreshing its own token.
"""
from __future__ import annotations
import atexit
import functools
import os
import re
import secrets
import shlex
import shutil
import signal
import socket
import subprocess
import tempfile
import threading
import time
import urllib.parse
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from types import MappingProxyType
from typing import Final
from pydantic import Field
from pydantic_settings import BaseSettings, SettingsConfigDict
from litellm._logging import verbose_proxy_logger
from litellm.proxy.db.token_auth import (
DatabaseTokenAuth,
IAMEndpoint,
mint_database_token,
parse_database_token_expiration,
parse_iam_endpoint_from_url,
)
PGBOUNCER_ENV_PREFIX: Final = "LITELLM_PGBOUNCER_"
PGBOUNCER_POOLED_ENV_VAR: Final = "LITELLM_PGBOUNCER_POOLED_DATABASE_URL"
PGBOUNCER_LISTEN_ADDR: Final = "127.0.0.1"
PGBOUNCER_POOL_USER: Final = "litellm_pgbouncer"
PGBOUNCER_INI_NAME: Final = "pgbouncer.ini"
PGBOUNCER_USERLIST_NAME: Final = "userlist.txt"
PGBOUNCER_CA_NAME: Final = "server-ca.pem"
PGBOUNCER_RESTART_DELAY_SECONDS: Final = 1.0
PGBOUNCER_READY_TIMEOUT_SECONDS: Final = 15.0
PGBOUNCER_STOP_GRACE_SECONDS: Final = 10.0
PGBOUNCER_UNPRIVILEGED_USER: Final = "nobody"
PGBOUNCER_MIN_VERSION: Final = (1, 19)
PGBOUNCER_MAX_PASSWORD_BYTES: Final = 2048
PGBOUNCER_VERSION_PATTERN: Final = re.compile(r"PgBouncer (\d+)\.(\d+)")
PGBOUNCER_TOKEN_REFRESH_BUFFER_SECONDS: Final = 180.0
PGBOUNCER_TOKEN_FALLBACK_REFRESH_SECONDS: Final = 600.0
PGBOUNCER_TOKEN_RETRY_SECONDS: Final = 30.0
# Prisma's client-side TLS params describe the hop to Postgres, which becomes
# PgBouncer's server side. They move into ``server_tls_*`` and must not stay on
# the loopback URL: the listener speaks plain TCP and Prisma would refuse it
# under ``sslmode=require`` or ``channel_binding=require``.
PRISMA_TLS_PARAM_KEYS: Final[frozenset[str]] = frozenset(
{"sslmode", "sslcert", "sslaccept", "sslidentity", "sslpassword", "channel_binding", "gssencmode"}
)
POOLED_URL_DROPPED_KEYS: Final[frozenset[str]] = PRISMA_TLS_PARAM_KEYS | frozenset(("options", "pgbouncer"))
PGBOUNCER_SSLMODES: Final[frozenset[str]] = frozenset(
{"disable", "allow", "prefer", "require", "verify-ca", "verify-full"}
)
class PgBouncerSettings(BaseSettings):
"""``LITELLM_PGBOUNCER_*`` env vars, read once in the supervisor."""
model_config = SettingsConfigDict(
env_prefix=PGBOUNCER_ENV_PREFIX, case_sensitive=False, extra="ignore", frozen=True
)
enabled: bool = False
port: int = Field(default=6432, ge=1, le=65535)
max_db_connections: int = Field(default=20, ge=1)
max_client_conn: int = Field(default=1000, ge=1)
binary: str = "pgbouncer"
@dataclass(frozen=True, slots=True)
class PgBouncerPlan:
ini: str
pooled_url: str
upstream_user: str
upstream_password: str | None
pool_password: str
ca_source: str | None = None
def userlist(self, upstream_password: str) -> str:
return "".join(
f"{_userlist_quote(user)} {_userlist_quote(password)}\n"
for user, password in ((self.upstream_user, upstream_password), (PGBOUNCER_POOL_USER, self.pool_password))
)
@dataclass(frozen=True, slots=True)
class PgBouncerError:
reason: str
def _single_quoted(value: str) -> str:
"""Quote for SQL and for PgBouncer's ``[databases]`` connection string: both double a literal ``'``."""
return "'" + value.replace("'", "''") + "'"
def _userlist_quote(value: str) -> str:
return '"' + value.replace('"', '""') + '"'
def _option_settings(tokens: Sequence[str]) -> tuple[str, ...] | None:
"""The ``name=value`` settings in a libpq ``options`` string, or None if it holds anything else.
Accepts ``-c name=value``, ``-cname=value`` and ``--name=value``; a
detached ``-c`` is folded into the token that follows it first.
"""
folded: Final = tuple(
f"-c{tokens[index + 1]}" if token == "-c" and index + 1 < len(tokens) else token
for index, token in enumerate(tokens)
if index == 0 or tokens[index - 1] != "-c"
)
settings: Final = tuple(token[2:] for token in folded if token.startswith(("-c", "--")) and "=" in token[2:])
return settings if len(settings) == len(folded) else None
def _connect_query(options: str) -> str | PgBouncerError:
"""Turn Prisma's ``options=-c name=value ...`` startup param into ``SET`` statements.
PgBouncer rejects any ``-c`` setting in ``options`` that is not one of the
handful it tracks (``statement_timeout`` and ``lock_timeout`` are not), so
the settings are applied to each new server connection instead. Every
client shares them, which is what the single ``DATABASE_URL`` gave anyway.
"""
settings: Final = _option_settings(tuple(shlex.split(options)))
if settings is None:
return PgBouncerError(f"cannot translate the DATABASE_URL options {options!r} into PgBouncer settings")
return "; ".join(
f"SET {name.strip()} TO {_single_quoted(value.strip())}"
for name, value in (setting.split("=", 1) for setting in settings)
)
def _server_tls_settings(sslmode: str, sslcert: str, sslaccept: str, ca_path: Path) -> tuple[str, ...] | PgBouncerError:
"""``server_tls_*`` lines naming ``ca_path``, the runtime-dir copy of the bundle: the original (or the
0600 root pinned by ``pin_bundle_root``) is often unreadable for the user PgBouncer drops to."""
if sslmode not in PGBOUNCER_SSLMODES:
return PgBouncerError(f"unsupported sslmode {sslmode!r} on DATABASE_URL")
verify: Final = sslmode in ("verify-ca", "verify-full") or (sslmode == "require" and sslaccept == "strict")
if verify and not sslcert:
return PgBouncerError(
"DATABASE_URL asks for a verified TLS connection but names no CA bundle; "
"add sslcert=<ca.pem> (or sslrootcert=) so the in-container PgBouncer can verify Postgres"
)
mode: Final = "verify-full" if verify else sslmode
return (f"server_tls_sslmode = {mode}", *((f"server_tls_ca_file = {ca_path}",) if sslcert else ()))
def plan_pgbouncer(
upstream_url: str,
settings: PgBouncerSettings,
runtime_dir: Path,
run_as_user: str | None,
) -> PgBouncerPlan | PgBouncerError:
"""Render the PgBouncer config for ``upstream_url`` and the loopback URL Prisma uses instead.
Params describing Prisma's own pool (``connection_limit``, ``pool_timeout``,
...) stay on the pooled URL; the TLS params and ``options`` describe the hop
to Postgres and move into the PgBouncer config. The upstream password is
left out of the config on purpose: PgBouncer then takes it from the auth
file, which can be rewritten while it runs. ``run_as_user`` is the
unprivileged user PgBouncer drops to when the proxy runs as root, which
PgBouncer itself refuses to do.
"""
parsed: Final = urllib.parse.urlsplit(upstream_url)
params: Final[Mapping[str, str]] = MappingProxyType(
dict(urllib.parse.parse_qsl(parsed.query, keep_blank_values=True))
)
dbname: Final = urllib.parse.unquote(parsed.path.lstrip("/"))
username: Final = urllib.parse.unquote(parsed.username or "")
password: Final = None if parsed.password is None else urllib.parse.unquote(parsed.password)
if not parsed.hostname or not username or not dbname:
return PgBouncerError("DATABASE_URL must carry a host, user and database name for the in-container PgBouncer")
if username == PGBOUNCER_POOL_USER:
return PgBouncerError(
f"the database user cannot be named {PGBOUNCER_POOL_USER!r}: that is the user the workers log in to the "
"in-container PgBouncer as, and PgBouncer keeps one password per user"
)
if "sslidentity" in params:
return PgBouncerError("client certificates (sslidentity) are not supported with the in-container PgBouncer")
tls: Final = _server_tls_settings(
params.get("sslmode", "prefer"),
params.get("sslcert", ""),
params.get("sslaccept", ""),
runtime_dir / PGBOUNCER_CA_NAME,
)
if isinstance(tls, PgBouncerError):
return tls
connect_query: Final = _connect_query(params["options"]) if params.get("options") else ""
if isinstance(connect_query, PgBouncerError):
return connect_query
upstream: Final = " ".join(
(
f"host={_single_quoted(parsed.hostname)}",
f"port={parsed.port or 5432}",
f"dbname={_single_quoted(dbname)}",
f"user={_single_quoted(username)}",
*((f"connect_query={_single_quoted(connect_query)}",) if connect_query else ()),
)
)
ini: Final = "\n".join(
(
"[databases]",
f"{dbname} = {upstream}",
"",
"[pgbouncer]",
f"listen_addr = {PGBOUNCER_LISTEN_ADDR}",
f"listen_port = {settings.port}",
f"unix_socket_dir = {runtime_dir}",
f"auth_file = {runtime_dir / PGBOUNCER_USERLIST_NAME}",
"auth_type = scram-sha-256",
f"stats_users = {PGBOUNCER_POOL_USER}",
"pool_mode = transaction",
f"max_client_conn = {settings.max_client_conn}",
f"default_pool_size = {settings.max_db_connections}",
f"max_db_connections = {settings.max_db_connections}",
"ignore_startup_parameters = extra_float_digits",
*tls,
*((f"user = {run_as_user}",) if run_as_user else ()),
"",
)
)
pooled_query: Final = urllib.parse.urlencode(
(*((key, value) for key, value in params.items() if key not in POOLED_URL_DROPPED_KEYS), ("pgbouncer", "true"))
)
pool_password: Final = secrets.token_urlsafe(32)
pooled_url: Final = urllib.parse.urlunsplit(
parsed._replace(
netloc=f"{PGBOUNCER_POOL_USER}:{pool_password}@{PGBOUNCER_LISTEN_ADDR}:{settings.port}", query=pooled_query
)
)
return PgBouncerPlan(
ini=ini,
pooled_url=pooled_url,
upstream_user=username,
upstream_password=password,
pool_password=pool_password,
ca_source=params.get("sslcert") or None,
)
def pooled_database_url(upstream_url: str, settings: PgBouncerSettings) -> str | PgBouncerError:
"""The loopback URL of a PgBouncer another container in the pod already runs for ``upstream_url``.
Only the container that started PgBouncer knows the pool user's password, so
this logs in as the upstream user, whom the auth file lists as well.
"""
plan: Final = plan_pgbouncer(upstream_url, settings, runtime_dir=Path("/nonexistent"), run_as_user=None)
if isinstance(plan, PgBouncerError):
return plan
password: Final = urllib.parse.urlsplit(upstream_url).password or ""
credentials: Final = f"{urllib.parse.quote(plan.upstream_user, safe='')}:{password}"
return urllib.parse.urlunsplit(
urllib.parse.urlsplit(plan.pooled_url)._replace(netloc=f"{credentials}@{PGBOUNCER_LISTEN_ADDR}:{settings.port}")
)
def _write_private(path: Path, content: str, run_as_user: str | None) -> None:
with open(os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600), "w", encoding="utf-8") as handle:
handle.write(content)
if run_as_user is not None:
shutil.chown(path, user=run_as_user)
def write_pgbouncer_ini(plan: PgBouncerPlan, runtime_dir: Path, run_as_user: str | None) -> Path | PgBouncerError:
"""Write the ini (mode 0600) and the CA copy, and return the ini path.
``run_as_user`` is the user PgBouncer drops to when started as root; it has
to own the files it re-reads on reload and the socket directory.
"""
ini_path: Final = runtime_dir / PGBOUNCER_INI_NAME
ca_path: Final = runtime_dir / PGBOUNCER_CA_NAME
if plan.ca_source is not None:
try:
shutil.copyfile(plan.ca_source, ca_path)
except OSError as error:
return PgBouncerError(f"cannot read the CA bundle {plan.ca_source!r} named by sslcert: {error}")
_write_private(ini_path, plan.ini, run_as_user)
if run_as_user is not None:
runtime_dir.chmod(0o700)
for path in (runtime_dir, *((ca_path,) if plan.ca_source is not None else ())):
shutil.chown(path, user=run_as_user)
return ini_path
def write_userlist(userlist: str, runtime_dir: Path, run_as_user: str | None) -> Path:
"""Replace the auth file in one step, so a PgBouncer starting or reloading meanwhile reads the old or the new one whole."""
userlist_path: Final = runtime_dir / PGBOUNCER_USERLIST_NAME
staged_path: Final = runtime_dir / f".{PGBOUNCER_USERLIST_NAME}.next"
_write_private(staged_path, userlist, run_as_user)
os.replace(staged_path, userlist_path)
return userlist_path
def export_pooled_database_url(pooled_url: str) -> None:
os.environ["DATABASE_URL"] = pooled_url
os.environ[PGBOUNCER_POOLED_ENV_VAR] = "true"
def database_url_is_pooled(environ: Mapping[str, str] = os.environ) -> bool:
return environ.get(PGBOUNCER_POOLED_ENV_VAR) == "true"
@dataclass(frozen=True, slots=True)
class PgBouncerTokenSource:
auth: DatabaseTokenAuth
endpoint: IAMEndpoint
def mint(self) -> str:
"""The token as Postgres expects it: ``mint_database_token`` returns it percent-encoded for a URL."""
return urllib.parse.unquote(mint_database_token(self.auth, self.endpoint))
def expires_at(self, token: str) -> datetime | None:
return parse_database_token_expiration(self.auth, token)
def _utcnow() -> datetime:
return datetime.now(timezone.utc).replace(tzinfo=None)
class PgBouncerTokenRefresher:
"""Keeps the token in PgBouncer's auth file current from a daemon thread.
``install`` gets each fresh token and is expected to rewrite the auth file
and reload PgBouncer. The next refresh is due ``buffer_seconds`` before the
token expires, or ``fallback_seconds`` later when the expiry cannot be read.
A refresh that fails leaves the previous auth file in place and is retried
after ``retry_seconds``: the old token stays good until it expires, so a
transient credential-provider error costs nothing unless it persists.
"""
def __init__(
self,
source: PgBouncerTokenSource,
install: Callable[[str], None],
*,
buffer_seconds: float = PGBOUNCER_TOKEN_REFRESH_BUFFER_SECONDS,
fallback_seconds: float = PGBOUNCER_TOKEN_FALLBACK_REFRESH_SECONDS,
retry_seconds: float = PGBOUNCER_TOKEN_RETRY_SECONDS,
now: Callable[[], datetime] = _utcnow,
) -> None:
self._source: Final = source
self._install: Final = install
self._buffer_seconds: Final = buffer_seconds
self._fallback_seconds: Final = fallback_seconds
self._retry_seconds: Final = retry_seconds
self._now: Final = now
self._stopping: Final = threading.Event()
self._delay: float = 0.0
self._thread: threading.Thread | None = None
def refresh(self) -> float | PgBouncerError:
label: Final = self._source.auth.label
try:
token: Final = self._source.mint()
except Exception as mint_error:
return PgBouncerError(f"could not mint a {label} for the in-container pgbouncer: {mint_error!r}")
if len(token.encode()) >= PGBOUNCER_MAX_PASSWORD_BYTES:
return PgBouncerError(
f"the {label} is {len(token.encode())} bytes long, but PgBouncer's auth file holds passwords of at "
f"most {PGBOUNCER_MAX_PASSWORD_BYTES - 1} bytes"
)
try:
self._install(token)
except OSError as install_error:
return PgBouncerError(f"could not install the {label} into the pgbouncer auth file: {install_error}")
expires_at: Final = self._source.expires_at(token)
if expires_at is None:
return self._fallback_seconds
return max(self._retry_seconds, (expires_at - self._now()).total_seconds() - self._buffer_seconds)
def start(self) -> PgBouncerError | None:
primed: Final = self.refresh()
if isinstance(primed, PgBouncerError):
return primed
self._delay = primed
self._thread = threading.Thread(target=self._run, daemon=True, name="litellm-pgbouncer-token-refresh")
self._thread.start()
return None
def _run(self) -> None:
while not self._stopping.wait(self._delay):
self._delay = self._refresh_and_report()
def _refresh_and_report(self) -> float:
outcome: Final = self.refresh()
if isinstance(outcome, PgBouncerError):
verbose_proxy_logger.error(
"In-container pgbouncer keeps its current %s (%s); retrying in %.0fs.",
self._source.auth.label,
outcome.reason,
self._retry_seconds,
)
return self._retry_seconds
verbose_proxy_logger.info(
"In-container pgbouncer picked up a fresh %s; the next one is due in %.0fs.",
self._source.auth.label,
outcome,
)
return outcome
def stop(self) -> None:
self._stopping.set()
if self._thread is not None:
self._thread.join()
def _port_open(port: int) -> bool:
try:
with socket.create_connection((PGBOUNCER_LISTEN_ADDR, port), timeout=0.5):
return True
except OSError:
return False
def _unix_socket_open(path: Path) -> bool:
with socket.socket(socket.AF_UNIX) as probe:
probe.settimeout(0.5)
try:
probe.connect(str(path))
except OSError:
return False
return True
def unix_socket_path(runtime_dir: Path, port: int) -> Path:
return runtime_dir / f".s.PGSQL.{port}"
def pgbouncer_version(binary: str) -> tuple[int, int] | PgBouncerError:
"""``(major, minor)`` from ``<binary> --version``.
Readiness relies on PgBouncer exiting when it cannot bind its TCP port,
which it does from 1.19 on. Older releases log a warning and serve the unix
socket alone, so their socket would vouch for a port held by someone else.
"""
try:
output: Final = subprocess.run(
(binary, "--version"), capture_output=True, text=True, check=False, timeout=10
).stdout
except (OSError, subprocess.TimeoutExpired) as run_error:
return PgBouncerError(f"could not run {binary!r} --version: {run_error}")
found: Final = PGBOUNCER_VERSION_PATTERN.search(output)
if found is None:
return PgBouncerError(f"{binary!r} --version did not report a PgBouncer version: {output.strip()!r}")
return int(found[1]), int(found[2])
def _end(process: subprocess.Popen[bytes]) -> None:
if process.poll() is not None:
return
process.terminate()
try:
process.wait(timeout=PGBOUNCER_STOP_GRACE_SECONDS)
except subprocess.TimeoutExpired:
process.kill()
process.wait()
class PgBouncerProcess:
"""Runs ``argv`` as a foreground child and restarts it whenever it exits on its own.
Prisma reconnects by itself after a failed query, so a PgBouncer crash
costs the requests in flight plus one failed query per idle pooled
connection the crash severed, and nothing else once the replacement is
listening again. A replacement that cannot be spawned, finds its port
taken, exits again or never starts listening is retried every
``restart_delay_seconds`` until ``stop`` is called.
A connect probe of ``port`` cannot tell the child from another process
that grabbed the port after the availability check, so readiness also
needs ``socket_path``: the unix socket PgBouncer creates in the private
runtime directory, which it only does once every TCP listener is bound
(PgBouncer 1.19 or newer, see ``pgbouncer_version``).
"""
def __init__(
self,
argv: Sequence[str],
port: int,
socket_path: Path,
restart_delay_seconds: float = PGBOUNCER_RESTART_DELAY_SECONDS,
ready_timeout_seconds: float = PGBOUNCER_READY_TIMEOUT_SECONDS,
) -> None:
self.argv: Final = tuple(argv)
self.port: Final = port
self.socket_path: Final = socket_path
self.restart_delay_seconds: Final = restart_delay_seconds
self.ready_timeout_seconds: Final = ready_timeout_seconds
self._stopping: Final = threading.Event()
self._lock: Final = threading.Lock()
self._process: subprocess.Popen[bytes] | None = None
@property
def pid(self) -> int | None:
with self._lock:
return None if self._process is None else self._process.pid
def _spawn(self) -> subprocess.Popen[bytes] | PgBouncerError | None:
"""Start a child, or None once ``stop`` ran; both take the lock so no child can slip in after a stop.
The port has to be free first: a listener that is already there would
pass the readiness check while the child fails to bind.
"""
with self._lock:
if self._stopping.is_set():
return None
if _port_open(self.port):
return PgBouncerError(f"{PGBOUNCER_LISTEN_ADDR}:{self.port} is already in use by another process")
try:
process: Final = subprocess.Popen(self.argv)
except OSError as spawn_error:
return PgBouncerError(f"could not start {self.argv[0]!r}: {spawn_error}")
self._process = process
return process
def _wait_ready(self, process: subprocess.Popen[bytes]) -> PgBouncerError | None:
deadline: Final = time.monotonic() + self.ready_timeout_seconds
while time.monotonic() < deadline:
if process.poll() is not None:
return PgBouncerError(f"pgbouncer exited with status {process.returncode} during startup")
if _port_open(self.port) and _unix_socket_open(self.socket_path):
return None
time.sleep(0.1)
if _port_open(self.port):
return PgBouncerError(
f"{PGBOUNCER_LISTEN_ADDR}:{self.port} is served by another process, not the pgbouncer that was started"
)
return PgBouncerError(
f"pgbouncer did not start listening on {PGBOUNCER_LISTEN_ADDR}:{self.port} "
f"within {self.ready_timeout_seconds:.0f}s"
)
def start(self) -> PgBouncerError | None:
"""Spawn PgBouncer, wait until it listens on port and unix socket, then supervise it from a daemon thread."""
process: Final = self._spawn()
if process is None:
return PgBouncerError("pgbouncer was stopped before it started")
if isinstance(process, PgBouncerError):
return process
not_ready: Final = self._wait_ready(process)
if not_ready is not None:
self.stop()
return not_ready
self._watch(process)
return None
def _watch(self, process: subprocess.Popen[bytes]) -> None:
threading.Thread(
target=self._supervise, args=(process,), daemon=True, name="litellm-pgbouncer-supervisor"
).start()
def _supervise(self, process: subprocess.Popen[bytes]) -> None:
status: Final = process.wait()
if self._stopping.is_set():
return
verbose_proxy_logger.error(
"In-container pgbouncer (pid %s) exited with status %s; restarting in %.1fs.",
process.pid,
status,
self.restart_delay_seconds,
)
self._restart_after_delay()
def _restart_after_delay(self) -> None:
time.sleep(self.restart_delay_seconds)
process: Final = self._spawn()
if process is None:
return
if isinstance(process, PgBouncerError):
self._retry_restart(process.reason)
return
not_ready: Final = self._wait_ready(process)
if not_ready is None:
self._watch(process)
return
_end(process)
self._retry_restart(not_ready.reason)
def _retry_restart(self, reason: str) -> None:
if self._stopping.is_set():
return
verbose_proxy_logger.error(
"In-container pgbouncer could not be restarted (%s); retrying in %.1fs.", reason, self.restart_delay_seconds
)
threading.Thread(target=self._restart_after_delay, daemon=True, name="litellm-pgbouncer-supervisor").start()
def reload(self) -> None:
with self._lock:
if self._process is not None:
self._process.send_signal(signal.SIGHUP)
def stop(self) -> None:
with self._lock:
self._stopping.set()
process: Final = self._process
if process is not None:
_end(process)
def install_pgbouncer_token(
plan: PgBouncerPlan, runtime_dir: Path, run_as_user: str | None, pooler: PgBouncerProcess, token: str
) -> None:
write_userlist(plan.userlist(token), runtime_dir, run_as_user)
pooler.reload()
def _install_upstream_password(
plan: PgBouncerPlan,
runtime_dir: Path,
run_as_user: str | None,
pooler: PgBouncerProcess,
token_auth: DatabaseTokenAuth | None,
upstream_url: str,
) -> PgBouncerTokenRefresher | None | PgBouncerError:
if token_auth is None:
if plan.upstream_password is None:
return PgBouncerError(
"DATABASE_URL carries no password and neither IAM_TOKEN_DB_AUTH nor AZURE_POSTGRESQL_AUTH is on, "
"so the in-container PgBouncer has nothing to authenticate to Postgres with"
)
write_userlist(plan.userlist(plan.upstream_password), runtime_dir, run_as_user)
return None
refresher: Final = PgBouncerTokenRefresher(
PgBouncerTokenSource(auth=token_auth, endpoint=parse_iam_endpoint_from_url(upstream_url)),
functools.partial(install_pgbouncer_token, plan, runtime_dir, run_as_user, pooler),
)
failed: Final = refresher.start()
if failed is not None:
return failed
return refresher
def _only_in_this_process(action: Callable[[], None]) -> Callable[[], None]:
"""An exit hook that does nothing in a forked child, which inherits the parent's ``atexit`` table."""
owner_pid: Final = os.getpid()
def run() -> None:
if os.getpid() == owner_pid:
action()
return run
def start_in_container_pgbouncer(
settings: PgBouncerSettings,
upstream_url: str,
token_auth: DatabaseTokenAuth | None = None,
register_exit_hook: Callable[[Callable[[], None]], object] = atexit.register,
) -> str | PgBouncerError:
"""Start the pooler for ``upstream_url`` and return the loopback URL the workers must use.
The pooler lives as long as this process: it is stopped from the exit hooks
once the worker manager has returned, and only by the process that started
it (gunicorn forks its workers, so they carry the hooks too). PgBouncer
refuses to run as root, so a root proxy (the default image) has it drop to
``nobody``. With ``token_auth`` the password on ``upstream_url`` is ignored:
the pooler mints its own tokens and renews them for as long as it runs.
"""
version: Final = pgbouncer_version(settings.binary)
if isinstance(version, PgBouncerError):
return version
if version < PGBOUNCER_MIN_VERSION:
return PgBouncerError(
f"PgBouncer {version[0]}.{version[1]} keeps running after failing to bind its TCP port, so the proxy "
f"cannot tell it apart from another listener; {PGBOUNCER_MIN_VERSION[0]}.{PGBOUNCER_MIN_VERSION[1]} "
"or newer is required"
)
runtime_dir: Final = Path(tempfile.mkdtemp(prefix="litellm-pgbouncer-"))
register_exit_hook(_only_in_this_process(lambda: shutil.rmtree(runtime_dir, ignore_errors=True)))
run_as_user: Final = PGBOUNCER_UNPRIVILEGED_USER if os.geteuid() == 0 else None
plan: Final = plan_pgbouncer(upstream_url, settings, runtime_dir, run_as_user)
if isinstance(plan, PgBouncerError):
return plan
ini_path: Final = write_pgbouncer_ini(plan, runtime_dir, run_as_user)
if isinstance(ini_path, PgBouncerError):
return ini_path
pooler: Final = PgBouncerProcess(
argv=(settings.binary, str(ini_path)),
port=settings.port,
socket_path=unix_socket_path(runtime_dir, settings.port),
)
refresher: Final = _install_upstream_password(plan, runtime_dir, run_as_user, pooler, token_auth, upstream_url)
if isinstance(refresher, PgBouncerError):
return refresher
failed: Final = pooler.start()
if failed is not None:
if refresher is not None:
refresher.stop()
return failed
register_exit_hook(_only_in_this_process(pooler.stop))
if refresher is not None:
register_exit_hook(_only_in_this_process(refresher.stop))
verbose_proxy_logger.info(
"In-container pgbouncer (pid %s) listening on %s:%s; capping this pod at %s upstream database connections%s.",
pooler.pid,
PGBOUNCER_LISTEN_ADDR,
settings.port,
settings.max_db_connections,
"" if token_auth is None else f" and renewing its {token_auth.label} before each one expires",
)
return plan.pooled_url

105
litellm/proxy/health_endpoints/_health_endpoints.py
View file

@ -7,7 +7,7 @@ import secrets
import time
import traceback
from collections.abc import Iterable, Mapping
from datetime import datetime, timedelta
from datetime import datetime, timedelta, timezone
from typing import Any, Final, Literal, TypedDict, cast
import fastapi
@ -41,6 +41,7 @@ from litellm.proxy.auth.auth_utils import (
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
from litellm.proxy.db.health_check_latest import LatestHealthCheckRow
from litellm.proxy.db.proxy_worker_heartbeat import count_live_proxy_workers
from litellm.proxy.health_check import (
ADMIN_ONLY_HEALTH_DISPLAY_PARAMS,
@ -747,13 +748,42 @@ def _aggregate_health_check_results(
return model_results
class _AggregatedHealthResult(TypedDict):
"""One entry of ``_aggregate_health_check_results``: a model's counts for this cycle."""
model_name: ReadOnly[str]
model_id: ReadOnly[str | None]
healthy_count: ReadOnly[int]
unhealthy_count: ReadOnly[int]
error_message: ReadOnly[str | None]
def _new_health_status(result: _AggregatedHealthResult) -> str:
return "healthy" if result["healthy_count"] > 0 else "unhealthy"
def _should_persist_health_check_result(
result: _AggregatedHealthResult, latest_checks_map: Mapping[str, LatestHealthCheckRow]
) -> bool:
"""
True when this result has to be written: no previous row, the status changed, or the
previous row is older than one hour (periodic refresh while the status is stable).
"""
lookup_key: Final = result["model_id"] if result["model_id"] else result["model_name"]
last_check: Final = latest_checks_map.get(lookup_key)
if last_check is None or last_check.status != _new_health_status(result):
return True
time_since_last_check: Final = (datetime.now(timezone.utc) - last_check.checked_at).total_seconds()
return time_since_last_check >= 3600 # 1 hour threshold
async def _save_health_check_results_if_changed(
prisma_client,
model_results: dict,
latest_checks_map: dict,
start_time: float,
checked_by: str | None = None,
):
) -> bool:
"""
Save health check results to database, but only if status changed or >1 hour since last save.
@ -764,47 +794,39 @@ async def _save_health_check_results_if_changed(
- Status changes: Immediate write (no delay)
- Result: ~92% reduction in DB writes for stable systems, while maintaining real-time updates on changes
The writes are awaited rather than detached so the caller learns whether this cycle's
persistence completed.
Args:
prisma_client: Database client
model_results: Dictionary of aggregated health check results per model
latest_checks_map: Dictionary mapping model_id/model_name to latest health check
start_time: Start time of health check for calculating response time
checked_by: Identifier for who/what performed the check
Returns:
True when every row that needed writing was written (including when nothing needed
writing); False when any write failed.
"""
for result in model_results.values():
new_status = "healthy" if result["healthy_count"] > 0 else "unhealthy"
# Check if we should save this result
should_save = True
lookup_key = result["model_id"] if result["model_id"] else result["model_name"]
if lookup_key in latest_checks_map:
last_check = latest_checks_map[lookup_key]
# Only save if status changed or if it's been a while since last check
if last_check.status == new_status:
# Check if last check was recent (within 1 hour)
if last_check.checked_at:
from datetime import datetime, timezone
time_since_last_check = (datetime.now(timezone.utc) - last_check.checked_at).total_seconds()
# Only skip if status unchanged AND checked recently (within 1 hour)
# This ensures we still get periodic updates even if status is stable
if time_since_last_check < 3600: # 1 hour threshold
should_save = False
if should_save:
asyncio.create_task(
prisma_client.save_health_check_result(
model_name=result["model_name"],
model_id=result["model_id"],
status=new_status,
healthy_count=result["healthy_count"],
unhealthy_count=result["unhealthy_count"],
error_message=result["error_message"],
response_time_ms=(time.time() - start_time) * 1000,
details=None,
checked_by=checked_by,
)
)
to_write: Final = tuple(
result for result in model_results.values() if _should_persist_health_check_result(result, latest_checks_map)
)
writes: Final = tuple(
prisma_client.save_health_check_result(
model_name=result["model_name"],
model_id=result["model_id"],
status=_new_health_status(result),
healthy_count=result["healthy_count"],
unhealthy_count=result["unhealthy_count"],
error_message=result["error_message"],
response_time_ms=(time.time() - start_time) * 1000,
details=None,
checked_by=checked_by,
)
for result in to_write
)
rows: Final = await asyncio.gather(*writes)
return all(row is not None for row in rows)
async def _save_background_health_checks_to_db(
@ -814,7 +836,7 @@ async def _save_background_health_checks_to_db(
unhealthy_endpoints: list,
start_time: float,
checked_by: str | None = None,
):
) -> bool:
"""
Save background health check results to database for each model.
@ -823,9 +845,13 @@ async def _save_background_health_checks_to_db(
OPTIMIZATION: Only saves to database if the status has changed from the last saved check.
This dramatically reduces database writes when health status remains stable.
Returns:
True when this cycle's persistence completed; False when it was skipped or any step
failed. Never raises: a database failure must not break the health check loop.
"""
if prisma_client is None:
return
return False
try:
# Step 1: Build mapping from model parameter to model info
@ -848,7 +874,7 @@ async def _save_background_health_checks_to_db(
latest_checks_map[key] = check
# Step 4: Save aggregated results, but only if status changed
await _save_health_check_results_if_changed(
return await _save_health_check_results_if_changed(
prisma_client,
model_results,
latest_checks_map,
@ -858,6 +884,7 @@ async def _save_background_health_checks_to_db(
except Exception as db_error:
verbose_proxy_logger.warning("Failed to save background health checks to database: %s", db_error)
# Continue execution - don't let database save failure break health checks
return False
_PROXY_ADMIN_ROLES: Final = frozenset(

23
litellm/proxy/hooks/batch_enqueued_tokens.py
View file

@ -9,6 +9,7 @@ the reservation is refunded when the batch reaches a terminal state
"""
import asyncio
import logging
import math
import time
import uuid
@ -19,6 +20,7 @@ from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeAlias
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError
from litellm._logging import verbose_proxy_logger
from litellm.caching.redis_cache import log_redis_failure
from litellm.constants import BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY, BATCH_ENQUEUED_TOKEN_TTL_SECONDS
from litellm.proxy._types import UserAPIKeyAuth
@ -233,8 +235,11 @@ class BatchEnqueuedTokenStore:
try:
return await self._reserve_via_redis(reserve_script, refund_script, tokens=tokens, scopes=scopes)
except Exception as e: # noqa: BLE001 # any Redis failure must fall back to the in-memory counters
verbose_proxy_logger.warning(
"Redis enqueued-token reserve failed, falling back to in-memory: %s", str(e)
log_redis_failure(
verbose_proxy_logger,
logging.WARNING,
"Redis enqueued-token reserve failed, falling back to in-memory",
e,
)
return await self._reserve_in_memory(tokens=tokens, scopes=scopes, span=litellm_parent_otel_span)
@ -374,8 +379,11 @@ class BatchEnqueuedTokenStore:
(serialized, ttl),
)
except Exception as e: # noqa: BLE001 # any Redis failure must fall back to the in-memory record
verbose_proxy_logger.warning(
"Redis enqueued-token reservation save failed, falling back to in-memory: %s", str(e)
log_redis_failure(
verbose_proxy_logger,
logging.WARNING,
"Redis enqueued-token reservation save failed, falling back to in-memory",
e,
)
else:
return
@ -421,8 +429,11 @@ class BatchEnqueuedTokenStore:
await pop_script((self._record_key(batch_id),), (BATCH_ENQUEUED_TOKEN_TTL_SECONDS,))
)
except Exception as e: # noqa: BLE001 # any Redis failure must fall back to the in-memory record
verbose_proxy_logger.warning(
"Redis enqueued-token reservation pop failed, falling back to in-memory: %s", str(e)
log_redis_failure(
verbose_proxy_logger,
logging.WARNING,
"Redis enqueued-token reservation pop failed, falling back to in-memory",
e,
)
return None

18
litellm/proxy/hooks/max_budget_per_session_limiter.py
View file

@ -14,11 +14,13 @@ Works across multiple proxy instances via DualCache (in-memory + Redis).
Follows the same pattern as max_iterations_limiter.py.
"""
import logging
import os
from typing import TYPE_CHECKING, Any, Final
from litellm import DualCache
from litellm._logging import verbose_proxy_logger
from litellm.caching.redis_cache import log_redis_failure
from litellm.exceptions import RateLimitType
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import UserAPIKeyAuth
@ -215,9 +217,11 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger):
return float(result)
return 0.0
except Exception as e:
verbose_proxy_logger.warning(
"MaxBudgetPerSessionHandler: Redis GET failed, falling back to in-memory: %s",
str(e),
log_redis_failure(
verbose_proxy_logger,
logging.WARNING,
"MaxBudgetPerSessionHandler: Redis GET failed, falling back to in-memory",
e,
)
result = await self.internal_usage_cache.async_get_cache(
@ -239,9 +243,11 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger):
)
return float(result)
except Exception as e:
verbose_proxy_logger.warning(
"MaxBudgetPerSessionHandler: Redis INCRBYFLOAT failed, falling back to in-memory: %s",
str(e),
log_redis_failure(
verbose_proxy_logger,
logging.WARNING,
"MaxBudgetPerSessionHandler: Redis INCRBYFLOAT failed, falling back to in-memory",
e,
)
return await self._in_memory_increment_spend(cache_key, amount)

58
litellm/proxy/hooks/parallel_request_limiter_v3.py
View file

@ -6,6 +6,7 @@ This is currently in development and not yet ready for production.
import asyncio
import binascii
import logging
import os
import uuid
from collections.abc import Awaitable, Callable, Mapping, Sequence, Set
@ -26,7 +27,7 @@ from typing_extensions import NotRequired, ReadOnly
from litellm import DualCache
from litellm._logging import verbose_proxy_logger
from litellm.caching.redis_cache import RedisCircuitBreakerOpenError
from litellm.caching.redis_cache import log_redis_failure
from litellm.constants import DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE, INTERNAL_CALL_ORIGIN_METADATA_KEY
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.prompt_templates.common_utils import (
@ -1229,7 +1230,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
)
all_cache_values.extend(group_cache_values)
except Exception as e:
verbose_proxy_logger.warning("Redis Lua script failed for hash tag %s: %s", hash_tag, e)
log_redis_failure(
verbose_proxy_logger, logging.WARNING, f"Redis Lua script failed for hash tag {hash_tag}", e
)
# Fallback to in-memory cache for this group
group_cache_values = await self.in_memory_cache_sliding_window(
keys=group_keys,
@ -1476,7 +1479,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
)
counts = [max(0, int(value)) for value in raw_counts]
except Exception as e: # noqa: BLE001 - any Redis/Lua failure degrades to the local mirror, never a 500
verbose_proxy_logger.warning("parallel_count_script failed, using local mirror: %s", e)
log_redis_failure(
verbose_proxy_logger, logging.WARNING, "parallel_count_script failed, using local mirror", e
)
counts = await self._read_local_gauge_counts(gauge_keys, parent_otel_span)
else:
counts = await self._read_local_gauge_counts(gauge_keys, parent_otel_span)
@ -1506,7 +1511,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
],
)
except Exception as e: # noqa: BLE001 - any Redis/Lua failure degrades to in-memory enforcement, never a 500
verbose_proxy_logger.warning("parallel_acquire_script failed, falling back to in-memory gauge: %s", e)
log_redis_failure(
verbose_proxy_logger,
logging.WARNING,
"parallel_acquire_script failed, falling back to in-memory gauge",
e,
)
async with self._check_and_increment_lock:
return await self._acquire_parallel_slots_in_memory(gauges, slot_id, parent_otel_span)
if int(raw[0]) == 1:
@ -1632,7 +1642,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
)
return
except Exception as e: # noqa: BLE001 - any Redis/Lua failure degrades to the in-memory release, never a 500
verbose_proxy_logger.warning("parallel_release_script failed, falling back to in-memory release: %s", e)
log_redis_failure(
verbose_proxy_logger,
logging.WARNING,
"parallel_release_script failed, falling back to in-memory release",
e,
)
async with self._check_and_increment_lock:
for counter_key in counter_keys:
@ -1815,12 +1830,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
# state ambiguous. Refund any prior groups so Redis returns
# to its pre-call state, then fall back to in-memory for the
# whole call (counters there are independent of Redis).
verbose_proxy_logger.error(
"atomic_check_and_increment_by_n: Redis Lua execution failed (%s: %s). Refunding %s prior descriptors and falling back to in-memory enforcement — counters will diverge from Redis until window expires (window_size=%ss).",
type(e).__name__,
log_redis_failure(
verbose_proxy_logger,
logging.ERROR,
f"atomic_check_and_increment_by_n: Redis Lua execution failed ({type(e).__name__}). Refunding "
f"{len(applied)} prior descriptors and falling back to in-memory enforcement, counters will "
f"diverge from Redis until window expires (window_size={self.window_size}s)",
e,
len(applied),
self.window_size,
)
await self._refund_applied_descriptor_groups(applied)
flat_meta: list[AtomicCounterMeta] = [m for _k, _a, group_meta in descriptor_groups for m in group_meta]
@ -1867,8 +1883,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
value=-entry["increment"],
)
except Exception as e:
verbose_proxy_logger.warning(
"Failed to refund %s on cross-descriptor rollback: %s", entry["counter_key"], e
log_redis_failure(
verbose_proxy_logger,
logging.WARNING,
f"Failed to refund {entry['counter_key']} on cross-descriptor rollback",
e,
)
def _build_atomic_response(
@ -3857,12 +3876,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
)
except Exception as e:
log: Final = (
verbose_proxy_logger.debug
if isinstance(e, RedisCircuitBreakerOpenError)
else verbose_proxy_logger.warning
log_redis_failure(
verbose_proxy_logger, logging.WARNING, "TTL preservation failed, falling back to regular pipeline", e
)
log("TTL preservation failed, falling back to regular pipeline: %s", e)
# Fallback to regular pipeline on error
await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline(
increment_list=pipeline_operations,
litellm_parent_otel_span=parent_otel_span,
@ -3927,9 +3944,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
)
continue
except Exception as e: # noqa: BLE001 # Redis failures use the plain increment fallback
verbose_proxy_logger.warning(
"Window-guarded token adjustment failed for %s: %s",
operation["key"],
log_redis_failure(
verbose_proxy_logger,
logging.WARNING,
f"Window-guarded token adjustment failed for {operation['key']}",
e,
)
if operation["increment_value"] > 0:

54
litellm/proxy/hooks/proxy_track_cost_callback.py
View file

@ -27,6 +27,16 @@ from litellm.proxy.db.db_spend_update_writer import (
get_llm_router,
)
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
from litellm.proxy.spend_tracking.spend_event import (
ObjectMapping,
SpendEventBuildError,
SpendEventDecodeError,
build_spend_event,
decode_spend_event,
is_offloadable_success,
spend_event_callback_args,
)
from litellm.proxy.spend_tracking.spend_event_producer import SpendEventProducer
from litellm.proxy.spend_tracking.spend_log_error_logger import (
should_suppress_spend_log_tracebacks,
spend_log_error,
@ -34,6 +44,7 @@ from litellm.proxy.spend_tracking.spend_log_error_logger import (
from litellm.proxy.spend_tracking.spend_tracking_utils import (
_sanitize_error_information_for_spend_logs,
get_request_model_access_groups,
should_store_prompts_and_responses_in_spend_logs,
)
from litellm.proxy.utils import ProxyUpdateSpend
from litellm.types.utils import (
@ -71,8 +82,43 @@ _CAPTURED_IDENTITY_CALL_TYPES: Final[frozenset[str]] = frozenset(
class _ProxyDBLogger(CustomLogger):
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
await self._PROXY_track_cost_callback(kwargs, response_obj, start_time, end_time)
def __init__(
self,
spend_event_producer: SpendEventProducer | None = None,
*,
turn_off_message_logging: bool = False,
message_logging: bool = True,
) -> None:
super().__init__(turn_off_message_logging=turn_off_message_logging, message_logging=message_logging)
self.spend_event_producer = spend_event_producer
async def async_log_success_event(
self, kwargs: ObjectMapping, response_obj: object, start_time: datetime, end_time: datetime
) -> None:
if self.spend_event_producer is None or not is_offloadable_success(response_obj):
await self._PROXY_track_cost_callback(kwargs, response_obj, start_time, end_time)
return
event: Final = build_spend_event(
kwargs,
response_obj,
start_time,
end_time,
store_bodies=should_store_prompts_and_responses_in_spend_logs(),
)
if isinstance(event, SpendEventBuildError):
verbose_proxy_logger.warning("collector: tracking cost in-process, event not buildable: %s", event.reason)
await self._PROXY_track_cost_callback(kwargs, response_obj, start_time, end_time)
return
await self.spend_event_producer.publish(event)
async def run_spend_event(self, line: bytes) -> None:
"""Run the unchanged cost pipeline on a serialized spend event (sidecar consumer and in-process fallback)."""
event: Final = decode_spend_event(line)
if isinstance(event, SpendEventDecodeError):
verbose_proxy_logger.error("collector: discarding undecodable spend event: %s", event.reason)
return
args: Final = spend_event_callback_args(event)
await self._PROXY_track_cost_callback(args.kwargs, args.response_obj, args.start_time, args.end_time)
async def async_post_call_failure_hook(
self,
@ -503,6 +549,10 @@ def _write_spend_metadata_to_kwargs(kwargs: dict, metadata: dict) -> None:
bucket[key] = value
async def run_spend_event(line: bytes) -> None:
await _ProxyDBLogger().run_spend_event(line)
def _is_unbilled_interaction_response(completion_response: object) -> bool:
from litellm.interactions.background_cost_polling import missing_usage_is_expected
from litellm.types.interactions import InteractionsAPIResponse

18
litellm/proxy/hooks/sensitive_data_routing.py
View file

@ -10,11 +10,13 @@ this hook manages:
Works across multiple proxy instances via DualCache (in-memory + Redis).
"""
import logging
import os
from typing import TYPE_CHECKING, Any, Final
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
from litellm.caching.redis_cache import log_redis_failure
from litellm.integrations.custom_guardrail import get_session_id_from_request_data
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import UserAPIKeyAuth
@ -96,9 +98,11 @@ class _PROXY_SensitiveDataRoutingHandler(CustomLogger):
)
return routed_model
except Exception as e:
verbose_proxy_logger.warning(
"SensitiveDataRoutingHandler: Redis GET failed, falling back to in-memory: %s",
str(e),
log_redis_failure(
verbose_proxy_logger,
logging.WARNING,
"SensitiveDataRoutingHandler: Redis GET failed, falling back to in-memory",
e,
)
result = await self.internal_usage_cache.async_get_cache(
@ -142,9 +146,11 @@ class _PROXY_SensitiveDataRoutingHandler(CustomLogger):
ttl=self.ttl,
)
except Exception as e:
verbose_proxy_logger.warning(
"SensitiveDataRoutingHandler: Redis SET failed, falling back to in-memory: %s",
str(e),
log_redis_failure(
verbose_proxy_logger,
logging.WARNING,
"SensitiveDataRoutingHandler: Redis SET failed, falling back to in-memory",
e,
)
await self.internal_usage_cache.async_set_cache(

2
litellm/proxy/management_endpoints/coordination_redis_endpoints.py
View file

@ -140,7 +140,7 @@ def _merge_over_saved(
def _validated_params(settings: Mapping[str, object]) -> CoordinationRedisParams:
"""Validate settings the way startup does: resolve env refs, then require a connection target."""
try:
params: Final = CoordinationRedisParams(**_resolve_env_refs(settings))
params: Final = CoordinationRedisParams.model_validate(_resolve_env_refs(settings))
except ValidationError as e:
invalid_fields: Final = sorted({str(error["loc"][0]) for error in e.errors() if error["loc"]})
raise HTTPException(

22
litellm/proxy/management_endpoints/mcp_management_endpoints.py
View file

@ -2278,6 +2278,28 @@ if MCP_AVAILABLE:
"""Persist the OAuth2 access token obtained by the calling user."""
prisma_client: Final = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy")
await _authorize_and_fetch_mcp_server(prisma_client, user_api_key_dict, server_id)
from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # keep manager import lazy
global_mcp_server_manager as _manager,
)
# This endpoint accepts an opaque token with no upstream identity validation, so it must be
# closed for identity-bound servers or it becomes a bypass of the token-relay binding check.
registry_server: Final = _manager.get_mcp_server_by_id(server_id)
binding: Final = registry_server.oauth_identity_binding if registry_server else None
if binding is not None and binding.mode == "enforce":
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail={ # mutable-ok: FastAPI exception detail requires a JSON-serializable dictionary
"error": "oauth_identity_binding_enforced",
"error_description": (
"Direct credential storage is disabled for this server: its OAuth identity "
"binding is enforced and this endpoint cannot validate the token's principal. "
"Complete the OAuth flow through the gateway instead."
),
"server_id": server_id,
"credential_stored": False,
},
)
user_id: Final = user_api_key_dict.user_id or ""
if not user_id:
raise HTTPException(

Some files were not shown because too many files have changed in this diff Show more