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

This commit is contained in:
mateo-berri 2026-09-10 15:47:55 -07:00
commit 3a0dce6bc4
77 changed files with 6649 additions and 407 deletions

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

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

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

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

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

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

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

71
gateway/launch.py Normal file
View file

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

View file

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

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"

View file

@ -355,6 +355,20 @@ db:
# only (e.g. when IAM_TOKEN_DB_AUTH supplies the token at runtime).
readReplicaUrl: ""
# In-container connection pool (PgBouncer, transaction mode) shared by every
# worker in the pod. Without it each --num_workers worker opens its own
# connection_limit connections to Postgres, so a pod's footprint against the
# database's connection ceiling is workers x connection_limit and grows with
# every replica. With it, the pod holds at most maxDbConnections upstream
# connections no matter how many workers run; the workers connect to the pool
# over loopback, with no extra network hop. Migrations still go straight to
# Postgres. Starting profile for numWorkers: 4 is maxDbConnections: 20, so
# a database with a 5000-connection ceiling fits roughly 200 replicas.
connectionPool:
enabled: false
maxDbConnections: 20
maxClientConn: 1000
# Use the Stackgres Helm chart to deploy an instance of a Stackgres cluster.
# The Stackgres Operator must already be installed within the target
# Kubernetes cluster.

View file

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

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

View file

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

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. The pool holds a static
# password, so it cannot be combined with `database.writer.useIAMAuth` or
# `useAzureEntraAuth` (rendering fails). Starting profile for
# `gateway.numWorkers: 4` is maxDbConnections: 20, so a database with a
# 5000-connection ceiling fits roughly 200 gateway replicas.
connectionPool:
enabled: false
maxDbConnections: 20
maxClientConn: 1000
# Optional Redis. Leave host empty to disable.
#
# This is the proxy's coordination store: cross-pod tpm/rpm rate limits, spend

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)

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

View file

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

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(

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(

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

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

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

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

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

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:

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)

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

View file

@ -703,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",
@ -932,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",
@ -2431,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))

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

View file

@ -0,0 +1,525 @@
"""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 pooler holds the database password from startup, so it cannot be combined
with ``IAM_TOKEN_DB_AUTH`` or ``AZURE_POSTGRESQL_AUTH``: those rotate the
password inside every worker on their own schedule, and PgBouncer would keep
authenticating upstream with the expired token.
"""
from __future__ import annotations
import atexit
import os
import re
import shlex
import shutil
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 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 AZURE_POSTGRESQL_AUTH_ENV_VAR, IAM_TOKEN_DB_AUTH_ENV_VAR
PGBOUNCER_ENV_PREFIX: Final = "LITELLM_PGBOUNCER_"
PGBOUNCER_LISTEN_ADDR: Final = "127.0.0.1"
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_VERSION_PATTERN: Final = re.compile(r"PgBouncer (\d+)\.(\d+)")
PGBOUNCER_LIST_DELIMITER_PATTERN: Final = re.compile(r"[,\s]")
PGBOUNCER_TOKEN_AUTH_CONFLICT: Final = (
f"the in-container pgbouncer cannot be combined with {IAM_TOKEN_DB_AUTH_ENV_VAR} or "
f"{AZURE_POSTGRESQL_AUTH_ENV_VAR}: each worker rotates the database password on its own schedule and the pooler "
"would keep using the expired token upstream. Disable the pooler or use a static database password"
)
# 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
userlist: str
pooled_url: str
ca_source: str | None = None
@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. ``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 password is None or not dbname:
return PgBouncerError(
"DATABASE_URL must carry a host, user, password and database name for the in-container PgBouncer"
)
if PGBOUNCER_LIST_DELIMITER_PATTERN.search(username):
return PgBouncerError(
f"the database user {username!r} cannot be named in PgBouncer's stats_users list: "
"PgBouncer splits list settings on commas and whitespace and has no quoting for them"
)
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"password={_single_quoted(password)}",
*((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 = {username}",
"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 ()),
"",
)
)
userlist: Final = f"{_userlist_quote(username)} {_userlist_quote(password)}\n"
pooled_query: Final = urllib.parse.urlencode(
(*((key, value) for key, value in params.items() if key not in POOLED_URL_DROPPED_KEYS), ("pgbouncer", "true"))
)
credentials: Final = f"{urllib.parse.quote(username, safe='')}:{urllib.parse.quote(password, safe='')}"
pooled_url: Final = urllib.parse.urlunsplit(
parsed._replace(netloc=f"{credentials}@{PGBOUNCER_LISTEN_ADDR}:{settings.port}", query=pooled_query)
)
return PgBouncerPlan(ini=ini, userlist=userlist, pooled_url=pooled_url, ca_source=params.get("sslcert") or None)
def write_pgbouncer_files(plan: PgBouncerPlan, runtime_dir: Path, run_as_user: str | None) -> Path | PgBouncerError:
"""Write the ini, userlist (both hold the password, so mode 0600) and 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
userlist_path: Final = runtime_dir / PGBOUNCER_USERLIST_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}")
for path, content in ((userlist_path, plan.userlist), (ini_path, plan.ini)):
path.touch(mode=0o600)
path.write_text(content, encoding="utf-8")
if run_as_user is not None:
runtime_dir.chmod(0o700)
for path in (runtime_dir, ini_path, userlist_path, *((ca_path,) if plan.ca_source is not None else ())):
shutil.chown(path, user=run_as_user)
return ini_path
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 stop(self) -> None:
with self._lock:
self._stopping.set()
process: Final = self._process
if process is not None:
_end(process)
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_enabled: bool = False,
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``.
"""
if token_auth_enabled:
return PgBouncerError(PGBOUNCER_TOKEN_AUTH_CONFLICT)
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_files(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),
)
failed: Final = pooler.start()
if failed is not None:
return failed
register_exit_hook(_only_in_this_process(pooler.stop))
verbose_proxy_logger.info(
"In-container pgbouncer (pid %s) listening on %s:%s; capping this pod at %s upstream database connections.",
pooler.pid,
PGBOUNCER_LISTEN_ADDR,
settings.port,
settings.max_db_connections,
)
return plan.pooled_url

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(

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(

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(

View file

@ -18,6 +18,7 @@ from pydantic import BaseModel, ConfigDict
import litellm
from litellm.constants import DEFAULT_NUM_WORKERS_LITELLM_PROXY
from litellm.proxy.db.pgbouncer import PgBouncerError, PgBouncerSettings, start_in_container_pgbouncer
from litellm.proxy.db.query_engine_reaper import start_query_engine_reaper
if TYPE_CHECKING:
@ -1377,6 +1378,21 @@ def run_server(
print(
f"Unable to connect to DB. DATABASE_URL found in environment, but prisma package not found." # noqa: F541
)
pgbouncer_settings: Final = PgBouncerSettings()
upstream_database_url: Final = os.getenv("DATABASE_URL")
if pgbouncer_settings.enabled and upstream_database_url is not None:
pooled_database_url: Final = start_in_container_pgbouncer(
pgbouncer_settings, upstream_database_url, token_auth_enabled=wants_rds_iam or wants_azure_entra
)
if isinstance(pooled_database_url, PgBouncerError):
print(
f"\033[1;31mLiteLLM Proxy: LITELLM_PGBOUNCER_ENABLED is set but the in-container pgbouncer "
f"could not start: {pooled_database_url.reason}\033[0m",
file=sys.stderr,
flush=True,
)
sys.exit(1)
os.environ["DATABASE_URL"] = pooled_database_url
if port == 4000 and ProxyInitializationHelpers._is_port_in_use(port):
port = random.randint(1024, 49152)
if prometheus_metrics_port == port:

View file

@ -16,7 +16,16 @@ import threading
import time
import traceback
import warnings
from collections.abc import AsyncGenerator, AsyncIterator, Callable, Collection, Mapping, MutableMapping, Sequence
from collections.abc import (
AsyncGenerator,
AsyncIterator,
Awaitable,
Callable,
Collection,
Mapping,
MutableMapping,
Sequence,
)
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from types import MappingProxyType, UnionType
@ -51,6 +60,7 @@ from litellm.constants import (
AIOHTTP_NEEDS_CLEANUP_CLOSED,
AIOHTTP_TTL_DNS_CACHE,
AUDIO_SPEECH_CHUNK_SIZE,
BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME,
BASE_MCP_ROUTE,
DAILY_TAG_SPEND_BATCH_MULTIPLIER,
DEFAULT_MAX_RECURSE_DEPTH,
@ -151,6 +161,7 @@ if TYPE_CHECKING:
from prisma import models as prisma_models
from litellm.integrations.opentelemetry import OpenTelemetry
from litellm.proxy.health_check_utils.shared_health_check_manager import SharedHealthCheckManager
Span = _Span | Any
else:
@ -233,7 +244,7 @@ def generate_feedback_box():
import contextlib
from collections import defaultdict
from contextlib import asynccontextmanager
from functools import lru_cache
from functools import lru_cache, partial
import litellm
import litellm._redis
@ -415,6 +426,7 @@ from litellm.proxy.config_resolvers.alerting import (
)
from litellm.proxy.container_endpoints.endpoints import router as container_router
from litellm.proxy.credential_endpoints.endpoints import router as credential_router
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import (
SPEND_LOG_CLEANUP_BOUND_SETTINGS,
SpendLogCleanup,
@ -3747,13 +3759,50 @@ async def _run_direct_health_check_with_instrumentation(
raise AssertionError("perform_health_check rejected every optional argument")
async def _window_gated_health_check_db_save(
save: Callable[[], Awaitable[bool]],
pod_lock_manager: PodLockManager | None,
lock_ttl: int | None,
) -> None:
"""
Persist at most once per window fleet-wide. A completed save keeps the lock as the
"this window's save is done" marker, so it is deliberately never released and expires
with the interval. A save that reports failure or is cancelled releases the lock so
another pod's cycle in the same window can retry, instead of the fleet going a whole
window without a write.
"""
if pod_lock_manager is None or pod_lock_manager.redis_cache is None:
await save()
return
acquired: Final = await pod_lock_manager.acquire_lock(
cronjob_id=BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME,
ttl=lock_ttl,
allow_reentrant=False,
)
if not acquired:
verbose_proxy_logger.debug("background_health_check_db_save_skipped another pod persisted this window")
return
try:
persisted: Final = await save()
except BaseException:
await pod_lock_manager.release_lock(cronjob_id=BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME)
raise
if not persisted:
verbose_proxy_logger.warning(
"background_health_check_db_save_incomplete released the window lock so another pod can retry"
)
await pod_lock_manager.release_lock(cronjob_id=BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME)
def _schedule_background_health_check_db_save(
prisma_client,
shared_health_manager,
prisma_client: PrismaClient | None,
shared_health_manager: "SharedHealthCheckManager | None",
model_list: list,
healthy_endpoints: list,
unhealthy_endpoints: list,
):
pod_lock_manager: PodLockManager | None = None,
lock_ttl: int | None = None,
) -> None:
"""Fire-and-forget: persist health check results to DB if prisma is available."""
if prisma_client is None:
return
@ -3765,16 +3814,16 @@ def _schedule_background_health_check_db_save(
checked_by: Final = shared_health_manager.pod_id if shared_health_manager is not None else "background_health_check"
start_time: Final = time_module.time()
asyncio.create_task(
_save_background_health_checks_to_db(
prisma_client,
model_list,
healthy_endpoints,
unhealthy_endpoints,
start_time,
checked_by=checked_by,
)
save: Final = partial(
_save_background_health_checks_to_db,
prisma_client,
model_list,
healthy_endpoints,
unhealthy_endpoints,
start_time,
checked_by=checked_by,
)
asyncio.create_task(_window_gated_health_check_db_save(save, pod_lock_manager, lock_ttl))
def _get_endpoint_exception_status(endpoint: dict, exceptions: dict) -> int:
@ -4081,6 +4130,8 @@ async def _run_background_health_check():
_llm_model_list,
healthy_endpoints,
unhealthy_endpoints,
pod_lock_manager=proxy_logging_obj.db_spend_update_writer.pod_lock_manager,
lock_ttl=health_check_interval,
)
# Write health state to router cache for health-check-driven routing
@ -4906,7 +4957,9 @@ class ProxyConfig:
if not isinstance(raw_params, dict):
raise ValueError("general_settings.coordination_redis must be a mapping of Redis connection params")
coordination_params: Final = CoordinationRedisParams(**_resolve_coordination_redis_env_refs(raw_params))
coordination_params: Final = CoordinationRedisParams.model_validate(
_resolve_coordination_redis_env_refs(raw_params)
)
if not coordination_params.has_connection_target():
raise ValueError(
"general_settings.coordination_redis needs a connection target: "
@ -9234,7 +9287,9 @@ class ProxyStartupEvent:
if persisted is None:
return None
coordination_params: Final = CoordinationRedisParams(**_resolve_coordination_redis_env_refs(persisted))
coordination_params: Final = CoordinationRedisParams.model_validate(
_resolve_coordination_redis_env_refs(persisted)
)
if not coordination_params.has_connection_target():
verbose_proxy_logger.warning(
"coordination_redis saved in the database names no connection target; ignoring it."

View file

@ -4633,16 +4633,16 @@ async def _assert_user_can_view_request_id(
Verify the requesting non-admin user is allowed to view this spend-log row.
Allowed when the log belongs to the user directly, or to one of their
permitted teams (admin or ``/spend/logs`` permission).
Raises HTTP 403 if not.
Raises HTTP 403 if not, including when no spend-log row exists for the
request_id (e.g. it was pruned by retention), so a missing row can't be
used to read a payload out of cold storage via the detail endpoint.
"""
row: Final = await _find_spend_log_row(prisma_client, request_id)
if row is None:
if row is not None and row.user is not None and row.user == user_api_key_dict.user_id:
return
if row.user is not None and row.user == user_api_key_dict.user_id:
return
if row.team_id:
if row is not None and row.team_id:
can_view: Final = await _can_team_member_view_log(
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,

View file

@ -125,6 +125,11 @@ from litellm.proxy.db.exception_handler import (
PrismaDBExceptionHandler,
call_with_db_reconnect_retry,
)
from litellm.proxy.db.health_check_latest import (
LatestHealthCheckRow,
fetch_latest_health_checks,
fetch_latest_health_checks_for_models,
)
from litellm.proxy.db.log_db_metrics import log_db_metrics
from litellm.proxy.db.prisma_client import (
PrismaWrapper,
@ -6462,48 +6467,13 @@ class PrismaClient:
verbose_proxy_logger.error("Error getting health check history: %s", e)
return []
async def get_all_latest_health_checks(self) -> "Sequence[prisma_models.LiteLLM_HealthCheckTable]":
"""
Get the latest health check for each model.
async def get_all_latest_health_checks(self) -> tuple[LatestHealthCheckRow, ...]:
"""Latest health check per (model_id, model_name), deduplicated in Postgres."""
return await fetch_latest_health_checks(self)
Uses DB-level DISTINCT ON (model_id, model_name) with ORDER BY checked_at DESC
(via Prisma ``distinct`` + ``order``) so we never load the full history into memory.
"""
try:
return await HealthCheckRepository(self).table.find_many(
distinct=["model_id", "model_name"],
order=[
{"model_id": "asc"},
{"model_name": "asc"},
{"checked_at": "desc"},
],
)
except Exception as e:
verbose_proxy_logger.error("Error getting all latest health checks: %s", e)
return []
async def get_latest_health_checks_for_models(
self, model_names: "Sequence[str]"
) -> "Sequence[prisma_models.LiteLLM_HealthCheckTable]":
"""
Get the latest health check for each of the named models.
Same DISTINCT ON as ``get_all_latest_health_checks``, bounded to the models asked
about, so a paged caller reads health for its page instead of for the whole table.
"""
if not model_names:
return ()
latest_first: Final = (("model_id", "asc"), ("model_name", "asc"), ("checked_at", "desc"))
order: Final = [{field: direction} for field, direction in latest_first] # mutable-ok: prisma order is a list
try:
return await HealthCheckRepository(self).table.find_many(
where={"model_name": {"in": list(model_names)}}, # mutable-ok: prisma filters are dicts and lists
distinct=["model_id", "model_name"], # mutable-ok: prisma distinct takes a list
order=order,
)
except Exception as e: # noqa: BLE001 # health decorates a list; a driver error must not fail the page
verbose_proxy_logger.error("Error getting latest health checks for models: %s", e)
return ()
async def get_latest_health_checks_for_models(self, model_names: Sequence[str]) -> tuple[LatestHealthCheckRow, ...]:
"""Same as ``get_all_latest_health_checks``, bounded to the named models."""
return await fetch_latest_health_checks_for_models(self, model_names)
### HELPER FUNCTIONS ###
@ -8023,6 +7993,88 @@ async def get_available_models_for_user(
return all_models
def _safe_get_model_info(model: str, get_model_info: Callable[[str], ModelInfo]) -> ModelInfo | None:
try:
return get_model_info(model)
except Exception as e:
verbose_proxy_logger.debug(
"create_model_info_response: cost map lookup failed for %s: %s",
model,
e,
)
return None
def _resolve_listing_model_info(
deployment_model: str | None,
listed_model: str,
listed_info: ModelInfo | None,
get_model_info: Callable[[str], ModelInfo],
) -> tuple[ModelInfo, ...]:
"""
Cost-map entries describing one deployment behind a listed model, best source first.
The name a model is listed under is an arbitrary public alias, so it often misses the
cost map and lands on a fallback-generalization rule that answers with a conservative
family baseline instead of the real model's limits; the deployment's underlying model
is what the request actually reaches. Both names are kept because either can
generalize, and because a deployment's own model is registered into the cost map as a
stub that carries no limits of its own. Exact entries are consulted before generalized
ones, and each field is then taken from the first entry that has it.
``listed_info`` is resolved once by the caller, since a group with several distinct
underlying models resolves the same alias for each of them.
"""
# Fast path, and the only one a wildcard-expanded name takes: with a single name
# there is nothing to order, so skip the generalization test entirely. This keeps
# the per-model cost of the listing on the hot path #33721 exists to protect.
if deployment_model is None or deployment_model == listed_model:
return () if listed_info is None else (listed_info,)
deployment_info: Final = _safe_get_model_info(deployment_model, get_model_info)
if deployment_info is None:
return () if listed_info is None else (listed_info,)
if listed_info is None:
return (deployment_info,)
from litellm.utils import is_generalized_model_info
# Both names resolved: the deployment's model leads unless it only generalized
# while the listed name is an exact cost-map entry.
if is_generalized_model_info(deployment_info) and not is_generalized_model_info(listed_info):
return (listed_info, deployment_info)
return (deployment_info, listed_info)
def _first_token_limit(candidates: tuple[ModelInfo, ...], field: str) -> int | None:
return next(
(limit for limit in (coerce_token_limit(info.get(field)) for info in candidates) if limit is not None),
None,
)
def _group_token_limit(candidate_sets: tuple[tuple[ModelInfo, ...], ...], field: str) -> int | None:
"""The widest limit any deployment behind the listed name declares for ``field``.
A model group is normally one model behind several interchangeable deployments, so
there is a single value to report and the choice of aggregate does not arise.
When a group genuinely mixes models no single number is right, and the widest is the
deliberate pick over the narrowest for two reasons. It is what ``/model_group/info``
has long reported to the Admin UI, so the two surfaces agree; disagreeing is the very
complaint this resolution path exists to fix. And of the two ways to be wrong,
under-advertising is worse: a client that trusts a narrowed window silently refuses
prompts the group would have served, while an over-long prompt that reaches a smaller
deployment comes back as a legible context-length error -- and does not reach one at
all when ``enable_pre_call_checks`` is set, which filters deployments the prompt does
not fit.
"""
limits: Final = tuple(
limit for limit in (_first_token_limit(candidates, field) for candidates in candidate_sets) if limit is not None
)
return max(limits) if limits else None
def create_model_info_response(
model_id: str,
provider: str,
@ -8047,31 +8099,48 @@ def create_model_info_response(
"owned_by": provider,
}
try:
model_cost_info: ModelInfo | None = get_model_info(model_id)
except Exception as e:
verbose_proxy_logger.debug(
"create_model_info_response: cost map lookup failed for %s: %s",
model_id,
e,
)
model_cost_info = None
listing_info: Final = llm_router.get_model_listing_info(model_id) if llm_router is not None else None
max_input_tokens: int | None = None
max_output_tokens: int | None = None
if model_cost_info is not None:
max_input_tokens = coerce_token_limit(model_cost_info.get("max_input_tokens"))
max_output_tokens = coerce_token_limit(model_cost_info.get("max_output_tokens"))
mode: Final = model_cost_info.get("mode")
if isinstance(mode, str):
base["mode"] = mode
# One entry per distinct model behind the listed name; (None,) when the router knows
# nothing about it, so the listed name is resolved on its own as before.
deployment_models: Final[tuple[str | None, ...]] = (
listing_info.cost_map_keys if listing_info is not None and listing_info.cost_map_keys else (None,)
)
listed_info: Final = _safe_get_model_info(model_id, get_model_info)
candidate_sets: Final = tuple(
_resolve_listing_model_info(
deployment_model=deployment_model,
listed_model=model_id,
listed_info=listed_info,
get_model_info=get_model_info,
)
for deployment_model in deployment_models
)
max_input_tokens: int | None = _group_token_limit(candidate_sets, "max_input_tokens")
max_output_tokens: int | None = _group_token_limit(candidate_sets, "max_output_tokens")
mode: Final = next(
(
m
for m in (
cast("Mapping[str, object]", info).get("mode") # cast-ok: an entry need not carry "mode"
for candidates in candidate_sets
for info in candidates
)
if isinstance(m, str)
),
None,
)
if mode is not None:
base["mode"] = mode
if listing_info is not None:
if listing_info.max_input_tokens is not None:
max_input_tokens = listing_info.max_input_tokens
if listing_info.max_output_tokens is not None:
max_output_tokens = listing_info.max_output_tokens
if llm_router is not None:
configured_input, configured_output = llm_router.get_configured_token_limits(model_id)
if configured_input is not None:
max_input_tokens = configured_input
if configured_output is not None:
max_output_tokens = configured_output
configured_mode: Final = llm_router.get_configured_mode(model_id)
if isinstance(configured_mode, str):
base["mode"] = configured_mode

View file

@ -230,6 +230,7 @@ from litellm.types.router import (
CredentialLiteLLMParams,
CustomRoutingStrategyBase,
Deployment,
DeploymentModelListingInfo,
DeploymentTypedDict,
FallbackAccessCheck,
GuardrailTypedDict,
@ -10214,15 +10215,71 @@ class Router:
return None
return Deployment(**first_usable) if isinstance(first_usable, dict) else first_usable
def get_model_listing_info(self, model_name: str) -> DeploymentModelListingInfo | None:
"""
Return what the concrete deployments behind model_name contribute to its
/v1/models entry: the cost-map keys for their underlying models, plus the widest
token limits explicitly configured in their model_info. Resolved via O(1) index
lookup.
Returns None for wildcard-expanded or unknown names, where the listed name is the
real model name and no deployment-specific information exists, and treats a
malformed configured limit as absent rather than failing the listing.
The whole group is read rather than just its first deployment, so a group that
mixes models does not advertise a window that depends on config order; the widest
one is reported, which is what get_model_group_info already shows the Admin UI.
Keys are deduplicated, so the ordinary group of interchangeable deployments of one
model still costs the caller a single cost-map lookup. Unlike get_model_group_info,
this never triggers pattern matching or deep copies, so it is safe to call per
listed model on the /v1/models hot path.
"""
indices: Final = self.model_name_to_deployment_indices.get(model_name)
if not indices:
return None
deployments: Final = tuple(self.model_list[index] for index in indices)
model_infos: Final = tuple(deployment.get("model_info") or MappingProxyType({}) for deployment in deployments)
params: Final = tuple(deployment.get("litellm_params") or MappingProxyType({}) for deployment in deployments)
# base_model resolution mirrors get_router_model_info: unset or blank means the
# deployment's own model name is the cost-map key.
cost_map_keys: Final = tuple(
dict.fromkeys( # deduplicates while preserving config order
key
for key in (
model_info.get("base_model") or litellm_params.get("base_model") or litellm_params.get("model")
for model_info, litellm_params in zip(model_infos, params)
)
if isinstance(key, str) and key
)
)
return DeploymentModelListingInfo(
cost_map_keys=cost_map_keys,
max_input_tokens=self._widest_configured_limit(model_infos, "max_input_tokens"),
max_output_tokens=self._widest_configured_limit(model_infos, "max_output_tokens"),
)
@staticmethod
def _widest_configured_limit(model_infos: Sequence[Mapping[str, Any]], field: str) -> int | None:
"""The largest usable value of ``field`` across a group's configured model_info blocks."""
limits: Final = tuple(
limit
for limit in (coerce_token_limit(model_info.get(field)) for model_info in model_infos)
if limit is not None
)
return max(limits) if limits else None
def get_configured_token_limits(self, model_name: str) -> "tuple[int | None, int | None]":
"""
Return (max_input_tokens, max_output_tokens) explicitly configured in a concrete
deployment's model_info for model_name, via O(1) index lookup.
Returns (None, None) for wildcard-expanded or unknown names, and treats a
malformed configured value as absent rather than failing the listing. Unlike
get_model_group_info, this never triggers pattern matching or deep copies, so it
is safe to call per listed model on the /v1/models hot path.
malformed configured value as absent rather than failing the caller.
Deliberately reads one deployment rather than aggregating the group the way
get_model_listing_info does: its caller truncates an embedding input to this
value, so the widest window in a mixed group would be the wrong answer there.
"""
deployment: Final = self.get_deployment_by_model_group_name(model_group_name=model_name)
if deployment is None:

View file

@ -237,4 +237,49 @@ CACHE_SETTINGS_FIELDS: Final[list[CacheSettingsField]] = [
ui_field_name="SSL Check Hostname",
redis_type=None,
),
CacheSettingsField(
field_name="aws_iam_auth",
field_type="Boolean",
field_value=None,
field_description="Enable AWS ElastiCache IAM authentication",
field_default=False,
ui_field_name="AWS IAM Authentication",
redis_type=None,
),
CacheSettingsField(
field_name="aws_iam_user_name",
field_type="String",
field_value=None,
field_description="AWS ElastiCache IAM user name",
field_default=None,
ui_field_name="AWS IAM User Name",
redis_type=None,
),
CacheSettingsField(
field_name="aws_iam_cache_name",
field_type="String",
field_value=None,
field_description="AWS ElastiCache cache name",
field_default=None,
ui_field_name="AWS IAM Cache Name",
redis_type=None,
),
CacheSettingsField(
field_name="aws_iam_region",
field_type="String",
field_value=None,
field_description="AWS region for ElastiCache IAM authentication",
field_default=None,
ui_field_name="AWS IAM Region",
redis_type=None,
),
CacheSettingsField(
field_name="aws_iam_serverless",
field_type="Boolean",
field_value=None,
field_description="The ElastiCache cache is serverless rather than a self-designed cluster",
field_default=False,
ui_field_name="AWS IAM Serverless Cache",
redis_type=None,
),
]

View file

@ -102,4 +102,41 @@ COORDINATION_REDIS_SETTINGS_FIELDS: Final[list[CoordinationRedisSettingsField]]
ui_field_name="Service Name",
section="sentinel",
),
CoordinationRedisSettingsField(
field_name="aws_iam_auth",
field_type="Boolean",
field_description="Enable AWS ElastiCache IAM authentication",
field_default=False,
ui_field_name="AWS IAM Authentication",
section="connection",
),
CoordinationRedisSettingsField(
field_name="aws_iam_user_name",
field_type="String",
field_description="AWS ElastiCache IAM user name",
ui_field_name="AWS IAM User Name",
section="connection",
),
CoordinationRedisSettingsField(
field_name="aws_iam_cache_name",
field_type="String",
field_description="AWS ElastiCache cache name",
ui_field_name="AWS IAM Cache Name",
section="connection",
),
CoordinationRedisSettingsField(
field_name="aws_iam_region",
field_type="String",
field_description="AWS region for ElastiCache IAM authentication",
ui_field_name="AWS IAM Region",
section="connection",
),
CoordinationRedisSettingsField(
field_name="aws_iam_serverless",
field_type="Boolean",
field_description="The ElastiCache cache is serverless rather than a self-designed cluster",
field_default=False,
ui_field_name="AWS IAM Serverless Cache",
section="connection",
),
]

View file

@ -1,7 +1,8 @@
from datetime import datetime
from typing import Any, Final, Literal
from pydantic import BaseModel, ConfigDict, field_validator
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
from typing_extensions import Self
from litellm.types.mcp import (
DEFAULT_SUBJECT_TOKEN_TYPE,
@ -38,6 +39,26 @@ class MCPOAuthMetadata(BaseModel):
usable in memory but must never be persisted as configuration."""
class MCPOAuthIdentityBinding(BaseModel):
"""Per-server policy binding stored per-user OAuth credentials to the authenticated LiteLLM caller.
When enabled for an interactive oauth2 server, the token relay validates the upstream OIDC
``id_token`` (signature via the pinned issuer's JWKS, issuer, audience, expiry, nonce) and compares its
principal claim to the LiteLLM caller's trusted identity before the token is returned, stored,
or cached. ``audit`` logs mismatches without changing behavior; ``enforce`` fails closed with
403 ``oauth_principal_mismatch`` and disables the direct ``oauth-user-credential`` POST, which
would otherwise bypass validation with an arbitrary opaque token.
"""
mode: Literal["disabled", "audit", "enforce"] = "disabled"
issuer: str
jwks_url: str | None = None
audiences: list[str] = Field(min_length=1) # mutable-ok: public Pydantic schema requires list values
principal_claim: str = "email"
caller_field: Literal["user_email", "user_id"] = "user_email"
require_email_verified: bool = True
class MCPServer(BaseModel):
server_id: str
name: str
@ -173,6 +194,7 @@ class MCPServer(BaseModel):
# response (supports dot-notation for nested fields, e.g. "team.enterprise_id").
# Tokens that fail validation are rejected before storage.
token_validation: dict[str, Any] | None = None
oauth_identity_binding: MCPOAuthIdentityBinding | None = None
# Optional TTL override (seconds) for the Redis per-user token cache, capped
# at the token's expires_in minus the expiry buffer so a cached entry never
# outlives the token. Defaults to the token's expires_in minus the expiry
@ -225,6 +247,14 @@ class MCPServer(BaseModel):
"""
return self.oauth2_flow == "client_credentials"
@model_validator(mode="after")
def validate_identity_binding_mode(self) -> Self:
binding: Final = self.oauth_identity_binding
if binding is not None and binding.mode != "disabled":
if not self.needs_user_oauth_token or self.delegate_auth_to_upstream:
raise ValueError("oauth_identity_binding requires gateway-managed per-user OAuth2 credentials")
return self
@property
def needs_user_oauth_token(self) -> bool:
"""True if this is an OAuth2 server that relies on per-user tokens (no client_credentials)."""

View file

@ -614,6 +614,24 @@ class Deployment(BaseModel):
setattr(self, key, value)
@dataclass(frozen=True, slots=True)
class DeploymentModelListingInfo:
"""What the deployments behind a model name contribute to its OpenAI-compatible listing entry.
``cost_map_keys`` are the names those deployments' underlying models are known by in
``litellm.model_cost`` (``base_model`` when set, else ``litellm_params.model``), which
is what a request actually reaches; the public model name they are listed under is an
arbitrary alias and often absent from the cost map. Keys are deduplicated in config
order, so the ordinary group -- several interchangeable deployments of one model --
carries exactly one. The token limits are the widest explicitly set in any
deployment's ``model_info``, which outrank anything the cost map says.
"""
cost_map_keys: tuple[str, ...]
max_input_tokens: int | None
max_output_tokens: int | None
class RouterErrors(enum.Enum):
"""
Enum for router specific errors with common codes

View file

@ -2971,24 +2971,33 @@ def _resolve_builtin_model_cost_entry(key: str, provider: str) -> dict[str, obje
return None
def is_generalized_model_info(model_info: ModelInfo) -> bool:
"""Whether ``model_info`` came from a fallback-generalization capability rule.
Detected as the resolved key missing ``litellm.model_cost`` while matching a
capability rule. A rule-derived entry carries no pricing and only a conservative
family-baseline context window, so callers holding a second candidate name should
prefer an exact cost-map entry from that name over this one.
"""
key: Final = cast("Mapping[str, object]", model_info).get("key") # cast-ok: partial dicts may omit "key"
if not isinstance(key, str):
return False
return key not in litellm.model_cost and match_capability_generalizations(key) is not None
def _get_builtin_model_info_for_registration(model: str) -> ModelInfo | None:
"""Resolve ``model`` to its built-in cost-map entry for registration merging.
Returns ``None`` when the lookup raises or when it resolved via a
fallback-generalization capability rule, detected as the resolved key missing
``litellm.model_cost`` while matching a capability rule. A rule-derived entry
carries no pricing, so treating it as a hit would skip the built-in
cache-pricing inheritance for prefix-mangled keys.
fallback-generalization capability rule. A rule-derived entry carries no
pricing, so treating it as a hit would skip the built-in cache-pricing
inheritance for prefix-mangled keys.
"""
try:
info: Final = get_model_info(model=model)
except Exception:
return None
if info["key"] in litellm.model_cost:
return info
if match_capability_generalizations(info["key"]) is None:
return info
return None
return None if is_generalized_model_info(info) else info
_runtime_registered_model_cost: Final[dict[str, dict[str, object]]] = {} # mutable-ok: replayed on reload

View file

@ -258,6 +258,39 @@ gateway_metrics_port = 4001
gateway_metrics_scrape_cidrs = ["10.0.0.0/16"]
```
### In-container connection pool
Each of the `gateway_num_workers` uvicorn workers opens its own Prisma pool
straight to Postgres, so one task holds `workers x connection_limit`
connections and the fleet's footprint against the database ceiling grows with
every task. `gateway_connection_pool_enabled` runs a PgBouncer (transaction
mode, loopback) inside the gateway container that all workers share, capping
the task at `gateway_pool_max_db_connections` upstream connections however
many workers it runs. `gateway_pool_max_client_conn` bounds the worker-side
connections the pooler accepts. The module sets
`LITELLM_PGBOUNCER_ENABLED`, `LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS` and
`LITELLM_PGBOUNCER_MAX_CLIENT_CONN` on the gateway container only; the backend
and the migration task keep the direct connection.
```hcl
create_database = false
database_url = "postgresql://litellm:<password>@db.internal:5432/litellm"
gateway_num_workers = 4
gateway_connection_pool_enabled = true
gateway_pool_max_db_connections = 20
gateway_pool_max_client_conn = 1000
```
The pool needs a static database password, so it is only valid with an
existing database via `database_url`. The module-created Aurora authenticates
with rotating IAM tokens (see [Aurora + IAM auth](#aurora--iam-auth)), which
the pooler cannot follow, and `terraform plan` rejects that combination.
The componentized `gateway_image` starts through `python -m gateway.launch`,
which reads these variables, starts the pooler once per task and hands the
workers its loopback URL; the classic `litellm` image honours them the same
way.
### Scaling the gateway on requests and tokens
By default the gateway service target-tracks CPU (`gateway_cpu_target`) and

View file

@ -213,6 +213,12 @@ locals {
# otherwise we keep the image's ENTRYPOINT and only override `command`.
gateway_uvicorn_args = "--host 0.0.0.0 --port 4000 --workers ${var.gateway_num_workers}"
gateway_pool_env = var.gateway_connection_pool_enabled ? [
{ name = "LITELLM_PGBOUNCER_ENABLED", value = "true" },
{ name = "LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS", value = tostring(var.gateway_pool_max_db_connections) },
{ name = "LITELLM_PGBOUNCER_MAX_CLIENT_CONN", value = tostring(var.gateway_pool_max_client_conn) },
] : []
metrics_enabled = var.gateway_metrics_port != null
metrics_multiproc_dir = "/tmp/litellm_prometheus_multiproc"
metrics_volume = "prometheus-multiproc"
@ -253,7 +259,7 @@ locals {
backend_uvicorn_args = "--host 0.0.0.0 --port 4001"
gateway_launch_cmd = "case \"$USE_DDTRACE\" in [Tt][Rr][Uu][Ee]) export DD_TRACE_OPENAI_ENABLED=\"False\"; exec ddtrace-run uvicorn gateway.main:app ${local.gateway_uvicorn_args};; *) exec uvicorn gateway.main:app ${local.gateway_uvicorn_args};; esac"
gateway_launch_cmd = "case \"$USE_DDTRACE\" in [Tt][Rr][Uu][Ee]) export DD_TRACE_OPENAI_ENABLED=\"False\"; exec ddtrace-run python -m gateway.launch ${local.gateway_uvicorn_args};; *) exec python -m gateway.launch ${local.gateway_uvicorn_args};; esac"
backend_launch_cmd = "case \"$USE_DDTRACE\" in [Tt][Rr][Uu][Ee]) export DD_TRACE_OPENAI_ENABLED=\"False\"; exec ddtrace-run uvicorn backend.main:app ${local.backend_uvicorn_args};; *) exec uvicorn backend.main:app ${local.backend_uvicorn_args};; esac"
gateway_proxy_overrides = local.proxy_config_enabled ? {
@ -298,6 +304,11 @@ resource "aws_ecs_task_definition" "gateway" {
)
error_message = "billing_metrics_client_cert_pem and billing_metrics_client_key_pem are both required when billing_metrics_endpoint is set."
}
precondition {
condition = !var.gateway_connection_pool_enabled || local.byo_database
error_message = "gateway_connection_pool_enabled requires an existing database via database_url with create_database = false: the module-created Aurora authenticates with IAM tokens, which the in-container pgbouncer cannot follow because it holds a static database password."
}
}
family = "${local.name}-gateway"
@ -323,6 +334,7 @@ resource "aws_ecs_task_definition" "gateway" {
local.gateway_extra_env_list,
local.proxy_config_env,
local.metrics_env,
local.gateway_pool_env,
)
secrets = concat(local.shared_secrets, local.gateway_extra_secrets_list)
mountPoints = local.metrics_mount_points

View file

@ -0,0 +1,123 @@
# Plan-only coverage for the in-container PgBouncer knobs on the gateway task.
# `mock_provider` keeps this offline: no AWS credentials, no API calls, no
# resources. Run from terraform/litellm/aws with `terraform test`.
mock_provider "aws" {
mock_data "aws_iam_policy_document" {
defaults = {
json = "{\"Version\":\"2012-10-17\",\"Statement\":[]}"
}
}
}
mock_provider "random" {}
variables {
region = "us-east-1"
tenant = "acme"
env = "test"
azs = ["us-east-1a", "us-east-1b"]
allow_plaintext_alb = true
}
run "pool_off_by_default" {
command = plan
assert {
condition = length(local.gateway_pool_env) == 0
error_message = "The gateway must get no LITELLM_PGBOUNCER_* env unless gateway_connection_pool_enabled is set."
}
}
run "pool_enabled_renders_the_three_vars_with_configured_sizes" {
command = plan
variables {
create_database = false
database_url = "postgresql://litellm:pw@db.internal:5432/litellm"
gateway_num_workers = 4
gateway_connection_pool_enabled = true
gateway_pool_max_db_connections = 8
gateway_pool_max_client_conn = 250
}
assert {
condition = alltrue([
length(local.gateway_pool_env) == 3,
local.gateway_pool_env[0].name == "LITELLM_PGBOUNCER_ENABLED" && local.gateway_pool_env[0].value == "true",
local.gateway_pool_env[1].name == "LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS" && local.gateway_pool_env[1].value == "8",
local.gateway_pool_env[2].name == "LITELLM_PGBOUNCER_MAX_CLIENT_CONN" && local.gateway_pool_env[2].value == "250",
])
error_message = "The pool env must carry the enabled flag and the configured sizes as strings."
}
}
run "pool_enabled_uses_the_module_default_sizes" {
command = plan
variables {
create_database = false
database_url = "postgresql://litellm:pw@db.internal:5432/litellm"
gateway_connection_pool_enabled = true
}
assert {
condition = alltrue([
length(local.gateway_pool_env) == 3,
local.gateway_pool_env[1].name == "LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS" && local.gateway_pool_env[1].value == "20",
local.gateway_pool_env[2].name == "LITELLM_PGBOUNCER_MAX_CLIENT_CONN" && local.gateway_pool_env[2].value == "1000",
])
error_message = "The pool env must fall back to the module defaults of 20 upstream and 1000 client connections."
}
}
run "gateway_starts_through_the_pool_aware_launcher" {
command = plan
variables {
gateway_num_workers = 4
}
assert {
condition = alltrue([
strcontains(local.gateway_launch_cmd, "exec python -m gateway.launch --host 0.0.0.0 --port 4000 --workers 4"),
strcontains(local.gateway_launch_cmd, "exec ddtrace-run python -m gateway.launch --host 0.0.0.0 --port 4000 --workers 4"),
!strcontains(local.gateway_launch_cmd, "uvicorn gateway.main:app"),
local.gateway_proxy_overrides.command[0] == local.gateway_launch_cmd,
])
error_message = "The gateway must start through gateway.launch (with and without ddtrace) so the pooler starts once before uvicorn forks the workers."
}
}
run "pool_with_module_created_iam_aurora_fails_at_plan" {
command = plan
variables {
gateway_connection_pool_enabled = true
}
expect_failures = [
aws_ecs_task_definition.gateway,
]
}
run "pool_without_any_database_fails_at_plan" {
command = plan
variables {
create_database = false
gateway_connection_pool_enabled = true
}
expect_failures = [
aws_ecs_task_definition.gateway,
]
}
run "module_created_iam_aurora_without_the_pool_still_plans" {
command = plan
assert {
condition = contains(local.managed_db_env, { name = "IAM_TOKEN_DB_AUTH", value = "true" })
error_message = "Without the pool the module-created Aurora must keep IAM token auth."
}
}

View file

@ -200,6 +200,45 @@ variable "gateway_num_workers" {
}
}
variable "gateway_connection_pool_enabled" {
description = <<-EOT
Run an in-container PgBouncer (transaction mode, loopback) in each gateway
task, shared by every uvicorn worker. Without it each of the
`gateway_num_workers` workers opens its own Prisma pool straight to
Postgres, so a task's footprint against the database connection ceiling is
workers x connection_limit and grows with every task. Sets
LITELLM_PGBOUNCER_ENABLED / LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS /
LITELLM_PGBOUNCER_MAX_CLIENT_CONN on the gateway container only. Requires
an existing database via `database_url`: the module-created Aurora
authenticates with IAM tokens, which the pooler cannot follow because it
holds one static password for the life of the task.
EOT
type = bool
default = false
}
variable "gateway_pool_max_db_connections" {
description = "Upstream Postgres connections one gateway task may hold when gateway_connection_pool_enabled is set, regardless of gateway_num_workers. 20 suits 4 workers; a 5000-connection database then fits roughly 200 tasks."
type = number
default = 20
validation {
condition = var.gateway_pool_max_db_connections >= 1
error_message = "gateway_pool_max_db_connections must be >= 1."
}
}
variable "gateway_pool_max_client_conn" {
description = "Client connections the in-container PgBouncer accepts from the gateway workers when gateway_connection_pool_enabled is set."
type = number
default = 1000
validation {
condition = var.gateway_pool_max_client_conn >= 1
error_message = "gateway_pool_max_client_conn must be >= 1."
}
}
variable "backend_cpu" {
description = "Fargate CPU units for the backend task (1024 = 1 vCPU)."
type = number

View file

@ -287,6 +287,40 @@ the gateway on GKE with the Helm chart's `targetTokensPerSecond` (see
"Dependencies only" below) rather than wiring the counter into Cloud
Monitoring, which the autoscaler would ignore
### In-container connection pool
Each of the `gateway_num_workers` uvicorn workers opens its own Prisma pool
straight to Cloud SQL, so one instance holds `workers x connection_limit`
connections and the fleet's footprint against the database ceiling grows with
every instance Cloud Run adds. `gateway_connection_pool_enabled` runs a
PgBouncer (transaction mode, loopback) inside the gateway container that all
workers share, capping the instance at `gateway_pool_max_db_connections`
upstream connections however many workers it runs.
`gateway_pool_max_client_conn` bounds the worker-side connections the pooler
accepts. The module sets `LITELLM_PGBOUNCER_ENABLED`,
`LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS` and `LITELLM_PGBOUNCER_MAX_CLIENT_CONN`
on the gateway service only; the backend service and the migrations job keep
the direct connection
```hcl
gateway_num_workers = 4
gateway_connection_pool_enabled = true
gateway_pool_max_db_connections = 20
gateway_pool_max_client_conn = 1000
```
The pooler holds one static database password for the life of the instance.
This stack authenticates to Cloud SQL with the Secret Manager password (see
[Database authentication](#database-authentication)), so nothing else is
needed; a Cloud SQL Auth Proxy sidecar with IAM auth would not work with the
pool
The gateway container starts through `python -m gateway.launch` (the
componentized image's own entrypoint) rather than `uvicorn` directly. The
launcher reads these variables, starts the pooler once per instance before
uvicorn forks the workers and hands them its loopback `DATABASE_URL`. It also
honours `KEEPALIVE_TIMEOUT` from `gateway_extra_env` the way the image does
## Tenant deployment
Every resource the stack creates is named `${tenant}-litellm-${env}` (or

View file

@ -138,10 +138,16 @@ locals {
"export DATABASE_URL_READ_REPLICA=\"postgresql://$${DATABASE_USER}:$${DATABASE_PASSWORD}@$${DATABASE_HOST_READ_REPLICA}:$${DATABASE_PORT_READ_REPLICA}/$${DATABASE_NAME}\"",
]
gateway_pool_env = var.gateway_connection_pool_enabled ? [
{ name = "LITELLM_PGBOUNCER_ENABLED", value = "true" },
{ name = "LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS", value = tostring(var.gateway_pool_max_db_connections) },
{ name = "LITELLM_PGBOUNCER_MAX_CLIENT_CONN", value = tostring(var.gateway_pool_max_client_conn) },
] : []
gateway_uvicorn_args = "--host 0.0.0.0 --port 4000 --workers ${var.gateway_num_workers}"
backend_uvicorn_args = "--host 0.0.0.0 --port 4001"
gateway_launch_cmd = "case \"$USE_DDTRACE\" in [Tt][Rr][Uu][Ee]) export DD_TRACE_OPENAI_ENABLED=\"False\"; exec ddtrace-run uvicorn gateway.main:app ${local.gateway_uvicorn_args};; *) exec uvicorn gateway.main:app ${local.gateway_uvicorn_args};; esac"
gateway_launch_cmd = "case \"$USE_DDTRACE\" in [Tt][Rr][Uu][Ee]) export DD_TRACE_OPENAI_ENABLED=\"False\"; exec ddtrace-run python -m gateway.launch ${local.gateway_uvicorn_args};; *) exec python -m gateway.launch ${local.gateway_uvicorn_args};; esac"
backend_launch_cmd = "case \"$USE_DDTRACE\" in [Tt][Rr][Uu][Ee]) export DD_TRACE_OPENAI_ENABLED=\"False\"; exec ddtrace-run uvicorn backend.main:app ${local.backend_uvicorn_args};; *) exec uvicorn backend.main:app ${local.backend_uvicorn_args};; esac"
gateway_args = join(" && ", concat(
@ -229,7 +235,7 @@ resource "google_cloud_run_v2_service" "gateway" {
}
dynamic "env" {
for_each = concat(local.shared_env_kv, local.gateway_otel_env_kv, local.billing_metrics_env_kv, local.gateway_extra_env_kv, local.proxy_config_env, local.metrics_env_kv)
for_each = concat(local.shared_env_kv, local.gateway_otel_env_kv, local.billing_metrics_env_kv, local.gateway_extra_env_kv, local.proxy_config_env, local.metrics_env_kv, local.gateway_pool_env)
content {
name = env.value.name
value = env.value.value

View file

@ -0,0 +1,135 @@
# Plan-only coverage for the in-container PgBouncer knobs on the gateway
# service. `mock_provider` keeps this offline: no GCP credentials, no API
# calls, no resources. Run from terraform/litellm/gcp with `terraform test`.
mock_provider "google" {
mock_resource "google_redis_instance" {
defaults = {
host = "10.0.0.4"
port = 6379
server_ca_certs = [{
cert = "-----BEGIN CERTIFICATE-----\nmock\n-----END CERTIFICATE-----"
}]
}
}
}
mock_provider "google-beta" {}
mock_provider "random" {}
variables {
project_id = "test-project"
tenant = "tenant"
env = "test"
allow_plaintext_lb = true
image_registry = "us-central1-docker.pkg.dev/test-project/litellm"
}
run "pool_off_by_default" {
command = plan
assert {
condition = length(local.gateway_pool_env) == 0
error_message = "The gateway must get no LITELLM_PGBOUNCER_* env unless gateway_connection_pool_enabled is set."
}
assert {
condition = !anytrue([
for e in google_cloud_run_v2_service.gateway[0].template[0].containers[0].env : startswith(e.name, "LITELLM_PGBOUNCER_")
])
error_message = "The gateway service must carry no LITELLM_PGBOUNCER_* env by default."
}
}
run "pool_enabled_renders_the_three_vars_with_configured_sizes" {
command = plan
variables {
gateway_num_workers = 4
gateway_connection_pool_enabled = true
gateway_pool_max_db_connections = 8
gateway_pool_max_client_conn = 250
}
assert {
condition = alltrue([
length(local.gateway_pool_env) == 3,
local.gateway_pool_env[0].name == "LITELLM_PGBOUNCER_ENABLED" && local.gateway_pool_env[0].value == "true",
local.gateway_pool_env[1].name == "LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS" && local.gateway_pool_env[1].value == "8",
local.gateway_pool_env[2].name == "LITELLM_PGBOUNCER_MAX_CLIENT_CONN" && local.gateway_pool_env[2].value == "250",
])
error_message = "The pool env must carry the enabled flag and the configured sizes as strings."
}
assert {
condition = alltrue([
contains([for e in google_cloud_run_v2_service.gateway[0].template[0].containers[0].env : e.name], "LITELLM_PGBOUNCER_ENABLED"),
contains([for e in google_cloud_run_v2_service.gateway[0].template[0].containers[0].env : e.name], "LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS"),
contains([for e in google_cloud_run_v2_service.gateway[0].template[0].containers[0].env : e.name], "LITELLM_PGBOUNCER_MAX_CLIENT_CONN"),
])
error_message = "The gateway service must receive all three LITELLM_PGBOUNCER_* env vars."
}
assert {
condition = !anytrue(concat(
[for e in google_cloud_run_v2_service.backend[0].template[0].containers[0].env : startswith(e.name, "LITELLM_PGBOUNCER_")],
[for e in google_cloud_run_v2_job.migrations[0].template[0].template[0].containers[0].env : startswith(e.name, "LITELLM_PGBOUNCER_")],
))
error_message = "The backend service and the migrations job must keep their direct database connection."
}
}
run "pool_enabled_uses_the_module_default_sizes" {
command = plan
variables {
gateway_connection_pool_enabled = true
}
assert {
condition = alltrue([
length(local.gateway_pool_env) == 3,
local.gateway_pool_env[1].name == "LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS" && local.gateway_pool_env[1].value == "20",
local.gateway_pool_env[2].name == "LITELLM_PGBOUNCER_MAX_CLIENT_CONN" && local.gateway_pool_env[2].value == "1000",
])
error_message = "The pool env must fall back to the module defaults of 20 upstream and 1000 client connections."
}
}
run "gateway_starts_through_the_pool_aware_launcher" {
command = plan
variables {
gateway_num_workers = 4
}
assert {
condition = alltrue([
strcontains(local.gateway_launch_cmd, "exec python -m gateway.launch --host 0.0.0.0 --port 4000 --workers 4"),
strcontains(local.gateway_launch_cmd, "exec ddtrace-run python -m gateway.launch --host 0.0.0.0 --port 4000 --workers 4"),
!strcontains(local.gateway_launch_cmd, "uvicorn gateway.main:app"),
endswith(google_cloud_run_v2_service.gateway[0].template[0].containers[0].args[0], local.gateway_launch_cmd),
])
error_message = "The gateway must start through gateway.launch (with and without ddtrace) so the pooler starts once before uvicorn forks the workers."
}
assert {
condition = strcontains(local.backend_launch_cmd, "uvicorn backend.main:app")
error_message = "The backend has no workers to share a pooler and keeps starting uvicorn directly."
}
}
run "pool_sizes_below_one_fail_at_plan" {
command = plan
variables {
gateway_connection_pool_enabled = true
gateway_pool_max_db_connections = 0
gateway_pool_max_client_conn = 0
}
expect_failures = [
var.gateway_pool_max_db_connections,
var.gateway_pool_max_client_conn,
]
}

View file

@ -206,6 +206,45 @@ variable "gateway_num_workers" {
}
}
variable "gateway_connection_pool_enabled" {
description = <<-EOT
Run an in-container PgBouncer (transaction mode, loopback) in each gateway
instance, shared by every uvicorn worker. Without it each of the
`gateway_num_workers` workers opens its own Prisma pool straight to
Cloud SQL, so an instance's footprint against the database connection
ceiling is workers x connection_limit and grows with every instance. Sets
LITELLM_PGBOUNCER_ENABLED / LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS /
LITELLM_PGBOUNCER_MAX_CLIENT_CONN on the gateway container only. The
module's Cloud SQL authenticates with the static password in Secret
Manager, which is what the pooler needs. Mirrors the AWS stack's
gateway_connection_pool_enabled.
EOT
type = bool
default = false
}
variable "gateway_pool_max_db_connections" {
description = "Upstream Cloud SQL connections one gateway instance may hold when gateway_connection_pool_enabled is set, regardless of gateway_num_workers. 20 suits 4 workers."
type = number
default = 20
validation {
condition = var.gateway_pool_max_db_connections >= 1
error_message = "gateway_pool_max_db_connections must be >= 1."
}
}
variable "gateway_pool_max_client_conn" {
description = "Client connections the in-container PgBouncer accepts from the gateway workers when gateway_connection_pool_enabled is set."
type = number
default = 1000
validation {
condition = var.gateway_pool_max_client_conn >= 1
error_message = "gateway_pool_max_client_conn must be >= 1."
}
}
# Cloud Run autoscales out of the box (request-rate driven). The min/max
# bounds mirror the HPA replica bounds in helm/litellm/values.yaml so each
# stack scales over the same range. Cloud Run has no direct CPU-utilization

View file

@ -0,0 +1,155 @@
import os
import socket
import sys
import textwrap
import urllib.parse
from pathlib import Path
from typing import Final, cast
import pytest
from uvicorn.importer import import_from_string
from uvicorn.main import main as uvicorn_main
import gateway.main
from gateway.launch import GATEWAY_APP, main, pool_database_url, uvicorn_argv
from litellm.proxy.db.db_url_settings import DatabaseURLSettings
from litellm.proxy.db.pgbouncer import PgBouncerError, PgBouncerSettings
DB_ENV: Final = {
"DATABASE_HOST": "db.internal",
"DATABASE_PORT": "5432",
"DATABASE_USER": "litellm_pool",
"DATABASE_NAME": "litellm",
"DATABASE_PASSWORD": "p@ss",
}
def _free_port() -> int:
with socket.socket() as probe:
probe.bind(("127.0.0.1", 0))
return cast(tuple[str, int], probe.getsockname())[1]
def _fake_pooler(tmp_path: Path) -> Path:
script: Final = tmp_path / "fake-pgbouncer"
script.write_text(
textwrap.dedent(
f"""\
#!{sys.executable}
import configparser, select, socket, sys
if sys.argv[1:] == ["--version"]:
print("PgBouncer 1.25.2")
sys.exit(0)
ini = configparser.ConfigParser()
ini.read(sys.argv[1])
port = ini.getint("pgbouncer", "listen_port")
tcp = socket.socket()
tcp.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
tcp.bind(("127.0.0.1", port))
tcp.listen()
unix = socket.socket(socket.AF_UNIX)
unix.bind(ini.get("pgbouncer", "unix_socket_dir") + f"/.s.PGSQL.{{port}}")
unix.listen()
while True:
for ready in select.select([tcp, unix], [], [])[0]:
ready.accept()[0].close()
"""
)
)
script.chmod(0o700)
return script
def _query(url: str) -> dict[str, str]:
return dict(urllib.parse.parse_qsl(urllib.parse.urlsplit(url).query))
@pytest.fixture
def password_env(monkeypatch: pytest.MonkeyPatch) -> dict[str, str]:
for var in ("DATABASE_URL", "IAM_TOKEN_DB_AUTH", "AZURE_POSTGRESQL_AUTH", "DATABASE_HOST_READ_REPLICA"):
monkeypatch.setenv(var, "")
monkeypatch.delenv(var)
for var, value in DB_ENV.items():
monkeypatch.setenv(var, value)
return dict(DB_ENV)
def _uvicorn_params(argv: tuple[str, ...]) -> dict[str, object]:
return uvicorn_main.make_context("uvicorn", list(argv)).params
class TestUvicornArgv:
def test_keepalive_env_reaches_uvicorn(self):
params: Final = _uvicorn_params(uvicorn_argv(("--workers", "4"), {"KEEPALIVE_TIMEOUT": "75"}))
assert params["app"] == GATEWAY_APP
assert params["workers"] == 4
assert params["timeout_keep_alive"] == 75
def test_unset_env_keeps_the_uvicorn_default(self):
assert _uvicorn_params(uvicorn_argv(("--workers", "4"), {}))["timeout_keep_alive"] == 5
def test_an_explicit_flag_wins_over_the_env(self):
argv: Final = uvicorn_argv(("--timeout-keep-alive", "30"), {"KEEPALIVE_TIMEOUT": "75"})
assert _uvicorn_params(argv)["timeout_keep_alive"] == 30
def test_the_app_uvicorn_is_told_to_serve_is_the_trimmed_gateway(self):
assert import_from_string(cast(str, _uvicorn_params(uvicorn_argv((), {}))["app"])) is gateway.main.app
class TestPoolDatabaseUrl:
def test_a_disabled_pooler_yields_no_url_to_install(self, password_env: dict[str, str]):
settings: Final = DatabaseURLSettings.from_env()
settings.apply_to_env()
environ: Final = {"DATABASE_URL": "postgresql://litellm_pool:p%40ss@db.internal:5432/litellm"}
assert pool_database_url(settings, PgBouncerSettings(enabled=False), environ) is None
def test_a_missing_upstream_url_is_reported(self, password_env: dict[str, str]):
environ: Final[dict[str, str]] = {}
outcome: Final = pool_database_url(DatabaseURLSettings.from_env(), PgBouncerSettings(enabled=True), environ)
assert isinstance(outcome, PgBouncerError)
assert "DATABASE_URL" in outcome.reason
def test_token_auth_is_refused(self, password_env: dict[str, str], monkeypatch: pytest.MonkeyPatch, tmp_path: Path):
monkeypatch.setenv("IAM_TOKEN_DB_AUTH", "true")
environ: Final = {"DATABASE_URL": "postgresql://litellm:token@db.internal:5432/litellm"}
outcome: Final = pool_database_url(
DatabaseURLSettings.from_env(),
PgBouncerSettings(enabled=True, port=_free_port(), binary=str(_fake_pooler(tmp_path))),
environ,
)
assert isinstance(outcome, PgBouncerError)
assert "IAM_TOKEN_DB_AUTH" in outcome.reason
class TestMain:
def test_workers_inherit_the_loopback_url_the_supervisor_installed(
self, password_env: dict[str, str], monkeypatch: pytest.MonkeyPatch, tmp_path: Path
):
port: Final = _free_port()
monkeypatch.setenv("LITELLM_PGBOUNCER_ENABLED", "true")
monkeypatch.setenv("LITELLM_PGBOUNCER_PORT", str(port))
monkeypatch.setenv("LITELLM_PGBOUNCER_BINARY", str(_fake_pooler(tmp_path)))
monkeypatch.setenv("KEEPALIVE_TIMEOUT", "75")
served: Final[list[tuple[str, ...]]] = []
main(("--workers", "4"), serve=lambda argv: served.append(tuple(argv)))
pooled: Final = os.environ["DATABASE_URL"]
assert urllib.parse.urlsplit(pooled).netloc == f"litellm_pool:p%40ss@127.0.0.1:{port}"
assert _query(pooled)["pgbouncer"] == "true"
assert _uvicorn_params(served[0])["timeout_keep_alive"] == 75
DatabaseURLSettings.from_env().apply_to_env()
worker_url: Final = os.environ["DATABASE_URL"]
assert urllib.parse.urlsplit(worker_url).netloc == f"litellm_pool:p%40ss@127.0.0.1:{port}"
assert _query(worker_url)["pgbouncer"] == "true"
def test_a_pooler_that_cannot_start_stops_the_gateway_before_uvicorn(
self, password_env: dict[str, str], monkeypatch: pytest.MonkeyPatch, tmp_path: Path
):
monkeypatch.setenv("LITELLM_PGBOUNCER_ENABLED", "true")
monkeypatch.setenv("LITELLM_PGBOUNCER_BINARY", str(tmp_path / "missing-pgbouncer"))
served: Final[list[tuple[str, ...]]] = []
with pytest.raises(SystemExit) as stopped:
main(("--workers", "4"), serve=lambda argv: served.append(tuple(argv)))
assert "missing-pgbouncer" in str(stopped.value)
assert served == []

View file

@ -22,6 +22,7 @@ class _Server:
server_id="srv",
configured_token_url=None,
):
self.oauth_identity_binding = None
self.token_url = token_url
self.configured_token_url = configured_token_url
self.client_id = client_id
@ -287,3 +288,40 @@ async def test_refresh_uses_admin_entered_token_url_when_issuer_yield_empties_re
assert token is not None
assert token.access_token == "new-at"
assert posted[0][0] == "https://idp.example.com/token"
@pytest.mark.asyncio
async def test_identity_rejection_never_persists_or_returns_refreshed_token():
from unittest.mock import AsyncMock
from fastapi import HTTPException
validator = AsyncMock(side_effect=HTTPException(status_code=403, detail="oauth_principal_mismatch"))
persist = AsyncMock()
refresher = AuthorizationCodeRefresher(
_lookup(_Server()),
_endpoint({"access_token": "foreign-token", "id_token": "foreign-identity"}),
persist,
identity_validator=validator,
)
assert await refresher.refresh("alice", "srv", OAuthToken(access_token="old", refresh_token="old-rt")) is None
persist.assert_not_awaited()
@pytest.mark.asyncio
async def test_verified_refresh_preserves_binding_proof_in_storage():
from unittest.mock import AsyncMock
validator = AsyncMock(return_value="verified-binding")
persist = AsyncMock()
refresher = AuthorizationCodeRefresher(
_lookup(_Server()),
_endpoint({"access_token": "new", "refresh_token": "rotated"}),
persist,
identity_validator=validator,
)
token = await refresher.refresh("alice", "srv", OAuthToken(access_token="old", refresh_token="old-rt"))
assert token.access_token == "new"
assert token.refresh_token == "rotated"
assert token.identity_binding_proof == "verified-binding"
assert persist.await_args.kwargs["identity_binding_proof"] == "verified-binding"

View file

@ -246,3 +246,94 @@ async def test_lazy_store_invalidate_works_after_redis_chain_is_built() -> None:
assert build_calls == 1
assert redis_store.invalidations == [("u", "s")]
@pytest.mark.asyncio
async def test_enforcement_invalidates_cached_legacy_credentials_before_use():
from litellm.types.mcp import MCPAuth, MCPTransport
from litellm.types.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer
server = MCPServer(
server_id="srv",
name="srv",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
oauth_identity_binding=MCPOAuthIdentityBinding(
mode="enforce",
issuer="https://idp.example.com",
audiences=["client"],
),
)
cached = _RecordingStore("belongs-to-bob")
store = LazyPerUserOAuthTokenStore(
lambda server_id: server,
store_builder=lambda lookup: (cached, False),
redis_available=lambda: False,
)
assert await store.fetch("alice", "srv") is None
assert cached.calls == [("alice", "srv")]
assert cached.invalidations == [("alice", "srv")]
@pytest.mark.asyncio
async def test_enforced_cache_hit_avoids_credential_read_and_rejects_changed_policy(monkeypatch):
from unittest.mock import AsyncMock
from litellm.proxy._experimental.mcp_server.oauth_identity_binding import current_binding_proof
from litellm.proxy._experimental.mcp_server.outbound_credentials import per_user_oauth_store as module
from litellm.types.mcp_server.mcp_server_manager import MCPServer
server = MCPServer(
server_id="srv",
name="srv",
transport="http",
auth_type="oauth2",
oauth_identity_binding={
"mode": "enforce",
"issuer": "https://idp.example",
"audiences": ["client"],
"caller_field": "user_id",
"principal_claim": "sub",
},
)
proof = await current_binding_proof(server.oauth_identity_binding, "alice", "srv")
read = AsyncMock(return_value={"access_token": "alice-token", "identity_binding_proof": proof})
monkeypatch.setattr(module, "_read_credential", read)
monkeypatch.setattr(module, "_runtime_backend_and_coordinator", lambda: (None, None, False))
store = LazyPerUserOAuthTokenStore(lambda _: server, redis_available=lambda: False)
assert (await store.fetch("alice", "srv")).access_token == "alice-token"
assert (await store.fetch("alice", "srv")).access_token == "alice-token"
read.assert_awaited_once_with("alice", "srv")
server.oauth_identity_binding = server.oauth_identity_binding.model_copy(update={"audiences": ["changed"]})
assert await store.fetch("alice", "srv") is None
read.assert_awaited_once()
assert await store.fetch("alice", "srv") is None
assert read.await_count == 2
@pytest.mark.asyncio
async def test_expired_unverified_credential_never_reaches_refresh(monkeypatch):
from unittest.mock import AsyncMock
from litellm.proxy._experimental.mcp_server.outbound_credentials import per_user_oauth_store as module
from litellm.types.mcp_server.mcp_server_manager import MCPServer
server = MCPServer(
server_id="srv",
name="srv",
transport="http",
auth_type="oauth2",
token_url="https://idp.example/token",
oauth_identity_binding={"mode": "enforce", "issuer": "https://idp.example", "audiences": ["client"]},
)
read = AsyncMock(
return_value={"access_token": "bob", "refresh_token": "bob-refresh", "expires_at": "2000-01-01T00:00:00Z"}
)
post = AsyncMock()
monkeypatch.setattr(module, "_read_credential", read)
monkeypatch.setattr(module, "_post_token_endpoint", post)
monkeypatch.setattr(module, "_runtime_backend_and_coordinator", lambda: (None, None, False))
store = LazyPerUserOAuthTokenStore(lambda _: server, redis_available=lambda: False)
assert await store.fetch("alice", "srv") is None
post.assert_not_awaited()

View file

@ -52,3 +52,17 @@ def test_undecryptable_blob_is_a_miss():
def test_empty_plaintext_is_a_miss():
codec = OAuthTokenCacheCodec(encrypt=lambda s: s, decrypt=lambda b: b)
assert codec.decode("") is None
def test_bound_token_round_trip_preserves_proof_without_refresh_secret():
codec = _wrapping_codec()
blob = codec.encode(OAuthToken("alice-token", refresh_token="private-refresh", identity_binding_proof="proof"))
assert "private-refresh" not in blob
decoded = codec.decode(blob)
assert decoded == OAuthToken("alice-token", identity_binding_proof="proof")
def test_malformed_bound_entries_fail_closed():
codec = _wrapping_codec()
for payload in ("not-json", "{}", '{"access_token":"at"}', '{"access_token":1,"identity_binding_proof":"p"}'):
assert codec.decode("enc:litellm-bound-oauth-v1:" + payload) is None

View file

@ -628,6 +628,7 @@ async def test_oauth_round_trip_returns_payload():
access_token,
refresh_token="rfr-xyz",
scopes=["a", "b"],
identity_binding_proof="verified-proof",
)
stored = _stored_value(prisma)
@ -642,6 +643,7 @@ async def test_oauth_round_trip_returns_payload():
assert result["access_token"] == access_token
assert result["refresh_token"] == "rfr-xyz"
assert result["scopes"] == ["a", "b"]
assert result["identity_binding_proof"] == "verified-proof"
@pytest.mark.asyncio
@ -1184,6 +1186,7 @@ class _RefreshResponse:
def _refresh_server(**overrides):
base = dict(
oauth_identity_binding=None,
token_url="https://idp.example.com/token",
server_id="srv-1",
client_id="cid",
@ -1309,7 +1312,7 @@ async def test_refresh_user_oauth_token_uses_client_secret_basic(monkeypatch):
sends HTTP Basic and keeps the secret out of the body."""
import litellm.proxy._experimental.mcp_server.db as db_mod
server = MagicMock()
server = MagicMock(oauth_identity_binding=None)
server.token_url = "https://idp.example.com/oauth2/token"
server.server_id = "srv"
server.client_id = "cid"
@ -1348,7 +1351,7 @@ async def test_refresh_user_oauth_token_defaults_to_client_secret_post(monkeypat
the body (client_secret_post) and sends no Authorization header."""
import litellm.proxy._experimental.mcp_server.db as db_mod
server = MagicMock()
server = MagicMock(oauth_identity_binding=None)
server.token_url = "https://idp.example.com/oauth2/token"
server.server_id = "srv"
server.client_id = "cid"
@ -1609,3 +1612,95 @@ def test_partial_update_defers_omitted_eligibility_fields_to_the_stored_row():
request = UpdateMCPServerRequest(server_id="relay-update", per_server_oauth_discovery=True)
assert request.per_server_oauth_discovery is True
@pytest.mark.asyncio
async def test_enforcement_rejects_preexisting_unverified_credential():
from litellm.types.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer
server = MCPServer(
server_id="srv-1", name="srv-1", url="https://mcp.example.com", transport=MCPTransport.http, auth_type=MCPAuth.oauth2,
oauth_identity_binding=MCPOAuthIdentityBinding(
mode="enforce", issuer="https://idp.example.com", audiences=["client"],
),
)
result = await resolve_valid_user_oauth_token(
user_id="alice", server=server,
cred={"access_token": "belongs-to-bob", "refresh_token": "bobs-refresh-token"},
)
assert result is None
@pytest.mark.asyncio
async def test_refresh_identity_rejection_returns_reauthentication_without_persisting(monkeypatch):
from fastapi import HTTPException
from litellm.proxy._experimental.mcp_server import db as module
validator = AsyncMock(side_effect=HTTPException(status_code=403, detail="oauth_principal_mismatch"))
monkeypatch.setattr(module, "enforce_oauth_identity_binding", validator)
result, captured = await _run_refresh(
monkeypatch, _refresh_server(), {"access_token": "bob", "refresh_token": "rotated"}
)
assert result is None
assert captured["data"]["grant_type"] == "refresh_token"
module.store_user_oauth_credential.assert_not_awaited()
@pytest.mark.asyncio
async def test_verified_legacy_cache_reads_avoid_database_and_reject_policy_changes(monkeypatch):
from litellm.caching.caching import DualCache
from litellm.proxy import proxy_server
from litellm.proxy._experimental.mcp_server import db as module
from litellm.proxy._experimental.mcp_server.oauth2_token_cache import mcp_per_user_token_cache
from litellm.proxy._experimental.mcp_server.oauth_identity_binding import current_binding_proof
from litellm.types.mcp_server.mcp_server_manager import MCPServer
server = MCPServer(
server_id="srv",
name="srv",
transport="http",
auth_type="oauth2",
oauth_identity_binding={
"mode": "enforce",
"issuer": "https://idp.example",
"audiences": ["client"],
"caller_field": "user_id",
"principal_claim": "sub",
},
)
proof = await current_binding_proof(server.oauth_identity_binding, "alice", "srv")
monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache())
read = AsyncMock(return_value=None)
monkeypatch.setattr(module, "get_user_oauth_credential", read)
monkeypatch.setattr(proxy_server, "prisma_client", MagicMock())
await mcp_per_user_token_cache.set("alice", "srv", "alice-token", 60, identity_binding_proof=proof)
assert await module.resolve_user_oauth_access_token("alice", server) == "alice-token"
assert await module.resolve_user_oauth_access_token("alice", server) == "alice-token"
read.assert_not_awaited()
server.oauth_identity_binding = server.oauth_identity_binding.model_copy(update={"audiences": ["changed"]})
assert await module.resolve_user_oauth_access_token("alice", server) is None
read.assert_awaited_once()
assert await mcp_per_user_token_cache.get_token("alice", "srv") is None
@pytest.mark.asyncio
async def test_unverified_legacy_cache_cannot_bypass_enforcement(monkeypatch):
from litellm.caching.caching import DualCache
from litellm.proxy import proxy_server
from litellm.proxy._experimental.mcp_server import db as module
from litellm.proxy._experimental.mcp_server.oauth2_token_cache import mcp_per_user_token_cache
from litellm.types.mcp_server.mcp_server_manager import MCPServer
server = MCPServer(
server_id="srv",
name="srv",
transport="http",
auth_type="oauth2",
oauth_identity_binding={"mode": "enforce", "issuer": "https://idp.example", "audiences": ["client"]},
)
monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache())
monkeypatch.setattr(proxy_server, "prisma_client", MagicMock())
monkeypatch.setattr(module, "get_user_oauth_credential", AsyncMock(return_value={"access_token": "bob"}))
await mcp_per_user_token_cache.set("alice", "srv", "bob", 60)
assert await module.resolve_user_oauth_access_token("alice", server) is None
assert await mcp_per_user_token_cache.get("alice", "srv") is None

View file

@ -8123,6 +8123,137 @@ async def test_token_exchange_pairs_client_secret_with_server_client_id():
assert "client_secret" not in sent
@pytest.mark.asyncio
async def test_token_exchange_refresh_passes_presented_refresh_ownership():
from fastapi import Request
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
exchange_token_with_server,
)
from litellm.proxy._experimental.mcp_server.oauth_identity_binding import RefreshTokenPresented
from litellm.proxy._types import MCPTransport
from litellm.types.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer
server = MCPServer(
server_id="srv-1",
name="srv-1",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
client_id="cid",
token_url="https://provider.example/token",
oauth_identity_binding=MCPOAuthIdentityBinding(
mode="enforce",
issuer="https://provider.example",
audiences=["cid"],
),
)
request = MagicMock(spec=Request)
request.base_url = "https://litellm.example.com/"
request.headers = {}
response = MagicMock()
response.json.return_value = {"access_token": "at"}
response.raise_for_status = MagicMock()
client = MagicMock()
client.post = AsyncMock(return_value=response)
enforce = AsyncMock()
with (
patch( # test-quality-ok: no injection seam exists for the exchange's HTTP and identity collaborators
"litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client",
return_value=client,
),
patch( # test-quality-ok: no injection seam exists for request identity extraction
"litellm.proxy._experimental.mcp_server.discoverable_endpoints._extract_user_id_from_request",
new=AsyncMock(return_value="user-a"),
),
patch( # test-quality-ok: captures the ownership value at the exchange boundary
"litellm.proxy._experimental.mcp_server.discoverable_endpoints.enforce_oauth_identity_binding",
new=enforce,
),
):
await exchange_token_with_server(
request=request,
mcp_server=server,
grant_type="refresh_token",
code=None,
redirect_uri=None,
client_id="cid",
client_secret=None,
code_verifier=None,
refresh_token="rt-1",
)
ownership = enforce.await_args.kwargs["refresh_ownership"]
assert isinstance(ownership, RefreshTokenPresented)
assert ownership.refresh_token == "rt-1"
@pytest.mark.asyncio
async def test_token_exchange_authorization_code_passes_no_refresh_ownership(monkeypatch):
from fastapi import Request
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
exchange_token_with_server,
seal_bridge_authorization_code,
)
from litellm.proxy._types import MCPTransport
from litellm.types.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer
monkeypatch.setenv("LITELLM_SALT_KEY", "identity-binding-test-salt")
server = MCPServer(
server_id="srv-1",
name="srv-1",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
client_id="cid",
token_url="https://provider.example/token",
oauth_identity_binding=MCPOAuthIdentityBinding(
mode="enforce",
issuer="https://provider.example",
audiences=["cid"],
),
)
request = MagicMock(spec=Request)
request.base_url = "https://litellm.example.com/"
request.headers = {}
response = MagicMock()
response.json.return_value = {"access_token": "at"}
response.raise_for_status = MagicMock()
client = MagicMock()
client.post = AsyncMock(return_value=response)
enforce = AsyncMock()
with (
patch( # test-quality-ok: no injection seam exists for the exchange's HTTP and identity collaborators
"litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client",
return_value=client,
),
patch( # test-quality-ok: no injection seam exists for request identity extraction
"litellm.proxy._experimental.mcp_server.discoverable_endpoints._extract_user_id_from_request",
new=AsyncMock(return_value="user-a"),
),
patch( # test-quality-ok: captures the ownership value at the exchange boundary
"litellm.proxy._experimental.mcp_server.discoverable_endpoints.enforce_oauth_identity_binding",
new=enforce,
),
):
result = await exchange_token_with_server(
request=request,
mcp_server=server,
grant_type="authorization_code",
code=seal_bridge_authorization_code("auth-code", "user-a", "srv-1", "login-nonce"),
redirect_uri="https://litellm.example.com/callback",
client_id="cid",
client_secret=None,
code_verifier="test-verifier",
)
assert result.status_code == 200
assert json.loads(result.body)["access_token"] == "at"
assert enforce.await_args.kwargs["refresh_ownership"] is None
assert enforce.await_args.kwargs["expected_nonce"] == "login-nonce"
def _upstream_token_response(status_code: int, *, json_body: object = None, text_body: str = "") -> "httpx.Response":
import httpx
@ -10916,3 +11047,97 @@ def test_introspect_route_answers_for_authenticated_caller(monkeypatch):
assert active.status_code == 200
assert active.json()["active"] is True
assert active.json()["sub"] == "u1"
@pytest.mark.asyncio
async def test_identity_bound_authorization_carries_nonce_and_caller_through_callback(monkeypatch):
from http.cookies import SimpleCookie
from urllib.parse import parse_qs, urlparse
from fastapi import Request
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
_oauth_state_cookie_name, authorize_with_server, callback, open_bridge_authorization_code,
)
from litellm.types.mcp import MCPTransport
from litellm.types.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer
monkeypatch.setenv("LITELLM_SALT_KEY", "identity-binding-test-salt")
server = MCPServer(
server_id="srv", name="srv", transport=MCPTransport.http, auth_type=MCPAuth.oauth2,
client_id="client", authorization_url="https://idp.example.com/authorize",
token_url="https://idp.example.com/token",
oauth_identity_binding=MCPOAuthIdentityBinding(
mode="enforce", issuer="https://idp.example.com", audiences=["client"],
),
)
request = Request({"type": "http", "scheme": "https", "server": ("proxy.example.com", 443),
"path": "/authorize", "query_string": b"", "headers": []})
with (
patch( # test-quality-ok: isolate authenticated request resolution from the real encrypted OAuth round trip
"litellm.proxy._experimental.mcp_server.discoverable_endpoints._extract_user_id_from_request",
new=AsyncMock(return_value="alice")),
patch( # test-quality-ok: isolate user access lookup while testing nonce and caller preservation
"litellm.proxy._experimental.mcp_server.discoverable_endpoints._bridge_authorize_access_denial",
new=AsyncMock(return_value=None)),
):
authorized = await authorize_with_server(
request, server, "client", "http://127.0.0.1:6274/callback", state="client-state",
code_challenge="pkce-challenge", code_challenge_method="S256",
)
query = parse_qs(urlparse(authorized.headers["location"]).query)
assert len(query["nonce"][0]) >= 32
cookies = SimpleCookie()
cookies.load(authorized.headers["set-cookie"])
name = _oauth_state_cookie_name(query["state"][0])
callback_request = Request({**request.scope, "path": "/callback",
"headers": [(b"cookie", f"{name}={cookies[name].value}".encode())]})
completed = await callback(callback_request, code="upstream-code", state=query["state"][0])
returned = parse_qs(urlparse(completed.headers["location"]).query)
sealed = open_bridge_authorization_code(returned["code"][0])
assert sealed.litellm_user_id == "alice"
assert sealed.mcp_server_id == "srv"
assert sealed.upstream_code == "upstream-code"
assert sealed.oauth_nonce == query["nonce"][0]
assert returned["state"] == ["client-state"]
@pytest.mark.asyncio
async def test_enforced_login_warms_verified_token_readable_without_database_lookup(monkeypatch):
from types import SimpleNamespace
from litellm.caching.caching import DualCache
from litellm.proxy import proxy_server
from litellm.proxy._experimental.mcp_server import db, mcp_server_manager
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import _store_per_user_token_server_side
from litellm.proxy._experimental.mcp_server.oauth2_token_cache import mcp_per_user_token_cache
from litellm.proxy._experimental.mcp_server.oauth_identity_binding import current_binding_proof
from litellm.proxy._experimental.mcp_server.outbound_credentials.dual_cache_token_backend import DualCacheTokenCacheBackend
from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import CachedOAuthTokenStore
from litellm.proxy._experimental.mcp_server.outbound_credentials.per_user_oauth_store import LazyPerUserOAuthTokenStore
from litellm.proxy._experimental.mcp_server.outbound_credentials.v2_token_store import V2PerUserTokenStore
from litellm.types.mcp_server.mcp_server_manager import MCPServer
server = MCPServer(
server_id="srv", name="srv", transport="http", auth_type="oauth2",
oauth_identity_binding={"mode": "enforce", "issuer": "https://idp.example", "audiences": ["client"],
"caller_field": "user_id", "principal_claim": "sub"},
)
proof = await current_binding_proof(server.oauth_identity_binding, "alice", "srv")
cache = DualCache()
monkeypatch.setenv("LITELLM_SALT_KEY", "test-cache-warm-encryption-salt")
monkeypatch.setattr(proxy_server, "user_api_key_cache", cache)
monkeypatch.setattr(proxy_server, "prisma_client", MagicMock())
monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", SimpleNamespace(invalidate_user_oauth_token_cache=AsyncMock()))
monkeypatch.setattr(db, "store_user_oauth_credential", AsyncMock())
await _store_per_user_token_server_side(
server=server, user_id="alice", token_response={"access_token": "alice-token", "refresh_token": "private", "expires_in": 3600},
identity_binding_proof=proof,
)
read = AsyncMock(return_value=None)
cached = CachedOAuthTokenStore(
V2PerUserTokenStore(read), default_ttl_seconds=300, backend=DualCacheTokenCacheBackend(cache, mcp_per_user_token_cache._codec()),
)
store = LazyPerUserOAuthTokenStore(lambda _: server, store_builder=lambda _: (cached, True), redis_available=lambda: True)
token = await store.fetch("alice", "srv")
assert token is not None and token.access_token == "alice-token"
assert token.identity_binding_proof == proof
assert token.refresh_token is None
read.assert_not_awaited()

View file

@ -0,0 +1,667 @@
import time
from collections.abc import Mapping
from types import SimpleNamespace
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import jwt
import pytest
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from fastapi import HTTPException
from pydantic import ValidationError
from litellm.proxy._experimental.mcp_server.oauth_identity_binding import (
RefreshOwnershipProven,
RefreshTokenPresented,
VerifiedRefreshToken,
_discover_jwks_url,
_fetch_issuer_jwks,
_load_caller_principal,
_load_stored_refresh_token,
_select_signing_key,
current_binding_proof,
enforce_oauth_identity_binding,
)
from litellm.types.mcp import MCPAuth, MCPTransport
from litellm.types.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer
ISSUER: Final = "https://idp.example.com"
AUDIENCE: Final = "litellm-client"
KID: Final = "test-key"
_PRIVATE_KEY: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048)
_PRIVATE_PEM: Final = _PRIVATE_KEY.private_bytes(
encoding=serialization.Encoding.PEM,
format=serialization.PrivateFormat.PKCS8,
encryption_algorithm=serialization.NoEncryption(),
)
_PUBLIC_JWK: Final = {
**jwt.algorithms.RSAAlgorithm.to_jwk(_PRIVATE_KEY.public_key(), as_dict=True),
"kid": KID,
"alg": "RS256",
"use": "sig",
}
def _sign_id_token(claims: Mapping[str, object]) -> str:
payload: Final = {
"iss": ISSUER,
"aud": AUDIENCE,
"exp": int(time.time()) + 300,
"iat": int(time.time()),
"nonce": "test-nonce",
"sub": "upstream-user",
**claims,
}
return jwt.encode(payload, _PRIVATE_PEM, algorithm="RS256", headers={"kid": KID})
async def _jwks_fetcher(_binding: MCPOAuthIdentityBinding) -> list[Mapping[str, object]]:
return [_PUBLIC_JWK]
def _caller_loader(email: str | None):
async def load(_user_id: str, _binding: MCPOAuthIdentityBinding) -> str | None:
return email
return load
def _stored_refresh_token_loader(refresh_token: str | None):
async def load(_user_id: str, _server_id: str, _binding: MCPOAuthIdentityBinding) -> VerifiedRefreshToken | None:
return VerifiedRefreshToken(refresh_token, "verified-binding") if refresh_token else None
return load
def _server(mode: str = "enforce", **binding_overrides: object) -> MCPServer:
return MCPServer(
server_id="srv-1",
name="srv-1",
url="https://mcp.example.com",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
oauth_identity_binding=MCPOAuthIdentityBinding(
mode=mode,
issuer=ISSUER,
audiences=[AUDIENCE],
**binding_overrides,
),
)
@pytest.mark.asyncio
async def test_fetch_issuer_jwks_fetches_and_caches_keys():
binding: Final = _server(jwks_url="https://idp.example.com/jwks-cache").oauth_identity_binding
assert binding is not None
response: Final = MagicMock()
response.json.return_value = {"keys": [_PUBLIC_JWK]}
client: Final = MagicMock()
client.get = AsyncMock(return_value=response)
with patch( # test-quality-ok: no HTTP dependency injection seam exists for JWKS fetching
"litellm.proxy._experimental.mcp_server.oauth_identity_binding.get_async_httpx_client",
return_value=client,
):
first: Final = await _fetch_issuer_jwks(binding)
second: Final = await _fetch_issuer_jwks(binding)
assert first == second == [_PUBLIC_JWK]
client.get.assert_awaited_once_with("https://idp.example.com/jwks-cache")
@pytest.mark.asyncio
async def test_fetch_issuer_jwks_rejects_malformed_document():
binding: Final = _server(jwks_url="https://idp.example.com/jwks-invalid").oauth_identity_binding
assert binding is not None
response: Final = MagicMock()
response.json.return_value = {}
client: Final = MagicMock()
client.get = AsyncMock(return_value=response)
with patch( # test-quality-ok: no HTTP dependency injection seam exists for JWKS fetching
"litellm.proxy._experimental.mcp_server.oauth_identity_binding.get_async_httpx_client",
return_value=client,
):
with pytest.raises(TypeError, match="has no 'keys' array"):
await _fetch_issuer_jwks(binding)
@pytest.mark.asyncio
async def test_discover_jwks_url_returns_provider_uri():
response: Final = MagicMock()
response.json.return_value = {"jwks_uri": "https://idp.example.com/jwks"}
client: Final = MagicMock()
client.get = AsyncMock(return_value=response)
with patch( # test-quality-ok: no HTTP dependency injection seam exists for OIDC discovery
"litellm.proxy._experimental.mcp_server.oauth_identity_binding.get_async_httpx_client",
return_value=client,
):
result: Final = await _discover_jwks_url("https://idp.example.com/")
assert result == "https://idp.example.com/jwks"
client.get.assert_awaited_once_with("https://idp.example.com/.well-known/openid-configuration")
@pytest.mark.asyncio
async def test_discover_jwks_url_rejects_missing_provider_uri():
response: Final = MagicMock()
response.json.return_value = {}
client: Final = MagicMock()
client.get = AsyncMock(return_value=response)
with patch( # test-quality-ok: no HTTP dependency injection seam exists for OIDC discovery
"litellm.proxy._experimental.mcp_server.oauth_identity_binding.get_async_httpx_client",
return_value=client,
):
with pytest.raises(ValueError, match="returned no jwks_uri"):
await _discover_jwks_url("https://idp.example.com")
def test_select_signing_key_returns_matching_key_or_rejection():
token: Final = _sign_id_token({"email": "alice@example.com", "email_verified": True})
selected: Final = _select_signing_key(token, [_PUBLIC_JWK])
rejected: Final = _select_signing_key(token, [])
assert selected.__class__.__name__ == "PyJWK"
assert rejected.code == "oauth_identity_binding_failed"
@pytest.mark.asyncio
async def test_load_caller_principal_supports_user_id_and_database_email():
user_id_binding: Final = _server(caller_field="user_id").oauth_identity_binding
assert user_id_binding is not None
assert await _load_caller_principal("user-a", user_id_binding) == "user-a"
with patch( # test-quality-ok: caller loading is a lazy database boundary without injection
"litellm.proxy._experimental.mcp_server.bridge_token_flow.load_active_user_by_id",
new=AsyncMock(side_effect=["no_active_key", SimpleNamespace(user_email="alice@example.com")]),
):
assert await _load_caller_principal("user-a", _server().oauth_identity_binding) is None
assert await _load_caller_principal("user-a", _server().oauth_identity_binding) == "alice@example.com"
@pytest.mark.asyncio
async def test_load_stored_refresh_token_returns_credential_and_fails_closed():
binding: Final = _server(caller_field="user_id", principal_claim="sub").oauth_identity_binding
proof: Final = await current_binding_proof(binding, "user-a", "srv-1")
get_credential: Final = AsyncMock(return_value={"refresh_token": "rt-1", "identity_binding_proof": proof})
with (
patch( # test-quality-ok: stored-token loading is a lazy database boundary without injection
"litellm.proxy._experimental.mcp_server.db.get_user_oauth_credential",
new=get_credential,
),
patch( # test-quality-ok: stored-token loading is a lazy database boundary without injection
"litellm.proxy.utils.get_prisma_client_or_throw",
return_value="prisma",
),
):
assert await _load_stored_refresh_token("user-a", "srv-1", binding) == VerifiedRefreshToken("rt-1", proof)
get_credential.return_value = {"refresh_token": "rt-1"}
assert await _load_stored_refresh_token("user-a", "srv-1", binding) is None
with patch( # test-quality-ok: stored-token loading is a lazy database boundary without injection
"litellm.proxy.utils.get_prisma_client_or_throw",
side_effect=RuntimeError("database unavailable"),
):
assert await _load_stored_refresh_token("user-a", "srv-1", binding) is None
@pytest.mark.asyncio
async def test_matching_principal_passes():
token: Final = _sign_id_token({"email": "Alice@Example.com", "email_verified": True})
result: Final = await enforce_oauth_identity_binding(
server=_server(),
token_response={"access_token": "at", "id_token": token},
litellm_user_id="user-a",
grant_type="authorization_code",
expected_nonce="test-nonce",
refresh_ownership=None,
jwks_fetcher=_jwks_fetcher,
caller_principal_loader=_caller_loader("alice@example.com"),
)
assert result is not None
@pytest.mark.asyncio
async def test_mismatched_principal_rejected():
token: Final = _sign_id_token({"email": "mallory@example.com", "email_verified": True})
with pytest.raises(HTTPException) as exc_info:
await enforce_oauth_identity_binding(
server=_server(),
token_response={"access_token": "at", "id_token": token},
litellm_user_id="user-a",
grant_type="authorization_code",
expected_nonce="test-nonce",
refresh_ownership=None,
jwks_fetcher=_jwks_fetcher,
caller_principal_loader=_caller_loader("alice@example.com"),
)
assert exc_info.value.status_code == 403
assert exc_info.value.detail["error"] == "oauth_principal_mismatch"
assert exc_info.value.detail["credential_stored"] is False
@pytest.mark.asyncio
async def test_missing_upstream_principal_rejected():
token: Final = _sign_id_token({"email_verified": True})
with pytest.raises(HTTPException) as exc_info:
await enforce_oauth_identity_binding(
server=_server(),
token_response={"access_token": "at", "id_token": token},
litellm_user_id="user-a",
grant_type="authorization_code",
expected_nonce="test-nonce",
refresh_ownership=None,
jwks_fetcher=_jwks_fetcher,
caller_principal_loader=_caller_loader("alice@example.com"),
)
assert exc_info.value.detail["error"] == "oauth_identity_binding_failed"
assert "no usable" in exc_info.value.detail["error_description"]
@pytest.mark.asyncio
async def test_jwks_fetch_failure_is_rejected():
async def fail(_binding: MCPOAuthIdentityBinding) -> list[Mapping[str, object]]:
raise RuntimeError("jwks unavailable")
token: Final = _sign_id_token({"email": "alice@example.com", "email_verified": True})
with pytest.raises(HTTPException) as exc_info:
await enforce_oauth_identity_binding(
server=_server(),
token_response={"access_token": "at", "id_token": token},
litellm_user_id="user-a",
grant_type="authorization_code",
expected_nonce="test-nonce",
refresh_ownership=None,
jwks_fetcher=fail,
caller_principal_loader=_caller_loader("alice@example.com"),
)
assert exc_info.value.detail["error"] == "oauth_identity_binding_failed"
assert "jwks unavailable" in exc_info.value.detail["error_description"]
@pytest.mark.asyncio
async def test_missing_signing_key_is_rejected():
token: Final = _sign_id_token({"email": "alice@example.com", "email_verified": True})
async def no_keys(_binding: MCPOAuthIdentityBinding) -> tuple[Mapping[str, object], ...]:
return ()
with pytest.raises(HTTPException) as exc_info:
await enforce_oauth_identity_binding(
server=_server(),
token_response={"access_token": "at", "id_token": token},
litellm_user_id="user-a",
grant_type="authorization_code",
expected_nonce="test-nonce",
refresh_ownership=None,
jwks_fetcher=no_keys,
caller_principal_loader=_caller_loader("alice@example.com"),
)
assert exc_info.value.detail["error"] == "oauth_identity_binding_failed"
assert "signing key" in exc_info.value.detail["error_description"]
@pytest.mark.asyncio
async def test_missing_caller_principal_is_rejected():
token: Final = _sign_id_token({"email": "alice@example.com", "email_verified": True})
with pytest.raises(HTTPException) as exc_info:
await enforce_oauth_identity_binding(
server=_server(),
token_response={"access_token": "at", "id_token": token},
litellm_user_id="user-a",
grant_type="authorization_code",
expected_nonce="test-nonce",
refresh_ownership=None,
jwks_fetcher=_jwks_fetcher,
caller_principal_loader=_caller_loader(None),
)
assert exc_info.value.detail["error"] == "oauth_identity_binding_failed"
assert "has no" in exc_info.value.detail["error_description"]
@pytest.mark.asyncio
async def test_refresh_without_id_token_requires_litellm_identity_for_presented_token():
with pytest.raises(HTTPException) as exc_info:
await enforce_oauth_identity_binding(
server=_server(),
token_response={"access_token": "at"},
litellm_user_id=None,
grant_type="refresh_token",
refresh_ownership=RefreshTokenPresented("rt-1"),
jwks_fetcher=_jwks_fetcher,
caller_principal_loader=_caller_loader("alice@example.com"),
stored_refresh_token_loader=_stored_refresh_token_loader("rt-1"),
)
assert exc_info.value.detail["error"] == "oauth_identity_binding_failed"
assert "no resolvable LiteLLM user identity" in exc_info.value.detail["error_description"]
@pytest.mark.asyncio
async def test_user_id_principal_matching_uses_exact_comparison():
token: Final = _sign_id_token({"sub": "user-a"})
result: Final = await enforce_oauth_identity_binding(
server=_server(principal_claim="sub", caller_field="user_id"),
token_response={"access_token": "at", "id_token": token},
litellm_user_id="user-a",
grant_type="authorization_code",
expected_nonce="test-nonce",
refresh_ownership=None,
jwks_fetcher=_jwks_fetcher,
caller_principal_loader=_caller_loader("user-a"),
)
assert result is not None
@pytest.mark.asyncio
async def test_missing_id_token_rejected_on_authorization_code():
with pytest.raises(HTTPException) as exc_info:
await enforce_oauth_identity_binding(
server=_server(),
token_response={"access_token": "at"},
litellm_user_id="user-a",
grant_type="authorization_code",
expected_nonce="test-nonce",
refresh_ownership=None,
jwks_fetcher=_jwks_fetcher,
caller_principal_loader=_caller_loader("alice@example.com"),
)
assert exc_info.value.status_code == 403
assert exc_info.value.detail["error"] == "oauth_identity_binding_failed"
@pytest.mark.asyncio
async def test_refresh_without_id_token_allowed_when_presented_token_matches_stored_credential():
result: Final = await enforce_oauth_identity_binding(
server=_server(),
token_response={"access_token": "at"},
litellm_user_id="user-a",
grant_type="refresh_token",
refresh_ownership=RefreshTokenPresented("rt-1"),
jwks_fetcher=_jwks_fetcher,
caller_principal_loader=_caller_loader("alice@example.com"),
stored_refresh_token_loader=_stored_refresh_token_loader("rt-1"),
)
assert result is not None
@pytest.mark.asyncio
async def test_refresh_without_id_token_rejects_different_presented_token():
with pytest.raises(HTTPException) as exc_info:
await enforce_oauth_identity_binding(
server=_server(),
token_response={"access_token": "at"},
litellm_user_id="user-a",
grant_type="refresh_token",
refresh_ownership=RefreshTokenPresented("rt-stolen"),
jwks_fetcher=_jwks_fetcher,
caller_principal_loader=_caller_loader("alice@example.com"),
stored_refresh_token_loader=_stored_refresh_token_loader("rt-1"),
)
assert exc_info.value.status_code == 403
assert exc_info.value.detail["error"] == "oauth_identity_binding_failed"
@pytest.mark.asyncio
async def test_refresh_without_id_token_rejects_missing_stored_token():
with pytest.raises(HTTPException) as exc_info:
await enforce_oauth_identity_binding(
server=_server(),
token_response={"access_token": "at"},
litellm_user_id="user-a",
grant_type="refresh_token",
refresh_ownership=RefreshTokenPresented("rt-1"),
jwks_fetcher=_jwks_fetcher,
caller_principal_loader=_caller_loader("alice@example.com"),
stored_refresh_token_loader=_stored_refresh_token_loader(None),
)
assert exc_info.value.status_code == 403
assert exc_info.value.detail["error"] == "oauth_identity_binding_failed"
@pytest.mark.asyncio
async def test_identity_envelope_does_not_prove_upstream_binding():
with pytest.raises(HTTPException) as error:
await enforce_oauth_identity_binding(
server=_server(),
token_response={"access_token": "at"},
litellm_user_id="user-a",
grant_type="refresh_token",
refresh_ownership=RefreshOwnershipProven(),
jwks_fetcher=_jwks_fetcher,
caller_principal_loader=_caller_loader("alice@example.com"),
)
assert error.value.status_code == 403
@pytest.mark.asyncio
async def test_refresh_without_id_token_rejects_without_ownership_proof():
with pytest.raises(HTTPException) as exc_info:
await enforce_oauth_identity_binding(
server=_server(),
token_response={"access_token": "at"},
litellm_user_id="user-a",
grant_type="refresh_token",
refresh_ownership=None,
jwks_fetcher=_jwks_fetcher,
caller_principal_loader=_caller_loader("alice@example.com"),
stored_refresh_token_loader=_stored_refresh_token_loader("rt-1"),
)
assert exc_info.value.status_code == 403
assert exc_info.value.detail["error"] == "oauth_identity_binding_failed"
@pytest.mark.asyncio
async def test_refresh_with_mismatched_id_token_rejected():
token: Final = _sign_id_token({"email": "mallory@example.com", "email_verified": True})
with pytest.raises(HTTPException):
await enforce_oauth_identity_binding(
server=_server(),
token_response={"access_token": "at", "id_token": token},
litellm_user_id="user-a",
grant_type="refresh_token",
refresh_ownership=None,
jwks_fetcher=_jwks_fetcher,
caller_principal_loader=_caller_loader("alice@example.com"),
)
@pytest.mark.asyncio
async def test_audit_mode_logs_but_does_not_reject(caplog):
token: Final = _sign_id_token({"email": "mallory@example.com", "email_verified": True})
result: Final = await enforce_oauth_identity_binding(
server=_server(mode="audit"),
token_response={"access_token": "at", "id_token": token},
litellm_user_id="user-a",
grant_type="authorization_code",
refresh_ownership=None,
jwks_fetcher=_jwks_fetcher,
caller_principal_loader=_caller_loader("alice@example.com"),
)
assert result is None
assert "oauth_principal_mismatch" in caplog.text
assert "nonce" not in caplog.text
@pytest.mark.asyncio
async def test_unverified_email_rejected():
token: Final = _sign_id_token({"email": "alice@example.com", "email_verified": False})
with pytest.raises(HTTPException) as exc_info:
await enforce_oauth_identity_binding(
server=_server(),
token_response={"access_token": "at", "id_token": token},
litellm_user_id="user-a",
grant_type="authorization_code",
expected_nonce="test-nonce",
refresh_ownership=None,
jwks_fetcher=_jwks_fetcher,
caller_principal_loader=_caller_loader("alice@example.com"),
)
assert exc_info.value.detail["error"] == "oauth_identity_binding_failed"
@pytest.mark.asyncio
async def test_wrong_issuer_rejected():
payload: Final = {
"iss": "https://evil.example.com",
"aud": AUDIENCE,
"exp": int(time.time()) + 300,
"email": "alice@example.com",
"email_verified": True,
}
token: Final = jwt.encode(payload, _PRIVATE_PEM, algorithm="RS256", headers={"kid": KID})
with pytest.raises(HTTPException) as exc_info:
await enforce_oauth_identity_binding(
server=_server(),
token_response={"access_token": "at", "id_token": token},
litellm_user_id="user-a",
grant_type="authorization_code",
expected_nonce="test-nonce",
refresh_ownership=None,
jwks_fetcher=_jwks_fetcher,
caller_principal_loader=_caller_loader("alice@example.com"),
)
assert exc_info.value.detail["error"] == "oauth_identity_binding_failed"
@pytest.mark.asyncio
async def test_no_litellm_identity_rejected():
token: Final = _sign_id_token({"email": "alice@example.com", "email_verified": True})
with pytest.raises(HTTPException) as exc_info:
await enforce_oauth_identity_binding(
server=_server(),
token_response={"access_token": "at", "id_token": token},
litellm_user_id=None,
grant_type="authorization_code",
expected_nonce="test-nonce",
refresh_ownership=None,
jwks_fetcher=_jwks_fetcher,
caller_principal_loader=_caller_loader("alice@example.com"),
)
assert exc_info.value.status_code == 403
@pytest.mark.asyncio
async def test_disabled_binding_is_noop():
result: Final = await enforce_oauth_identity_binding(
server=_server(mode="disabled"),
token_response={"access_token": "at"},
litellm_user_id=None,
grant_type="authorization_code",
expected_nonce="test-nonce",
refresh_ownership=None,
jwks_fetcher=_jwks_fetcher,
caller_principal_loader=_caller_loader(None),
)
assert result is None
@pytest.mark.asyncio
async def test_no_binding_is_noop():
server: Final = MCPServer(
server_id="srv-2",
name="srv-2",
url="https://mcp.example.com",
transport=MCPTransport.http,
)
result: Final = await enforce_oauth_identity_binding(
server=server,
token_response={"access_token": "at"},
litellm_user_id=None,
grant_type="authorization_code",
expected_nonce="test-nonce",
refresh_ownership=None,
jwks_fetcher=_jwks_fetcher,
caller_principal_loader=_caller_loader(None),
)
assert result is None
def test_identity_binding_requires_non_empty_audiences():
with pytest.raises(ValidationError):
MCPOAuthIdentityBinding(mode="enforce", issuer=ISSUER, audiences=[])
with pytest.raises(ValidationError):
MCPOAuthIdentityBinding(mode="enforce", issuer=ISSUER)
@pytest.mark.asyncio
async def test_wrong_audience_rejected():
token: Final = _sign_id_token({"aud": "other-client", "email": "alice@example.com", "email_verified": True})
with pytest.raises(HTTPException) as exc_info:
await enforce_oauth_identity_binding(
server=_server(),
token_response={"access_token": "at", "id_token": token},
litellm_user_id="user-a",
grant_type="authorization_code",
expected_nonce="test-nonce",
refresh_ownership=None,
jwks_fetcher=_jwks_fetcher,
caller_principal_loader=_caller_loader("alice@example.com"),
)
assert exc_info.value.status_code == 403
assert exc_info.value.detail["error"] == "oauth_identity_binding_failed"
@pytest.mark.asyncio
@pytest.mark.parametrize("nonce", [None, "another-login"])
async def test_authorization_code_rejects_missing_or_foreign_nonce(nonce):
token = _sign_id_token({"email": "alice@example.com", "email_verified": True, "nonce": nonce})
with pytest.raises(HTTPException) as error:
await enforce_oauth_identity_binding(
server=_server(),
token_response={"access_token": "at", "id_token": token},
litellm_user_id="user-a",
grant_type="authorization_code",
refresh_ownership=None,
expected_nonce="this-login",
jwks_fetcher=_jwks_fetcher,
caller_principal_loader=_caller_loader("alice@example.com"),
)
assert error.value.status_code == 403
assert error.value.detail["error"] == "oauth_identity_binding_failed"
@pytest.mark.asyncio
async def test_binding_proof_rejects_changed_user_or_policy():
from litellm.proxy._experimental.mcp_server.oauth_identity_binding import credential_binding_matches
binding = _server(caller_field="user_id", principal_claim="sub").oauth_identity_binding
proof = await current_binding_proof(binding, "alice", "srv-1")
credential = {"identity_binding_proof": proof}
assert await credential_binding_matches(binding, "alice", "srv-1", credential)
assert not await credential_binding_matches(binding, "bob", "srv-1", credential)
assert not await credential_binding_matches(binding, "alice", "other-server", credential)
changed = binding.model_copy(update={"audiences": ["different-client"]})
assert not await credential_binding_matches(changed, "alice", "srv-1", credential)
@pytest.mark.parametrize("auth_type", [MCPAuth.oauth_delegate, MCPAuth.true_passthrough])
def test_identity_binding_rejects_modes_without_gateway_credential_custody(auth_type):
with pytest.raises(ValidationError, match="gateway-managed per-user"):
MCPServer(
server_id="srv",
name="srv",
transport=MCPTransport.http,
auth_type=auth_type,
oauth_identity_binding=_server().oauth_identity_binding,
)
@pytest.mark.asyncio
async def test_audit_matching_login_without_nonce_does_not_report_failure(caplog):
token: Final = _sign_id_token({"email": "alice@example.com", "email_verified": True})
result: Final = await enforce_oauth_identity_binding(
server=_server(mode="audit"),
token_response={"access_token": "at", "id_token": token},
litellm_user_id="user-a",
grant_type="authorization_code",
refresh_ownership=None,
jwks_fetcher=_jwks_fetcher,
caller_principal_loader=_caller_loader("alice@example.com"),
)
assert result is None
assert "oauth_identity_binding audit" not in caplog.text

View file

@ -481,7 +481,7 @@ async def test_archive_source_registers_and_is_served_verbatim_in_marketplace():
assert response.action == "created"
assert response.plugin.source == _ARCHIVE_SOURCE
marketplace = json.loads((await get_marketplace()).body)
marketplace = json.loads((await get_marketplace(request=MagicMock())).body)
assert marketplace["plugins"] == [{"name": "s3-skill", "source": _ARCHIVE_SOURCE, "version": "1.0.0"}]

View file

@ -2033,6 +2033,82 @@ def test_proxy_admin_viewer_can_access_logs_page_endpoints(route):
)
@pytest.mark.parametrize(
"user_role",
[LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY],
)
def test_internal_user_can_access_logs_drawer_detail_route(user_role):
"""
The Logs drawer detail fetch (GET /spend/logs/ui/{request_id}) must pass
route_checks for plain internal users, not just admins the handler
itself already self-authorizes row ownership via
_assert_user_can_view_request_id.
"""
route = "/spend/logs/ui/abc-request-id"
user_obj = LiteLLM_UserTable(
user_id="internal_user",
user_email="user@example.com",
user_role=user_role.value,
)
valid_token = UserAPIKeyAuth(
user_id="internal_user",
user_role=user_role.value,
)
request = MagicMock(spec=Request)
request.query_params = {}
try:
RouteChecks.non_proxy_admin_allowed_routes_check(
user_obj=user_obj,
_user_role=user_role.value,
route=route,
request=request,
valid_token=valid_token,
request_data={},
)
except Exception as e:
pytest.fail(f"{user_role.value} should be able to access {route}. Got error: {str(e)}")
@pytest.mark.parametrize(
"route_group_name",
[
"spend_tracking_routes",
"internal_user_routes",
"internal_user_view_only_routes",
"admin_viewer_routes",
"org_admin_allowed_routes",
],
)
def test_logs_drawer_detail_route_in_every_route_group(route_group_name):
"""
/spend/logs/ui/{request_id} must be reachable through
RouteChecks.check_route_access under each role's own route group, so a
partial revert (removing the route from `spend_tracking_routes` while
leaving `non_proxy_admin_allowed_routes_check` alone) is also caught.
"""
from litellm.proxy._types import LiteLLMRoutes
allowed_routes = getattr(LiteLLMRoutes, route_group_name).value
assert RouteChecks.check_route_access(
route="/spend/logs/ui/req-34099", allowed_routes=allowed_routes
)
def test_logs_drawer_detail_route_allowed_for_scoped_virtual_key():
"""
A virtual key scoped to `allowed_routes=["spend_tracking_routes"]` must be
able to reach the Logs drawer detail route.
"""
valid_token = UserAPIKeyAuth(
user_id="scoped_key_user",
allowed_routes=["spend_tracking_routes"],
)
assert RouteChecks.is_virtual_key_allowed_to_call_route(
route="/spend/logs/ui/req-34099", valid_token=valid_token
)
@pytest.mark.parametrize("route", ADMIN_VIEWER_LOGS_PAGE_ROUTES)
def test_internal_user_blocked_from_admin_viewer_logs_routes(route):
"""

View file

@ -0,0 +1,111 @@
from datetime import datetime, timezone
from unittest.mock import AsyncMock, MagicMock
import pytest
from litellm.proxy.db.health_check_latest import (
LATEST_HEALTH_CHECKS_FOR_MODELS_SQL,
LATEST_HEALTH_CHECKS_SQL,
fetch_latest_health_checks,
fetch_latest_health_checks_for_models,
)
def _prisma(rows):
prisma = MagicMock()
prisma.db.query_raw = AsyncMock(return_value=rows)
return prisma
def _raw_row(**overrides):
row = {
"health_check_id": "hc-1",
"model_name": "gpt-4",
"model_id": "deployment-abc",
"status": "healthy",
"healthy_count": 1,
"unhealthy_count": 0,
"error_message": None,
"response_time_ms": 12.5,
"details": None,
"checked_by": "pod-1",
"checked_at": "2026-08-25T00:00:00+00:00",
"created_at": "2026-08-25T00:00:00+00:00",
"updated_at": "2026-08-25T00:00:00+00:00",
}
return {**row, **overrides}
@pytest.mark.asyncio
async def test_fetch_all_runs_one_distinct_on_query_with_no_parameters():
"""The dedup must be in the SQL: prisma find_many(distinct=...) streams the whole history table."""
prisma = _prisma([])
assert await fetch_latest_health_checks(prisma) == ()
assert prisma.db.query_raw.await_args.args == (LATEST_HEALTH_CHECKS_SQL,)
assert 'DISTINCT ON ("model_id", "model_name")' in LATEST_HEALTH_CHECKS_SQL
assert '"checked_at" DESC' in LATEST_HEALTH_CHECKS_SQL
@pytest.mark.asyncio
async def test_raw_datetimes_come_back_tz_aware_with_or_without_an_offset():
"""The save path subtracts checked_at from datetime.now(timezone.utc); a naive value would TypeError."""
naive = _raw_row(health_check_id="hc-naive", model_id=None, checked_at="2026-08-25T00:00:00")
aware = _raw_row(health_check_id="hc-aware", checked_at="2026-08-25T01:00:00+02:00")
rows = await fetch_latest_health_checks(_prisma([naive, aware]))
assert {row.health_check_id: (row.model_id, row.checked_at) for row in rows} == {
"hc-naive": (None, datetime(2026, 8, 25, 0, 0, tzinfo=timezone.utc)),
"hc-aware": ("deployment-abc", datetime(2026, 8, 24, 23, 0, tzinfo=timezone.utc)),
}
@pytest.mark.asyncio
async def test_json_details_decode_from_text_and_pass_through_as_dict():
rows = await fetch_latest_health_checks(
_prisma(
[
_raw_row(health_check_id="text", details='{"region": "eu"}'),
_raw_row(health_check_id="dict", details={"region": "us"}),
_raw_row(health_check_id="none", details=None),
]
)
)
assert {row.health_check_id: row.details for row in rows} == {
"text": {"region": "eu"},
"dict": {"region": "us"},
"none": None,
}
@pytest.mark.asyncio
async def test_fetch_all_degrades_to_no_rows_when_the_query_fails():
prisma = _prisma([])
prisma.db.query_raw.side_effect = RuntimeError("db down")
assert await fetch_latest_health_checks(prisma) == ()
@pytest.mark.asyncio
async def test_fetch_all_degrades_to_no_rows_for_a_malformed_row():
assert await fetch_latest_health_checks(_prisma([{"unexpected": "shape"}])) == ()
@pytest.mark.asyncio
async def test_fetch_for_models_binds_the_page_as_the_only_parameter():
prisma = _prisma([_raw_row()])
rows = await fetch_latest_health_checks_for_models(prisma, ("gpt-4", "claude-opus"))
assert prisma.db.query_raw.await_args.args == (LATEST_HEALTH_CHECKS_FOR_MODELS_SQL, ["gpt-4", "claude-opus"])
assert [row.model_name for row in rows] == ["gpt-4"]
assert 'WHERE "model_name" = ANY($1)' in LATEST_HEALTH_CHECKS_FOR_MODELS_SQL
@pytest.mark.asyncio
async def test_fetch_for_models_skips_the_database_for_an_empty_page():
prisma = _prisma([])
assert await fetch_latest_health_checks_for_models(prisma, ()) == ()
prisma.db.query_raw.assert_not_awaited()
@pytest.mark.asyncio
async def test_fetch_for_models_degrades_to_no_rows_when_the_query_fails():
prisma = _prisma([])
prisma.db.query_raw.side_effect = RuntimeError("db down")
assert await fetch_latest_health_checks_for_models(prisma, ("gpt-4",)) == ()

View file

@ -0,0 +1,624 @@
import configparser
import logging
import os
import signal
import socket
import stat
import sys
import tempfile
import textwrap
import time
import urllib.parse
from collections.abc import Callable
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Final, cast
import pytest
from litellm._logging import verbose_proxy_logger
from litellm.proxy.db.pgbouncer import (
PgBouncerError,
PgBouncerPlan,
PgBouncerProcess,
PgBouncerSettings,
pgbouncer_version,
plan_pgbouncer,
start_in_container_pgbouncer,
unix_socket_path,
write_pgbouncer_files,
)
UPSTREAM: Final = (
"postgresql://app:p%40ss%27w@db.internal:5433/litellm"
"?schema=public&connection_limit=10&pool_timeout=20"
"&sslmode=require&sslaccept=strict&sslcert=/certs/ca.pem"
"&options=-c%20statement_timeout%3D7000%20-c%20lock_timeout%3D3000"
)
SETTINGS: Final = PgBouncerSettings(enabled=True, port=6543, max_db_connections=8, max_client_conn=400)
def _plan(url: str = UPSTREAM, run_as_user: str | None = None) -> PgBouncerPlan:
plan: Final = plan_pgbouncer(url, SETTINGS, Path("/run/pgb"), run_as_user)
assert isinstance(plan, PgBouncerPlan), plan
return plan
def _ini(plan: PgBouncerPlan) -> configparser.ConfigParser:
parser: Final = configparser.ConfigParser(interpolation=None)
parser.read_string(plan.ini)
return parser
def _query(url: str) -> dict[str, str]:
return dict(urllib.parse.parse_qsl(urllib.parse.urlsplit(url).query, keep_blank_values=True))
class TestPlanPgBouncer:
def test_upstream_credentials_and_timeouts_move_into_the_pgbouncer_config(self):
ini: Final = _ini(_plan())
assert ini["databases"]["litellm"] == (
"host='db.internal' port=5433 dbname='litellm' user='app' password='p@ss''w' "
"connect_query='SET statement_timeout TO ''7000''; SET lock_timeout TO ''3000'''"
)
assert _plan().userlist == '"app" "p@ss\'w"\n'
def test_an_upstream_without_a_port_is_reached_on_the_postgres_default(self):
ini: Final = _ini(_plan("postgresql://app:pw@db/litellm"))
assert ini["databases"]["litellm"] == "host='db' port=5432 dbname='litellm' user='app' password='pw'"
def test_pool_is_sized_from_settings_in_transaction_mode(self):
pgb: Final = _ini(_plan())["pgbouncer"]
assert pgb["pool_mode"] == "transaction"
assert pgb["max_db_connections"] == "8"
assert pgb["default_pool_size"] == "8"
assert pgb["max_client_conn"] == "400"
assert pgb["auth_type"] == "scram-sha-256"
assert pgb["listen_addr"] == "127.0.0.1"
assert pgb["listen_port"] == "6543"
assert pgb["auth_file"] == "/run/pgb/userlist.txt"
assert pgb["unix_socket_dir"] == "/run/pgb"
def test_the_app_user_can_read_the_pgbouncer_console(self):
assert _ini(_plan())["pgbouncer"]["stats_users"] == "app"
@pytest.mark.parametrize("user", ["app,admin", "app%20admin", "app%09admin"])
def test_a_user_pgbouncer_would_split_into_several_console_users_is_refused(self, user: str):
outcome: Final = plan_pgbouncer(f"postgresql://{user}:pw@db/litellm", SETTINGS, Path("/run/pgb"), None)
assert isinstance(outcome, PgBouncerError)
assert "stats_users" in outcome.reason
def test_pooled_url_points_prisma_at_loopback_without_prepared_statements(self):
pooled: Final = urllib.parse.urlsplit(_plan().pooled_url)
assert (pooled.hostname, pooled.port, pooled.path) == ("127.0.0.1", 6543, "/litellm")
assert (pooled.username, pooled.password) == ("app", "p%40ss%27w")
assert _query(_plan().pooled_url) == {
"schema": "public",
"connection_limit": "10",
"pool_timeout": "20",
"pgbouncer": "true",
}
@pytest.mark.parametrize("hop_param", ["channel_binding=require", "gssencmode=require"])
def test_transport_params_for_the_postgres_hop_stay_off_the_plain_tcp_loopback_url(self, hop_param: str):
pooled: Final = _plan(f"postgresql://app:pw@db/litellm?connection_limit=5&{hop_param}").pooled_url
assert _query(pooled) == {"connection_limit": "5", "pgbouncer": "true"}
def test_verified_tls_becomes_server_side_verify_full_with_a_ca_copy_in_the_runtime_dir(self):
plan: Final = _plan()
pgb: Final = _ini(plan)["pgbouncer"]
assert pgb["server_tls_sslmode"] == "verify-full"
assert pgb["server_tls_ca_file"] == "/run/pgb/server-ca.pem"
assert plan.ca_source == "/certs/ca.pem"
def test_unverified_require_stays_require_without_a_ca_file(self):
plan: Final = _plan("postgresql://app:pw@db/litellm?sslmode=require")
pgb: Final = _ini(plan)["pgbouncer"]
assert pgb["server_tls_sslmode"] == "require"
assert "server_tls_ca_file" not in pgb
assert plan.ca_source is None
def test_no_tls_params_default_to_prefer(self):
assert _ini(_plan("postgresql://app:pw@db/litellm"))["pgbouncer"]["server_tls_sslmode"] == "prefer"
def test_verification_without_a_ca_bundle_is_refused(self):
outcome: Final = plan_pgbouncer(
"postgresql://app:pw@db/litellm?sslmode=require&sslaccept=strict", SETTINGS, Path("/run/pgb"), None
)
assert isinstance(outcome, PgBouncerError)
assert "sslcert" in outcome.reason
def test_client_certificates_are_refused(self):
outcome: Final = plan_pgbouncer(
"postgresql://app:pw@db/litellm?sslidentity=/certs/client.p12", SETTINGS, Path("/run/pgb"), None
)
assert isinstance(outcome, PgBouncerError)
assert "sslidentity" in outcome.reason
@pytest.mark.parametrize(
"url",
[
"postgresql://app@db/litellm",
"postgresql://app:pw@db",
"postgresql://:pw@db/litellm",
],
)
def test_urls_missing_forwardable_credentials_are_refused(self, url: str):
outcome: Final = plan_pgbouncer(url, SETTINGS, Path("/run/pgb"), None)
assert isinstance(outcome, PgBouncerError)
def test_every_options_spelling_becomes_a_set_statement(self):
options: Final = urllib.parse.quote("-c a=1 -cb=2 --c=3")
ini: Final = _ini(_plan(f"postgresql://app:pw@db/litellm?options={options}"))
assert ini["databases"]["litellm"].endswith("connect_query='SET a TO ''1''; SET b TO ''2''; SET c TO ''3'''")
def test_options_that_are_not_settings_are_refused(self):
outcome: Final = plan_pgbouncer(
"postgresql://app:pw@db/litellm?options=-c%20search_path", SETTINGS, Path("/run/pgb"), None
)
assert isinstance(outcome, PgBouncerError)
assert "options" in outcome.reason
def test_run_as_user_is_only_written_when_given(self):
assert _ini(_plan(run_as_user="nobody"))["pgbouncer"]["user"] == "nobody"
assert "user" not in _ini(_plan())["pgbouncer"]
class TestWritePgBouncerFiles:
def test_files_hold_the_plan_and_are_private_to_the_owner(self, tmp_path: Path):
plan: Final = _plan("postgresql://app:pw@db/litellm")
ini_path: Final = write_pgbouncer_files(plan, tmp_path, None)
assert isinstance(ini_path, Path), ini_path
assert ini_path == tmp_path / "pgbouncer.ini"
assert ini_path.read_text() == plan.ini
assert (tmp_path / "userlist.txt").read_text() == plan.userlist
for path in (ini_path, tmp_path / "userlist.txt"):
assert stat.S_IMODE(path.stat().st_mode) == 0o600
assert not (tmp_path / "server-ca.pem").exists()
def test_the_ca_bundle_is_copied_next_to_the_ini_pgbouncer_reads(self, tmp_path: Path):
bundle: Final = tmp_path / "rds-root.pem"
bundle.write_text("-----BEGIN CERTIFICATE-----\nMIIB\n-----END CERTIFICATE-----\n")
runtime_dir: Final = tmp_path / "run"
runtime_dir.mkdir()
plan: Final = plan_pgbouncer(
f"postgresql://app:pw@db/litellm?sslmode=verify-full&sslcert={bundle}", SETTINGS, runtime_dir, None
)
assert isinstance(plan, PgBouncerPlan), plan
ini_path: Final = write_pgbouncer_files(plan, runtime_dir, None)
assert isinstance(ini_path, Path), ini_path
ca_file: Final = Path(_ini(plan)["pgbouncer"]["server_tls_ca_file"])
assert ca_file.parent == runtime_dir
assert ca_file.read_text() == bundle.read_text()
def test_an_unreadable_ca_bundle_is_reported(self, tmp_path: Path):
plan: Final = plan_pgbouncer(
f"postgresql://app:pw@db/litellm?sslmode=verify-full&sslcert={tmp_path / 'missing.pem'}",
SETTINGS,
tmp_path,
None,
)
assert isinstance(plan, PgBouncerPlan), plan
outcome: Final = write_pgbouncer_files(plan, tmp_path, None)
assert isinstance(outcome, PgBouncerError)
assert "missing.pem" in outcome.reason
assert not (tmp_path / "pgbouncer.ini").exists()
def _bound_port(sock: socket.socket) -> int:
return cast(tuple[str, int], sock.getsockname())[1]
def _free_port() -> int:
with socket.socket() as probe:
probe.bind(("127.0.0.1", 0))
return _bound_port(probe)
def _fake_pooler(
tmp_path: Path,
port: int,
exit_immediately: bool = False,
port_file: Path | None = None,
bind_delay_seconds: float = 0.0,
version_banner: str = "PgBouncer 1.25.2\nlibevent 2.1.13-stable",
) -> Path:
"""An executable that listens like PgBouncer: on the TCP port first, then on ``.s.PGSQL.<port>`` in the socket dir.
Port and socket dir come from the ini it is given, else from ``port`` and
``tmp_path``. With ``port_file`` each start reads the port from that file
instead. ``bind_delay_seconds`` holds the bind back, like a slow start.
``--version`` prints ``version_banner``.
"""
script: Final = tmp_path / "fake-pgbouncer"
script.write_text(
textwrap.dedent(
f"""\
#!{sys.executable}
import configparser, os, pathlib, select, socket, sys, time
if sys.argv[1:] == ["--version"]:
print({version_banner!r})
sys.exit(0)
if {exit_immediately!r}:
sys.exit(3)
ini = configparser.ConfigParser()
ini.read(sys.argv[1:2])
port = ini.getint("pgbouncer", "listen_port", fallback={port})
if not {port_file is None!r}:
port = int(pathlib.Path({str(port_file)!r}).read_text())
socket_dir = ini.get("pgbouncer", "unix_socket_dir", fallback={str(tmp_path)!r})
socket_path = f"{{socket_dir}}/.s.PGSQL.{{port}}"
time.sleep({bind_delay_seconds!r})
listener = socket.socket()
listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
listener.bind(("127.0.0.1", port))
listener.listen()
if os.path.exists(socket_path):
os.unlink(socket_path)
unix_listener = socket.socket(socket.AF_UNIX)
unix_listener.bind(socket_path)
unix_listener.listen()
while True:
for ready in select.select([listener, unix_listener], [], [])[0]:
conn, _ = ready.accept()
conn.close()
"""
)
)
script.chmod(0o700)
return script
def _listening(port: int) -> bool:
try:
with socket.create_connection(("127.0.0.1", port), timeout=0.5):
return True
except OSError:
return False
def _wait_until(condition: Callable[[], bool], timeout_seconds: float = 5.0) -> bool:
deadline: Final = time.monotonic() + timeout_seconds
while time.monotonic() < deadline:
if condition():
return True
time.sleep(0.05)
return False
class TestPgBouncerProcess:
def test_start_waits_for_the_listener_and_stop_ends_it(self, tmp_path: Path):
port: Final = _free_port()
pooler: Final = PgBouncerProcess(
argv=(str(_fake_pooler(tmp_path, port)),), port=port, socket_path=unix_socket_path(tmp_path, port)
)
assert pooler.start() is None
assert _listening(port)
pid: Final = pooler.pid
assert pid is not None
pooler.stop()
assert _wait_until(lambda: not _listening(port))
with pytest.raises(ProcessLookupError):
os.kill(pid, 0)
def test_a_crashed_pooler_is_restarted_with_a_new_pid(self, tmp_path: Path):
port: Final = _free_port()
pooler: Final = PgBouncerProcess(
argv=(str(_fake_pooler(tmp_path, port)),),
port=port,
socket_path=unix_socket_path(tmp_path, port),
restart_delay_seconds=0.1,
)
assert pooler.start() is None
first_pid: Final = pooler.pid
assert first_pid is not None
os.kill(first_pid, signal.SIGKILL)
assert _wait_until(lambda: pooler.pid not in (None, first_pid) and _listening(port))
pooler.stop()
assert _wait_until(lambda: not _listening(port))
def test_a_failed_restart_is_retried_until_the_pooler_is_back(
self, tmp_path: Path, caplog: pytest.LogCaptureFixture
):
port: Final = _free_port()
script: Final = _fake_pooler(tmp_path, port)
pooler: Final = PgBouncerProcess(
argv=(str(script),), port=port, socket_path=unix_socket_path(tmp_path, port), restart_delay_seconds=0.1
)
assert pooler.start() is None
first_pid: Final = pooler.pid
assert first_pid is not None
hidden: Final = script.rename(tmp_path / "hidden")
with caplog.at_level(logging.ERROR, logger=verbose_proxy_logger.name):
os.kill(first_pid, signal.SIGKILL)
assert _wait_until(lambda: any("could not be restarted" in record.message for record in caplog.records))
assert not _listening(port)
hidden.rename(script)
assert _wait_until(lambda: pooler.pid not in (None, first_pid) and _listening(port))
pooler.stop()
assert _wait_until(lambda: not _listening(port))
def test_a_replacement_that_never_listens_is_replaced_again(self, tmp_path: Path, caplog: pytest.LogCaptureFixture):
port: Final = _free_port()
port_file: Final = tmp_path / "port"
port_file.write_text(str(port))
script: Final = _fake_pooler(tmp_path, port, port_file=port_file)
pooler: Final = PgBouncerProcess(
argv=(str(script),),
port=port,
socket_path=unix_socket_path(tmp_path, port),
restart_delay_seconds=0.1,
ready_timeout_seconds=0.3,
)
assert pooler.start() is None
first_pid: Final = pooler.pid
assert first_pid is not None
wrong_port: Final = _free_port()
port_file.write_text(str(wrong_port))
with caplog.at_level(logging.ERROR, logger=verbose_proxy_logger.name):
os.kill(first_pid, signal.SIGKILL)
assert _wait_until(lambda: _listening(wrong_port))
assert _wait_until(lambda: any("did not start listening" in record.message for record in caplog.records))
port_file.write_text(str(port))
assert _wait_until(lambda: _listening(port))
assert _wait_until(lambda: not _listening(wrong_port))
pooler.stop()
assert _wait_until(lambda: not _listening(port))
def test_stopping_during_the_restart_delay_leaves_no_pooler_behind(self, tmp_path: Path):
port: Final = _free_port()
pooler: Final = PgBouncerProcess(
argv=(str(_fake_pooler(tmp_path, port)),),
port=port,
socket_path=unix_socket_path(tmp_path, port),
restart_delay_seconds=0.3,
)
assert pooler.start() is None
first_pid: Final = pooler.pid
assert first_pid is not None
os.kill(first_pid, signal.SIGKILL)
assert _wait_until(lambda: not _listening(port))
pooler.stop()
time.sleep(1.0)
assert not _listening(port)
assert pooler.pid == first_pid
def test_a_stopped_pooler_is_not_restarted(self, tmp_path: Path, caplog: pytest.LogCaptureFixture):
port: Final = _free_port()
pooler: Final = PgBouncerProcess(
argv=(str(_fake_pooler(tmp_path, port)),),
port=port,
socket_path=unix_socket_path(tmp_path, port),
restart_delay_seconds=0.1,
)
assert pooler.start() is None
with caplog.at_level(logging.ERROR, logger=verbose_proxy_logger.name):
pooler.stop()
time.sleep(0.5)
assert not _listening(port)
assert caplog.records == []
def test_a_pooler_that_exits_during_startup_is_reported(self, tmp_path: Path):
port: Final = _free_port()
pooler: Final = PgBouncerProcess(
argv=(str(_fake_pooler(tmp_path, port, exit_immediately=True)),),
port=port,
socket_path=unix_socket_path(tmp_path, port),
)
outcome: Final = pooler.start()
assert isinstance(outcome, PgBouncerError)
assert "status 3" in outcome.reason
def test_a_missing_binary_is_reported(self, tmp_path: Path):
outcome: Final = PgBouncerProcess(
argv=("/nonexistent/pgbouncer",), port=_free_port(), socket_path=tmp_path / "sock"
).start()
assert isinstance(outcome, PgBouncerError)
assert "/nonexistent/pgbouncer" in outcome.reason
def test_a_port_owned_by_someone_else_is_refused_before_spawning(self, tmp_path: Path):
with socket.socket() as squatter:
squatter.bind(("127.0.0.1", 0))
squatter.listen()
port: Final = _bound_port(squatter)
pooler: Final = PgBouncerProcess(
argv=(str(_fake_pooler(tmp_path, port)),), port=port, socket_path=unix_socket_path(tmp_path, port)
)
outcome: Final = pooler.start()
assert isinstance(outcome, PgBouncerError)
assert f"127.0.0.1:{port} is already in use" in outcome.reason
assert pooler.pid is None
def test_a_replacement_waits_until_a_squatter_leaves_the_port(
self, tmp_path: Path, caplog: pytest.LogCaptureFixture
):
port: Final = _free_port()
pooler: Final = PgBouncerProcess(
argv=(str(_fake_pooler(tmp_path, port)),),
port=port,
socket_path=unix_socket_path(tmp_path, port),
restart_delay_seconds=0.5,
)
assert pooler.start() is None
first_pid: Final = pooler.pid
assert first_pid is not None
os.kill(first_pid, signal.SIGKILL)
assert _wait_until(lambda: not _listening(port))
with socket.socket() as squatter, caplog.at_level(logging.ERROR, logger=verbose_proxy_logger.name):
squatter.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
squatter.bind(("127.0.0.1", port))
squatter.listen()
assert _wait_until(lambda: any("already in use" in record.message for record in caplog.records))
assert pooler.pid == first_pid
assert _wait_until(lambda: pooler.pid not in (None, first_pid) and _listening(port))
pooler.stop()
assert _wait_until(lambda: not _listening(port))
def test_a_listener_that_grabs_the_port_after_the_spawn_is_not_taken_for_the_pooler(self, tmp_path: Path):
port: Final = _free_port()
pooler: Final = PgBouncerProcess(
argv=(str(_fake_pooler(tmp_path, port, bind_delay_seconds=0.5)),),
port=port,
socket_path=unix_socket_path(tmp_path, port),
ready_timeout_seconds=3.0,
)
with socket.socket() as squatter, ThreadPoolExecutor(max_workers=1) as starter:
squatter.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
starting: Final = starter.submit(pooler.start)
assert _wait_until(lambda: pooler.pid is not None)
squatter.bind(("127.0.0.1", port))
squatter.listen()
outcome: Final = starting.result()
assert isinstance(outcome, PgBouncerError)
assert "exited with status 1" in outcome.reason
def test_a_port_served_by_a_stranger_while_the_pooler_is_still_starting_is_reported(self, tmp_path: Path):
port: Final = _free_port()
pooler: Final = PgBouncerProcess(
argv=(str(_fake_pooler(tmp_path, port, bind_delay_seconds=30.0)),),
port=port,
socket_path=unix_socket_path(tmp_path, port),
ready_timeout_seconds=0.5,
)
with socket.socket() as squatter, ThreadPoolExecutor(max_workers=1) as starter:
squatter.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
starting: Final = starter.submit(pooler.start)
assert _wait_until(lambda: pooler.pid is not None)
squatter.bind(("127.0.0.1", port))
squatter.listen()
outcome: Final = starting.result()
assert isinstance(outcome, PgBouncerError)
assert f"127.0.0.1:{port} is served by another process" in outcome.reason
pid: Final = pooler.pid
assert pid is not None
with pytest.raises(ProcessLookupError):
os.kill(pid, 0)
def test_a_pooler_that_never_listens_times_out(self, tmp_path: Path):
port: Final = _free_port()
pooler: Final = PgBouncerProcess(
argv=(str(_fake_pooler(tmp_path, _free_port())),),
port=port,
socket_path=unix_socket_path(tmp_path, port),
ready_timeout_seconds=0.5,
)
outcome: Final = pooler.start()
assert isinstance(outcome, PgBouncerError)
assert "did not start listening" in outcome.reason
pid: Final = pooler.pid
assert pid is not None
with pytest.raises(ProcessLookupError):
os.kill(pid, 0)
def _runtime_dir_listening_on(port: int) -> Path:
matches: Final = tuple(
ini.parent
for ini in Path(tempfile.gettempdir()).glob("litellm-pgbouncer-*/pgbouncer.ini")
if f"listen_port = {port}\n" in ini.read_text()
)
assert len(matches) == 1, matches
return matches[0]
class TestStartInContainerPgBouncer:
def test_returns_the_loopback_url_once_the_pooler_listens(self, tmp_path: Path):
port: Final = _free_port()
settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(_fake_pooler(tmp_path, port)))
pooled: Final = start_in_container_pgbouncer(settings, "postgresql://app:pw@db/litellm?connection_limit=5")
assert pooled == f"postgresql://app:pw@127.0.0.1:{port}/litellm?connection_limit=5&pgbouncer=true"
assert _listening(port)
@pytest.mark.filterwarnings("ignore:This process .* is multi-threaded:DeprecationWarning")
def test_a_forked_worker_exiting_leaves_the_pooler_and_its_files_to_the_parent(self, tmp_path: Path):
port: Final = _free_port()
settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(_fake_pooler(tmp_path, port)))
exit_hooks: Final[list[Callable[[], None]]] = []
pooled: Final = start_in_container_pgbouncer(
settings, "postgresql://app:pw@db/litellm", register_exit_hook=exit_hooks.append
)
assert isinstance(pooled, str), pooled
runtime_dir: Final = _runtime_dir_listening_on(port)
worker: Final = os.fork()
if worker == 0:
try:
for hook in exit_hooks:
hook()
finally:
os._exit(0)
if not _wait_until(lambda: os.waitpid(worker, os.WNOHANG) != (0, 0)):
os.kill(worker, signal.SIGKILL)
pytest.fail("the forked worker did not exit: an exit hook blocked on state inherited from the parent")
assert _listening(port)
assert (runtime_dir / "pgbouncer.ini").exists()
for hook in exit_hooks:
hook()
assert _wait_until(lambda: not _listening(port))
assert not runtime_dir.exists()
def test_a_bad_upstream_url_is_reported_without_starting_anything(self, tmp_path: Path):
port: Final = _free_port()
settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(_fake_pooler(tmp_path, port)))
outcome: Final = start_in_container_pgbouncer(settings, "postgresql://app@db/litellm")
assert isinstance(outcome, PgBouncerError)
assert not _listening(port)
def test_token_auth_is_refused_without_starting_anything(self, tmp_path: Path):
port: Final = _free_port()
settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(_fake_pooler(tmp_path, port)))
outcome: Final = start_in_container_pgbouncer(
settings, "postgresql://app:pw@db/litellm", token_auth_enabled=True
)
assert isinstance(outcome, PgBouncerError)
assert "IAM_TOKEN_DB_AUTH" in outcome.reason
assert "AZURE_POSTGRESQL_AUTH" in outcome.reason
assert not _listening(port)
def test_a_pgbouncer_that_survives_a_failed_tcp_bind_is_refused_without_starting(self, tmp_path: Path):
port: Final = _free_port()
binary: Final = _fake_pooler(tmp_path, port, version_banner="PgBouncer 1.18.1\nlibevent 2.1.12-stable")
settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(binary))
outcome: Final = start_in_container_pgbouncer(settings, "postgresql://app:pw@db/litellm")
assert isinstance(outcome, PgBouncerError)
assert "PgBouncer 1.18" in outcome.reason
assert "1.19" in outcome.reason
assert not _listening(port)
def test_the_first_version_that_dies_on_a_failed_tcp_bind_is_accepted(self, tmp_path: Path):
port: Final = _free_port()
binary: Final = _fake_pooler(tmp_path, port, version_banner="PgBouncer 1.19.0")
settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(binary))
assert start_in_container_pgbouncer(settings, "postgresql://app:pw@db/litellm") == (
f"postgresql://app:pw@127.0.0.1:{port}/litellm?pgbouncer=true"
)
assert _listening(port)
class TestPgBouncerVersion:
def test_reads_major_and_minor_from_the_banner(self, tmp_path: Path):
assert pgbouncer_version(str(_fake_pooler(tmp_path, _free_port()))) == (1, 25)
def test_a_binary_that_cannot_run_is_reported(self, tmp_path: Path):
outcome: Final = pgbouncer_version(str(tmp_path / "missing-pgbouncer"))
assert isinstance(outcome, PgBouncerError)
assert "missing-pgbouncer" in outcome.reason
def test_a_banner_without_a_version_is_reported(self, tmp_path: Path):
outcome: Final = pgbouncer_version(str(_fake_pooler(tmp_path, _free_port(), version_banner="something else")))
assert isinstance(outcome, PgBouncerError)
assert "something else" in outcome.reason
class TestPgBouncerSettings:
def test_reads_the_litellm_pgbouncer_env_vars(self, monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv("LITELLM_PGBOUNCER_ENABLED", "true")
monkeypatch.setenv("LITELLM_PGBOUNCER_PORT", "7000")
monkeypatch.setenv("LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS", "12")
settings: Final = PgBouncerSettings()
assert (settings.enabled, settings.port, settings.max_db_connections) == (True, 7000, 12)
def test_defaults_are_off(self, monkeypatch: pytest.MonkeyPatch):
monkeypatch.delenv("LITELLM_PGBOUNCER_ENABLED", raising=False)
assert PgBouncerSettings().enabled is False

View file

@ -4802,6 +4802,72 @@ async def test_store_mcp_oauth_user_credential_returns_status():
assert result.expires_at == "2099-01-01T00:00:00+00:00"
@pytest.mark.asyncio
async def test_store_mcp_oauth_user_credential_blocked_when_identity_binding_enforced():
"""The direct opaque-token POST must be closed for enforce-mode identity-bound servers,
otherwise it bypasses the token-relay principal check."""
from litellm.proxy._types import MCPOAuthUserCredentialRequest
from litellm.types.mcp import MCPAuth, MCPTransport
from litellm.types.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer
if not mgmt_endpoints.MCP_AVAILABLE:
pytest.skip("MCP module not installed")
from litellm.proxy._experimental.mcp_server import mcp_server_manager as manager_module
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
store_mcp_oauth_user_credential,
)
server_id = "srv-binding-1"
bound_server = MCPServer(
server_id=server_id,
name=server_id,
url="https://mcp.example.com",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
oauth_identity_binding=MCPOAuthIdentityBinding(
mode="enforce",
issuer="https://idp.example.com",
audiences=["litellm-client"],
),
)
store_mock = AsyncMock(return_value=None)
with (
patch( # test-quality-ok: mirrors the existing store-credential tests in this file
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
return_value=_make_prisma_client(),
),
patch( # test-quality-ok: mirrors the existing store-credential tests in this file
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server",
new=AsyncMock(return_value=generate_mock_mcp_server_db_record(server_id=server_id)),
),
patch( # test-quality-ok: mirrors the existing store-credential tests in this file
"litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view",
return_value=True,
),
patch.object( # test-quality-ok: registry is a module-level singleton; injecting it would change the endpoint signature
manager_module.global_mcp_server_manager,
"get_mcp_server_by_id",
return_value=bound_server,
),
patch( # test-quality-ok: asserting the DB write is never reached is the point of the test
"litellm.proxy.management_endpoints.mcp_management_endpoints.store_user_oauth_credential",
new=store_mock,
),
):
with pytest.raises(HTTPException) as exc_info:
await store_mcp_oauth_user_credential(
server_id=server_id,
payload=MCPOAuthUserCredentialRequest(access_token="opaque-tok", expires_in=3600),
user_api_key_dict=_make_user_auth("user-123"),
)
assert exc_info.value.status_code == 403
assert exc_info.value.detail["error"] == "oauth_identity_binding_enforced"
store_mock.assert_not_called()
@pytest.mark.asyncio
async def test_delete_mcp_oauth_user_credential_only_deletes_oauth():
"""delete_mcp_oauth_user_credential only deletes OAuth2 credentials, not BYOK."""

View file

@ -2169,7 +2169,7 @@ class TestModelInfoEndpoint:
):
mock_router.get_fully_blocked_model_names.return_value = set()
mock_router.get_model_list.return_value = []
mock_router.get_configured_token_limits.return_value = (None, None)
mock_router.get_model_listing_info.return_value = None
mock_router.get_deployment_by_model_group_name.return_value = Deployment(
model_name="gpt-4",
litellm_params=LiteLLM_Params(model="openai/gpt-4"),
@ -2246,7 +2246,7 @@ class TestModelInfoEndpoint:
):
mock_router.get_fully_blocked_model_names.return_value = set()
mock_router.get_model_list.return_value = []
mock_router.get_configured_token_limits.return_value = (None, None)
mock_router.get_model_listing_info.return_value = None
mock_router.get_deployment_by_model_group_name.return_value = Deployment(
model_name="team-model-1",
litellm_params=LiteLLM_Params(model="custom/team-model-1"),

View file

@ -20,6 +20,7 @@ from unittest.mock import AsyncMock, MagicMock
import pytest
import litellm.proxy.proxy_server as proxy_server
from litellm.constants import BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME
from litellm.proxy.proxy_server import (
_adaptive_router_flusher_loop,
_get_endpoint_exception_status,
@ -111,13 +112,11 @@ async def test_run_direct_health_check_with_instrumentation_returns_results(
lambda _gs: {},
)
healthy, unhealthy, exceptions = (
await _run_direct_health_check_with_instrumentation(
model_list=[{"model_name": "gpt-4"}],
details=False,
max_concurrency=1,
instrumentation_context={"source": "test"},
)
healthy, unhealthy, exceptions = await _run_direct_health_check_with_instrumentation(
model_list=[{"model_name": "gpt-4"}],
details=False,
max_concurrency=1,
instrumentation_context={"source": "test"},
)
assert normalize(
@ -245,6 +244,137 @@ async def test_schedule_background_health_check_db_save_invalid_no_event_loop_ra
)
def _lock_manager(redis_cache, acquired):
manager = MagicMock()
manager.redis_cache = redis_cache
manager.acquire_lock = AsyncMock(return_value=acquired)
manager.release_lock = AsyncMock()
return manager
def _capture_saves(monkeypatch, persisted=True):
saves = []
async def _fake_save(*_args, **kwargs):
saves.append(kwargs)
return persisted
import litellm.proxy.health_endpoints._health_endpoints as he
monkeypatch.setattr(he, "_save_background_health_checks_to_db", _fake_save)
return saves
def _cancel_during_save(monkeypatch):
async def _fake_save(*_args, **_kwargs):
raise asyncio.CancelledError()
import litellm.proxy.health_endpoints._health_endpoints as he
monkeypatch.setattr(he, "_save_background_health_checks_to_db", _fake_save)
def _schedule_with(lock_manager):
_schedule_background_health_check_db_save(
prisma_client=MagicMock(),
shared_health_manager=None,
model_list=[],
healthy_endpoints=[],
unhealthy_endpoints=[],
pod_lock_manager=lock_manager,
lock_ttl=300,
)
@pytest.mark.asyncio
async def test_schedule_background_health_check_db_save_skips_a_window_another_pod_persisted(monkeypatch):
saves = _capture_saves(monkeypatch)
lock_manager = _lock_manager(redis_cache=MagicMock(), acquired=False)
_schedule_with(lock_manager)
await asyncio.sleep(0)
assert saves == []
@pytest.mark.asyncio
async def test_schedule_background_health_check_db_save_holds_the_window_lock_for_the_whole_interval(monkeypatch):
"""The lock is the "saved this window" marker: never reentrant, TTL = interval, and never released."""
saves = _capture_saves(monkeypatch)
lock_manager = _lock_manager(redis_cache=MagicMock(), acquired=True)
_schedule_with(lock_manager)
await asyncio.sleep(0)
assert normalize(
{
"saves": len(saves),
"lock_request": lock_manager.acquire_lock.await_args.kwargs,
"released": lock_manager.release_lock.await_count,
}
) == {
"saves": 1,
"lock_request": {
"cronjob_id": BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME,
"ttl": 300,
"allow_reentrant": False,
},
"released": 0,
}
@pytest.mark.asyncio
async def test_schedule_background_health_check_db_save_releases_the_window_lock_when_the_save_reports_failure(
monkeypatch,
):
"""A failed save must not burn the window: release the lock so another pod's cycle can retry."""
saves = _capture_saves(monkeypatch, persisted=False)
lock_manager = _lock_manager(redis_cache=MagicMock(), acquired=True)
_schedule_with(lock_manager)
await asyncio.sleep(0)
assert normalize(
{
"saves": len(saves),
"release_request": lock_manager.release_lock.await_args.kwargs,
"release_count": lock_manager.release_lock.await_count,
}
) == {
"saves": 1,
"release_request": {"cronjob_id": BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME},
"release_count": 1,
}
@pytest.mark.asyncio
async def test_schedule_background_health_check_db_save_releases_the_window_lock_when_the_save_is_cancelled(
monkeypatch,
):
"""A pod shutting down mid-save releases the lock instead of holding it until the TTL."""
_cancel_during_save(monkeypatch)
lock_manager = _lock_manager(redis_cache=MagicMock(), acquired=True)
_schedule_with(lock_manager)
await asyncio.sleep(0)
assert (
lock_manager.release_lock.await_args.kwargs,
lock_manager.release_lock.await_count,
) == ({"cronjob_id": BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME}, 1)
@pytest.mark.asyncio
async def test_schedule_background_health_check_db_save_runs_ungated_without_redis(monkeypatch):
saves = _capture_saves(monkeypatch, persisted=False)
lock_manager = _lock_manager(redis_cache=None, acquired=True)
_schedule_with(lock_manager)
await asyncio.sleep(0)
assert (len(saves), lock_manager.acquire_lock.await_count, lock_manager.release_lock.await_count) == (1, 0, 0)
# ---------------------------------------------------------------------------
# _get_endpoint_exception_status
# ---------------------------------------------------------------------------
@ -319,13 +449,9 @@ def test_write_health_state_to_router_cache_sets_states(monkeypatch):
_write_health_state_to_router_cache(healthy, unhealthy, exceptions)
fake_router.health_state_cache.set_deployment_health_states.assert_called_once_with(
fake_states
)
fake_router.health_state_cache.set_deployment_health_states.assert_called_once_with(fake_states)
call_args = fake_router.health_state_cache.set_deployment_health_states.call_args[
0
][0]
call_args = fake_router.health_state_cache.set_deployment_health_states.call_args[0][0]
assert normalize(
{
"states_keys": sorted(call_args.keys()),
@ -367,9 +493,7 @@ def test_write_health_state_to_router_cache_populates_for_listing_filter(monkeyp
fake_router.cooldown_time = 30
monkeypatch.setattr(proxy_server, "llm_router", fake_router)
monkeypatch.setattr(
proxy_server, "general_settings", {"model_list_healthy_only": True}
)
monkeypatch.setattr(proxy_server, "general_settings", {"model_list_healthy_only": True})
fake_states = {"m1": {"is_healthy": True}, "m2": {"is_healthy": False}}
@ -403,9 +527,7 @@ def test_write_health_state_to_router_cache_populates_for_listing_filter(monkeyp
{"m2": SimpleNamespace(status_code=500)},
)
fake_router.health_state_cache.set_deployment_health_states.assert_called_once_with(
fake_states
)
fake_router.health_state_cache.set_deployment_health_states.assert_called_once_with(fake_states)
assert cooldowns == []
assert failures == []
@ -415,9 +537,7 @@ def test_write_health_state_to_router_cache_swallows_internal_failures(monkeypat
fake_router = MagicMock()
fake_router.enable_health_check_routing = True
fake_router.health_check_ignore_transient_errors = False
fake_router.health_state_cache.set_deployment_health_states.side_effect = (
RuntimeError("cache exploded")
)
fake_router.health_state_cache.set_deployment_health_states.side_effect = RuntimeError("cache exploded")
monkeypatch.setattr(proxy_server, "llm_router", fake_router)
@ -447,9 +567,7 @@ async def test_adaptive_router_flusher_loop_flushes_each_router(monkeypatch):
from litellm.types.router import TaggedPreRoutingStrategy
fake_router = MagicMock()
fake_router.adaptive_routers = {
"alpha": [TaggedPreRoutingStrategy(tags=(), strategy=fake_ar)]
}
fake_router.adaptive_routers = {"alpha": [TaggedPreRoutingStrategy(tags=(), strategy=fake_ar)]}
monkeypatch.setattr(proxy_server, "llm_router", fake_router)
monkeypatch.setattr(proxy_server, "prisma_client", MagicMock())
@ -547,12 +665,8 @@ async def test_run_background_health_check_runs_one_cycle_then_cancels(monkeypat
"_run_direct_health_check_with_instrumentation",
_fake_direct,
)
monkeypatch.setattr(
proxy_server, "_schedule_background_health_check_db_save", lambda *a, **kw: None
)
monkeypatch.setattr(
proxy_server, "_write_health_state_to_router_cache", lambda *a, **kw: None
)
monkeypatch.setattr(proxy_server, "_schedule_background_health_check_db_save", lambda *a, **kw: None)
monkeypatch.setattr(proxy_server, "_write_health_state_to_router_cache", lambda *a, **kw: None)
monkeypatch.setattr(
proxy_server,
"health_check_filter_kwargs_from_general_settings",
@ -630,12 +744,8 @@ async def test_run_background_health_check_probes_only_listed_model_groups(monke
"_run_direct_health_check_with_instrumentation",
_fake_direct,
)
monkeypatch.setattr(
proxy_server, "_schedule_background_health_check_db_save", lambda *a, **kw: None
)
monkeypatch.setattr(
proxy_server, "_write_health_state_to_router_cache", lambda *a, **kw: None
)
monkeypatch.setattr(proxy_server, "_schedule_background_health_check_db_save", lambda *a, **kw: None)
monkeypatch.setattr(proxy_server, "_write_health_state_to_router_cache", lambda *a, **kw: None)
monkeypatch.setattr(
proxy_server,
"health_check_filter_kwargs_from_general_settings",

View file

@ -17,6 +17,7 @@ import litellm
from litellm.proxy import proxy_server
from litellm.proxy import utils as proxy_utils
from litellm.proxy.utils import create_model_info_response
from litellm.types.router import DeploymentModelListingInfo
from .conftest import normalize # type: ignore[import-not-found]
@ -166,13 +167,16 @@ def test_anthropic_format_carries_router_configured_token_limits(client, auth_as
built from another entry's lookup shows up as the wrong numbers."""
def _configured(model_name):
return (300000, 32000) if model_name == "gpt-4" else (500000, 4096)
max_input, max_output = (300000, 32000) if model_name == "gpt-4" else (500000, 4096)
return DeploymentModelListingInfo(
cost_map_keys=(model_name,), max_input_tokens=max_input, max_output_tokens=max_output
)
def _cost_map_lookup(model_id):
max_input, max_output = (200000, 64000) if model_id == "gpt-4" else (100000, 8000)
return {"max_input_tokens": max_input, "max_output_tokens": max_output, "mode": "chat"}
patched_models.get_configured_token_limits = MagicMock(side_effect=_configured)
patched_models.get_model_listing_info = MagicMock(side_effect=_configured)
def _resolved(**kwargs):
return create_model_info_response(**kwargs, get_model_info=_cost_map_lookup)

View file

@ -880,7 +880,7 @@ async def test_v1_models_translates_team_model_for_access_group_key(monkeypatch)
router.get_model_names.return_value = ["model_name_teamX_uuid9"]
router.get_model_access_groups.return_value = {"grp-a": ["model_name_teamX_uuid9"]}
router.get_fully_blocked_model_names.return_value = set()
router.get_configured_token_limits.return_value = (None, None)
router.get_model_listing_info.return_value = None
router.model_list = [team_dep]
router.get_model_list.return_value = [team_dep]
@ -922,7 +922,7 @@ async def test_v1_models_keeps_internal_names_when_public_name_flag_disabled(
router.get_model_names.return_value = ["model_name_teamX_uuid9"]
router.get_model_access_groups.return_value = {"grp-a": ["model_name_teamX_uuid9"]}
router.get_fully_blocked_model_names.return_value = set()
router.get_configured_token_limits.return_value = (None, None)
router.get_model_listing_info.return_value = None
router.model_list = [team_dep]
router.get_model_list.return_value = [team_dep]
@ -957,7 +957,7 @@ async def test_v1_models_translates_team_model_with_metadata(monkeypatch):
router.get_model_names.return_value = ["model_name_teamX_uuid9"]
router.get_model_access_groups.return_value = {"grp-a": ["model_name_teamX_uuid9"]}
router.get_fully_blocked_model_names.return_value = set()
router.get_configured_token_limits.return_value = (None, None)
router.get_model_listing_info.return_value = None
router.model_list = [team_dep]
router.get_model_list.return_value = [team_dep]
router.get_model_group_info.return_value = None
@ -1003,7 +1003,7 @@ async def test_v1_models_metadata_fallbacks_use_internal_routing_key(monkeypatch
router.get_model_names.return_value = ["model_name_teamX_uuid9"]
router.get_model_access_groups.return_value = {"grp-a": ["model_name_teamX_uuid9"]}
router.get_fully_blocked_model_names.return_value = set()
router.get_configured_token_limits.return_value = (None, None)
router.get_model_listing_info.return_value = None
router.model_list = [team_dep]
router.get_model_list.return_value = [team_dep]
# Fallbacks are keyed on the internal routing name, as the router stores them.
@ -1060,7 +1060,7 @@ async def test_v1_models_metadata_does_not_leak_other_team_fallbacks(monkeypatch
router.get_model_names.return_value = ["model_name_teamX_uuid9"]
router.get_model_access_groups.return_value = {"grp-a": ["model_name_teamX_uuid9"]}
router.get_fully_blocked_model_names.return_value = set()
router.get_configured_token_limits.return_value = (None, None)
router.get_model_listing_info.return_value = None
router.model_list = [team_x, team_y]
router.get_model_list.return_value = [team_x, team_y]
router.fallbacks = [
@ -1315,7 +1315,7 @@ def test_translate_team_model_names_for_listing_respects_legacy_flag():
def _public_named_router(*team_rows: dict) -> MagicMock:
router = MagicMock()
router.get_model_list.return_value = list(team_rows)
router.get_configured_token_limits.return_value = (None, None)
router.get_model_listing_info.return_value = None
return router

View file

@ -455,6 +455,34 @@ async def test_assert_user_can_view_request_id_rejects_both_users_none():
assert exc_info.value.status_code == 403
@pytest.mark.asyncio
async def test_assert_user_can_view_request_id_rejects_missing_row():
"""
A request_id with no spend-log row (e.g. pruned by retention) must not
authorize reading the payload from cold storage; a missing row is not
the same as an owned row.
"""
class MockSpendLogs:
async def find_unique(self, where, include=None):
return None
class MockDB:
def __init__(self):
self.litellm_spendlogs = MockSpendLogs()
class MockPrisma:
def __init__(self):
self.db = MockDB()
auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_1")
with pytest.raises(HTTPException) as exc_info:
await spend_management_endpoints._assert_user_can_view_request_id(
MockPrisma(), auth, "req-missing-row"
)
assert exc_info.value.status_code == 403
def test_ui_view_request_response_forbids_non_admin_without_db(client, monkeypatch):
"""
Without prisma, non-admins cannot be authorized to read request/response
@ -6777,3 +6805,161 @@ async def test_ui_view_spend_logs_search_returns_flat_rows_when_grouping_by_sess
assert "next_session_cursor" not in data
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
def _fake_prisma_with_owned_spend_log(owner_user_id, messages_json, response_json):
class _Row:
user = owner_user_id
team_id = None
class _SpendLogs:
async def find_unique(self, where, include=None):
return _Row()
class _DB:
def __init__(self):
self.litellm_spendlogs = _SpendLogs()
async def query_raw(self, _sql, *_args):
return [
{
"messages": messages_json,
"response": response_json,
"proxy_server_request": "{}",
"metadata": "{}",
}
]
class _Prisma:
def __init__(self):
self.db = _DB()
return _Prisma()
def test_ui_view_request_response_internal_user_owner_gets_payload(client, monkeypatch):
"""
An internal_user who owns the spend-log row can fetch the Logs drawer
detail payload for their own request (regression for #34099, where the
route was blocked for INTERNAL_USER before reaching this ownership check).
"""
messages_json = json.dumps([{"role": "user", "content": "hi"}])
response_json = json.dumps({"choices": [{"message": {"content": "hello"}}]})
monkeypatch.setattr(
"litellm.proxy.proxy_server.prisma_client",
_fake_prisma_with_owned_spend_log("user_a", messages_json, response_json),
)
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_a"
)
try:
response = client.get(
"/spend/logs/ui/req-owned-by-user-a",
headers={"Authorization": "Bearer sk-test"},
)
assert response.status_code == 200
body = response.json()
assert json.loads(body["messages"]) == [{"role": "user", "content": "hi"}]
assert json.loads(body["response"]) == {
"choices": [{"message": {"content": "hello"}}]
}
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
class _RecordingAdditionalLoggingUtils:
"""Injectable custom logger that records every request_id it's asked for."""
def __init__(self, payload):
self._payload = payload
self.requested_ids = []
async def get_request_response_payload(self, request_id, start_time_utc, end_time_utc):
self.requested_ids.append(request_id)
return self._payload
def test_ui_view_request_response_internal_user_non_owner_forbidden(client, monkeypatch):
"""
A different internal_user requesting someone else's row is forbidden;
guards against _assert_user_can_view_request_id being skipped in the
detail-drawer handler. Also proves the handler stops before it ever asks
a custom logger or the DB for the payload.
"""
messages_json = json.dumps([{"role": "user", "content": "hi"}])
response_json = json.dumps({"choices": [{"message": {"content": "hello"}}]})
fake_prisma = _fake_prisma_with_owned_spend_log("user_a", messages_json, response_json)
original_query_raw = fake_prisma.db.query_raw
query_raw_calls = []
async def _spy_query_raw(*args, **kwargs):
query_raw_calls.append((args, kwargs))
return await original_query_raw(*args, **kwargs)
fake_prisma.db.query_raw = _spy_query_raw
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", fake_prisma)
custom_logger = _RecordingAdditionalLoggingUtils({"messages": "should-not-be-returned"})
monkeypatch.setattr(
litellm.logging_callback_manager,
"get_active_additional_logging_utils_from_custom_logger",
lambda: [custom_logger],
)
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_b"
)
try:
response = client.get(
"/spend/logs/ui/req-owned-by-user-a",
headers={"Authorization": "Bearer sk-test"},
)
assert response.status_code == 403
assert custom_logger.requested_ids == []
assert query_raw_calls == []
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
def test_ui_view_request_response_internal_user_missing_row_forbidden(client, monkeypatch):
"""
Regression for the fail-open in _assert_user_can_view_request_id: a
request_id with no spend-log row (e.g. pruned by retention) must be
denied before the handler ever consults a custom logger, otherwise a
non-admin who guesses/obtains a request_id could read another tenant's
payload out of cold storage. Fails if `if row is None: return` is
reintroduced.
"""
class _SpendLogs:
async def find_unique(self, where, include=None):
return None
class _DB:
def __init__(self):
self.litellm_spendlogs = _SpendLogs()
from types import SimpleNamespace
fake_prisma = SimpleNamespace(db=_DB())
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", fake_prisma)
custom_logger = _RecordingAdditionalLoggingUtils({"messages": "should-not-be-returned"})
monkeypatch.setattr(
litellm.logging_callback_manager,
"get_active_additional_logging_utils_from_custom_logger",
lambda: [custom_logger],
)
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_a"
)
try:
response = client.get(
"/spend/logs/ui/req-pruned",
headers={"Authorization": "Bearer sk-test"},
)
assert response.status_code == 403
assert custom_logger.requested_ids == []
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)

View file

@ -6,6 +6,8 @@ from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.db.health_check_latest import LatestHealthCheckRow
from litellm.proxy.health_endpoints._health_endpoints import (
_aggregate_health_check_results,
_build_model_param_to_info_mapping,
@ -13,6 +15,7 @@ from litellm.proxy.health_endpoints._health_endpoints import (
_save_background_health_checks_to_db,
_save_health_check_results_if_changed,
_save_health_check_to_db,
latest_health_checks_endpoint,
)
from litellm.proxy.utils import PrismaClient
@ -21,12 +24,8 @@ from litellm.proxy.utils import PrismaClient
def mock_prisma():
"""Simplified mock PrismaClient with bound methods"""
client = MagicMock()
client.db.litellm_healthchecktable.create = AsyncMock(
return_value={"id": "test-id"}
)
client.db.litellm_healthchecktable.find_many = AsyncMock(
return_value=[{"id": "1", "model_name": "test"}]
)
client.db.litellm_healthchecktable.create = AsyncMock(return_value={"id": "test-id"})
client.db.litellm_healthchecktable.find_many = AsyncMock(return_value=[{"id": "1", "model_name": "test"}])
# Bind actual methods
import types
@ -52,14 +51,10 @@ def mock_prisma():
("healthy", 1, 0, False), # Database error case
],
)
async def test_save_health_check_result(
mock_prisma, status, healthy, unhealthy, should_succeed
):
async def test_save_health_check_result(mock_prisma, status, healthy, unhealthy, should_succeed):
"""Test health check result saving with various scenarios"""
if not should_succeed:
mock_prisma.db.litellm_healthchecktable.create.side_effect = Exception(
"DB Error"
)
mock_prisma.db.litellm_healthchecktable.create.side_effect = Exception("DB Error")
result = await mock_prisma.save_health_check_result(
model_name="test-model",
@ -187,9 +182,7 @@ def test_aggregate_health_check_results():
{"model": "gpt-4", "error": "Rate limit exceeded"},
]
result = _aggregate_health_check_results(
model_param_to_info, healthy_endpoints, unhealthy_endpoints
)
result = _aggregate_health_check_results(model_param_to_info, healthy_endpoints, unhealthy_endpoints)
# Check gpt-3.5-turbo is healthy
gpt35_key = ("model-123", "gpt-3.5-turbo")
@ -220,9 +213,7 @@ def test_aggregate_health_check_results_multiple_endpoints():
]
unhealthy_endpoints = []
result = _aggregate_health_check_results(
model_param_to_info, healthy_endpoints, unhealthy_endpoints
)
result = _aggregate_health_check_results(model_param_to_info, healthy_endpoints, unhealthy_endpoints)
key = ("model-123", "gpt-3.5-turbo")
assert result[key]["healthy_count"] == 2
@ -398,7 +389,7 @@ async def test_save_background_health_checks_to_db():
start_time = 1234567890.0
await _save_background_health_checks_to_db(
persisted = await _save_background_health_checks_to_db(
mock_prisma,
model_list,
healthy_endpoints,
@ -407,7 +398,8 @@ async def test_save_background_health_checks_to_db():
"background_health_check",
)
# Should call get_all_latest_health_checks and save_health_check_result
# Should call get_all_latest_health_checks and save_health_check_result, and report completion
assert persisted is True
mock_prisma.get_all_latest_health_checks.assert_called_once()
mock_prisma.save_health_check_result.assert_called_once()
@ -418,22 +410,112 @@ async def test_save_background_health_checks_to_db():
assert call_kwargs["checked_by"] == "background_health_check"
def _two_model_results():
return {
("model-1", "gpt-4"): {
"model_name": "gpt-4",
"model_id": "model-1",
"healthy_count": 1,
"unhealthy_count": 0,
"error_message": None,
},
("model-2", "gpt-4o"): {
"model_name": "gpt-4o",
"model_id": "model-2",
"healthy_count": 0,
"unhealthy_count": 1,
"error_message": "boom",
},
}
@pytest.mark.asyncio
async def test_save_health_check_results_if_changed_awaits_every_write_and_reports_success():
"""Writes are awaited, not detached, so the caller can tell the cycle's persistence completed."""
mock_prisma = MagicMock()
mock_prisma.save_health_check_result = AsyncMock(return_value={"id": "row"})
persisted = await _save_health_check_results_if_changed(
mock_prisma, _two_model_results(), {}, 1234567890.0, "background_health_check"
)
assert (persisted, mock_prisma.save_health_check_result.await_count) == (True, 2)
@pytest.mark.asyncio
async def test_save_health_check_results_if_changed_reports_failure_when_a_write_returns_none():
"""save_health_check_result swallows DB errors and returns None; that must surface as False."""
mock_prisma = MagicMock()
mock_prisma.save_health_check_result = AsyncMock(side_effect=[{"id": "row"}, None])
persisted = await _save_health_check_results_if_changed(
mock_prisma, _two_model_results(), {}, 1234567890.0, "background_health_check"
)
assert (persisted, mock_prisma.save_health_check_result.await_count) == (False, 2)
@pytest.mark.asyncio
async def test_save_health_check_results_if_changed_reports_success_when_nothing_needed_writing():
mock_prisma = MagicMock()
mock_prisma.save_health_check_result = AsyncMock()
model_results = {
("model-1", "gpt-4"): {
"model_name": "gpt-4",
"model_id": "model-1",
"healthy_count": 1,
"unhealthy_count": 0,
"error_message": None,
},
}
latest_checks_map = {
"model-1": MagicMock(status="healthy", checked_at=datetime.now(timezone.utc) - timedelta(minutes=5)),
}
persisted = await _save_health_check_results_if_changed(
mock_prisma, model_results, latest_checks_map, 1234567890.0, "background_health_check"
)
assert (persisted, mock_prisma.save_health_check_result.await_count) == (True, 0)
def _one_model_setup():
model_list = [
{
"model_name": "gpt-3.5-turbo",
"model_info": {"id": "model-123"},
"litellm_params": {"model": "gpt-3.5-turbo"},
},
]
return model_list, [{"model": "gpt-3.5-turbo"}], []
@pytest.mark.asyncio
async def test_save_background_health_checks_to_db_returns_false_when_a_write_fails():
mock_prisma = MagicMock()
mock_prisma.get_all_latest_health_checks = AsyncMock(return_value=[])
mock_prisma.save_health_check_result = AsyncMock(return_value=None)
model_list, healthy_endpoints, unhealthy_endpoints = _one_model_setup()
persisted = await _save_background_health_checks_to_db(
mock_prisma, model_list, healthy_endpoints, unhealthy_endpoints, 1234567890.0, "background_health_check"
)
assert (persisted, mock_prisma.save_health_check_result.await_count) == (False, 1)
@pytest.mark.asyncio
async def test_save_background_health_checks_to_db_no_prisma():
"""Test graceful handling when no prisma client"""
result = await _save_background_health_checks_to_db(
None, [], [], [], 0.0, "background_health_check"
)
assert result is None
result = await _save_background_health_checks_to_db(None, [], [], [], 0.0, "background_health_check")
assert result is False
@pytest.mark.asyncio
async def test_save_background_health_checks_to_db_exception_handling():
"""Test exception handling in background health check save"""
mock_prisma = MagicMock()
mock_prisma.get_all_latest_health_checks = AsyncMock(
side_effect=Exception("DB Error")
)
mock_prisma.get_all_latest_health_checks = AsyncMock(side_effect=Exception("DB Error"))
model_list = [
{
@ -443,104 +525,134 @@ async def test_save_background_health_checks_to_db_exception_handling():
},
]
# Should not raise exception, should handle gracefully
await _save_background_health_checks_to_db(
# Must not raise (the health check loop has to survive a DB outage) but must report
# the failure, so the window lock can be released for another pod to retry
persisted = await _save_background_health_checks_to_db(
mock_prisma, model_list, [], [], 0.0, "background_health_check"
)
# Function should complete without raising
assert persisted is False
def _raw_latest_row(model_name: str, model_id, checked_at: datetime) -> dict:
return {
"health_check_id": f"hc-{model_id or 'no-id'}-{model_name}",
"model_name": model_name,
"model_id": model_id,
"status": "healthy",
"healthy_count": 1,
"unhealthy_count": 0,
"error_message": None,
"response_time_ms": 10.0,
"details": None,
"checked_by": "pod-1",
"checked_at": checked_at.isoformat(),
"created_at": checked_at.isoformat(),
"updated_at": checked_at.isoformat(),
}
@pytest.mark.asyncio
async def test_get_all_latest_health_checks_with_model_id(mock_prisma):
"""Test get_all_latest_health_checks properly groups by model_id"""
mock_check2 = MagicMock()
mock_check2.model_id = "model-456"
mock_check2.model_name = "gpt-3.5-turbo"
mock_check2.checked_at = datetime.now(timezone.utc) - timedelta(minutes=5)
mock_check3 = MagicMock()
mock_check3.model_id = "model-123"
mock_check3.model_name = "gpt-3.5-turbo"
mock_check3.checked_at = datetime.now(timezone.utc) - timedelta(
minutes=1
) # Latest for model-123
# Order by checked_at desc
mock_prisma.db.litellm_healthchecktable.find_many = AsyncMock(
return_value=[mock_check3, mock_check2]
)
result = await mock_prisma.get_all_latest_health_checks()
# Should return 2 unique models (by model_id)
assert len(result) == 2
# Should have latest check for each model_id
model_ids = {check.model_id for check in result}
assert "model-123" in model_ids
assert "model-456" in model_ids
# model-123 should have the latest check (1 minute ago)
model123_check = next(c for c in result if c.model_id == "model-123")
assert model123_check.checked_at == mock_check3.checked_at
@pytest.mark.asyncio
async def test_get_all_latest_health_checks_without_model_id(mock_prisma):
"""Test get_all_latest_health_checks groups by model_name when model_id is None"""
mock_check2 = MagicMock()
mock_check2.model_id = None
mock_check2.model_name = "gpt-3.5-turbo"
mock_check2.checked_at = datetime.now(timezone.utc) - timedelta(minutes=1) # Latest
mock_prisma.db.litellm_healthchecktable.find_many = AsyncMock(
return_value=[mock_check2]
)
result = await mock_prisma.get_all_latest_health_checks()
# Should return 1 unique model (by model_name)
assert len(result) == 1
assert result[0].model_name == "gpt-3.5-turbo"
assert result[0].checked_at == mock_check2.checked_at # Latest
@pytest.mark.asyncio
async def test_get_all_latest_health_checks_same_name_with_and_without_model_id(
mock_prisma,
):
async def test_get_all_latest_health_checks_keeps_every_distinct_group_with_its_own_checked_at(mock_prisma):
"""
Same model_name can appear twice after DISTINCT ON: once keyed by (model_id, name)
and once by (NULL, name) different Postgres groups than a single row with id.
Postgres owns the dedup. (id, name), (other id, name) and (NULL, name) are distinct groups and each row
must arrive typed, with its own checked_at, for the 1h re-save compare and the id-or-name lookup key.
"""
now = datetime.now(timezone.utc)
with_id = MagicMock()
with_id.model_id = "deployment-abc"
with_id.model_name = "gpt-4"
with_id.checked_at = now - timedelta(minutes=2)
without_id = MagicMock()
without_id.model_id = None
without_id.model_name = "gpt-4"
without_id.checked_at = now - timedelta(minutes=1)
mock_prisma.db.litellm_healthchecktable.find_many = AsyncMock(
return_value=[without_id, with_id]
mock_prisma.db.query_raw = AsyncMock(
return_value=[
_raw_latest_row("gpt-3.5-turbo", "model-123", now - timedelta(minutes=1)),
_raw_latest_row("gpt-3.5-turbo", "model-456", now - timedelta(minutes=5)),
_raw_latest_row("gpt-4", "deployment-abc", now - timedelta(minutes=2)),
_raw_latest_row("gpt-4", None, now - timedelta(minutes=3)),
]
)
result = await mock_prisma.get_all_latest_health_checks()
assert len(result) == 2
names = {r.model_name for r in result}
assert names == {"gpt-4"}
ids = {r.model_id for r in result}
assert "deployment-abc" in ids
assert None in ids
assert {(check.model_id, check.model_name): check.checked_at for check in result} == {
("model-123", "gpt-3.5-turbo"): now - timedelta(minutes=1),
("model-456", "gpt-3.5-turbo"): now - timedelta(minutes=5),
("deployment-abc", "gpt-4"): now - timedelta(minutes=2),
(None, "gpt-4"): now - timedelta(minutes=3),
}
by_key = {(r.model_id, r.model_name): r for r in result}
assert by_key[("deployment-abc", "gpt-4")].checked_at == with_id.checked_at
assert by_key[(None, "gpt-4")].checked_at == without_id.checked_at
@pytest.mark.asyncio
async def test_save_background_health_checks_compares_raw_checked_at_against_utc_now(mock_prisma):
"""
Raw rows carry ISO strings and the engine may omit the offset. A naive checked_at would TypeError
inside the 1h compare, be swallowed, and silently stop every save; a stale row must still re-save.
"""
stale = (datetime.now(timezone.utc) - timedelta(hours=2)).replace(tzinfo=None)
fresh = datetime.now(timezone.utc) - timedelta(minutes=5)
mock_prisma.db.query_raw = AsyncMock(
return_value=[
_raw_latest_row("stale-model", "stale-id", stale),
_raw_latest_row("fresh-model", "fresh-id", fresh),
]
)
mock_prisma.save_health_check_result = AsyncMock()
model_list = [
{"model_name": "stale-model", "model_info": {"id": "stale-id"}, "litellm_params": {"model": "openai/stale"}},
{"model_name": "fresh-model", "model_info": {"id": "fresh-id"}, "litellm_params": {"model": "openai/fresh"}},
]
await _save_background_health_checks_to_db(
mock_prisma,
model_list,
[{"model": "openai/stale"}, {"model": "openai/fresh"}],
[],
time.time(),
"pod-1",
)
await asyncio.sleep(0)
assert [call.kwargs["model_id"] for call in mock_prisma.save_health_check_result.await_args_list] == ["stale-id"]
@pytest.mark.asyncio
async def test_latest_health_checks_endpoint_serialises_raw_rows(monkeypatch):
row = LatestHealthCheckRow(
health_check_id="hc-1",
model_name="gpt-4",
model_id="deployment-abc",
status="healthy",
healthy_count=1,
unhealthy_count=0,
error_message=None,
response_time_ms=12.5,
details='{"region": "eu"}',
checked_by="pod-1",
checked_at=datetime(2026, 8, 25),
created_at=datetime(2026, 8, 25, tzinfo=timezone.utc),
updated_at=datetime(2026, 8, 25, tzinfo=timezone.utc),
)
prisma = MagicMock()
prisma.get_all_latest_health_checks = AsyncMock(return_value=(row,))
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma)
response = await latest_health_checks_endpoint(user_api_key_dict=UserAPIKeyAuth())
assert response == {
"latest_health_checks": {
"deployment-abc": {
"health_check_id": "hc-1",
"model_name": "gpt-4",
"model_id": "deployment-abc",
"status": "healthy",
"healthy_count": 1,
"unhealthy_count": 0,
"error_message": None,
"response_time_ms": 12.5,
"details": {"region": "eu"},
"checked_by": "pod-1",
"checked_at": "2026-08-25T00:00:00+00:00",
"created_at": "2026-08-25T00:00:00+00:00",
}
},
"total_models": 1,
}
@pytest.mark.asyncio
@ -628,12 +740,7 @@ def test_parse_background_health_check_model_groups_unset_returns_none():
assert parse_background_health_check_model_groups(None) is None
assert parse_background_health_check_model_groups({}) is None
assert (
parse_background_health_check_model_groups(
{"background_health_check_model_groups": None}
)
is None
)
assert parse_background_health_check_model_groups({"background_health_check_model_groups": None}) is None
def test_parse_background_health_check_model_groups_list_returns_frozenset():
@ -650,9 +757,7 @@ def test_parse_background_health_check_model_groups_malformed_raises(bad_value):
from litellm.proxy.health_check import parse_background_health_check_model_groups
with pytest.raises(ValueError, match="must be a list of model group names"):
parse_background_health_check_model_groups(
{"background_health_check_model_groups": bad_value}
)
parse_background_health_check_model_groups({"background_health_check_model_groups": bad_value})
def test_filter_deployments_to_model_groups():
@ -665,9 +770,7 @@ def test_filter_deployments_to_model_groups():
]
assert filter_deployments_to_model_groups(model_list, None) == tuple(model_list)
assert filter_deployments_to_model_groups(
model_list, frozenset({"prod-openai"})
) == (model_list[0], model_list[2])
assert filter_deployments_to_model_groups(model_list, frozenset({"prod-openai"})) == (model_list[0], model_list[2])
assert filter_deployments_to_model_groups(model_list, frozenset()) == ()

View file

@ -887,6 +887,7 @@ from typing import cast
import litellm
from litellm.proxy.utils import create_model_info_response
from litellm.types.router import DeploymentModelListingInfo
from litellm.types.utils import ModelInfo
@ -914,7 +915,7 @@ def test_create_model_info_response_includes_max_tokens_from_lookup():
def test_create_model_info_response_does_not_call_router_group_info():
router = MagicMock()
router.get_configured_token_limits.return_value = (None, None)
router.get_model_listing_info.return_value = None
response = create_model_info_response(
model_id="some-model",
@ -929,7 +930,9 @@ def test_create_model_info_response_does_not_call_router_group_info():
def test_create_model_info_response_uses_deployment_limits_when_not_in_cost_map():
router = MagicMock()
router.get_configured_token_limits.return_value = (32000, 8000)
router.get_model_listing_info.return_value = DeploymentModelListingInfo(
cost_map_keys=("my-custom-deployment",), max_input_tokens=32000, max_output_tokens=8000
)
response = create_model_info_response(
model_id="my-custom-deployment",
@ -986,7 +989,9 @@ def test_create_model_info_response_uses_deployment_mode_for_auto_router():
def test_create_model_info_response_deployment_limits_override_cost_map():
router = MagicMock()
router.get_configured_token_limits.return_value = (200000, None)
router.get_model_listing_info.return_value = DeploymentModelListingInfo(
cost_map_keys=("gpt-4o",), max_input_tokens=200000, max_output_tokens=None
)
response = create_model_info_response(
model_id="gpt-4o",
@ -999,6 +1004,54 @@ def test_create_model_info_response_deployment_limits_override_cost_map():
assert response["max_output_tokens"] == 16384
def test_create_model_info_response_reports_widest_window_in_a_mixed_group():
"""A group mixing models advertises the widest window, not whichever is listed first."""
limits = {
"small-model": _fake_model_info(max_input_tokens=200000, max_output_tokens=4096, mode="chat"),
"large-model": _fake_model_info(max_input_tokens=1000000, max_output_tokens=128000, mode="chat"),
}
for keys in (("small-model", "large-model"), ("large-model", "small-model")):
router = MagicMock()
router.get_model_listing_info.return_value = DeploymentModelListingInfo(
cost_map_keys=keys, max_input_tokens=None, max_output_tokens=None
)
response = create_model_info_response(
model_id="house-claude",
provider="openai",
llm_router=router,
get_model_info=lambda model: limits[model],
)
assert response["max_input_tokens"] == 1000000, keys
assert response["max_output_tokens"] == 128000, keys
def test_create_model_info_response_resolves_alias_once_per_listing():
"""The alias is the same for every deployment in the group, so it is looked up once."""
seen: list[str] = []
def _tracking_get_model_info(model: str) -> ModelInfo:
seen.append(model)
return _fake_model_info(max_input_tokens=128000)
router = MagicMock()
router.get_model_listing_info.return_value = DeploymentModelListingInfo(
cost_map_keys=("model-a", "model-b"), max_input_tokens=None, max_output_tokens=None
)
create_model_info_response(
model_id="house-model",
provider="openai",
llm_router=router,
get_model_info=_tracking_get_model_info,
)
assert seen.count("house-model") == 1
assert sorted(seen) == ["house-model", "model-a", "model-b"]
def test_create_model_info_response_survives_malformed_configured_limits():
from litellm import Router
@ -1960,6 +2013,122 @@ async def test_proxy_only_error_5xx_keeps_traceback_and_runs_sync_callbacks(monk
assert "test_proxy_utils" in captured["async_traceback"]
def test_create_model_info_response_resolves_alias_to_deployment_model():
"""A public model name that is not itself a cost-map key must not be resolved through
the fallback-generalization rules: `bedrock-claude-opus-5` matches the generic
claude-family baseline (200k/64k) by substring, while the deployment it fronts really
accepts 1M/128k. Regression for the /v1/models alias resolution introduced in v1.94.0."""
from litellm import Router
saved_model_cost = dict(litellm.model_cost)
try:
router = Router(
model_list=[
{
"model_name": "bedrock-claude-opus-5",
"litellm_params": {
"custom_llm_provider": "bedrock",
"model": "bedrock/eu.anthropic.claude-opus-5",
},
"model_info": {"base_model": "eu.anthropic.claude-opus-5"},
}
]
)
response = create_model_info_response(
model_id="bedrock-claude-opus-5", provider="openai", llm_router=router
)
finally:
litellm.model_cost.clear()
litellm.model_cost.update(saved_model_cost)
assert response["max_input_tokens"] == 1000000
assert response["max_output_tokens"] == 128000
def test_create_model_info_response_keeps_exact_alias_over_generalized_deployment_model():
"""Mirror of the alias bug: when the deployment points at a custom backend name that
only matches a generalization rule, the listed name's exact cost-map entry is the
better answer and must win."""
from litellm import Router
saved_model_cost = dict(litellm.model_cost)
try:
router = Router(
model_list=[
{
"model_name": "claude-opus-5",
"litellm_params": {
"custom_llm_provider": "bedrock",
"model": "bedrock/my-claude-opus-5-provisioned",
},
}
]
)
response = create_model_info_response(
model_id="claude-opus-5", provider="openai", llm_router=router
)
finally:
litellm.model_cost.clear()
litellm.model_cost.update(saved_model_cost)
assert response["max_input_tokens"] == 1000000
def test_create_model_info_response_falls_back_to_alias_for_opaque_deployment_name():
"""An Azure deployment named after the resource rather than the model has no cost-map
entry; the listed name still does, and must keep answering."""
from litellm import Router
saved_model_cost = dict(litellm.model_cost)
try:
router = Router(
model_list=[
{
"model_name": "gpt-4o",
"litellm_params": {"model": "azure/my-gpt4o-deployment"},
}
]
)
response = create_model_info_response(
model_id="gpt-4o", provider="openai", llm_router=router
)
finally:
litellm.model_cost.clear()
litellm.model_cost.update(saved_model_cost)
assert response["max_input_tokens"] == 128000
assert response["max_output_tokens"] == 16384
def test_create_model_info_response_resolves_mode_through_deployment_model():
"""`mode` is derived from the same lookup, so an aliased embedding deployment
currently reports no mode at all; it must report `embedding`."""
from litellm import Router
saved_model_cost = dict(litellm.model_cost)
try:
router = Router(
model_list=[
{
"model_name": "my-embeddings",
"litellm_params": {"model": "openai/text-embedding-3-small"},
}
]
)
response = create_model_info_response(
model_id="my-embeddings", provider="openai", llm_router=router
)
finally:
litellm.model_cost.clear()
litellm.model_cost.update(saved_model_cost)
assert response["mode"] == "embedding"
@pytest.mark.parametrize(
"key_metadata, team_metadata, expected_to_run",
[

View file

@ -22,6 +22,10 @@ from unittest.mock import AsyncMock, MagicMock
import pytest
from litellm.proxy.db.health_check_latest import (
LATEST_HEALTH_CHECKS_FOR_MODELS_SQL,
LATEST_HEALTH_CHECKS_SQL,
)
from litellm.proxy.utils import PrismaClient
@ -261,36 +265,49 @@ async def test_get_health_check_history_db_error_returns_empty_list(
assert await prisma_client.get_health_check_history() == []
def _raw_health_check_row(model_name: str = "gpt-4", model_id: str | None = "deployment-abc") -> dict[str, Any]:
return {
"health_check_id": f"hc-{model_name}",
"model_name": model_name,
"model_id": model_id,
"status": "healthy",
"healthy_count": 1,
"unhealthy_count": 0,
"error_message": None,
"response_time_ms": 12.5,
"details": None,
"checked_by": "pod-1",
"checked_at": "2026-08-25T00:00:00+00:00",
"created_at": "2026-08-25T00:00:00+00:00",
"updated_at": "2026-08-25T00:00:00+00:00",
}
@pytest.mark.asyncio
async def test_get_all_latest_health_checks_uses_distinct(
async def test_get_all_latest_health_checks_dedups_in_postgres(
prisma_client: PrismaClient,
) -> None:
rows = [MagicMock(name=f"row-{i}") for i in range(3)]
prisma_client.db.litellm_healthchecktable.find_many = AsyncMock(return_value=rows)
"""A revert to prisma find_many(distinct=...) streams the whole history table; the SQL must own the DISTINCT."""
prisma_client.db.query_raw = AsyncMock(return_value=[_raw_health_check_row()])
result = await prisma_client.get_all_latest_health_checks()
kwargs = prisma_client.db.litellm_healthchecktable.find_many.await_args.kwargs
actual = {
"len": len(result),
"distinct": kwargs["distinct"],
"order_len": len(kwargs["order"]),
"first_order": kwargs["order"][0],
"query": prisma_client.db.query_raw.await_args.args,
"rows": [(row.model_id, row.model_name, row.status) for row in result],
"row_type": type(result[0]).__name__,
}
assert actual == {
"len": 3,
"distinct": ["model_id", "model_name"],
"order_len": 3,
"first_order": {"model_id": "asc"},
"query": (LATEST_HEALTH_CHECKS_SQL,),
"rows": [("deployment-abc", "gpt-4", "healthy")],
"row_type": "LatestHealthCheckRow",
}
@pytest.mark.asyncio
async def test_get_all_latest_health_checks_db_error_returns_empty_list(
async def test_get_all_latest_health_checks_db_error_returns_no_rows(
prisma_client: PrismaClient,
) -> None:
prisma_client.db.litellm_healthchecktable.find_many = AsyncMock(
side_effect=RuntimeError("oops")
)
assert await prisma_client.get_all_latest_health_checks() == []
prisma_client.db.query_raw = AsyncMock(side_effect=RuntimeError("oops"))
assert await prisma_client.get_all_latest_health_checks() == ()
@pytest.mark.asyncio
@ -298,18 +315,15 @@ async def test_get_latest_health_checks_for_models_bounds_the_query_to_those_mod
prisma_client: PrismaClient,
) -> None:
"""A paged caller reads health for its page; an unbounded read is the bug this exists to avoid."""
prisma_client.db.litellm_healthchecktable.find_many = AsyncMock(return_value=[])
await prisma_client.get_latest_health_checks_for_models(["gpt-5", "claude-opus"])
kwargs = prisma_client.db.litellm_healthchecktable.find_many.await_args.kwargs
prisma_client.db.query_raw = AsyncMock(return_value=[_raw_health_check_row(model_name="gpt-5")])
result = await prisma_client.get_latest_health_checks_for_models(["gpt-5", "claude-opus"])
actual = {
"where": kwargs["where"],
"distinct": kwargs["distinct"],
"order": kwargs["order"],
"query": prisma_client.db.query_raw.await_args.args,
"rows": [row.model_name for row in result],
}
assert actual == {
"where": {"model_name": {"in": ["gpt-5", "claude-opus"]}},
"distinct": ["model_id", "model_name"],
"order": [{"model_id": "asc"}, {"model_name": "asc"}, {"checked_at": "desc"}],
"query": (LATEST_HEALTH_CHECKS_FOR_MODELS_SQL, ["gpt-5", "claude-opus"]),
"rows": ["gpt-5"],
}
@ -317,14 +331,14 @@ async def test_get_latest_health_checks_for_models_bounds_the_query_to_those_mod
async def test_get_latest_health_checks_for_models_does_not_query_for_an_empty_page(
prisma_client: PrismaClient,
) -> None:
prisma_client.db.litellm_healthchecktable.find_many = AsyncMock(return_value=[])
prisma_client.db.query_raw = AsyncMock(return_value=[])
assert await prisma_client.get_latest_health_checks_for_models([]) == ()
assert prisma_client.db.litellm_healthchecktable.find_many.await_count == 0
assert prisma_client.db.query_raw.await_count == 0
@pytest.mark.asyncio
async def test_get_latest_health_checks_for_models_db_error_returns_empty_list(
async def test_get_latest_health_checks_for_models_db_error_returns_no_rows(
prisma_client: PrismaClient,
) -> None:
prisma_client.db.litellm_healthchecktable.find_many = AsyncMock(side_effect=RuntimeError("oops"))
prisma_client.db.query_raw = AsyncMock(side_effect=RuntimeError("oops"))
assert await prisma_client.get_latest_health_checks_for_models(["gpt-5"]) == ()

View file

@ -38,12 +38,16 @@ _STUB_TEMPLATE = """#!/bin/sh
_ENTRYPOINT_RE = re.compile(r"^ENTRYPOINT\s+(\[.*\])\s*$", re.MULTILINE)
_CMD_RE = re.compile(r"^CMD\s+(\[.*\])\s*$", re.MULTILINE)
_COPY_RE = re.compile(r"^COPY\s+(?!--from)(\S+)\s+(\S+)\s*$", re.MULTILINE)
_APP_TARGET_RE = re.compile(r"(?:gateway|backend)\.main:app")
_APP_TARGET_RE = re.compile(r"(?:gateway|backend)\.main:app|gateway\.launch")
_TF_STRING_LOCAL_RE = re.compile(r'^\s*(\w+)\s*=\s*"((?:[^"\\]|\\.)*)"\s*$', re.MULTILINE)
_TF_INTERPOLATION_RE = re.compile(r"\$\{(local|var)\.(\w+)\}")
TERRAFORM_LAUNCH_SITES = {TERRAFORM_ECS: 2, TERRAFORM_CLOUDRUN: 2}
TERRAFORM_VAR_STUBS = {"gateway_num_workers": "2"}
COMPONENT_LAUNCHERS = {
"gateway": ("python", "-m", "gateway.launch"),
"backend": ("uvicorn", "backend.main:app"),
}
_MAX_INTERPOLATION_PASSES = 5
@ -63,7 +67,7 @@ def _run_entrypoint(
"""Run `script` with stubbed executables on PATH and return the recorded lines."""
bin_dir = tmp_path / "bin"
bin_dir.mkdir(parents=True)
_write_stubs(bin_dir, ("ddtrace-run", "uvicorn", "litellm"))
_write_stubs(bin_dir, ("ddtrace-run", "uvicorn", "python", "litellm"))
record = tmp_path / "record.txt"
env = {
@ -330,22 +334,63 @@ def test_entrypoint_script_has_no_carriage_returns() -> None:
@pytest.mark.parametrize(
"dockerfile, app_target",
"dockerfile, launcher",
[
(GATEWAY_DOCKERFILE, "gateway.main:app"),
(BACKEND_DOCKERFILE, "backend.main:app"),
(GATEWAY_DOCKERFILE, "python -m gateway.launch"),
(BACKEND_DOCKERFILE, "uvicorn backend.main:app"),
],
)
def test_component_images_launch_uvicorn_through_the_entrypoint(dockerfile: Path, app_target: str) -> None:
def test_component_images_launch_uvicorn_through_the_entrypoint(dockerfile: Path, launcher: str) -> None:
entrypoint = " ".join(_entrypoint_argv(dockerfile))
assert IMAGE_ENTRYPOINT_PATH in entrypoint, f"{dockerfile} bypasses the ddtrace-aware entrypoint"
assert app_target in entrypoint
assert entrypoint.index(IMAGE_ENTRYPOINT_PATH) < entrypoint.index("uvicorn"), (
assert launcher in entrypoint
assert entrypoint.index(IMAGE_ENTRYPOINT_PATH) < entrypoint.index(launcher), (
f"{dockerfile} must invoke uvicorn through the entrypoint, not the other way around"
)
@pytest.mark.parametrize(
"use_ddtrace, num_workers, expected_exec, expected_args",
[
(None, "4", "exec=python", "args=-m gateway.launch --workers 4 --host 0.0.0.0 --port 4000"),
(None, None, "exec=python", "args=-m gateway.launch --workers 1 --host 0.0.0.0 --port 4000"),
("true", "4", "exec=ddtrace-run", "args=python -m gateway.launch --workers 4 --host 0.0.0.0 --port 4000"),
],
)
def test_gateway_image_execs_the_supervisor_with_its_worker_count(
use_ddtrace: str | None, num_workers: str | None, expected_exec: str, expected_args: str, tmp_path: Path
) -> None:
"""Run the gateway image's ENTRYPOINT + CMD and record what the container execs.
The Dockerfile's `/app/...` script path is resolved to the checked-in script and `python`
is stubbed on PATH, so the assertion is on the argv `gateway.launch` receives, not on the
Dockerfile text.
"""
entrypoint = tuple(
part.replace(IMAGE_ENTRYPOINT_PATH, str(COMPONENT_ENTRYPOINT)) for part in _entrypoint_argv(GATEWAY_DOCKERFILE)
)
bin_dir = tmp_path / "bin"
bin_dir.mkdir(parents=True)
_write_stubs(bin_dir, ("ddtrace-run", "python", "uvicorn"))
record = tmp_path / "record.txt"
overrides = {"USE_DDTRACE": use_ddtrace, "NUM_WORKERS": num_workers}
env = {
**{k: v for k, v in os.environ.items() if k not in ("DD_TRACE_OPENAI_ENABLED", *overrides)},
**{k: v for k, v in overrides.items() if v is not None},
"PATH": f"{bin_dir}{os.pathsep}{os.environ['PATH']}",
"RECORD": str(record),
"PYTHONPATH": PYTHONPATH_SENTINEL,
}
result = subprocess.run(
[*entrypoint, *_cmd_argv(GATEWAY_DOCKERFILE)], env=env, capture_output=True, text=True, check=False
)
assert result.returncode == 0, f"stdout={result.stdout} stderr={result.stderr}"
assert tuple(record.read_text().splitlines())[:2] == (expected_exec, expected_args)
@pytest.mark.parametrize("dockerfile", [GATEWAY_DOCKERFILE, BACKEND_DOCKERFILE])
def test_component_images_make_the_entrypoint_executable(dockerfile: Path) -> None:
body = dockerfile.read_text()
@ -367,17 +412,18 @@ def test_terraform_launch_command_matches_the_script_contract(
implementations under the same environment and asserts they agree on which binary is exec'd
and on whether the openai integration is disabled.
"""
app_target = f"{component}.main:app"
launcher = COMPONENT_LAUNCHERS[component]
app_target = " ".join(launcher[1:])
command = _resolve_tf_local(terraform_file, f"{component}_launch_cmd")
bin_dir = tmp_path / "bin"
bin_dir.mkdir(parents=True)
_write_stubs(bin_dir, ("ddtrace-run", "uvicorn"))
_write_stubs(bin_dir, ("ddtrace-run", "uvicorn", "python"))
from_terraform = _run_shell_command(command, bin_dir, tmp_path / "terraform.txt", use_ddtrace)
from_script = _run_entrypoint(
COMPONENT_ENTRYPOINT,
("uvicorn", app_target),
launcher,
use_ddtrace=use_ddtrace,
tmp_path=tmp_path / "script",
)
@ -387,12 +433,14 @@ def test_terraform_launch_command_matches_the_script_contract(
)
assert from_terraform[2] == from_script[2], f"{terraform_file} disagrees with the script on the openai integration"
assert app_target in from_terraform[1]
assert "gateway.main:app" not in from_terraform[1], f"{terraform_file} bypasses the gateway.launch supervisor"
if use_ddtrace in TRUTHY_USE_DDTRACE:
assert from_terraform[0] == "exec=ddtrace-run"
assert from_terraform[1].startswith(f"args={launcher[0]} ")
assert from_terraform[2] == "DD_TRACE_OPENAI_ENABLED=False"
else:
assert from_terraform[0] == "exec=uvicorn"
assert from_terraform[0] == f"exec={launcher[0]}"
assert from_terraform[2] == "DD_TRACE_OPENAI_ENABLED=<unset>"

View file

@ -10,13 +10,16 @@ from redis.credentials import CredentialProvider
import litellm
from litellm._redis import (
_AWS_IAM_KWARG_NAMES,
_async_auth_kwargs,
_coerce_redis_kwargs_types,
_get_redis_client_logic,
_get_redis_cluster_kwargs,
_get_redis_env_kwarg_mapping,
_get_redis_kwargs,
_get_redis_url_kwargs,
_pretty_print_redis_config,
_uses_tls,
get_redis_async_client,
get_redis_client,
get_redis_connection_pool,
@ -24,12 +27,14 @@ from litellm._redis import (
)
from litellm._redis_credential_provider import (
AzureADCredentialProvider,
ElastiCacheIAMCredentialProvider,
GCPIAMCredentialProvider,
_token_cache,
)
from litellm.caching.redis_cache import RedisCache
from litellm.caching.redis_cluster_cache import RedisClusterCache
from litellm.constants import REDIS_CLUSTER_HEALTH_CHECK_INTERVAL
from litellm.proxy._types import CoordinationRedisParams
class _StubCredentialProvider(CredentialProvider):
@ -78,6 +83,8 @@ def clean_redis_environment(monkeypatch):
"REDIS_URL",
"REDIS_CLUSTER_NODES",
"REDIS_SENTINEL_NODES",
"AWS_REGION",
"AWS_DEFAULT_REGION",
*_get_redis_env_kwarg_mapping(),
):
monkeypatch.delenv(var, raising=False)
@ -110,6 +117,29 @@ def test_credential_provider_is_not_environment_derived():
assert "credential_provider" not in mapping.values()
_AWS_IAM_SETTINGS = {
"aws_iam_auth",
"aws_iam_user_name",
"aws_iam_cache_name",
"aws_iam_region",
"aws_iam_serverless",
}
def test_aws_iam_settings_are_environment_derived():
allowed = _get_redis_kwargs()
mapping = _get_redis_env_kwarg_mapping()
assert _AWS_IAM_SETTINGS <= allowed
assert set(_AWS_IAM_KWARG_NAMES) == _AWS_IAM_SETTINGS
assert {f for f in CoordinationRedisParams.model_fields if f.startswith("aws_iam_")} == _AWS_IAM_SETTINGS
assert mapping["REDIS_AWS_IAM_AUTH"] == "aws_iam_auth"
assert mapping["REDIS_AWS_IAM_USER_NAME"] == "aws_iam_user_name"
assert mapping["REDIS_AWS_IAM_CACHE_NAME"] == "aws_iam_cache_name"
assert mapping["REDIS_AWS_IAM_REGION"] == "aws_iam_region"
assert mapping["REDIS_AWS_IAM_SERVERLESS"] == "aws_iam_serverless"
def test_sync_direct_preserves_credential_provider_identity(clean_redis_environment):
provider = _StubCredentialProvider()
@ -300,6 +330,333 @@ def test_gcp_kwargs_never_survive_client_logic(clean_redis_environment, override
assert "gcp_ssl_ca_certs" not in redis_kwargs
def test_aws_iam_environment_settings_install_provider(clean_redis_environment, monkeypatch):
monkeypatch.setenv("REDIS_AWS_IAM_AUTH", "true")
monkeypatch.setenv("REDIS_AWS_IAM_USER_NAME", "iam-user")
monkeypatch.setenv("REDIS_AWS_IAM_CACHE_NAME", "cache.example.com")
monkeypatch.setenv("REDIS_AWS_IAM_REGION", "us-east-1")
monkeypatch.setenv("REDIS_AWS_IAM_SERVERLESS", "1")
monkeypatch.setenv("REDIS_SSL", "1")
redis_kwargs = _get_redis_client_logic(host="cache.example.com", port=6379)
provider = redis_kwargs["credential_provider"]
assert isinstance(provider, ElastiCacheIAMCredentialProvider)
assert provider._is_serverless is True
assert not _AWS_IAM_SETTINGS & redis_kwargs.keys()
@pytest.mark.parametrize(
"transport",
[
pytest.param({"host": "cache.example.com", "port": 6379}, id="host_without_ssl"),
pytest.param({"host": "cache.example.com", "port": 6379, "ssl": False}, id="host_ssl_false"),
pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "false"}, id="host_ssl_false_string"),
pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "0"}, id="host_ssl_zero_string"),
pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "no"}, id="host_ssl_no_string"),
pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "off"}, id="host_ssl_off_string"),
pytest.param({"url": "redis://cache.example.com:6379", "ssl": True}, id="plaintext_url"),
pytest.param(
{"startup_nodes": [{"host": "cache.example.com", "port": 6379}]},
id="cluster_without_ssl",
),
pytest.param(
{
"url": "rediss://cache.example.com:6379",
"startup_nodes": [{"host": "cache.example.com", "port": 6379}],
},
id="cluster_without_ssl_ignores_url_scheme",
),
pytest.param(
{
"sentinel_nodes": [("sentinel.example.com", 26379)],
"service_name": "cache",
},
id="sentinel_without_ssl",
),
],
)
def test_aws_iam_auth_rejects_non_tls_connections(clean_redis_environment, transport):
with pytest.raises(ValueError, match="requires TLS"):
_get_redis_client_logic(
**transport,
aws_iam_auth=True,
aws_iam_user_name="iam-user",
aws_iam_cache_name="cache.example.com",
aws_iam_region="us-east-1",
)
@pytest.mark.parametrize(
"transport",
[
pytest.param({"host": "cache.example.com", "port": 6379, "ssl": True}, id="host"),
pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "true"}, id="host_ssl_true_string"),
pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "True"}, id="host_ssl_true_capitalized"),
pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "1"}, id="host_ssl_one_string"),
pytest.param({"host": "cache.example.com", "port": 6379, "ssl": "yes"}, id="host_ssl_yes_string"),
pytest.param(
{"startup_nodes": [{"host": "cache.example.com", "port": 6379}], "ssl": "1"},
id="cluster_ssl_one_string",
),
pytest.param({"url": "rediss://cache.example.com:6379"}, id="url"),
pytest.param(
{
"startup_nodes": [{"host": "cache.example.com", "port": 6379}],
"ssl": True,
},
id="cluster",
),
pytest.param(
{
"sentinel_nodes": [("sentinel.example.com", 26379)],
"service_name": "cache",
"ssl": True,
},
id="sentinel",
),
],
)
def test_aws_iam_auth_accepts_tls_connections(clean_redis_environment, transport):
redis_kwargs = _get_redis_client_logic(
**transport,
aws_iam_auth=True,
aws_iam_user_name="iam-user",
aws_iam_cache_name="cache.example.com",
aws_iam_region="us-east-1",
)
assert isinstance(redis_kwargs["credential_provider"], ElastiCacheIAMCredentialProvider)
@pytest.mark.parametrize("ssl", ["true", "True", "TRUE", "1", "yes", "YES", "false", "0", "no", "off", "", "maybe"])
def test_tls_detection_agrees_with_the_ssl_kwarg_coercion(ssl):
assert _uses_tls({"ssl": ssl}) is _coerce_redis_kwargs_types({"ssl": ssl})["ssl"]
def test_aws_iam_settings_are_removed_for_url_and_static_credentials(clean_redis_environment):
redis_kwargs = _get_redis_client_logic(
url="rediss://url-user:url-pass@cache.example.com:6380",
aws_iam_auth=True,
aws_iam_user_name="iam-user",
aws_iam_cache_name="cache.example.com",
aws_iam_region="us-east-1",
username="static-user",
password="static-password",
)
assert isinstance(redis_kwargs["credential_provider"], ElastiCacheIAMCredentialProvider)
assert redis_kwargs["url"] == "rediss://cache.example.com:6380"
assert "username" not in redis_kwargs
assert "password" not in redis_kwargs
assert not _AWS_IAM_SETTINGS & redis_kwargs.keys()
@pytest.mark.parametrize("missing", ["aws_iam_user_name", "aws_iam_cache_name", "aws_iam_region"])
def test_aws_iam_missing_setting_fails_closed(clean_redis_environment, missing):
settings = {
"aws_iam_auth": True,
"aws_iam_user_name": "iam-user",
"aws_iam_cache_name": "cache.example.com",
"aws_iam_region": "us-east-1",
"ssl": True,
}
settings[missing] = None
with pytest.raises(ValueError, match=missing):
_get_redis_client_logic(host="cache.example.com", port=6379, **settings)
@pytest.mark.parametrize("region_var", ["AWS_REGION", "AWS_DEFAULT_REGION"])
def test_aws_iam_region_falls_back_to_environment(clean_redis_environment, monkeypatch, region_var):
monkeypatch.setenv(region_var, "sa-east-1")
redis_kwargs = _get_redis_client_logic(
host="cache.example.com",
port=6379,
ssl=True,
aws_iam_auth=True,
aws_iam_user_name="iam-user",
aws_iam_cache_name="cache.example.com",
)
provider = redis_kwargs["credential_provider"]
assert isinstance(provider, ElastiCacheIAMCredentialProvider)
assert provider._region == "sa-east-1"
def test_aws_iam_region_prefers_aws_region_over_default_region(clean_redis_environment, monkeypatch):
monkeypatch.setenv("AWS_REGION", "sa-east-1")
monkeypatch.setenv("AWS_DEFAULT_REGION", "eu-west-1")
redis_kwargs = _get_redis_client_logic(
host="cache.example.com",
port=6379,
ssl=True,
aws_iam_auth=True,
aws_iam_user_name="iam-user",
aws_iam_cache_name="cache.example.com",
)
assert redis_kwargs["credential_provider"]._region == "sa-east-1"
def test_aws_iam_region_prefers_explicit_over_environment(clean_redis_environment, monkeypatch):
monkeypatch.setenv("AWS_REGION", "sa-east-1")
redis_kwargs = _get_redis_client_logic(
host="cache.example.com",
port=6379,
ssl=True,
aws_iam_auth=True,
aws_iam_user_name="iam-user",
aws_iam_cache_name="cache.example.com",
aws_iam_region="explicit-region",
)
assert redis_kwargs["credential_provider"]._region == "explicit-region"
def test_aws_iam_settings_map_to_distinct_provider_fields(clean_redis_environment):
redis_kwargs = _get_redis_client_logic(
host="cache.example.com",
port=6379,
ssl=True,
aws_iam_auth=True,
aws_iam_user_name="iam-user-value",
aws_iam_cache_name="iam-cache-value",
aws_iam_region="iam-region-value",
)
provider = redis_kwargs["credential_provider"]
assert isinstance(provider, ElastiCacheIAMCredentialProvider)
assert provider._user_name == "iam-user-value"
assert provider._cache_name == "iam-cache-value"
assert provider._region == "iam-region-value"
@pytest.mark.parametrize("aws_iam_auth", [None, False, "", "false", "0", "no", "off"])
def test_aws_iam_auth_disabled_does_not_install_provider(clean_redis_environment, aws_iam_auth):
redis_kwargs = _get_redis_client_logic(
host="cache.example.com",
port=6379,
aws_iam_auth=aws_iam_auth,
aws_iam_user_name="iam-user",
aws_iam_cache_name="cache.example.com",
aws_iam_region="us-east-1",
)
assert "credential_provider" not in redis_kwargs
assert not _AWS_IAM_SETTINGS & redis_kwargs.keys()
@pytest.mark.parametrize("aws_iam_auth", [True, "true", "True", "TRUE", "1", "yes"])
def test_aws_iam_auth_enabled_by_any_truthy_flag(clean_redis_environment, aws_iam_auth):
redis_kwargs = _get_redis_client_logic(
host="cache.example.com",
port=6379,
ssl=True,
aws_iam_auth=aws_iam_auth,
aws_iam_user_name="iam-user",
aws_iam_cache_name="cache-name",
aws_iam_region="us-east-1",
)
assert isinstance(redis_kwargs["credential_provider"], ElastiCacheIAMCredentialProvider)
@pytest.mark.parametrize(
"aws_iam_serverless, expected",
[
pytest.param(None, False, id="unset"),
pytest.param(False, False, id="bool_false"),
pytest.param("false", False, id="string_false"),
pytest.param("0", False, id="string_zero"),
pytest.param(True, True, id="bool_true"),
pytest.param("true", True, id="string_true"),
pytest.param("1", True, id="string_one"),
],
)
def test_aws_iam_serverless_flag_reaches_the_provider(clean_redis_environment, aws_iam_serverless, expected):
redis_kwargs = _get_redis_client_logic(
host="cache.example.com",
port=6379,
ssl=True,
aws_iam_auth=True,
aws_iam_user_name="iam-user",
aws_iam_cache_name="cache-name",
aws_iam_region="us-east-1",
aws_iam_serverless=aws_iam_serverless,
)
provider = redis_kwargs["credential_provider"]
assert isinstance(provider, ElastiCacheIAMCredentialProvider)
assert provider._is_serverless is expected
assert "aws_iam_serverless" not in redis_kwargs
def test_explicit_provider_wins_over_aws_iam(clean_redis_environment):
provider = _StubCredentialProvider()
redis_kwargs = _get_redis_client_logic(
host="cache.example.com",
port=6379,
credential_provider=provider,
aws_iam_auth=True,
aws_iam_user_name="iam-user",
aws_iam_cache_name="cache.example.com",
aws_iam_region="us-east-1",
)
assert redis_kwargs["credential_provider"] is provider
assert "aws_iam_auth" not in redis_kwargs
def test_gcp_wins_over_aws_iam(clean_redis_environment):
redis_kwargs = _get_redis_client_logic(
host="cache.example.com",
port=6379,
gcp_service_account="sa@example.com",
aws_iam_auth=True,
aws_iam_user_name="iam-user",
aws_iam_cache_name="cache.example.com",
aws_iam_region="us-east-1",
)
assert "credential_provider" not in redis_kwargs
assert redis_kwargs["redis_connect_func"]._gcp_service_account == "sa@example.com"
def test_azure_wins_over_aws_iam(clean_redis_environment):
redis_kwargs = _get_redis_client_logic(
host="cache.example.com",
port=6379,
azure_redis_ad_token="true",
aws_iam_auth=True,
aws_iam_user_name="iam-user",
aws_iam_cache_name="cache.example.com",
aws_iam_region="us-east-1",
)
assert "credential_provider" not in redis_kwargs
assert redis_kwargs["redis_connect_func"]._azure_redis_ad_token is True
def test_async_cluster_installs_aws_iam_provider(clean_redis_environment):
startup_nodes = [{"host": "cluster-node", "port": 6379}]
client = get_redis_async_client(
startup_nodes=startup_nodes,
ssl=True,
aws_iam_auth=True,
aws_iam_user_name="iam-user",
aws_iam_cache_name="cache.example.com",
aws_iam_region="us-east-1",
)
assert isinstance(client.connection_kwargs["credential_provider"], ElastiCacheIAMCredentialProvider)
def test_provider_keeps_the_rest_of_the_url_intact(clean_redis_environment):
provider = _StubCredentialProvider()

View file

@ -0,0 +1,236 @@
import asyncio
import builtins
import sys
from types import SimpleNamespace
from urllib.parse import parse_qs, urlsplit
import pytest
from litellm._redis_credential_provider import ElastiCacheIAMCredentialProvider
class _FakeCredentials:
def __init__(self, access_key: str) -> None:
self.access_key = access_key
self.secret_key = "synthetic-secret"
self.token = "synthetic-session-token"
def get_frozen_credentials(self):
return self
class _RotatingFakeCredentials:
def __init__(self) -> None:
self.calls = 0
def __bool__(self) -> bool:
return False
def get_frozen_credentials(self):
self.calls += 1
return SimpleNamespace(
access_key=f"AKIA-SYNTHETIC-{self.calls}",
secret_key="synthetic-secret",
token="synthetic-session-token",
)
class _FakeResolver:
def __init__(self, credentials: _FakeCredentials | _RotatingFakeCredentials | None) -> None:
self.credentials = credentials
self.calls = 0
def __call__(self):
self.calls += 1
return self.credentials
def test_elasticache_provider_signs_expected_query():
resolver = _FakeResolver(_FakeCredentials("AKIA-SYNTHETIC"))
provider = ElastiCacheIAMCredentialProvider(
user_name="iam-user",
cache_name="cache.example.com",
region="us-east-1",
credentials_resolver=resolver,
)
user_name, token = provider.get_credentials()
parsed = urlsplit("https://" + token)
query = parse_qs(parsed.query)
assert user_name == "iam-user"
assert parsed.netloc == "cache.example.com"
assert query["Action"] == ["connect"]
assert query["User"] == ["iam-user"]
assert query["X-Amz-Expires"] == ["900"]
assert "elasticache" in query["X-Amz-Credential"][0]
assert query["X-Amz-Credential"][0].split("/")[2] == "us-east-1"
assert not token.startswith("https://")
def test_elasticache_provider_resolves_credentials_once_but_refreshes_signature():
rotating_credentials = _RotatingFakeCredentials()
resolver = _FakeResolver(rotating_credentials)
provider = ElastiCacheIAMCredentialProvider(
user_name="iam-user",
cache_name="cache.example.com",
region="us-east-1",
credentials_resolver=resolver,
)
first = provider.get_credentials()
second = provider.get_credentials()
async_result = asyncio.run(provider.get_credentials_async())
assert first[0] == second[0] == async_result[0] == "iam-user"
assert first[1] != second[1]
assert async_result[1] != second[1]
assert resolver.calls == 1
assert rotating_credentials.calls == 3
def test_elasticache_provider_uses_botocore_session_credentials(monkeypatch):
credentials = _FakeCredentials("AKIA-SYNTHETIC")
monkeypatch.setattr("botocore.session.get_session", lambda: SimpleNamespace(get_credentials=lambda: credentials))
provider = ElastiCacheIAMCredentialProvider(
user_name="iam-user",
cache_name="cache.example.com",
region="us-east-1",
)
user_name, token = provider.get_credentials()
assert user_name == "iam-user"
assert "AKIA-SYNTHETIC" in token
def test_elasticache_provider_reports_missing_botocore(monkeypatch):
original_import = builtins.__import__
def import_without_botocore(name, *args, **kwargs):
if name == "botocore.session":
raise ImportError("synthetic missing dependency")
return original_import(name, *args, **kwargs)
monkeypatch.delitem(sys.modules, "botocore.session", raising=False)
monkeypatch.setattr(builtins, "__import__", import_without_botocore)
provider = ElastiCacheIAMCredentialProvider(
user_name="iam-user",
cache_name="cache.example.com",
region="us-east-1",
)
with pytest.raises(ImportError, match="pip install boto3"):
provider.get_credentials()
def test_elasticache_provider_reports_missing_credentials():
provider = ElastiCacheIAMCredentialProvider(
user_name="iam-user",
cache_name="cache.example.com",
region="us-east-1",
credentials_resolver=_FakeResolver(None),
)
with pytest.raises(RuntimeError, match="Unable to resolve AWS credentials"):
provider.get_credentials()
def test_elasticache_provider_reports_missing_signing_dependency(monkeypatch):
original_import = builtins.__import__
def import_without_botocore_auth(name, *args, **kwargs):
if name == "botocore.auth":
raise ImportError("synthetic missing dependency")
return original_import(name, *args, **kwargs)
monkeypatch.delitem(sys.modules, "botocore.auth", raising=False)
monkeypatch.setattr(builtins, "__import__", import_without_botocore_auth)
provider = ElastiCacheIAMCredentialProvider(
user_name="iam-user",
cache_name="cache.example.com",
region="us-east-1",
credentials_resolver=_FakeResolver(_FakeCredentials("AKIA-SYNTHETIC")),
)
with pytest.raises(ImportError, match="pip install boto3"):
provider.get_credentials()
def test_elasticache_provider_recovers_after_a_failed_resolution():
resolver = _FakeResolver(None)
provider = ElastiCacheIAMCredentialProvider(
user_name="iam-user",
cache_name="cache.example.com",
region="us-east-1",
credentials_resolver=resolver,
)
with pytest.raises(RuntimeError, match="Unable to resolve AWS credentials"):
provider.get_credentials()
resolver.credentials = _FakeCredentials("AKIA-SYNTHETIC")
user_name, token = provider.get_credentials()
assert user_name == "iam-user"
assert token
assert resolver.calls == 2
@pytest.mark.parametrize(
"provider_kwargs, expected_operation_params",
[
pytest.param({}, frozenset({"Action", "User"}), id="default_is_self_designed"),
pytest.param({"is_serverless": False}, frozenset({"Action", "User"}), id="self_designed"),
pytest.param({"is_serverless": True}, frozenset({"Action", "User", "ResourceType"}), id="serverless"),
],
)
def test_elasticache_provider_signs_resource_type_only_for_serverless(provider_kwargs, expected_operation_params):
provider = ElastiCacheIAMCredentialProvider(
user_name="iam-user",
cache_name="cache-name",
region="us-east-1",
credentials_resolver=_FakeResolver(_FakeCredentials("AKIA-SYNTHETIC")),
**provider_kwargs,
)
_, token = provider.get_credentials()
query_string = urlsplit("https://" + token).query
param_names = tuple(pair.split("=", 1)[0] for pair in query_string.split("&"))
first_auth_param = next(i for i, name in enumerate(param_names) if name.startswith("X-Amz-"))
query = parse_qs(query_string)
assert frozenset(param_names[:first_auth_param]) == expected_operation_params
assert all(name.startswith("X-Amz-") for name in param_names[first_auth_param:])
assert query.get("ResourceType") == (["ServerlessCache"] if "ResourceType" in expected_operation_params else None)
assert query["X-Amz-Signature"]
def test_elasticache_provider_lowercases_the_cache_name():
provider = ElastiCacheIAMCredentialProvider(
user_name="iam-user",
cache_name="Mixed-Case-Cache",
region="us-east-1",
credentials_resolver=_FakeResolver(_FakeCredentials("AKIA-SYNTHETIC")),
)
_, token = provider.get_credentials()
assert urlsplit("https://" + token).netloc == "mixed-case-cache"
def test_elasticache_provider_encodes_reserved_characters_in_the_user_name():
user_name = "iam user/with+reserved&chars"
provider = ElastiCacheIAMCredentialProvider(
user_name=user_name,
cache_name="cache-name",
region="us-east-1",
credentials_resolver=_FakeResolver(_FakeCredentials("AKIA-SYNTHETIC")),
)
returned_user_name, token = provider.get_credentials()
query = parse_qs(urlsplit("https://" + token).query)
assert returned_user_name == user_name
assert query["User"] == [user_name]
assert query["Action"] == ["connect"]

View file

@ -8208,6 +8208,204 @@ def test_get_configured_token_limits_coerces_numeric_strings():
assert router.get_configured_token_limits("quoted-limits-model") == (32000, 8000)
def test_get_model_listing_info_prefers_base_model_over_litellm_params_model():
"""The cost-map key comes from base_model when set, so a deployment pointing at an
opaque backend name still resolves the real catalog entry."""
router = litellm.Router(
model_list=[
{
"model_name": "bedrock-claude-opus-5",
"litellm_params": {"model": "bedrock/eu.anthropic.claude-opus-5"},
"model_info": {"base_model": "eu.anthropic.claude-opus-5"},
}
]
)
info = router.get_model_listing_info("bedrock-claude-opus-5")
assert info is not None
assert info.cost_map_keys == ("eu.anthropic.claude-opus-5",)
def test_get_model_listing_info_falls_back_to_litellm_params_model():
router = litellm.Router(
model_list=[
{
"model_name": "bedrock-claude-opus-5",
"litellm_params": {"model": "bedrock/eu.anthropic.claude-opus-5"},
}
]
)
info = router.get_model_listing_info("bedrock-claude-opus-5")
assert info is not None
assert info.cost_map_keys == ("bedrock/eu.anthropic.claude-opus-5",)
def test_get_model_listing_info_ignores_blank_base_model():
"""A base_model set to an empty string is absent, not a cost-map key."""
router = litellm.Router(
model_list=[
{
"model_name": "bedrock-claude-opus-5",
"litellm_params": {"model": "bedrock/eu.anthropic.claude-opus-5"},
"model_info": {"base_model": ""},
}
]
)
info = router.get_model_listing_info("bedrock-claude-opus-5")
assert info is not None
assert info.cost_map_keys == ("bedrock/eu.anthropic.claude-opus-5",)
def test_get_model_listing_info_returns_none_for_unknown_name():
router = litellm.Router(
model_list=[
{
"model_name": "no-limits-model",
"litellm_params": {"model": "openai/some-unmapped-model"},
}
]
)
assert router.get_model_listing_info("not-a-real-model") is None
def test_get_model_listing_info_carries_configured_limits():
router = litellm.Router(
model_list=[
{
"model_name": "my-custom-model",
"litellm_params": {"model": "openai/some-unmapped-model"},
"model_info": {"max_input_tokens": 32000, "max_output_tokens": 8000},
}
]
)
info = router.get_model_listing_info("my-custom-model")
assert info is not None
assert (info.max_input_tokens, info.max_output_tokens) == (32000, 8000)
def test_widest_configured_limit_ignores_absent_and_malformed_values():
model_infos = (
{"max_input_tokens": 32000},
{},
{"max_input_tokens": "not-a-number"},
{"max_input_tokens": "128000"},
{"max_output_tokens": 4096},
)
assert litellm.Router._widest_configured_limit(model_infos, "max_input_tokens") == 128000
assert litellm.Router._widest_configured_limit(model_infos, "max_output_tokens") == 4096
assert litellm.Router._widest_configured_limit((), "max_input_tokens") is None
def test_get_model_listing_info_dedupes_interchangeable_deployments():
"""The ordinary group is N deployments of one model, so it yields exactly one key."""
router = litellm.Router(
model_list=[
{
"model_name": "gpt-4o",
"litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-a"},
},
{
"model_name": "gpt-4o",
"litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-b"},
},
]
)
info = router.get_model_listing_info("gpt-4o")
assert info is not None
assert info.cost_map_keys == ("openai/gpt-4o",)
def test_get_model_listing_info_collects_every_model_in_a_mixed_group():
router = litellm.Router(
model_list=[
{
"model_name": "house-claude",
"litellm_params": {"model": "anthropic/claude-3-haiku-20240307"},
},
{
"model_name": "house-claude",
"litellm_params": {"model": "bedrock/eu.anthropic.claude-opus-5"},
},
]
)
info = router.get_model_listing_info("house-claude")
assert info is not None
assert info.cost_map_keys == (
"anthropic/claude-3-haiku-20240307",
"bedrock/eu.anthropic.claude-opus-5",
)
def test_get_model_listing_info_reports_widest_configured_limits_in_a_mixed_group():
"""Matches how get_model_group_info aggregates for the Admin UI, so the two agree."""
router = litellm.Router(
model_list=[
{
"model_name": "house-model",
"litellm_params": {"model": "openai/some-unmapped-model"},
"model_info": {"max_input_tokens": 32000, "max_output_tokens": 4096},
},
{
"model_name": "house-model",
"litellm_params": {"model": "openai/another-unmapped-model"},
"model_info": {"max_input_tokens": 128000, "max_output_tokens": 16384},
},
]
)
info = router.get_model_listing_info("house-model")
assert info is not None
assert (info.max_input_tokens, info.max_output_tokens) == (128000, 16384)
def test_get_model_listing_info_reads_base_model_from_litellm_params():
"""base_model resolution mirrors get_router_model_info, which also accepts it there."""
router = litellm.Router(
model_list=[
{
"model_name": "azure-deployment",
"litellm_params": {
"model": "azure/my-azure-deployment-name",
"base_model": "azure/gpt-4o",
"api_key": "sk-a",
"api_base": "https://example.openai.azure.com",
},
}
]
)
info = router.get_model_listing_info("azure-deployment")
assert info is not None
assert info.cost_map_keys == ("azure/gpt-4o",)
def test_get_model_listing_info_skips_wildcard_pattern_matching():
router = litellm.Router(
model_list=[
{
"model_name": "bedrock/*",
"litellm_params": {"model": "bedrock/*"},
"model_info": {"max_input_tokens": 12345},
}
]
)
with patch.object(
router.pattern_router, "route", side_effect=AssertionError("pattern route called")
):
assert (
router.get_model_listing_info("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0")
is None
)
def test_get_configured_mode_reads_deployment_model_info():
router = litellm.Router(
model_list=[

View file

@ -26371,6 +26371,31 @@ export interface components {
* independently of the response-cache backend in `litellm_settings.cache_params`.
*/
CoordinationRedisParams: {
/**
* Aws Iam Auth
* @description enable AWS ElastiCache IAM authentication
*/
aws_iam_auth?: boolean | string | null;
/**
* Aws Iam Cache Name
* @description AWS ElastiCache cache name
*/
aws_iam_cache_name?: string | null;
/**
* Aws Iam Region
* @description AWS region for ElastiCache IAM authentication
*/
aws_iam_region?: string | null;
/**
* Aws Iam Serverless
* @description the ElastiCache cache is serverless rather than a self-designed cluster
*/
aws_iam_serverless?: boolean | string | null;
/**
* Aws Iam User Name
* @description AWS ElastiCache IAM user name
*/
aws_iam_user_name?: string | null;
/**
* Host
* @description Redis hostname