mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_sanitize_unknown_model_spend_rows
This commit is contained in:
commit
3a0dce6bc4
77 changed files with 6649 additions and 407 deletions
2
.github/ISSUE_TEMPLATE/config.yml
vendored
2
.github/ISSUE_TEMPLATE/config.yml
vendored
|
|
@ -1,4 +1,4 @@
|
|||
blank_issues_enabled: true
|
||||
blank_issues_enabled: false
|
||||
contact_links:
|
||||
- name: Schedule Demo
|
||||
url: https://enterprise.litellm.ai/demo
|
||||
|
|
|
|||
1
.github/workflows/test-unit.yml
vendored
1
.github/workflows/test-unit.yml
vendored
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -1 +1,3 @@
|
|||
Read @CLAUDE.md for coding guidelines
|
||||
|
||||
Before requesting maintainer review, verify the current PR tip passes required CI and code coverage, meets Greptile confidence of at least 4/5, and has acceptable Veria and Bugbot reviews. Inspect warnings and findings, fix actionable issues, and rerun the affected checks and reviewers after changes. Record evidence for any false positive or unavailable review; never treat a pending or missing bot result as a pass. Do not lower coverage thresholds or lint budgets to satisfy a check
|
||||
|
|
|
|||
19
Dockerfile
19
Dockerfile
|
|
@ -8,9 +8,25 @@ ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7
|
|||
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
|
||||
# Pinned by digest like the other base images; bump explicitly on Node upgrades.
|
||||
ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43
|
||||
# Checksum from https://www.pgbouncer.org/downloads/ (the Wolfi repo only carries 1.24.x)
|
||||
ARG PGBOUNCER_VERSION=1.25.2
|
||||
ARG PGBOUNCER_SHA256=924ad35113fd0a71c8e2dbe85b5d03445532e2b7b37a9f8a48983beea238b332
|
||||
|
||||
FROM $UV_IMAGE AS uvbin
|
||||
|
||||
FROM $LITELLM_BUILD_IMAGE AS pgbouncer-builder
|
||||
ARG PGBOUNCER_VERSION
|
||||
ARG PGBOUNCER_SHA256
|
||||
USER root
|
||||
RUN apk add --no-cache build-base pkgconf libevent-dev openssl-dev curl
|
||||
WORKDIR /build
|
||||
RUN curl -fsSL -o pgbouncer.tar.gz "https://www.pgbouncer.org/downloads/files/${PGBOUNCER_VERSION}/pgbouncer-${PGBOUNCER_VERSION}.tar.gz" && \
|
||||
echo "${PGBOUNCER_SHA256} pgbouncer.tar.gz" | sha256sum -c - && \
|
||||
tar xzf pgbouncer.tar.gz --strip-components=1 && \
|
||||
./configure --prefix=/usr/local --with-openssl=/usr && \
|
||||
make -j"$(nproc)" pgbouncer && \
|
||||
install -m 0755 pgbouncer /usr/local/bin/pgbouncer
|
||||
|
||||
# Admin UI builder. Pinned to the build platform so the architecture-independent
|
||||
# Next.js static export compiles once natively even in a multi-arch build,
|
||||
# instead of once per target arch under QEMU.
|
||||
|
|
@ -110,7 +126,8 @@ USER root
|
|||
RUN echo "https://packages.wolfi.dev/os" >> /etc/apk/repositories
|
||||
|
||||
# node (without npm) is required by the prisma CLI at runtime
|
||||
RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile
|
||||
RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile libevent
|
||||
COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer
|
||||
|
||||
WORKDIR /app
|
||||
ENV PATH="/app/.venv/bin:${PATH}" \
|
||||
|
|
|
|||
|
|
@ -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}" \
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
71
gateway/launch.py
Normal 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:])
|
||||
|
|
@ -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:
|
||||
|
|
|
|||
61
helm/litellm-helm/tests/connection_pool_tests.yaml
Normal file
61
helm/litellm-helm/tests/connection_pool_tests.yaml
Normal 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"
|
||||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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 }}
|
||||
|
|
|
|||
117
helm/litellm/tests/connection_pool_tests.yaml
Normal file
117
helm/litellm/tests/connection_pool_tests.yaml
Normal 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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -28,6 +28,8 @@ from litellm.proxy._experimental.mcp_server.oauth_utils import (
|
|||
build_upstream_oauth2_token_request,
|
||||
resolve_upstream_resource,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import OAuthToken
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.token_cache_codec import OAuthTokenCacheCodec
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
||||
decrypt_value_helper,
|
||||
encrypt_value_helper,
|
||||
|
|
@ -233,8 +235,17 @@ class MCPPerUserTokenCache:
|
|||
def _cache_key(self, user_id: str, server_id: str) -> str:
|
||||
return f"{MCP_PER_USER_TOKEN_REDIS_KEY_PREFIX}:{user_id}:{server_id}"
|
||||
|
||||
def _codec(self) -> OAuthTokenCacheCodec:
|
||||
return OAuthTokenCacheCodec(
|
||||
encrypt_value_helper,
|
||||
lambda blob: decrypt_value_helper(blob, key="mcp_per_user_token", exception_type="debug"),
|
||||
)
|
||||
|
||||
async def get(self, user_id: str, server_id: str) -> str | None:
|
||||
"""Return the plaintext access_token, or None on miss/error."""
|
||||
token: Final = await self.get_token(user_id, server_id)
|
||||
return token.access_token if token is not None else None
|
||||
|
||||
async def get_token(self, user_id: str, server_id: str) -> OAuthToken | None:
|
||||
try:
|
||||
from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415
|
||||
|
||||
|
|
@ -242,12 +253,7 @@ class MCPPerUserTokenCache:
|
|||
encrypted: Final = await user_api_key_cache.async_get_cache(key)
|
||||
if encrypted is None:
|
||||
return None
|
||||
plaintext: Final = decrypt_value_helper(
|
||||
encrypted,
|
||||
key="mcp_per_user_token",
|
||||
exception_type="debug",
|
||||
)
|
||||
return plaintext or None
|
||||
return self._codec().decode(encrypted)
|
||||
except Exception as exc:
|
||||
verbose_logger.debug(
|
||||
"MCPPerUserTokenCache.get failed for user=%s server=%s: %s",
|
||||
|
|
@ -263,13 +269,16 @@ class MCPPerUserTokenCache:
|
|||
server_id: str,
|
||||
access_token: str,
|
||||
ttl: int,
|
||||
identity_binding_proof: str | None = None,
|
||||
) -> None:
|
||||
"""Store NaCl-encrypted access_token in Redis with the given TTL."""
|
||||
try:
|
||||
from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415
|
||||
|
||||
key: Final = self._cache_key(user_id, server_id)
|
||||
encrypted: Final = encrypt_value_helper(access_token)
|
||||
encrypted: Final = self._codec().encode(
|
||||
OAuthToken(access_token=access_token, identity_binding_proof=identity_binding_proof)
|
||||
)
|
||||
await user_api_key_cache.async_set_cache(key, encrypted, ttl=ttl)
|
||||
verbose_logger.debug(
|
||||
"MCPPerUserTokenCache.set: cached token for user=%s server=%s ttl=%ds",
|
||||
|
|
|
|||
423
litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py
Normal file
423
litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py
Normal 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,
|
||||
},
|
||||
)
|
||||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
|
|||
96
litellm/proxy/db/health_check_latest.py
Normal file
96
litellm/proxy/db/health_check_latest.py
Normal 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 ()
|
||||
525
litellm/proxy/db/pgbouncer.py
Normal file
525
litellm/proxy/db/pgbouncer.py
Normal 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
|
||||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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."
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
),
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
),
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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)."""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
123
terraform/litellm/aws/tests/connection_pool.tftest.hcl
Normal file
123
terraform/litellm/aws/tests/connection_pool.tftest.hcl
Normal 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."
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
135
terraform/litellm/gcp/tests/connection_pool.tftest.hcl
Normal file
135
terraform/litellm/gcp/tests/connection_pool.tftest.hcl
Normal 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,
|
||||
]
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
|
|||
155
tests/test_gateway/test_launch.py
Normal file
155
tests/test_gateway/test_launch.py
Normal 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 == []
|
||||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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"}]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
"""
|
||||
|
|
|
|||
111
tests/test_litellm/proxy/db/test_health_check_latest.py
Normal file
111
tests/test_litellm/proxy/db/test_health_check_latest.py
Normal 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",)) == ()
|
||||
624
tests/test_litellm/proxy/db/test_pgbouncer.py
Normal file
624
tests/test_litellm/proxy/db/test_pgbouncer.py
Normal 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
|
||||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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"),
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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()) == ()
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
[
|
||||
|
|
|
|||
|
|
@ -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"]) == ()
|
||||
|
|
|
|||
|
|
@ -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>"
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
||||
|
|
|
|||
236
tests/test_litellm/test_redis_credential_provider.py
Normal file
236
tests/test_litellm/test_redis_credential_provider.py
Normal 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"]
|
||||
|
|
@ -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=[
|
||||
|
|
|
|||
25
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
25
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue