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

# Conflicts:
#	litellm/llms/bedrock/files/transformation.py
#	litellm/proxy/_lazy_features.py
#	litellm/proxy/analytics_endpoints/cache_activity.py

.circleci/config.yml

@@ -1084,9 +1084,7 @@ jobs:
           name: Run tests
           command: |
             mkdir -p test-results
-            TEST_FILES=$(printf "%s\n%s\n" \
-              "$(circleci tests glob "tests/ocr_tests/**/test_*.py")" \
-              "tests/test_litellm/ocr/test_rust_bridge.py")
+            TEST_FILES=$(circleci tests glob "tests/ocr_tests/**/test_*.py")
             echo "$TEST_FILES" | circleci tests run \
               --verbose \
               --command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
@@ -1440,6 +1438,7 @@ jobs:
             TEST_FILES=$(printf "%s\n" \
               tests/local_testing/test_dual_cache.py \
               tests/local_testing/test_redis_batch_optimizations.py \
+              tests/local_testing/test_redis_increment_with_floor.py \
               tests/local_testing/test_router_utils.py)
             echo "$TEST_FILES" | circleci tests run \
               --verbose \

.github/ISSUE_TEMPLATE/config.yml

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

.github/ci-coverage-allowlist.yml

@@ -106,11 +106,6 @@ dockerfiles:
       and lint workflows already exercise that output, so building the image adds no signal about it
     paths:
       - ui/Dockerfile
-  - reason: >-
-      The Rust gateway ships as its own chart and package with a separate release pipeline, so its
-      image is not part of this repo's Python image set
-    paths:
-      - litellm-rust/crates/ai-gateway/Dockerfile
   - reason: >-
       An example image under cookbook/ that is documentation rather than a shipped artifact
     paths:

.github/e2e-stack/down.sh

@@ -10,7 +10,7 @@ for pid_file in "${STACK_DIR}"/pids/*.pid; do
   rm -f "${pid_file}"
 done
 
-for container in e2e-nginx e2e-valkey e2e-jaeger e2e-postgres; do
+for container in e2e-nginx e2e-keycloak e2e-valkey e2e-jaeger e2e-postgres; do
   docker rm -f "${container}" >/dev/null 2>&1
 done
 

.github/e2e-stack/select_tests.py

@@ -11,6 +11,7 @@ UNSUPPORTED: Final = re.compile(
 )
 HARNESS: Final = re.compile(
     r"^tests/e2e/[A-Za-z0-9_.-]+\.(py|ini)$"
+    r"|^tests/e2e/idp_realm\.json$"
     r"|^tests/e2e/gateway/"
     r"|^\.github/e2e-stack/"
     r"|^\.github/workflows/test-e2e-changed\.yml$"

.github/e2e-stack/start-idp.sh

@@ -0,0 +1,45 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
+KEYCLOAK_IMAGE="${E2E_KEYCLOAK_IMAGE:-quay.io/keycloak/keycloak@sha256:ff4257d0d64efbe99ed1ddfaf07765cc3c36dc7518bf8324d41961327f441c54}"
+KEYCLOAK_PORT="${E2E_KEYCLOAK_PORT:-8081}"
+POSTGRES_IMAGE="${E2E_POSTGRES_IMAGE:-postgres:16.6}"
+: "${DATABASE_HOST:?}" "${DATABASE_PORT:?}" "${DATABASE_USER:?}" "${DATABASE_PASSWORD:?}" "${DATABASE_NAME:?}"
+
+DB_HOST="${DATABASE_HOST}"
+DB_NETWORK_ARGS=(--network bridge)
+IDP_NETWORK_ARGS=(-p "127.0.0.1:${KEYCLOAK_PORT}:${KEYCLOAK_PORT}")
+if [[ "$(uname)" == "Linux" ]]; then
+  DB_NETWORK_ARGS=(--network host)
+  IDP_NETWORK_ARGS=(--network host)
+elif [[ "${DB_HOST}" == "127.0.0.1" || "${DB_HOST}" == "localhost" ]]; then
+  DB_HOST=host.docker.internal
+fi
+
+docker run --rm "${DB_NETWORK_ARGS[@]}" -e "PGPASSWORD=${DATABASE_PASSWORD}" \
+  "${POSTGRES_IMAGE}" psql -h "${DB_HOST}" -p "${DATABASE_PORT}" \
+  -U "${DATABASE_USER}" -d "${DATABASE_NAME}" -v ON_ERROR_STOP=1 \
+  -c 'CREATE SCHEMA IF NOT EXISTS keycloak' >/dev/null
+
+docker rm -f e2e-keycloak >/dev/null 2>&1 || true
+docker run -d --name e2e-keycloak "${IDP_NETWORK_ARGS[@]}" --memory 1536m \
+  -v "${REPO_ROOT}/tests/e2e/idp_realm.json:/opt/keycloak/data/import/realm.json:ro" \
+  -e KC_DB=postgres -e "KC_DB_URL_HOST=${DB_HOST}" -e "KC_DB_URL_PORT=${DATABASE_PORT}" \
+  -e "KC_DB_URL_DATABASE=${DATABASE_NAME}" -e KC_DB_SCHEMA=keycloak \
+  -e "KC_DB_USERNAME=${DATABASE_USER}" -e "KC_DB_PASSWORD=${DATABASE_PASSWORD}" \
+  -e KC_DB_POOL_INITIAL_SIZE=2 -e KC_DB_POOL_MIN_SIZE=2 -e KC_DB_POOL_MAX_SIZE=10 \
+  -e "KC_HTTP_PORT=${KEYCLOAK_PORT}" -e KC_BOOTSTRAP_ADMIN_USERNAME=admin \
+  -e KC_BOOTSTRAP_ADMIN_PASSWORD=e2e-ephemeral-idp-not-a-secret \
+  "${KEYCLOAK_IMAGE}" start-dev --import-realm >/dev/null
+
+deadline=$((SECONDS + ${E2E_KEYCLOAK_STARTUP_TIMEOUT:-300}))
+until curl -fsS --connect-timeout 2 --max-time 3 \
+  "http://127.0.0.1:${KEYCLOAK_PORT}/realms/litellm-e2e/.well-known/openid-configuration" >/dev/null 2>&1; do
+  if ((SECONDS >= deadline)); then
+    echo 'e2e-stack: timed out waiting for the Keycloak realm' >&2
+    exit 1
+  fi
+  sleep 2
+done
+echo 'e2e-stack: Keycloak realm is up'

.github/e2e-stack/up.sh

@@ -25,6 +25,7 @@ DATABASE_PASSWORD="${E2E_DATABASE_PASSWORD:-dbpassword9090}"
 DATABASE_NAME="${E2E_DATABASE_NAME:-litellm}"
 JAEGER_OTLP_PORT="${E2E_JAEGER_OTLP_PORT:-4318}"
 JAEGER_QUERY_PORT="${E2E_JAEGER_QUERY_PORT:-16686}"
+KEYCLOAK_PORT="${E2E_KEYCLOAK_PORT:-8081}"
 
 MASTER_KEY="${LITELLM_MASTER_KEY:-sk-e2e-$(openssl rand -hex 16)}"
 
@@ -124,6 +125,9 @@ SERVER_ENV=(
   "OTEL_EXPORTER_OTLP_ENDPOINT=http://127.0.0.1:${JAEGER_OTLP_PORT}"
   "SSL_CERT_FILE=${CERTS_DIR}/ca-bundle.pem"
   "PYTHONPATH=${REPO_ROOT}"
+  "JWT_PUBLIC_KEY_URL=http://127.0.0.1:${KEYCLOAK_PORT}/realms/litellm-e2e/protocol/openid-connect/certs"
+  "JWT_ISSUER=http://127.0.0.1:${KEYCLOAK_PORT}/realms/litellm-e2e"
+  "JWT_AUDIENCE=litellm-e2e"
 )
 if [[ -n "${VERTEXAI_CREDENTIALS:-}" ]]; then
   printf '%s' "${VERTEXAI_CREDENTIALS}" > "${STACK_DIR}/vertex-adc.json"
@@ -132,6 +136,8 @@ fi
 
 cd "${REPO_ROOT}"
 
+env "${SERVER_ENV[@]}" "E2E_KEYCLOAK_PORT=${KEYCLOAK_PORT}" bash .github/e2e-stack/start-idp.sh
+
 log "running migrations"
 env "${SERVER_ENV[@]}" uv run --no-sync python migrations/run.py >"${LOGS_DIR}/migrations.log" 2>&1
 
@@ -200,6 +206,9 @@ LITELLM_MASTER_KEY=${MASTER_KEY}
 REDIS_HOST=127.0.0.1
 REDIS_PORT=${REDIS_PORT}
 E2E_OTEL_QUERY_URL=http://127.0.0.1:${JAEGER_QUERY_PORT}
+E2E_KEYCLOAK_URL=http://127.0.0.1:${KEYCLOAK_PORT}
+E2E_KEYCLOAK_ADMIN_USER=admin
+E2E_KEYCLOAK_ADMIN_PASSWORD=e2e-ephemeral-idp-not-a-secret
 SSL_CERT_FILE=${CERTS_DIR}/ca-bundle.pem
 DATABASE_URL=postgresql://${DATABASE_USER}:${DATABASE_PASSWORD}@${DATABASE_HOST}:${DATABASE_PORT}/${DATABASE_NAME}
 EOF

.github/pull_request_template.md

@@ -47,6 +47,10 @@ After: the same request comes back with real token counts, so the dashboard show
 
 <!-- e.g., "Fixes #000" -->
 
+## Affected release
+
+<!-- Only for a fix to a regression in a released or rc version (perf, memory, crash, or behavior): name the version it regressed in, e.g. "regression in v1.100.0" or "since v1.101.0-rc.1", and add the `backport-stable` label so the fix is cherry-picked onto the rc line before the stable is tagged. Leave the section blank otherwise -->
+
 ## Linear ticket
 
 <!-- if you are an internal contributor, add "Resolves " followed by the Linear ticket e.g., "Resolves LIT-1234" to link the Linear ticket to the GitHub PR. If you don't have one, leave the section blank rather than guessing -->
@@ -127,6 +131,8 @@ If you're seeing a delay in your PR being merged, ping the LiteLLM Team on [Slac
      - Low: anything else worth noting: naming, cleanup, an edge case nobody hits
      Nest bullets as deep as helps: hierarchy beats one long line when it makes things clearer to a
      human reader
+     If you assumed something instead of testing it, e.g. "only reproduces with X on" or "no
+     user-observable behavior difference", list it here too with what breaks if it is wrong
      Leave this section empty if there are none -->
 
 ## QA runbook
@@ -152,3 +158,4 @@ Example checklists:
 ## Final Attestation
 
 - [ ] The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR
+

.github/scripts/auto_update_price_and_context_window_file.py

@@ -1,6 +1,8 @@
 import asyncio
 import aiohttp
 import json
+import math
+from typing import Any
 
 # Asynchronously fetch data from a given URL
 async def fetch_data(url):
@@ -21,11 +23,157 @@ async def fetch_data(url):
         print("Error fetching data from URL:", e)
         return None
 
+
+FRIENDLI_API_URL = "https://api.friendli.ai/serverless/v1/models"
+FRIENDLI_PROVIDER = "friendliai"
+
+INHERITABLE_BASE_KEYS = (
+    "supports_pdf_input",
+    "supports_assistant_prefill",
+    "supports_adaptive_thinking",
+    "supports_output_config",
+)
+
+REASONING_EFFORT_LEVEL_ORDER = ("none", "minimal", "low", "medium", "high", "xhigh", "max")
+
+
+def _find_base_model_entry(base_model: str, local_data: dict) -> str | None:
+    if not base_model:
+        return None
+    bm_tail = base_model.split("/")[-1].lower()
+    if base_model in local_data:
+        return base_model
+    for key in local_data:
+        if key.startswith("sample_spec") or key == "fallback_generalizations":
+            continue
+        if key.split("/")[-1].lower() == bm_tail:
+            return key
+    return None
+
+
+def _reasoning_effort_levels(reasoning_options: list) -> list:
+    offered = {
+        val
+        for opt in reasoning_options or []
+        if opt.get("type") == "effort"
+        for val in opt.get("values", [])
+    }
+    return [level for level in REASONING_EFFORT_LEVEL_ORDER if level in offered]
+
+
+def _valid_token_price(value: object) -> bool:
+    try:
+        price = float(value)  # pyright: ignore[reportArgumentType]  # non-numeric values are rejected via the except
+    except (TypeError, ValueError):
+        return False
+    return math.isfinite(price) and price >= 0
+
+
+def _has_valid_token_prices(pricing: dict | None) -> bool:
+    prices = pricing or {}
+    return _valid_token_price(prices.get("input")) and _valid_token_price(prices.get("output"))
+
+
+def _pricing(pricing: dict) -> dict:
+    out: dict[str, Any] = {}
+    if not pricing:
+        return out
+    if "input" in pricing:
+        out["input_cost_per_token"] = float(pricing["input"])
+    if "output" in pricing:
+        out["output_cost_per_token"] = float(pricing["output"])
+    if "input_cache_read" in pricing and pricing["input_cache_read"] is not None:
+        out["cache_read_input_token_cost"] = float(pricing["input_cache_read"])
+    return out
+
+
+def _modality_flags(input_mods: list) -> dict:
+    mods = input_mods or []
+    has_image = "image" in mods
+    return {
+        "supports_vision": has_image,
+        "supports_image_input": has_image,
+        "supports_video_input": "video" in mods,
+    }
+
+
+def transform_friendli_data(data: list, local_data: dict) -> dict:
+    transformed: dict[str, dict] = {}
+    if not data:
+        return transformed
+    for model in data:
+        # An unpriced row must never wholesale-replace an already priced local entry:
+        # missing prices cost-calculate as zero, silently zeroing tracked spend
+        if not _has_valid_token_prices(model.get("pricing")):
+            continue
+        model_id = model["id"]
+        base_model = model.get("base_model") or ""
+        entry: dict[str, Any] = {
+            "litellm_provider": FRIENDLI_PROVIDER,
+        }
+
+        base_key = _find_base_model_entry(base_model, local_data)
+        if base_key:
+            base_entry = local_data[base_key]
+            for k in INHERITABLE_BASE_KEYS:
+                if k in base_entry:
+                    entry[k] = base_entry[k]
+
+        ctx = model.get("context_length")
+        if ctx is not None:
+            entry["max_input_tokens"] = int(ctx)
+        max_out = model.get("max_completion_tokens")
+        if max_out is not None:
+            entry["max_output_tokens"] = int(max_out)
+            entry["max_tokens"] = int(max_out)
+
+        pricing = _pricing(model.get("pricing", {}))
+        entry.update(pricing)
+        entry["supports_prompt_caching"] = "cache_read_input_token_cost" in pricing
+
+        reasoning = model.get("reasoning") is True
+        entry["supports_reasoning"] = reasoning
+        if reasoning:
+            entry["reasoning_effort_levels"] = _reasoning_effort_levels(
+                model.get("reasoning_options", [])
+            )
+
+        func = model.get("functionality", {})
+        entry["supports_function_calling"] = func.get("tool_call") is True
+        entry["supports_parallel_function_calling"] = func.get("parallel_tool_call") is True
+        is_struct = func.get("structured_output") is True
+        entry["supports_response_schema"] = is_struct
+        entry["supports_native_structured_output"] = is_struct
+        entry["supports_system_messages"] = func.get("system_messages") is True
+        entry["supports_tool_choice"] = func.get("tool_choice") is True
+
+        entry.update(_modality_flags(model.get("input_modalities", [])))
+
+        entry["mode"] = model.get("mode", "chat")
+
+        desc = model.get("description")
+        if desc:
+            entry["comment"] = desc
+
+        dep = model.get("deprecation_date")
+        if dep:
+            entry["deprecation_date"] = dep.split("T")[0]
+
+        entry["source"] = FRIENDLI_API_URL
+
+        transformed[f"{FRIENDLI_PROVIDER}/{model_id}"] = entry
+    return transformed
+
 # Synchronize local data with remote data
-def sync_local_data_with_remote(local_data, remote_data):
+def sync_local_data_with_remote(local_data, remote_data, replace_keys=frozenset()):
     # Update existing keys in local_data with values from remote_data
+    # (replace_keys entries are swapped wholesale so a field the remote catalog
+    #  dropped, e.g. cache pricing, cannot survive as a stale value)
     for key in (set(local_data) & set(remote_data)):
-        local_data[key].update(remote_data[key])
+        if key in replace_keys:
+            local_data[key] = remote_data[key]
+        else:
+            local_data[key].update(remote_data[key])
 
     # Add new keys from remote_data to local_data
     for key in (set(remote_data) - set(local_data)):
@@ -46,6 +194,8 @@ def write_to_file(file_path, data):
 # Update the existing models and add the missing models for OpenRouter
 def transform_openrouter_data(data):
     transformed = {}
+    if not data:
+        return transformed
     for row in data:
         # Add the fields 'max_tokens' and 'input_cost_per_token'
         obj = {
@@ -84,7 +234,14 @@ def transform_openrouter_data(data):
 # Update the existing models and add the missing models for Vercel AI Gateway
 def transform_vercel_ai_gateway_data(data):
     transformed = {}
+    if not data:
+        return transformed
     for row in data:
+        # Rows without token pricing or token limits (video/embedding models) previously KeyError'd the whole sync
+        if any(row.get(k) is None for k in ("context_window", "max_tokens")) or any(
+            row.get("pricing", {}).get(k) is None for k in ("input", "output")
+        ):
+            continue
         obj = {
             "max_tokens": row["context_window"],
             "input_cost_per_token": float(row["pricing"]["input"]),
@@ -143,13 +300,16 @@ def main():
     vercel_data = asyncio.run(fetch_data(vercel_ai_gateway_url))
     # Transform the fetched Vercel AI Gateway data
     vercel_data = transform_vercel_ai_gateway_data(vercel_data)
+
+    friendli_data = asyncio.run(fetch_data(FRIENDLI_API_URL))
+    friendli_data = transform_friendli_data(friendli_data, local_data)
     
     # Combine both datasets
-    all_remote_data = {**openrouter_data, **vercel_data}
+    all_remote_data = {**openrouter_data, **vercel_data, **friendli_data}
 
     # If both local and openrouter data are available, synchronize and save
     if local_data and all_remote_data:
-        sync_local_data_with_remote(local_data, all_remote_data)
+        sync_local_data_with_remote(local_data, all_remote_data, replace_keys=frozenset(friendli_data))
         write_to_file(local_file_path, local_data)
     else:
         print("Failed to fetch model data from either local file or URL.")

.github/scripts/verify_linux_native_wheel.py

@@ -205,7 +205,7 @@ def main(
     native_module: Final = load_native_module(native_path)
     native_module_loads: Final = native_module is not None
     panic_test_hook_absent: Final = native_module is not None and not hasattr(native_module, "_panic_for_test")
-    native_size_limit: Final = 20_000_000
+    native_size_limit: Final = 25_000_000
     native_size_within_limit: Final = native_member.file_size <= native_size_limit
     validations: Final = (
         (f"Python tag is {EXPECTED_PYTHON_TAG}", python_tag == EXPECTED_PYTHON_TAG),
@@ -222,7 +222,7 @@ def main(
         ("Python extension entry point is present", extension_entry_point_present),
         ("Native module loads", native_module_loads),
         ("Production module omits the panic test hook", panic_test_hook_absent),
-        ("Native extension does not exceed 20 MB", native_size_within_limit),
+        ("Native extension does not exceed 25 MB", native_size_within_limit),
         ("Wheel contents are valid", not unexpected_members),
     )
 

.github/workflows/_test-unit-base.yml

@@ -113,7 +113,7 @@ jobs:
         if: steps.changes.outputs.decision != 'skip'
         timeout-minutes: 8
         run: |
-          .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra mongodb
+          .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml
           uv run --no-sync python -c 'import os, sys; print(sys.version); assert f"{sys.version_info.major}.{sys.version_info.minor}" == os.environ["UV_PYTHON"]'
 
       - name: Cache Prisma binaries

.github/workflows/ai-gateway-image.yml

@@ -0,0 +1,73 @@
+name: ai-gateway image
+
+on:
+  push:
+    paths:
+      - "litellm-rust/**"
+      - "litellm/**"
+      - "enterprise/**"
+      - "litellm-proxy-extras/**"
+      - "pyproject.toml"
+      - "rust-toolchain.toml"
+      - ".github/workflows/ai-gateway-image.yml"
+  pull_request:
+    branches:
+      - main
+      - litellm_internal_staging
+      - litellm_oss_staging
+      - "litellm_**"
+    paths:
+      - "litellm-rust/**"
+      - "litellm/**"
+      - "enterprise/**"
+      - "litellm-proxy-extras/**"
+      - "pyproject.toml"
+      - "rust-toolchain.toml"
+      - ".github/workflows/ai-gateway-image.yml"
+
+permissions:
+  contents: read
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  ai-gateway-image:
+    name: ai-gateway release image
+    runs-on: ubuntu-latest
+    timeout-minutes: 60
+    permissions:
+      contents: read
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
+        with:
+          persist-credentials: false
+      - name: Build the release image
+        run: docker build -f litellm-rust/crates/ai-gateway/Dockerfile -t litellm-ai-gateway:${{ github.sha }} .
+      - name: Start the gateway and wait for readiness
+        env:
+          IMAGE: litellm-ai-gateway:${{ github.sha }}
+        run: |
+          docker run -d --name ai-gateway -p 4001:4001 \
+            -e LITELLM_MASTER_KEY=sk-ci-not-a-real-key \
+            -e OPENAI_API_KEY=sk-ci-not-a-real-key \
+            "$IMAGE"
+          for _ in $(seq 1 60); do
+            if curl -fsS http://127.0.0.1:4001/health/readiness; then
+              echo "gateway is serving readiness"
+              exit 0
+            fi
+            sleep 2
+          done
+          echo "gateway never became ready" >&2
+          docker logs ai-gateway >&2
+          exit 1
+      - name: Assert the gateway loaded the baked config
+        run: |
+          docker logs ai-gateway 2>&1 | tee gateway.log
+          grep 'via python config reader' gateway.log
+      - name: Stop the gateway
+        if: always()
+        run: docker rm -f ai-gateway || true

.github/workflows/guard-main-branch.yml

@@ -1,42 +0,0 @@
-name: Guard main branch
-
-on:
-  pull_request:
-    branches:
-      - main
-  merge_group:
-
-permissions: {}
-
-# DO NOT RENAME the job's `name:` — it is referenced by GitHub branch
-# protection as a required status check on `main`. Renaming silently
-# breaks the gate.
-jobs:
-  guard:
-    name: Verify PR source branch
-    runs-on: ubuntu-latest
-    timeout-minutes: 2
-    steps:
-      - name: Reject merge_group events
-        if: github.event_name == 'merge_group'
-        run: |
-          echo "::error::Merge queue is not supported for main. Disable merge queue or update this guard."
-          exit 1
-      - name: Check head branch name
-        env:
-          HEAD_REF: ${{ github.head_ref }}
-          HEAD_REPO: ${{ github.event.pull_request.head.repo.full_name }}
-          BASE_REPO: ${{ github.repository }}
-        run: |
-          echo "PR head repo: $HEAD_REPO"
-          echo "PR head branch: $HEAD_REF"
-          if [ "$HEAD_REPO" != "$BASE_REPO" ]; then
-            echo "::error::PRs to main must originate from the canonical repository ($BASE_REPO), not a fork ($HEAD_REPO). External contributors should open PRs against 'litellm_internal_staging' instead."
-            exit 1
-          fi
-          if [ "$HEAD_REF" = "litellm_internal_staging" ] || [[ "$HEAD_REF" == litellm_hotfix_?* ]]; then
-            echo "Allowed source branch."
-            exit 0
-          fi
-          echo "::error::PRs to main must originate from 'litellm_internal_staging' or a 'litellm_hotfix_*' branch. Got: '$HEAD_REF'. If this is a contribution, retarget the PR against 'litellm_internal_staging' instead."
-          exit 1

.github/workflows/test-code-quality.yml

@@ -81,7 +81,7 @@ jobs:
         run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_workflow_job_name_collisions.py
 
       - name: test_e2e_changed_gate
-        run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_e2e_changed_gate.py
+        run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_e2e_changed_gate.py tests/code_coverage_tests/test_e2e_idp_stack.py
 
       - name: router_code_coverage
         run: uv run --no-sync python ./tests/code_coverage_tests/router_code_coverage.py

.github/workflows/test-e2e-changed.yml

@@ -27,6 +27,8 @@ jobs:
           sparse-checkout: |
             .github/e2e-stack
             tests/e2e/access_control
+            tests/e2e/management/test_jwt_management_e2e.py
+            tests/e2e/other/test_jwt_auth_e2e.py
           persist-credentials: false
           ref: ${{ github.sha }}
 
@@ -45,7 +47,8 @@ jobs:
             --jq '.[] | select(.status != "removed") | .filename')"
           gh api "repos/${REPO}/pulls/${PR_NUMBER}" --jq '.head.sha' | grep -Fxq "${HEAD_SHA}"
           tests="$(printf '%s\n' "${files}" \
-            | python3 .github/e2e-stack/select_tests.py tests/e2e/access_control/test_*.py)"
+            | python3 .github/e2e-stack/select_tests.py tests/e2e/access_control/test_*.py \
+                tests/e2e/management/test_jwt_management_e2e.py tests/e2e/other/test_jwt_auth_e2e.py)"
           echo "tests=${tests}" >> "${GITHUB_OUTPUT}"
           if [ -n "${tests}" ]; then
             echo "any=true" >> "${GITHUB_OUTPUT}"

.github/workflows/test-e2e-redis-chaos.yml

@@ -0,0 +1,103 @@
+name: "Redis Chaos E2E"
+
+on:
+  workflow_dispatch:
+  workflow_call:
+    inputs:
+      ref:
+        description: "Commit SHA or ref to test. Defaults to the ref the workflow was triggered on"
+        required: false
+        type: string
+
+permissions:
+  contents: read
+
+jobs:
+  redis-chaos-e2e:
+    runs-on: ubuntu-latest-16-cores
+    timeout-minutes: 30
+    services:
+      postgres:
+        image: postgres:16.6@sha256:557fea37a744d5f4c8faab304b0a90858b53ab119735a88c131fd19dab802f36
+        env:
+          POSTGRES_USER: llmproxy
+          POSTGRES_PASSWORD: dbpassword9090
+          POSTGRES_DB: litellm
+        ports:
+          - 5432:5432
+        options: >-
+          --health-cmd "pg_isready -U llmproxy"
+          --health-interval 5s
+          --health-timeout 5s
+          --health-retries 10
+      valkey:
+        image: valkey/valkey:8.1.4@sha256:81db6d39e1bba3b3ff32bd3a1b19a6d69690f94a3954ec131277b9a26b95b3aa
+        ports:
+          - 6379:6379
+        options: >-
+          --health-cmd "valkey-cli ping"
+          --health-interval 5s
+          --health-timeout 5s
+          --health-retries 10
+    env:
+      DATABASE_URL: postgresql://llmproxy:dbpassword9090@localhost:5432/litellm
+      LITELLM_MASTER_KEY: sk-redis-chaos-e2e
+      LITELLM_LOG: WARNING
+      JSON_LOGS: "true"
+    steps:
+      - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
+        with:
+          persist-credentials: false
+          ref: ${{ inputs.ref || github.sha }}
+
+      - name: Set up Python
+        uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
+        with:
+          python-version: "3.12"
+
+      - name: Set up uv
+        uses: ./.github/actions/setup-uv-with-retries
+        with:
+          version: "0.10.9"
+
+      - name: Cache the Rust build
+        uses: ./.github/actions/cache-cargo-build
+
+      - name: Install dependencies
+        run: |
+          .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --group e2e-dev --extra proxy
+
+      - name: Cache Prisma binaries
+        uses: ./.github/actions/cache-prisma-binaries
+
+      - name: Generate Prisma client
+        run: |
+          uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
+
+      - name: Start a multi-worker proxy on the chaos config
+        run: |
+          nohup uv run --no-sync litellm --config tests/e2e/gateway/redis_chaos_ci_config.yml --port 4000 --num_workers 4 > proxy.log 2>&1 &
+          echo "E2E_PROXY_PID=$!" >> "$GITHUB_ENV"
+          echo "E2E_PROXY_LOG=$(pwd)/proxy.log" >> "$GITHUB_ENV"
+          for _ in $(seq 1 90); do
+            if curl -fs http://localhost:4000/health/liveliness > /dev/null; then
+              exit 0
+            fi
+            sleep 2
+          done
+          echo "proxy never became live"
+          tail -n 100 proxy.log
+          exit 1
+
+      - name: Run the Redis chaos load test
+        env:
+          E2E_REDIS_CHAOS: "1"
+          LITELLM_PROXY_URL: http://localhost:4000
+          REDIS_HOST: 127.0.0.1
+          REDIS_PORT: "6379"
+        run: |
+          uv run --no-sync pytest tests/e2e/load/test_redis_chaos_e2e.py -v --tb=short -rA -s
+
+      - name: Show proxy log on failure
+        if: failure()
+        run: tail -n 300 proxy.log

.github/workflows/test-linting.yml

@@ -113,7 +113,7 @@ jobs:
       - name: Check ruff format
         if: steps.changes.outputs.decision != 'skip'
         run: |
-          git diff --name-only --diff-filter=ACMR "$GATE_BASE_SHA" HEAD -- 'litellm/**/*.py' | grep -v '^litellm/enterprise/' > "$RUNNER_TEMP/ruff_format_files.txt" || true
+          git diff --name-only --diff-filter=ACMR "$GATE_BASE_SHA" HEAD -- ':(glob)litellm/**/*.py' | grep -v '^litellm/enterprise/' > "$RUNNER_TEMP/ruff_format_files.txt" || true
           if [ ! -s "$RUNNER_TEMP/ruff_format_files.txt" ]; then
             echo "No changed litellm Python files to check with ruff format."
             exit 0
@@ -172,7 +172,7 @@ jobs:
       - name: Check tests/e2e basedpyright (zero errors)
         if: steps.changes.outputs.decision != 'skip'
         run: |
-          if git diff --name-only --diff-filter=ACMRD "$GATE_BASE_SHA" HEAD -- 'tests/e2e/**/*.py' | grep -q .; then
+          if git diff --name-only --diff-filter=ACMRD "$GATE_BASE_SHA" HEAD -- ':(glob)tests/e2e/**/*.py' | grep -q .; then
             uv run --no-sync basedpyright tests/e2e
           else
             echo "No changed tests/e2e Python files; skipping."

.github/workflows/test-litellm-ui-unit.yml

@@ -66,7 +66,7 @@ jobs:
           BASE_SHA: ${{ github.event.pull_request.base.sha }}
           HEAD_SHA: ${{ github.event.pull_request.head.sha }}
         run: |
-          full_suite() { npm run test -- --run --pool forks --poolOptions.forks.maxForks=14; }
+          full_suite() { npm run test -- --run --pool forks --maxWorkers=14; }
 
           if [ -z "$BASE_SHA" ]; then
             echo "Push to $GITHUB_REF_NAME: running the full suite"
@@ -95,4 +95,4 @@ jobs:
 
           echo "Pull request: running tests related to ${#changed_files[@]} changed UI files"
           npm run test -- related "${changed_files[@]}" --run --passWithNoTests \
-            --pool forks --poolOptions.forks.maxForks=14
+            --pool forks --maxWorkers=14

.github/workflows/test-rust.yml

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

.github/workflows/test-terraform-modules.yml

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

.github/workflows/test-unit.yml

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

.gitignore

@@ -147,3 +147,6 @@ crash.*.log
 
 ui/litellm-dashboard/out/
 litellm.log
+
+.coverage-rust
+coverage-rust.xml

AGENTS.md

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

Dockerfile

@@ -8,9 +8,25 @@ ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7
 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
 # Pinned by digest like the other base images; bump explicitly on Node upgrades.
 ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43
+# Checksum from https://www.pgbouncer.org/downloads/ (the Wolfi repo only carries 1.24.x)
+ARG PGBOUNCER_VERSION=1.25.2
+ARG PGBOUNCER_SHA256=924ad35113fd0a71c8e2dbe85b5d03445532e2b7b37a9f8a48983beea238b332
 
 FROM $UV_IMAGE AS uvbin
 
+FROM $LITELLM_BUILD_IMAGE AS pgbouncer-builder
+ARG PGBOUNCER_VERSION
+ARG PGBOUNCER_SHA256
+USER root
+RUN apk add --no-cache build-base pkgconf libevent-dev openssl-dev curl
+WORKDIR /build
+RUN curl -fsSL -o pgbouncer.tar.gz "https://www.pgbouncer.org/downloads/files/${PGBOUNCER_VERSION}/pgbouncer-${PGBOUNCER_VERSION}.tar.gz" && \
+    echo "${PGBOUNCER_SHA256}  pgbouncer.tar.gz" | sha256sum -c - && \
+    tar xzf pgbouncer.tar.gz --strip-components=1 && \
+    ./configure --prefix=/usr/local --with-openssl=/usr && \
+    make -j"$(nproc)" pgbouncer && \
+    install -m 0755 pgbouncer /usr/local/bin/pgbouncer
+
 # Admin UI builder. Pinned to the build platform so the architecture-independent
 # Next.js static export compiles once natively even in a multi-arch build,
 # instead of once per target arch under QEMU.
@@ -67,7 +83,6 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr
     --extra semantic-router \
     --extra saml \
     --extra bedrock-realtime \
-    --extra mongodb \
     --python python3.13
 
 # Copy full source tree
@@ -90,7 +105,6 @@ RUN uv sync --frozen --no-default-groups --no-editable \
     --extra semantic-router \
     --extra saml \
     --extra bedrock-realtime \
-    --extra mongodb \
     --python python3.13
 
 RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
@@ -112,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}" \

Makefile

@@ -150,8 +150,8 @@ lint-install:
 # Diff-scoped format check, mirroring test-linting.yml's "Check ruff format" step:
 # only the litellm Python files changed vs the base are checked, so a pre-existing
 # format issue elsewhere doesn't block an unrelated commit. Git pathspecs match
-# recursively, so 'litellm/*.py' covers nested modules and the top-level files that
-# CI's 'litellm/**/*.py' skips, which makes this target a superset of the CI step.
+# recursively, so 'litellm/*.py' covers top-level files and nested modules alike,
+# the same set CI's ':(glob)litellm/**/*.py' selects.
 lint-format-check-changed: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
 	@base_ref=$$($(RESOLVE_BASE)) && \
 	changed=$$(git diff --name-only --diff-filter=ACMR "$$base_ref...HEAD" -- 'litellm/*.py') && \

README.md

@@ -262,6 +262,8 @@ curl -X POST 'http://0.0.0.0:4000/v1/chat/completions' \
 }
 ```
 
+For MCP OAuth, an upstream may advertise dynamic client registration but refuse requests with HTTP 401 or 403. If the provider requires a pre-registered OAuth app, configure its `credentials.client_id` and, when required, `credentials.client_secret` on the MCP server. This skips dynamic registration in the gateway sign-in flow. The provider must approve the app for MCP access; reaching its authorization page does not establish that login or tool calls will succeed
+
 [**Docs: MCP Gateway**](https://docs.litellm.ai/docs/mcp)
 
 </details>

basedpyright-code-budget.json

@@ -105,7 +105,7 @@
     "limit": 109
   },
   "reportUnknownMemberType": {
-    "limit": 38271
+    "limit": 38269
   },
   "reportUnknownParameterType": {
     "limit": 19584

docker/Dockerfile.database

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

docker/Dockerfile.non_root

@@ -7,9 +7,25 @@ ARG PROXY_EXTRAS_SOURCE=published
 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
 # Pinned by digest like the other base images; bump explicitly on Node upgrades.
 ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43
+# Checksum from https://www.pgbouncer.org/downloads/ (the Wolfi repo only carries 1.24.x)
+ARG PGBOUNCER_VERSION=1.25.2
+ARG PGBOUNCER_SHA256=924ad35113fd0a71c8e2dbe85b5d03445532e2b7b37a9f8a48983beea238b332
 
 FROM $UV_IMAGE AS uvbin
 
+FROM $LITELLM_BUILD_IMAGE AS pgbouncer-builder
+ARG PGBOUNCER_VERSION
+ARG PGBOUNCER_SHA256
+USER root
+RUN apk add --no-cache build-base pkgconf libevent-dev openssl-dev curl
+WORKDIR /build
+RUN curl -fsSL -o pgbouncer.tar.gz "https://www.pgbouncer.org/downloads/files/${PGBOUNCER_VERSION}/pgbouncer-${PGBOUNCER_VERSION}.tar.gz" && \
+    echo "${PGBOUNCER_SHA256}  pgbouncer.tar.gz" | sha256sum -c - && \
+    tar xzf pgbouncer.tar.gz --strip-components=1 && \
+    ./configure --prefix=/usr/local --with-openssl=/usr && \
+    make -j"$(nproc)" pgbouncer && \
+    install -m 0755 pgbouncer /usr/local/bin/pgbouncer
+
 # Admin UI builder. Pinned to the build platform so the architecture-independent
 # Next.js static export compiles once natively even in a multi-arch build,
 # instead of once per target arch under QEMU.
@@ -71,7 +87,6 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
     --extra semantic-router \
     --extra saml \
     --extra bedrock-realtime \
-    --extra mongodb \
     --python python3.13
 
 # Copy full source tree
@@ -100,7 +115,6 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
         --extra semantic-router \
         --extra saml \
         --extra bedrock-realtime \
-        --extra mongodb \
         --python python3.13 \
         --no-sources-package litellm-proxy-extras; \
     else \
@@ -111,7 +125,6 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
         --extra semantic-router \
         --extra saml \
         --extra bedrock-realtime \
-        --extra mongodb \
         --python python3.13; \
     fi
 
@@ -131,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

docker/component_entrypoint.sh

@@ -1,5 +1,11 @@
 #!/bin/sh
 
+# stale samples from a previous container incarnation would be summed into the aggregate
+if [ -n "$PROMETHEUS_MULTIPROC_DIR" ]; then
+    mkdir -p "$PROMETHEUS_MULTIPROC_DIR"
+    rm -f "$PROMETHEUS_MULTIPROC_DIR"/*.db
+fi
+
 case "$USE_DDTRACE" in
     [Tt][Rr][Uu][Ee])
         export DD_TRACE_OPENAI_ENABLED="False"

enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py

@@ -11,10 +11,14 @@ import sys
 sys.path.insert(
     0, os.path.abspath("../..")
 )  # Adds the parent directory to the system path
-import functools
+import configparser
+import contextlib
+import itertools
+import re
 import tempfile
+from collections.abc import Generator, Iterator, Sequence
 from contextvars import ContextVar
-from typing import TYPE_CHECKING, ClassVar, Literal, Optional
+from typing import TYPE_CHECKING, ClassVar, Final, Literal, Optional
 
 from litellm._logging import verbose_proxy_logger
 from litellm.caching.caching import DualCache
@@ -433,12 +437,101 @@ _default_detect_secrets_config = {
             "name": "ZendeskSecretKeyDetector",
             "path": _custom_plugins_path + "/zendesk_secret_key.py",
         },
+        {
+            "name": "CredentialKeywordDetector",
+            "path": _custom_plugins_path + "/credential_keyword.py",
+        },
         {"name": "Base64HighEntropyString", "limit": 4.5},
         {"name": "HexHighEntropyString", "limit": 3.0},
     ],
 }
 
 
+_CONFIG_SECTION: Final = "litellm-prompt"
+
+_ASSIGNMENT_LINE: Final = re.compile(r"[^\s\[#;:=][^:=]*[:=]")
+
+_SHELL_ASSIGNMENT: Final = re.compile(r"(?P<key>[^\s\[#;:=](?:[^:=]*[^\s:=])?)=(?P<value>\S+)")
+
+_SHELL_OPERATORS: Final = ";&|"
+
+_SHELL_TRAILER: Final = re.compile(r"\\|#.*|-*\w[\w.-]*=\S*")
+
+_SCAN_SUFFIX: Final = ".py"
+
+
+@contextlib.contextmanager
+def _temp_file(text: str) -> Generator[str, None, None]:
+    temp_file: Final = tempfile.NamedTemporaryFile(suffix=_SCAN_SUFFIX, delete=False)
+    try:
+        temp_file.write(text.encode("utf-8"))
+        temp_file.close()
+        yield temp_file.name
+    finally:
+        temp_file.close()
+        os.remove(temp_file.name)
+
+
+def _scan_lines(lines: Sequence[str]) -> frozenset[tuple[str, str]]:
+    from detect_secrets import SecretsCollection
+
+    secrets: Final = SecretsCollection()
+    with _temp_file("\n".join(lines)) as path:
+        secrets.scan_file(path)
+
+    return frozenset(
+        (found_secret.secret_value, found_secret.type)
+        for file in secrets.files
+        for found_secret in secrets[file]
+        if found_secret.secret_value is not None
+    )
+
+
+def _classify_line(state: tuple[bool, str | None], numbered: tuple[int, str]) -> tuple[bool, str | None]:
+    open_option: Final = state[0]
+    number, line = numbered
+    stripped: Final = line.strip()
+    if not stripped or stripped[0] in "#;":
+        return open_option, None
+    shell_assignment: Final = _SHELL_ASSIGNMENT.match(stripped)
+    if shell_assignment is not None:
+        return True, f"{shell_assignment['key']}_{number}={shell_assignment['value']}"
+    assignment: Final = _ASSIGNMENT_LINE.match(stripped)
+    if assignment is not None:
+        return True, f"{assignment.group()[:-1].strip()}_{number}{stripped[assignment.end() - 1 :]}"
+    if line[0].isspace() and open_option:
+        return True, line
+    return False, None
+
+
+def _parseable_lines(text: str) -> Iterator[str]:
+    states: Final = itertools.accumulate(enumerate(text.splitlines()), _classify_line, initial=(False, None))
+    return (line for _, line in states if line is not None)
+
+
+def _lone_value(line: str) -> str | None:
+    tokens: Final = line.split()
+    if not tokens or '"' in tokens[0]:
+        return None
+    value: Final = tokens[0].rstrip(_SHELL_OPERATORS)
+    if len(tokens) == 1 or value != tokens[0] or _SHELL_TRAILER.fullmatch(tokens[1]) is not None:
+        return value
+    return None
+
+
+def _quoted_assignments(text: str) -> tuple[str, ...]:
+    parser: Final = configparser.ConfigParser(interpolation=None)
+    parser.optionxform = str  # pyright: ignore[reportAttributeAccessIssue]  # configparser types optionxform as a method
+    parser.read_string(f"[{_CONFIG_SECTION}]\n" + "\n".join(_parseable_lines(text)))
+    return tuple(
+        f'{key} = "{value}"'
+        for section in parser
+        for key, values in parser.items(section)
+        for line in values.splitlines()
+        if (value := _lone_value(line)) is not None
+    )
+
+
 class _ENTERPRISE_SecretDetection(CustomGuardrail):
     # Keeps proxied traffic on async_pre_call_hook (the unified apply_guardrail
     # path skips should_run_check and never sees data["prompt"]).
@@ -449,35 +542,21 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail):
         super().__init__(**kwargs)
 
     def scan_message_for_secrets(self, message_content: str):
-        from detect_secrets import SecretsCollection
         from detect_secrets.settings import transient_settings
 
-        temp_file = tempfile.NamedTemporaryFile(delete=False)
-        temp_file.write(message_content.encode("utf-8"))
-        temp_file.close()
-
-        secrets = SecretsCollection()
-
         detect_secrets_config = (
             self.user_defined_detect_secrets_config or _default_detect_secrets_config
         )
         with transient_settings(detect_secrets_config):
-            secrets.scan_file(temp_file.name)
-
-        os.remove(temp_file.name)
+            found: Final = _scan_lines(
+                (*message_content.splitlines(), *_quoted_assignments(message_content))
+            )
 
         return [
-            {"type": found_secret.type, "value": found_secret.secret_value}
-            for file in sorted(secrets.files)
-            for found_secret in sorted(
-                secrets[file],
-                key=lambda secret: (
-                    -len(secret.secret_value or ""),
-                    secret.type,
-                    secret.secret_value or "",
-                ),
+            {"type": secret_type, "value": value}
+            for value, secret_type in sorted(
+                found, key=lambda pair: (-len(pair[0]), pair[1], pair[0])
             )
-            if found_secret.secret_value is not None
         ]
 
     def redact_text(self, text: str, source: str = "message") -> str:
@@ -490,15 +569,16 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail):
         if counts is not None:
             for secret in detected_secrets:
                 counts[secret["type"]] = counts.get(secret["type"], 0) + 1
-        secret_types = [secret["type"] for secret in detected_secrets]
+        secret_types: Final = sorted(
+            dict.fromkeys(secret["type"] for secret in detected_secrets)
+        )
         verbose_proxy_logger.warning(
-            f"Detected and redacted secrets in {source}: {secret_types}"
+            "Detected and redacted secrets in %s: %s", source, secret_types
         )
-        return functools.reduce(
-            lambda redacted, secret: redacted.replace(secret["value"], "[REDACTED]"),
-            detected_secrets,
-            text,
+        pattern: Final = re.compile(
+            "|".join(re.escape(secret["value"]) for secret in detected_secrets)
         )
+        return pattern.sub("[REDACTED]", text)
 
     async def should_run_check(self, user_api_key_dict: UserAPIKeyAuth) -> bool:
         if user_api_key_dict.permissions is not None:

enterprise/litellm_enterprise/enterprise_callbacks/secrets_plugins/credential_keyword.py

@@ -0,0 +1,63 @@
+import re
+from collections.abc import Generator, Mapping
+from string import punctuation
+from typing import Final
+
+from detect_secrets.plugins.keyword import (
+    QUOTES_REQUIRED_DENYLIST_REGEX_TO_GROUP,
+    KeywordDetector,
+)
+
+_CREDENTIAL_VALUE: Final = re.compile(r"[^\s()\[\]]+")
+_ENVIRONMENT_REFERENCE: Final = re.compile(r"os\.environ/\w+", re.IGNORECASE)
+_ENVIRONMENT_VARIABLE_NAME: Final = re.compile(r"[A-Z][A-Z0-9]*(?:_[A-Z0-9]+)+")
+_LOWERCASE_WORD_SEQUENCE: Final = re.compile(r"[a-z]+(?:[-._/][a-z]+)+")
+_ISO_8601_TIMESTAMP: Final = re.compile(
+    r"\d{4}-\d{2}-\d{2}(?:T\d{2}:\d{2}(?::\d{2}(?:\.\d+)?)?(?:Z|[+-]\d{2}:?\d{2})?)?"
+)
+_URL_WITHOUT_USERINFO_OR_QUERY: Final = re.compile(r"[A-Za-z][A-Za-z0-9+.-]*://[^\s@?]*")
+_BENIGN_VALUES: Final = (
+    _ENVIRONMENT_REFERENCE,
+    _ENVIRONMENT_VARIABLE_NAME,
+    _LOWERCASE_WORD_SEQUENCE,
+    _ISO_8601_TIMESTAMP,
+    _URL_WITHOUT_USERINFO_OR_QUERY,
+)
+
+
+class CredentialKeywordDetector(KeywordDetector):  # pyright: ignore[reportUntypedBaseClass]  # detect_secrets ships no type information
+    secret_type = "Credential Keyword"
+
+    def __init__(self, minimum_length: int = 12, keyword_exclude: str | None = None) -> None:
+        if (
+            not isinstance(minimum_length, int)  # pyright: ignore[reportUnnecessaryIsInstance]  # the value comes from an operator's YAML
+            or minimum_length < 1
+        ):
+            raise ValueError(f"minimum_length must be a positive integer, got {minimum_length!r}")
+        super().__init__(keyword_exclude=keyword_exclude)
+        self.minimum_length = minimum_length
+
+    def _is_credential(self, value: str) -> bool:
+        core: Final = value.strip(punctuation)
+        return (
+            len(value) >= self.minimum_length
+            and _CREDENTIAL_VALUE.fullmatch(value) is not None
+            and all(benign.fullmatch(core) is None for benign in _BENIGN_VALUES)
+        )
+
+    def analyze_string(
+        self,
+        string: str,
+        denylist_regex_to_group: Mapping[re.Pattern[str], int] | None = None,
+    ) -> Generator[str, None, None]:
+        if self.keyword_exclude is not None and self.keyword_exclude.search(string):
+            return
+        regex_to_group: Final = (
+            QUOTES_REQUIRED_DENYLIST_REGEX_TO_GROUP if denylist_regex_to_group is None else denylist_regex_to_group
+        )
+        yield from (
+            match.group(group)
+            for regex, group in regex_to_group.items()
+            for match in regex.finditer(string)
+            if self._is_credential(match.group(group))
+        )

enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py

@@ -142,6 +142,40 @@ class CheckBatchCost:
             verbose_proxy_logger.error(f"CheckBatchCost: could not look up team alias for team {team_id}: {e}")
             return None
 
+    async def _get_org_id(self, job: "LiteLLM_ManagedObjectTable", batch_id: str) -> str | None:
+        org_id = getattr(job, "org_id", None)
+        if org_id:
+            return org_id
+        api_key = getattr(job, "api_key", None)
+        team_id = getattr(job, "team_id", None)
+        if api_key:
+            try:
+                key_row: prisma_models.LiteLLM_VerificationToken | None = (
+                    await self.prisma_client.db.litellm_verificationtoken.find_unique(
+                        where={"token": api_key}
+                    )
+                )
+                key_org_id = getattr(key_row, "organization_id", None) if key_row is not None else None
+                if key_org_id:
+                    return key_org_id
+            except Exception as e:
+                verbose_proxy_logger.error(
+                    f"CheckBatchCost: could not resolve the key's org for batch {batch_id}, "
+                    f"still trying the team's: {e}"
+                )
+        if not team_id:
+            return None
+        try:
+            team_row: prisma_models.LiteLLM_TeamTable | None = (
+                await self.prisma_client.db.litellm_teamtable.find_unique(
+                    where={"team_id": team_id}
+                )
+            )
+            return getattr(team_row, "organization_id", None) if team_row is not None else None
+        except Exception as e:
+            verbose_proxy_logger.error(f"CheckBatchCost: could not resolve the team's org for batch {batch_id}: {e}")
+            return None
+
     async def _build_creator_attribution_metadata(
         self, job: "LiteLLM_ManagedObjectTable", batch_id: str
     ) -> dict[str, object]:
@@ -153,6 +187,10 @@ class CheckBatchCost:
         user_api_key_alias; when it has no alias, or the key has since been rotated or
         deleted, the field keeps the creating user's alias that _get_user_info filled in,
         because a resolvable name is more useful on the spend row than a null.
+
+        user_api_key_org_id must be resolved here too: the spend update writer reads it
+        off this metadata to increment organization spend, so leaving it out silently
+        drops batch cost from org accounting for keys and teams that belong to one.
         """
         api_key = getattr(job, "api_key", None)
         team_id = getattr(job, "team_id", None)
@@ -172,6 +210,9 @@ class CheckBatchCost:
         team_alias = await self._get_team_alias(team_id)
         if team_alias is not None:
             metadata["user_api_key_team_alias"] = team_alias
+        org_id: Final = await self._get_org_id(job, batch_id)
+        if org_id is not None:
+            metadata["user_api_key_org_id"] = org_id
         if isinstance(request_tags, list) and request_tags:
             metadata["tags"] = [tag for tag in request_tags if isinstance(tag, str)]
 
@@ -641,7 +682,7 @@ class CheckBatchCost:
         from litellm.files.main import afile_content
         from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
         from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
-        from litellm.litellm_core_utils.litellm_logging import deployment_pricing_model_info
+        from litellm.litellm_core_utils.litellm_logging import deployment_pricing_model_info, mask_api_base_credentials
         from litellm.proxy.openai_files_endpoints.common_utils import (
             _is_base64_encoded_unified_file_id,
         )
@@ -805,6 +846,7 @@ class CheckBatchCost:
             function_id=str(uuid.uuid4()),
         )
 
+        deployment_api_base: Final = deployment_info.litellm_params.api_base
         logging_obj.update_environment_variables(
             litellm_params={
                 # set the user-agent header so that S3 callback consumers can easily identify CheckBatchCost callbacks
@@ -813,9 +855,17 @@ class CheckBatchCost:
                         "user-agent": CHECK_BATCH_COST_USER_AGENT,
                     }
                 },
-                "metadata": await self._build_creator_attribution_metadata(job, batch_id),
+                **({"api_base": mask_api_base_credentials(deployment_api_base)} if deployment_api_base else {}),
+                "metadata": {
+                    **(await self._build_creator_attribution_metadata(job, batch_id)),
+                    # spend logs read the deployment identity off these metadata keys, so
+                    # without them the batch cost row carries no model_id or model_group
+                    "model_info": {"id": model_id},
+                    "model_group": deployment_info.model_name,
+                },
             },
             optional_params={},
+            custom_llm_provider=str(llm_provider) if llm_provider else None,
         )
 
         if not await self._claim_job_for_costing(job):
@@ -833,6 +883,8 @@ class CheckBatchCost:
                 batch_models=batch_result.models,
                 batch_successful_requests=batch_result.successful_requests,
                 batch_failed_requests=batch_result.failed_requests,
+                batch_prompt_cost=batch_result.prompt_cost,
+                batch_completion_cost=batch_result.completion_cost,
             )
         except Exception:
             await self._release_job_claim(job)

enterprise/litellm_enterprise/proxy/hooks/managed_files.py

@@ -27,10 +27,13 @@ import litellm
 from litellm import Router, verbose_logger
 from litellm._uuid import uuid
 from litellm.caching.caching import DualCache
+from litellm.constants import MAX_FILE_LIST_LIMIT
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.litellm_core_utils.prompt_templates.common_utils import (
     extract_file_metadata,
 )
+from openai.types.file_deleted import FileDeleted
+
 from litellm.llms.base_llm.files.transformation import BaseFileEndpoints
 from litellm.llms.base_llm.managed_resources.isolation import (
     build_list_page,
@@ -48,7 +51,6 @@ from litellm.proxy._types import (
 from litellm.proxy.openai_files_endpoints.common_utils import (
     BATCH_CREATE_HIDDEN_PARAM,
     FILE_LIST_CONTINUATION_CHUNK_SIZE,
-    MAX_FILE_LIST_LIMIT,
     _is_base64_encoded_unified_file_id,
     apply_unified_file_ids,
     decode_model_from_file_id,
@@ -280,6 +282,27 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         )
         verbose_logger.debug(f"LiteLLM Managed File object with id={file_id} stored in db: {result}")
 
+    async def _resolve_creator_org_id(self, user_api_key_dict: UserAPIKeyAuth) -> Optional[str]:
+        if user_api_key_dict.org_id:
+            return user_api_key_dict.org_id
+        if not user_api_key_dict.team_id:
+            return None
+        from litellm.proxy.auth.auth_checks import get_team_object
+        from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache
+
+        try:
+            team: Final = await get_team_object(
+                team_id=user_api_key_dict.team_id,
+                prisma_client=self.prisma_client,
+                user_api_key_cache=user_api_key_cache,
+                parent_otel_span=user_api_key_dict.parent_otel_span,
+                proxy_logging_obj=proxy_logging_obj,
+            )
+            return team.organization_id
+        except Exception as e:
+            verbose_logger.warning(f"could not resolve org for managed object attribution: {e}")
+            return None
+
     async def store_unified_object_id(
         self,
         unified_object_id: str,
@@ -352,6 +375,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
                     "file_purpose": file_purpose,
                     "created_by": resolve_resource_owner_id(user_api_key_dict),
                     "team_id": user_api_key_dict.team_id,
+                    "org_id": await self._resolve_creator_org_id(user_api_key_dict),
                     "updated_by": user_api_key_dict.user_id,
                     "status": file_object.status,
                     **attribution_columns,
@@ -1765,7 +1789,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         litellm_parent_otel_span: Optional[Span],
         llm_router: Router,
         **data: Dict,
-    ) -> OpenAIFileObject:
+    ) -> FileDeleted:
 
         # Check if file deletion should be blocked due to batch references
         await self._check_file_deletion_allowed(file_id)
@@ -1773,29 +1797,28 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         # file_id = convert_b64_uid_to_unified_uid(file_id)
         model_file_id_mapping = await self.get_model_file_id_mapping([file_id], litellm_parent_otel_span)
 
-        delete_response = None
         specific_model_file_id_mapping = model_file_id_mapping.get(file_id)
         if specific_model_file_id_mapping:
             # Remove conflicting keys from data to avoid duplicate keyword arguments
             filtered_data = {k: v for k, v in data.items() if k not in ("model", "file_id")}
             for model_id, model_file_id in specific_model_file_id_mapping.items():
-                delete_response = await llm_router.afile_delete(model=model_id, file_id=model_file_id, **filtered_data)  # type: ignore
+                credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_id)
+                delete_data = {
+                    **{k: v for k, v in filtered_data.items() if k != "_litellm_internal_model_credentials"},
+                    **(
+                        {"_litellm_internal_model_credentials": MappingProxyType(dict(credentials))}
+                        if credentials is not None
+                        else {}
+                    ),
+                }
+                await llm_router.afile_delete(model=model_id, file_id=model_file_id, **delete_data)
 
-        stored_file_object = await self.delete_unified_file_id(file_id, litellm_parent_otel_span)
+        await self.delete_unified_file_id(file_id, litellm_parent_otel_span)
 
-        # Record successful deletion metric only on actual success
-        if stored_file_object or delete_response:
-            prom_logger = self._get_prometheus_logger()
-            if prom_logger:
-                prom_logger.record_managed_file_deleted(result="success")
-
-        if stored_file_object:
-            return stored_file_object
-        elif delete_response:
-            delete_response.id = file_id
-            return delete_response
-        else:
-            raise Exception(f"LiteLLM Managed File object with id={file_id} not found")
+        prom_logger = self._get_prometheus_logger()
+        if prom_logger:
+            prom_logger.record_managed_file_deleted(result="success")
+        return FileDeleted(id=file_id, object="file", deleted=True)
 
     async def afile_content(
         self,

enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py

@@ -780,7 +780,10 @@ async def update_project(
 
         # Handle budget updates
         budget_fields = LiteLLM_BudgetTable.model_fields.keys()
-        budget_updates = {k: v for k, v in update_data.items() if k in budget_fields}
+        budget_updates = {
+            **{k: v for k, v in update_data.items() if k in budget_fields},
+            **({"max_budget": None} if "max_budget" in data.model_fields_set and data.max_budget is None else {}),
+        }
 
         if budget_updates and existing_project.budget_id:
             # Update existing budget

enterprise/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "litellm-enterprise"
-version = "0.1.65"
+version = "0.1.67"
 description = "Package for LiteLLM Enterprise features"
 readme = "README.md"
 requires-python = ">=3.9"
@@ -26,7 +26,7 @@ required-version = ">=0.10.9"
 module-root = ""
 
 [tool.commitizen]
-version = "0.1.65"
+version = "0.1.67"
 version_files = [
     "pyproject.toml:^version",
     "../pyproject.toml:litellm-enterprise==",

gateway/Dockerfile

@@ -1,9 +1,25 @@
 ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
 ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
+# Checksum from https://www.pgbouncer.org/downloads/ (the Wolfi repo only carries 1.24.x)
+ARG PGBOUNCER_VERSION=1.25.2
+ARG PGBOUNCER_SHA256=924ad35113fd0a71c8e2dbe85b5d03445532e2b7b37a9f8a48983beea238b332
 
 FROM $UV_IMAGE AS uvbin
 
+FROM $LITELLM_BUILD_IMAGE AS pgbouncer-builder
+ARG PGBOUNCER_VERSION
+ARG PGBOUNCER_SHA256
+USER root
+RUN apk add --no-cache build-base pkgconf libevent-dev openssl-dev curl
+WORKDIR /build
+RUN curl -fsSL -o pgbouncer.tar.gz "https://www.pgbouncer.org/downloads/files/${PGBOUNCER_VERSION}/pgbouncer-${PGBOUNCER_VERSION}.tar.gz" && \
+    echo "${PGBOUNCER_SHA256}  pgbouncer.tar.gz" | sha256sum -c - && \
+    tar xzf pgbouncer.tar.gz --strip-components=1 && \
+    ./configure --prefix=/usr/local --with-openssl=/usr && \
+    make -j"$(nproc)" pgbouncer && \
+    install -m 0755 pgbouncer /usr/local/bin/pgbouncer
+
 # ---------- Builder ----------
 FROM $LITELLM_BUILD_IMAGE AS builder
 
@@ -47,7 +63,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
         --extra extra_proxy \
         --extra semantic-router \
         --extra bedrock-realtime \
-        --extra mongodb \
         --python python3.13
 
 # Stage 2 — copy source and install the project + workspace members.
@@ -60,9 +75,12 @@ RUN --mount=type=cache,target=/root/.cache/uv \
         --extra extra_proxy \
         --extra semantic-router \
         --extra bedrock-realtime \
-        --extra mongodb \
         --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
@@ -75,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
@@ -92,15 +110,17 @@ ENV HOME=/home/nonroot \
 
 COPY --from=builder --chown=nonroot:nonroot /app /app
 COPY --from=builder /opt/prisma /opt/prisma
+COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer
 
 RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \
     find /app/.venv -type d -path "*/tornado/test" -delete && \
     chmod -R a+rX /opt/prisma && \
-    python -c "from prisma.client import BINARY_PATHS; paths = list(BINARY_PATHS.query_engine.values()); assert paths and all(p.startswith('/opt/prisma/') for p in paths), paths"
+    python -c "from prisma.client import BINARY_PATHS; paths = list(BINARY_PATHS.query_engine.values()); assert paths and all(p.startswith('/opt/prisma/') for p in paths), paths" && \
+    python -c "import litellm; from litellm.rust_bridge.loader import native_bridge_available; assert litellm.__file__ == '/app/litellm/__init__.py', litellm.__file__; assert native_bridge_available()"
 
 USER nonroot
 
 EXPOSE 4000/tcp
 
-ENTRYPOINT ["sh", "-c", "exec /app/docker/component_entrypoint.sh uvicorn gateway.main:app --workers \"${NUM_WORKERS:-1}\" \"$@\"", "--"]
+ENTRYPOINT ["sh", "-c", "exec /app/docker/component_entrypoint.sh python -m gateway.launch --workers \"${NUM_WORKERS:-1}\" \"$@\"", "--"]
 CMD ["--host", "0.0.0.0", "--port", "4000"]

gateway/launch.py

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

gateway/routes/allowlist.py

@@ -3,7 +3,8 @@
 The gateway exposes the LLM data-plane surface: chat/completions, embeddings,
 audio, batches, files, fine-tuning, rerank, ocr, rag, video, search, image,
 responses, vector stores, passthrough providers, realtime websockets, MCP
-tool-call endpoints, and operational endpoints (/health, /metrics).
+tool-call endpoints, and operational endpoints (/health, /metrics, and the
+/debug/memory/summary read of the serving worker's RSS).
 
 Any path not listed here is dropped from the gateway process so management/UI
 endpoints don't ride on the same pods.
@@ -121,6 +122,7 @@ GATEWAY_EXACT_PATHS: frozenset[str] = frozenset(
         "/docs/oauth2-redirect",
         "/redoc",
         "/test",
+        "/debug/memory/summary",
     }
 )
 

helm/litellm-helm/templates/_helpers.tpl

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

helm/litellm-helm/templates/deployment.yaml

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

helm/litellm-helm/templates/hpa.yaml

@@ -18,6 +18,15 @@ spec:
   {{- end }}
   metrics:
     {{- if .Values.autoscaling.targetCPUUtilizationPercentage }}
+    {{- if and .Values.collector.enabled .Values.collector.scaleOnProxyContainerCpu }}
+    - type: ContainerResource
+      containerResource:
+        name: cpu
+        container: {{ include "litellm.name" . }}
+        target:
+          type: Utilization
+          averageUtilization: {{ .Values.autoscaling.targetCPUUtilizationPercentage }}
+    {{- else }}
     - type: Resource
       resource:
         name: cpu
@@ -25,6 +34,7 @@ spec:
           type: Utilization
           averageUtilization: {{ .Values.autoscaling.targetCPUUtilizationPercentage }}
     {{- end }}
+    {{- end }}
     {{- if .Values.autoscaling.targetMemoryUtilizationPercentage }}
     - type: Resource
       resource:
@@ -33,4 +43,22 @@ spec:
           type: Utilization
           averageUtilization: {{ .Values.autoscaling.targetMemoryUtilizationPercentage }}
     {{- end }}
+    {{- with .Values.autoscaling.targetRequestsPerSecond }}
+    - type: Pods
+      pods:
+        metric:
+          name: litellm_requests_per_second
+        target:
+          type: AverageValue
+          averageValue: {{ toJson . | trimAll "\"" | quote }}
+    {{- end }}
+    {{- with .Values.autoscaling.targetTokensPerSecond }}
+    - type: Pods
+      pods:
+        metric:
+          name: litellm_tokens_per_second
+        target:
+          type: AverageValue
+          averageValue: {{ toJson . | trimAll "\"" | quote }}
+    {{- end }}
 {{- end }}

helm/litellm-helm/templates/keda.yaml

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

helm/litellm-helm/templates/service-metrics.yaml

@@ -0,0 +1,17 @@
+{{- if .Values.metricsServer.enabled }}
+apiVersion: v1
+kind: Service
+metadata:
+  name: {{ include "litellm.fullname" . }}-metrics
+  labels:
+    {{- include "litellm.labels" . | nindent 4 }}
+spec:
+  type: ClusterIP
+  ports:
+    - port: {{ .Values.metricsServer.port }}
+      targetPort: metrics
+      protocol: TCP
+      name: metrics
+  selector:
+    {{- include "litellm.selectorLabels" . | nindent 4 }}
+{{- end }}

helm/litellm-helm/templates/servicemonitor.yaml

@@ -26,7 +26,7 @@ spec:
     {{- toYaml .namespaceSelector.matchNames | nindent 4 }}
     {{- end }}
   endpoints:
-  - port: http
+  - port: {{ ternary "metrics" "http" $.Values.metricsServer.enabled }}
     path: /metrics/
     interval: {{ .interval }}
     scrapeTimeout: {{ .scrapeTimeout }}

helm/litellm-helm/tests/collector_tests.yaml

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

helm/litellm-helm/tests/connection_pool_tests.yaml

@@ -0,0 +1,112 @@
+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"
+
+  - it: should give the collector sidecar the same pool env as the proxy container
+    template: deployment.yaml
+    set:
+      collector.enabled: true
+      db.connectionPool.enabled: true
+      db.connectionPool.maxDbConnections: 8
+      db.connectionPool.maxClientConn: 400
+    asserts:
+      - equal:
+          path: spec.template.spec.containers[1].name
+          value: litellm-collector
+      - contains:
+          path: spec.template.spec.containers[1].env
+          content:
+            name: LITELLM_PGBOUNCER_ENABLED
+            value: "true"
+      - contains:
+          path: spec.template.spec.containers[1].env
+          content:
+            name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
+            value: "8"
+      - contains:
+          path: spec.template.spec.containers[1].env
+          content:
+            name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
+            value: "400"
+
+  - it: should give the collector sidecar no pool env when the pool is off
+    template: deployment.yaml
+    set:
+      collector.enabled: true
+    asserts:
+      - equal:
+          path: spec.template.spec.containers[1].name
+          value: litellm-collector
+      - notContains:
+          path: spec.template.spec.containers[1].env
+          content:
+            name: LITELLM_PGBOUNCER_ENABLED
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[1].env
+          content:
+            name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[1].env
+          content:
+            name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
+          any: true

helm/litellm-helm/tests/hpa_tests.yaml

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

helm/litellm-helm/tests/keda_tests.yaml

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

helm/litellm-helm/tests/metrics_server_tests.yaml

@@ -0,0 +1,106 @@
+suite: separate metrics server
+templates:
+  - configmap-litellm.yaml
+  - deployment.yaml
+  - service.yaml
+  - service-metrics.yaml
+  - servicemonitor.yaml
+tests:
+  - it: should not expose a metrics port or PROMETHEUS_METRICS_PORT by default
+    asserts:
+      - notContains:
+          path: spec.template.spec.containers[0].ports
+          content:
+            name: metrics
+          any: true
+        template: deployment.yaml
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: PROMETHEUS_METRICS_PORT
+          any: true
+        template: deployment.yaml
+      - lengthEqual:
+          path: spec.ports
+          count: 1
+        template: service.yaml
+      - hasDocuments:
+          count: 0
+        template: service-metrics.yaml
+
+  - it: should scrape the proxy port when the metrics server is disabled
+    template: servicemonitor.yaml
+    set:
+      serviceMonitor.enabled: true
+    asserts:
+      - equal:
+          path: spec.endpoints[0].port
+          value: http
+
+  - it: should wire the separate metrics server through container, a ClusterIP metrics service and servicemonitor
+    set:
+      metricsServer.enabled: true
+      metricsServer.port: 4101
+      serviceMonitor.enabled: true
+      service.type: LoadBalancer
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: PROMETHEUS_METRICS_PORT
+            value: "4101"
+        template: deployment.yaml
+      - contains:
+          path: spec.template.spec.containers[0].ports
+          content:
+            name: metrics
+            containerPort: 4101
+            protocol: TCP
+        template: deployment.yaml
+      - lengthEqual:
+          path: spec.ports
+          count: 1
+        template: service.yaml
+      - equal:
+          path: spec.type
+          value: LoadBalancer
+        template: service.yaml
+      - equal:
+          path: metadata.name
+          value: RELEASE-NAME-litellm-metrics
+        template: service-metrics.yaml
+      - equal:
+          path: spec.type
+          value: ClusterIP
+        template: service-metrics.yaml
+      - equal:
+          path: spec.ports
+          value:
+            - port: 4101
+              targetPort: metrics
+              protocol: TCP
+              name: metrics
+        template: service-metrics.yaml
+      - equal:
+          path: spec.selector
+          value:
+            app.kubernetes.io/name: litellm
+            app.kubernetes.io/instance: RELEASE-NAME
+        template: service-metrics.yaml
+      - equal:
+          path: spec.endpoints[0].port
+          value: metrics
+        template: servicemonitor.yaml
+      - equal:
+          path: spec.endpoints[0].path
+          value: /metrics/
+        template: servicemonitor.yaml
+
+  - it: should reject a metrics port equal to the proxy port
+    template: deployment.yaml
+    set:
+      metricsServer.enabled: true
+      metricsServer.port: 4000
+    asserts:
+      - failedTemplate:
+          errorMessage: metricsServer.port must differ from service.port

helm/litellm-helm/values.yaml

@@ -180,6 +180,58 @@ proxy_config:
   general_settings:
     master_key: os.environ/PROXY_MASTER_KEY
 
+# Serve Prometheus /metrics from a separate process (PROMETHEUS_METRICS_PORT)
+# so a scrape never runs on an inference worker. Adds a `metrics` port to the
+# container and a dedicated ClusterIP `<release>-metrics` Service, and the
+# ServiceMonitor scrapes it instead of the proxy port. The separate port has
+# no virtual-key auth: keep it off public ingress. Needs the proxy image
+# v1.101.0 or newer.
+metricsServer:
+  enabled: false
+  port: 4001
+
+# Opt-in sidecar that runs the post-response spend pipeline (cost calculation,
+# spend logs, spend counters, budget reservation reconciliation) so the proxy's
+# uvicorn workers only serialise a compact typed event and go back to serving
+# inference. Same image and tag as the proxy, second container in the same pod,
+# fed over loopback (a unix socket on a shared emptyDir, or 127.0.0.1 TCP). It
+# reuses the pod's in-container pgbouncer (db.connectionPool) and the same Redis
+# spend transaction buffer, so the per-pod DB connection budget is unchanged.
+# Delivery is at-most-once inside the pod: events already handed to the sidecar
+# are lost if it crashes before writing them; events the workers could not hand
+# over follow onUnavailable. Both containers drain on SIGTERM within
+# terminationGracePeriodSeconds
+collector:
+  enabled: false
+  # unix:///<dir>/<file>.sock (the <dir> becomes a shared emptyDir) or tcp://127.0.0.1:<port>
+  address: unix:///var/run/litellm/collector.sock
+  # Events each uvicorn worker holds in memory while the sidecar is slow or restarting
+  bufferSize: 1000
+  # fallback: run the pipeline in the worker when the sidecar is unreachable or the
+  # buffer is full (spend stays exact, that request costs proxy CPU again)
+  # drop: count and discard the event instead (spend under-reports)
+  onUnavailable: fallback
+  # How long the workers keep pushing buffered events on shutdown, and how long the
+  # sidecar keeps serving its open connections after SIGTERM
+  drainTimeoutSeconds: 10
+  command:
+    - python
+    - -m
+    - litellm.proxy.collector
+  # Sized independently of the proxy container; the pipeline is CPU bound
+  resources: {}
+  # requests:
+  #   cpu: 500m
+  #   memory: 1Gi
+  # limits:
+  #   cpu: "1"
+  #   memory: 2Gi
+  # When autoscaling.enabled, swap the pod-wide cpu Resource metric for an
+  # autoscaling/v2 ContainerResource metric on the proxy container only, so the
+  # sidecar's CPU never scales inference replicas. Needs Kubernetes 1.30+ (or the
+  # HPAContainerMetrics feature gate on 1.27 to 1.29)
+  scaleOnProxyContainerCpu: false
+
 resources:
   {}
   # Unset by default so the chart installs on small clusters such as Minikube, and so an
@@ -212,6 +264,25 @@ autoscaling:
   # Memory is a floor to provision under 'resources', not a signal to scale on.
   # targetMemoryUtilizationPercentage: 80
   # behavior: {}
+  # Opt-in per-pod workload targets, rendered as autoscaling/v2 `Pods` metrics
+  # named `litellm_requests_per_second` and `litellm_tokens_per_second` with an
+  # AverageValue target, alongside whichever resource targets are set (the HPA
+  # follows the metric asking for the most replicas). A Prometheus Adapter must
+  # serve those two names on custom.metrics.k8s.io from the proxy's counters,
+  # grouped by the scrape target's `pod` label (enable serviceMonitor below so
+  # every pod is scraped on its own):
+  #   litellm_requests_per_second:
+  #     sum(rate(litellm_proxy_total_requests_metric_total{<<.LabelMatchers>>}[1m])) by (<<.GroupBy>>)
+  #   litellm_tokens_per_second:
+  #     sum(rate(litellm_total_tokens_metric_total{<<.LabelMatchers>>}[1m])) by (<<.GroupBy>>)
+  # rate() over [1m] is already per second, so no `* 60`. How fast the HPA
+  # reacts is set by that window, the scrape interval and the HPA sync period
+  # (15s by default), not by the unit: keep serviceMonitor.interval at 15s or
+  # faster so a 1m window holds at least 4 samples. averageValue takes SI
+  # suffixes, so "6M" is six million tokens per second per pod. Tokens are
+  # counted when a response completes, so TPS trails long streams.
+  targetRequestsPerSecond: ""
+  targetTokensPerSecond: ""
 
 # Autoscaling with keda is mutually exclusive with hpa
 keda:
@@ -233,6 +304,23 @@ keda:
   #     metricName: http_requests_total
   #     threshold: '100'
   #     query: sum(rate(http_requests_total{deployment="my-deployment"}[2m]))
+  # First-class Prometheus triggers on the proxy's own request and token
+  # counters, appended to `triggers`. Each target is the per-second load one
+  # replica should carry: KEDA divides the release-wide
+  # `sum(rate(<counter>[1m]))` by it to pick the replica count. Thresholds
+  # are plain numbers (KEDA parses them as floats, no SI suffixes). The
+  # queries select samples by the release namespace and the `job` label the
+  # chart's ServiceMonitor produces (the metrics Service name), so enable
+  # serviceMonitor below together with metricsServer: the http port serves
+  # /metrics/ behind virtual-key auth and answers an unauthenticated scrape
+  # with 401. Reaction time comes from the [1m] window, the scrape interval
+  # and pollingInterval above, so keep both at 15s or faster. Tokens are
+  # counted at completion, so TPS trails long streams. serverAddress is
+  # required once either target is set.
+  prometheus:
+    serverAddress: ""
+    requestsPerSecond: ""
+    tokensPerSecond: ""
   behavior: {}
   # scaleDown:
   #   stabilizationWindowSeconds: 300
@@ -309,6 +397,20 @@ db:
   # only (e.g. when IAM_TOKEN_DB_AUTH supplies the token at runtime).
   readReplicaUrl: ""
 
+  # In-container connection pool (PgBouncer, transaction mode) shared by every
+  # worker in the pod. Without it each --num_workers worker opens its own
+  # connection_limit connections to Postgres, so a pod's footprint against the
+  # database's connection ceiling is workers x connection_limit and grows with
+  # every replica. With it, the pod holds at most maxDbConnections upstream
+  # connections no matter how many workers run; the workers connect to the pool
+  # over loopback, with no extra network hop. Migrations still go straight to
+  # Postgres. Starting profile for numWorkers: 4 is maxDbConnections: 20, so
+  # a database with a 5000-connection ceiling fits roughly 200 replicas.
+  connectionPool:
+    enabled: false
+    maxDbConnections: 20
+    maxClientConn: 1000
+
   # Use the Stackgres Helm chart to deploy an instance of a Stackgres cluster.
   #  The Stackgres Operator must already be installed within the target
   #  Kubernetes cluster.

helm/litellm/templates/_helpers.tpl

@@ -257,6 +257,14 @@ IAM_TOKEN_DB_AUTH / AZURE_POSTGRESQL_AUTH toggle that only the writer sets.
 - name: DATABASE_SCHEMA
   value: {{ .schema | quote }}
 {{- end }}
+{{- if .sslMode }}
+- name: DATABASE_SSLMODE
+  value: {{ .sslMode | quote }}
+{{- end }}
+{{- if .sslRootCert }}
+- name: DATABASE_SSLROOTCERT
+  value: {{ .sslRootCert | quote }}
+{{- end }}
 {{- if and .useIAMAuth .useAzureEntraAuth }}
 {{- fail "database.writer.useIAMAuth and database.writer.useAzureEntraAuth are mutually exclusive: the database password can only come from one token source" }}
 {{- end }}
@@ -360,6 +368,20 @@ harmless no-op for the Job and authoritative for the app pods.
 {{- end }}
 {{- end -}}
 
+{{/*
+In-container PgBouncer env for the gateway container. Under IAM or Entra auth the pooler mints and renews the database token itself.
+*/}}
+{{- define "litellm.connectionPoolEnv" -}}
+{{- with .Values.database.connectionPool -}}
+- name: LITELLM_PGBOUNCER_ENABLED
+  value: "true"
+- name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
+  value: {{ required "database.connectionPool.maxDbConnections is required when the pool is enabled" .maxDbConnections | quote }}
+- name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
+  value: {{ required "database.connectionPool.maxClientConn is required when the pool is enabled" .maxClientConn | quote }}
+{{- end }}
+{{- end -}}
+
 {{/*
 PodDisruptionBudget shared by gateway, backend, and ui.
 
@@ -441,3 +463,36 @@ ImplementationSpecific
 {{- .pathType -}}
 {{- end -}}
 {{- end -}}
+
+{{- define "litellm.gateway.prometheusMultiprocDir" -}}/tmp/litellm_prometheus_multiproc{{- end -}}
+
+{{/*
+Directory of the collector's unix socket, shared by the gateway and
+collector containers through an emptyDir. Empty when the sidecar is off
+or gateway.collector.address is a tcp://127.0.0.1:<port> address.
+*/}}
+{{- define "litellm.gateway.collectorSocketDir" -}}
+{{- if and .Values.gateway.collector.enabled (hasPrefix "unix://" .Values.gateway.collector.address) -}}
+{{- dir (trimPrefix "unix://" .Values.gateway.collector.address) -}}
+{{- end -}}
+{{- end -}}
+
+{{/*
+LITELLM_COLLECTOR_* env shared by the producer (gateway container) and the
+consumer (collector container), so both agree on the transport and the
+shutdown drain window.
+*/}}
+{{- define "litellm.gateway.collectorEnv" -}}
+{{- with .Values.gateway.collector }}
+- name: LITELLM_COLLECTOR_ENABLED
+  value: "true"
+- name: LITELLM_COLLECTOR_ADDRESS
+  value: {{ .address | quote }}
+- name: LITELLM_COLLECTOR_BUFFER_SIZE
+  value: {{ .bufferSize | quote }}
+- name: LITELLM_COLLECTOR_ON_UNAVAILABLE
+  value: {{ .onUnavailable | quote }}
+- name: LITELLM_COLLECTOR_DRAIN_TIMEOUT_SECONDS
+  value: {{ .drainTimeoutSeconds | quote }}
+{{- end }}
+{{- end -}}

helm/litellm/templates/gateway/deployment.yaml

@@ -61,17 +61,38 @@ 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 }}
+            {{- if .Values.gateway.metricsServer.enabled }}
+            {{- if eq (int .Values.gateway.metricsServer.port) 4000 }}
+            {{- fail "gateway.metricsServer.port must differ from the gateway port 4000" }}
+            {{- end }}
+            - name: PROMETHEUS_MULTIPROC_DIR
+              value: {{ include "litellm.gateway.prometheusMultiprocDir" . }}
+            {{- end }}
+            {{- if .Values.gateway.collector.enabled }}
+            {{- include "litellm.gateway.collectorEnv" . | nindent 12 }}
+            {{- end }}
           {{- include "litellm.envFrom" .Values.gateway | nindent 10 }}
-          {{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.billingMetrics.enabled }}
+          {{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled (include "litellm.gateway.collectorSocketDir" .) }}
           volumeMounts:
             {{- if .Values.gateway.config.create }}
             - name: gateway-config
               mountPath: /app/config/config.yaml
               subPath: config.yaml
             {{- end }}
+            {{- if .Values.gateway.metricsServer.enabled }}
+            - name: prometheus-multiproc
+              mountPath: {{ include "litellm.gateway.prometheusMultiprocDir" . }}
+            {{- end }}
+            {{- if include "litellm.gateway.collectorSocketDir" . }}
+            - name: collector-socket
+              mountPath: {{ include "litellm.gateway.collectorSocketDir" . }}
+            {{- end }}
             {{- if .Values.billingMetrics.enabled }}
             {{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }}
             {{- end }}
@@ -97,16 +118,102 @@ spec:
           {{- end }}
           resources:
             {{- toYaml .Values.gateway.resources | nindent 12 }}
+        {{- if .Values.gateway.metricsServer.enabled }}
+        - name: metrics
+          image: "{{ .Values.gateway.image.repository }}:{{ .Values.gateway.image.tag | default .Chart.AppVersion }}"
+          imagePullPolicy: {{ .Values.gateway.image.pullPolicy }}
+          {{- with .Values.gateway.securityContext }}
+          securityContext:
+            {{- toYaml . | nindent 12 }}
+          {{- end }}
+          command:
+            - python
+            - -m
+            - litellm.proxy.prometheus_metrics_server
+            - --port
+            - {{ .Values.gateway.metricsServer.port | quote }}
+          env:
+            - name: PROMETHEUS_MULTIPROC_DIR
+              value: {{ include "litellm.gateway.prometheusMultiprocDir" . }}
+          ports:
+            - name: metrics
+              containerPort: {{ .Values.gateway.metricsServer.port }}
+              protocol: TCP
+          volumeMounts:
+            - name: prometheus-multiproc
+              mountPath: {{ include "litellm.gateway.prometheusMultiprocDir" . }}
+          readinessProbe:
+            tcpSocket: { port: metrics }
+            periodSeconds: 10
+          livenessProbe:
+            tcpSocket: { port: metrics }
+            periodSeconds: 15
+            failureThreshold: 6
+          resources:
+            {{- toYaml .Values.gateway.metricsServer.resources | nindent 12 }}
+        {{- end }}
+        {{- if .Values.gateway.collector.enabled }}
+        - name: collector
+          image: "{{ .Values.gateway.image.repository }}:{{ .Values.gateway.image.tag | default .Chart.AppVersion }}"
+          imagePullPolicy: {{ .Values.gateway.image.pullPolicy }}
+          {{- with .Values.gateway.securityContext }}
+          securityContext:
+            {{- toYaml . | nindent 12 }}
+          {{- end }}
+          command:
+            - python
+            - -m
+            - litellm.proxy.collector
+          env:
+            {{- include "litellm.serverEnv" (dict "root" $ "component" .Values.gateway) | nindent 12 }}
+            {{- if .Values.gateway.config.create }}
+            - name: CONFIG_FILE_PATH
+              value: /app/config/config.yaml
+            {{- end }}
+            {{- if .Values.database.connectionPool.enabled }}
+            {{- include "litellm.connectionPoolEnv" $ | nindent 12 }}
+            {{- end }}
+            {{- include "litellm.gateway.collectorEnv" . | nindent 12 }}
+            - name: LITELLM_JOB_ROLE
+              value: collector
+          {{- include "litellm.envFrom" .Values.gateway | nindent 10 }}
+          {{- if or .Values.gateway.config.create .Values.gateway.volumeMounts (include "litellm.gateway.collectorSocketDir" .) }}
+          volumeMounts:
+            {{- if .Values.gateway.config.create }}
+            - name: gateway-config
+              mountPath: /app/config/config.yaml
+              subPath: config.yaml
+            {{- end }}
+            {{- if include "litellm.gateway.collectorSocketDir" . }}
+            - name: collector-socket
+              mountPath: {{ include "litellm.gateway.collectorSocketDir" . }}
+            {{- end }}
+            {{- with .Values.gateway.volumeMounts }}
+            {{- toYaml . | nindent 12 }}
+            {{- end }}
+          {{- end }}
+          resources:
+            {{- toYaml .Values.gateway.collector.resources | nindent 12 }}
+        {{- end }}
       {{- with .Values.gateway.extraContainers }}
         {{- tpl (toYaml .) $ | nindent 8 }}
       {{- end }}
-      {{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.billingMetrics.enabled }}
+      {{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled (include "litellm.gateway.collectorSocketDir" .) }}
       volumes:
         {{- if .Values.gateway.config.create }}
         - name: gateway-config
           configMap:
             name: {{ include "litellm.gateway.fullname" . }}-config
         {{- end }}
+        {{- if .Values.gateway.metricsServer.enabled }}
+        - name: prometheus-multiproc
+          emptyDir: {}
+        {{- end }}
+        {{- if include "litellm.gateway.collectorSocketDir" . }}
+        - name: collector-socket
+          emptyDir:
+            sizeLimit: 1Mi
+        {{- end }}
         {{- if .Values.billingMetrics.enabled }}
         {{- include "litellm.billingMetricsVolumes" . | nindent 8 }}
         {{- end }}

helm/litellm/templates/gateway/hpa.yaml

@@ -15,6 +15,15 @@ spec:
   maxReplicas: {{ .Values.gateway.hpa.maxReplicas }}
   metrics:
     {{- if .Values.gateway.hpa.targetCPUUtilizationPercentage }}
+    {{- if and .Values.gateway.collector.enabled .Values.gateway.collector.scaleOnGatewayContainerCpu }}
+    - type: ContainerResource
+      containerResource:
+        name: cpu
+        container: gateway
+        target:
+          type: Utilization
+          averageUtilization: {{ .Values.gateway.hpa.targetCPUUtilizationPercentage }}
+    {{- else }}
     - type: Resource
       resource:
         name: cpu
@@ -22,6 +31,7 @@ spec:
           type: Utilization
           averageUtilization: {{ .Values.gateway.hpa.targetCPUUtilizationPercentage }}
     {{- end }}
+    {{- end }}
     {{- if .Values.gateway.hpa.targetMemoryUtilizationPercentage }}
     - type: Resource
       resource:
@@ -30,6 +40,24 @@ spec:
           type: Utilization
           averageUtilization: {{ .Values.gateway.hpa.targetMemoryUtilizationPercentage }}
     {{- end }}
+    {{- with .Values.gateway.hpa.targetRequestsPerSecond }}
+    - type: Pods
+      pods:
+        metric:
+          name: litellm_requests_per_second
+        target:
+          type: AverageValue
+          averageValue: {{ toJson . | trimAll "\"" | quote }}
+    {{- end }}
+    {{- with .Values.gateway.hpa.targetTokensPerSecond }}
+    - type: Pods
+      pods:
+        metric:
+          name: litellm_tokens_per_second
+        target:
+          type: AverageValue
+          averageValue: {{ toJson . | trimAll "\"" | quote }}
+    {{- end }}
   {{- with .Values.gateway.hpa.behavior }}
   behavior:
     {{- toYaml . | nindent 4 }}

helm/litellm/templates/gateway/service-metrics.yaml

@@ -0,0 +1,18 @@
+{{- if and .Values.gateway.enabled .Values.gateway.metricsServer.enabled }}
+apiVersion: v1
+kind: Service
+metadata:
+  name: {{ include "litellm.gateway.fullname" . }}-metrics
+  labels:
+    {{- include "litellm.commonLabels" . | nindent 4 }}
+    app.kubernetes.io/component: gateway
+spec:
+  type: ClusterIP
+  ports:
+    - port: {{ .Values.gateway.metricsServer.port }}
+      targetPort: metrics
+      protocol: TCP
+      name: metrics
+  selector:
+    {{- include "litellm.gateway.selectorLabels" . | nindent 4 }}
+{{- end }}

helm/litellm/templates/gateway/servicemonitor.yaml

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

helm/litellm/templates/ingress.yaml

@@ -89,7 +89,7 @@
   at "/" Prefix would swallow the whole backend management API) instead of
   adding to it.
 */}}
-{{- $builtinPathKeys := list "/test|Exact" "/|Prefix" -}}
+{{- $builtinPathKeys := list "/test|Exact" "/debug/memory/summary|Exact" "/|Prefix" -}}
 apiVersion: networking.k8s.io/v1
 kind: Ingress
 metadata:
@@ -129,6 +129,8 @@ spec:
           # --- Gateway data plane ---
           # Exact /test only (see the $gatewayPrefixes comment above);
           # /test/* MCP management endpoints fall to the backend catch-all.
+          # Exact /debug/memory/summary reads a serving worker's RSS (the e2e memory
+          # gate); the rest of /debug/* stays on the backend.
           - path: /test
             pathType: Exact
             backend:
@@ -136,6 +138,13 @@ spec:
                 name: {{ $gatewayName }}
                 port:
                   number: {{ $gatewayPort }}
+          - path: /debug/memory/summary
+            pathType: Exact
+            backend:
+              service:
+                name: {{ $gatewayName }}
+                port:
+                  number: {{ $gatewayPort }}
           {{- range $gatewayPrefixes }}
           {{- $pathType := include "litellm.ingress.pathType" (dict "controller" $controller "path" . "pathType" "Prefix") }}
           {{- $builtinPathKeys = append $builtinPathKeys (printf "%s|%s" . $pathType) }}

helm/litellm/tests/collector_tests.yaml

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

helm/litellm/tests/connection_pool_tests.yaml

@@ -0,0 +1,204 @@
+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: collector sidecar gets the same pool env as the gateway container, the metrics sidecar none
+    template: gateway/deployment.yaml
+    set:
+      gateway.collector.enabled: true
+      gateway.metricsServer.enabled: true
+      database.connectionPool.enabled: true
+      database.connectionPool.maxDbConnections: 8
+      database.connectionPool.maxClientConn: 250
+    asserts:
+      - equal:
+          path: spec.template.spec.containers[1].name
+          value: metrics
+      - notContains:
+          path: spec.template.spec.containers[1].env
+          content:
+            name: LITELLM_PGBOUNCER_ENABLED
+          any: true
+      - equal:
+          path: spec.template.spec.containers[2].name
+          value: collector
+      - contains:
+          path: spec.template.spec.containers[2].env
+          content:
+            name: LITELLM_PGBOUNCER_ENABLED
+            value: "true"
+      - contains:
+          path: spec.template.spec.containers[2].env
+          content:
+            name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
+            value: "8"
+      - contains:
+          path: spec.template.spec.containers[2].env
+          content:
+            name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
+            value: "250"
+
+  - it: collector sidecar gets no pool env when the pool is off
+    template: gateway/deployment.yaml
+    set:
+      gateway.collector.enabled: true
+    asserts:
+      - equal:
+          path: spec.template.spec.containers[1].name
+          value: collector
+      - notContains:
+          path: spec.template.spec.containers[1].env
+          content:
+            name: LITELLM_PGBOUNCER_ENABLED
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[1].env
+          content:
+            name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[1].env
+          content:
+            name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
+          any: true
+
+  - it: pool with IAM auth renders both the pool and the token auth flag
+    template: gateway/deployment.yaml
+    set:
+      gateway.collector.enabled: true
+      database.connectionPool.enabled: true
+      database.writer.useIAMAuth: true
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: LITELLM_PGBOUNCER_ENABLED
+            value: "true"
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: IAM_TOKEN_DB_AUTH
+            value: "true"
+      - contains:
+          path: spec.template.spec.containers[1].env
+          content:
+            name: LITELLM_PGBOUNCER_ENABLED
+            value: "true"
+      - contains:
+          path: spec.template.spec.containers[1].env
+          content:
+            name: IAM_TOKEN_DB_AUTH
+            value: "true"
+
+  - it: pool with Entra auth renders both the pool and the token auth flag
+    template: gateway/deployment.yaml
+    set:
+      database.connectionPool.enabled: true
+      database.writer.useAzureEntraAuth: true
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: LITELLM_PGBOUNCER_ENABLED
+            value: "true"
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: AZURE_POSTGRESQL_AUTH
+            value: "true"
+
+  - it: IAM auth without the pool still renders
+    template: gateway/deployment.yaml
+    set:
+      database.writer.useIAMAuth: true
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: IAM_TOKEN_DB_AUTH
+            value: "true"
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: LITELLM_PGBOUNCER_ENABLED
+          any: true

helm/litellm/tests/database_auth_tests.yaml

@@ -4,6 +4,7 @@ templates:
   - gateway/configmap.yaml
   - backend/deployment.yaml
   - backend/configmap.yaml
+  - migrations-job.yaml
 values:
   - ./values/required.yaml
 tests:
@@ -67,6 +68,82 @@ tests:
             value: "true"
           any: true
 
+  - it: emits no TLS env by default
+    template: gateway/deployment.yaml
+    asserts:
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_SSLMODE
+          any: true
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_SSLROOTCERT
+          any: true
+
+  - it: writer sslMode and sslRootCert reach gateway and backend as DATABASE_SSLMODE and DATABASE_SSLROOTCERT
+    templates:
+      - gateway/deployment.yaml
+      - backend/deployment.yaml
+    set:
+      database.writer.useIAMAuth: true
+      database.writer.sslMode: verify-full
+      database.writer.sslRootCert: /etc/ssl/certs/ca-certificates.crt
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_SSLMODE
+            value: verify-full
+          any: true
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_SSLROOTCERT
+            value: /etc/ssl/certs/ca-certificates.crt
+          any: true
+
+  - it: writer sslMode and sslRootCert reach the collector sidecar and the migrations job, which dial Postgres themselves
+    set:
+      gateway.collector.enabled: true
+      database.connectionPool.enabled: true
+      database.writer.sslMode: verify-full
+      database.writer.sslRootCert: /etc/ssl/certs/ca-certificates.crt
+    asserts:
+      - equal:
+          path: spec.template.spec.containers[1].name
+          value: collector
+        template: gateway/deployment.yaml
+      - contains:
+          path: spec.template.spec.containers[1].env
+          content:
+            name: DATABASE_SSLMODE
+            value: verify-full
+          any: true
+        template: gateway/deployment.yaml
+      - contains:
+          path: spec.template.spec.containers[1].env
+          content:
+            name: DATABASE_SSLROOTCERT
+            value: /etc/ssl/certs/ca-certificates.crt
+          any: true
+        template: gateway/deployment.yaml
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_SSLMODE
+            value: verify-full
+          any: true
+        template: migrations-job.yaml
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: DATABASE_SSLROOTCERT
+            value: /etc/ssl/certs/ca-certificates.crt
+          any: true
+        template: migrations-job.yaml
+
   - it: writer rejects both token sources at once
     template: gateway/deployment.yaml
     set:

helm/litellm/tests/hpa_workload_metrics_tests.yaml

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

helm/litellm/tests/ingress_controller_tests.yaml

@@ -97,6 +97,16 @@ tests:
                 name: RELEASE-NAME-litellm-gateway
                 port:
                   number: 4000
+      - contains:
+          path: spec.rules[0].http.paths
+          content:
+            path: /debug/memory/summary
+            pathType: Exact
+            backend:
+              service:
+                name: RELEASE-NAME-litellm-gateway
+                port:
+                  number: 4000
       - equal:
           path: spec.rules[0].http.paths[-1]
           value:

helm/litellm/tests/ingress_extra_paths_tests.yaml

@@ -288,6 +288,17 @@ tests:
       - failedTemplate:
           errorMessage: "ingress.extraPaths[0]: path /test with pathType Exact is already routed by this chart, and a duplicate would take it over rather than add to it"
 
+  - it: rejects an entry that would take over the exact /debug/memory/summary route
+    set:
+      ingress.enabled: true
+      ingress.extraPaths:
+        - path: /debug/memory/summary
+          pathType: Exact
+          service: backend
+    asserts:
+      - failedTemplate:
+          errorMessage: "ingress.extraPaths[0]: path /debug/memory/summary with pathType Exact is already routed by this chart, and a duplicate would take it over rather than add to it"
+
   - it: allows a built-in path under a different pathType, which is a distinct rule
     set:
       ingress.enabled: true

helm/litellm/tests/metrics_server_tests.yaml

@@ -0,0 +1,148 @@
+suite: test gateway metrics sidecar
+templates:
+  - gateway/configmap.yaml
+  - gateway/deployment.yaml
+  - gateway/service.yaml
+  - gateway/service-metrics.yaml
+values:
+  - ./values/required.yaml
+tests:
+  - it: adds no sidecar, volume, env or service port when the metrics server is off
+    asserts:
+      - lengthEqual:
+          path: spec.template.spec.containers
+          count: 1
+        template: gateway/deployment.yaml
+      - notContains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: PROMETHEUS_MULTIPROC_DIR
+          any: true
+        template: gateway/deployment.yaml
+      - notContains:
+          path: spec.template.spec.volumes
+          content:
+            name: prometheus-multiproc
+          any: true
+        template: gateway/deployment.yaml
+      - lengthEqual:
+          path: spec.ports
+          count: 1
+        template: gateway/service.yaml
+      - hasDocuments:
+          count: 0
+        template: gateway/service-metrics.yaml
+
+  - it: runs the metrics server as a sidecar over a shared multiproc dir and exposes it on a ClusterIP metrics service
+    set:
+      gateway.metricsServer.enabled: true
+      gateway.metricsServer.port: 4101
+      gateway.service.type: LoadBalancer
+      gateway.image.tag: v1.101.0
+    asserts:
+      - contains:
+          path: spec.template.spec.containers[0].env
+          content:
+            name: PROMETHEUS_MULTIPROC_DIR
+            value: /tmp/litellm_prometheus_multiproc
+        template: gateway/deployment.yaml
+      - contains:
+          path: spec.template.spec.containers[0].volumeMounts
+          content:
+            name: prometheus-multiproc
+            mountPath: /tmp/litellm_prometheus_multiproc
+        template: gateway/deployment.yaml
+      - equal:
+          path: spec.template.spec.containers[1].name
+          value: metrics
+        template: gateway/deployment.yaml
+      - equal:
+          path: spec.template.spec.containers[1].image
+          value: ghcr.io/berriai/litellm-gateway:v1.101.0
+        template: gateway/deployment.yaml
+      - equal:
+          path: spec.template.spec.containers[1].command
+          value:
+            - python
+            - -m
+            - litellm.proxy.prometheus_metrics_server
+            - --port
+            - "4101"
+        template: gateway/deployment.yaml
+      - equal:
+          path: spec.template.spec.containers[1].env
+          value:
+            - name: PROMETHEUS_MULTIPROC_DIR
+              value: /tmp/litellm_prometheus_multiproc
+        template: gateway/deployment.yaml
+      - equal:
+          path: spec.template.spec.containers[1].ports
+          value:
+            - name: metrics
+              containerPort: 4101
+              protocol: TCP
+        template: gateway/deployment.yaml
+      - equal:
+          path: spec.template.spec.containers[1].volumeMounts
+          value:
+            - name: prometheus-multiproc
+              mountPath: /tmp/litellm_prometheus_multiproc
+        template: gateway/deployment.yaml
+      - equal:
+          path: spec.template.spec.containers[1].readinessProbe.tcpSocket.port
+          value: metrics
+        template: gateway/deployment.yaml
+      - equal:
+          path: spec.template.spec.containers[1].livenessProbe.tcpSocket.port
+          value: metrics
+        template: gateway/deployment.yaml
+      - equal:
+          path: spec.template.spec.containers[1].resources.requests.cpu
+          value: 50m
+        template: gateway/deployment.yaml
+      - contains:
+          path: spec.template.spec.volumes
+          content:
+            name: prometheus-multiproc
+            emptyDir: {}
+        template: gateway/deployment.yaml
+      - lengthEqual:
+          path: spec.ports
+          count: 1
+        template: gateway/service.yaml
+      - equal:
+          path: spec.type
+          value: LoadBalancer
+        template: gateway/service.yaml
+      - equal:
+          path: metadata.name
+          value: RELEASE-NAME-litellm-gateway-metrics
+        template: gateway/service-metrics.yaml
+      - equal:
+          path: spec.type
+          value: ClusterIP
+        template: gateway/service-metrics.yaml
+      - equal:
+          path: spec.ports
+          value:
+            - port: 4101
+              targetPort: metrics
+              protocol: TCP
+              name: metrics
+        template: gateway/service-metrics.yaml
+      - equal:
+          path: spec.selector
+          value:
+            app.kubernetes.io/name: litellm
+            app.kubernetes.io/instance: RELEASE-NAME
+            app.kubernetes.io/component: gateway
+        template: gateway/service-metrics.yaml
+
+  - it: rejects a metrics port equal to the gateway port
+    template: gateway/deployment.yaml
+    set:
+      gateway.metricsServer.enabled: true
+      gateway.metricsServer.port: 4000
+    asserts:
+      - failedTemplate:
+          errorMessage: gateway.metricsServer.port must differ from the gateway port 4000

helm/litellm/values.yaml

@@ -208,6 +208,11 @@ database:
       name: litellm-writer-secret
       usernameKey: username
       passwordKey: password
+    # libpq sslmode / sslrootcert applied to the writer and reader URLs (Prisma and the
+    # in-container PgBouncer); e.g. verify-full with /etc/ssl/certs/ca-certificates.crt for AWS RDS.
+    # sslRootCert on its own implies sslMode verify-full
+    sslMode: ""
+    sslRootCert: ""
 
   # Optional read-replica routing. When `reader.host` is set, the proxy routes
   # reads (find_*, count, group_by, query_raw/_first) to this endpoint while
@@ -225,6 +230,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 and its collector sidecar only: the backend runs a
+  # single worker and the migrations Job must keep a direct connection. With
+  # `database.writer.useIAMAuth` or `useAzureEntraAuth` the pool mints and
+  # renews the database token itself, so the workers never see it. Starting profile for
+  # `gateway.numWorkers: 4` is maxDbConnections: 20, so a database with a
+  # 5000-connection ceiling fits roughly 200 gateway replicas.
+  connectionPool:
+    enabled: false
+    maxDbConnections: 20
+    maxClientConn: 1000
+
 # Optional Redis. Leave host empty to disable.
 #
 # This is the proxy's coordination store: cross-pod tpm/rpm rate limits, spend
@@ -268,6 +293,68 @@ gateway:
   config:
     create: true
     proxy_config: {}
+  # Serve Prometheus /metrics from a `metrics` sidecar container (same image,
+  # `python -m litellm.proxy.prometheus_metrics_server`) that aggregates the
+  # workers' PROMETHEUS_MULTIPROC_DIR samples over a shared emptyDir, so a
+  # scrape never runs on an inference worker. Adds a `metrics` port to the pod
+  # and a dedicated ClusterIP `<gateway>-metrics` Service; point your scrape
+  # config at it. The port has no virtual-key auth: keep it off public ingress.
+  # Needs the gateway image v1.101.0 or newer.
+  metricsServer:
+    enabled: false
+    port: 4001
+    resources:
+      requests:
+        cpu: 50m
+        memory: 128Mi
+      limits:
+        memory: 512Mi
+  # Prometheus Operator ServiceMonitor for the gateway pods. Scrapes the
+  # `<gateway>-metrics` Service, so it requires metricsServer above (the http
+  # port serves /metrics/ behind virtual-key auth). Every pod is its own scrape
+  # target, so the samples carry the `pod` label the per-pod autoscaling
+  # queries below group by.
+  serviceMonitor:
+    enabled: false
+    labels: {}
+    interval: 15s
+    scrapeTimeout: 10s
+  # Opt-in `collector` sidecar (same image, `python -m litellm.proxy.collector`)
+  # that runs the post-response spend pipeline (cost calculation, spend logs,
+  # spend counters, budget reservation reconciliation) so the uvicorn workers
+  # only serialise a compact event over loopback and go back to serving
+  # requests. It shares the pod's env, proxy config, in-container pgbouncer and
+  # Redis spend buffer, so the per-pod DB connection budget is unchanged.
+  # Delivery is at-most-once inside the pod: events already handed over are
+  # lost if the sidecar dies before writing them; events the workers cannot
+  # hand over follow `onUnavailable`.
+  collector:
+    enabled: false
+    # unix:///<dir>/<file>.sock (the <dir> becomes a shared emptyDir) or
+    # tcp://127.0.0.1:<port>
+    address: unix:///var/run/litellm/collector.sock
+    # Events each uvicorn worker holds in memory while the sidecar is slow or
+    # restarting.
+    bufferSize: 1000
+    # fallback: run the pipeline in the worker when the sidecar is unreachable
+    # or the buffer is full (spend stays exact, that request costs gateway CPU
+    # again). drop: count and discard the event instead (spend under-reports).
+    onUnavailable: fallback
+    # How long the workers keep pushing buffered events on shutdown, and how
+    # long the sidecar keeps serving open connections after SIGTERM.
+    drainTimeoutSeconds: 10
+    # Sized independently of the gateway container; the pipeline is CPU bound.
+    resources:
+      requests:
+        cpu: 500m
+        memory: 1Gi
+      limits:
+        cpu: "1"
+        memory: 2Gi
+    # With hpa.targetCPUUtilizationPercentage set, scale on an autoscaling/v2
+    # ContainerResource metric of the `gateway` container only, so the
+    # sidecar's CPU never drives inference replicas. Needs Kubernetes 1.30+.
+    scaleOnGatewayContainerCpu: true
   image:
     repository: ghcr.io/berriai/litellm-gateway
     tag: ""              # defaults to .Chart.AppVersion
@@ -324,6 +411,25 @@ gateway:
     #     policies:
     #       - { type: Percent, value: 100, periodSeconds: 30 }
     behavior: {}
+    # Opt-in per-pod workload targets, rendered as autoscaling/v2 `Pods` metrics
+    # named `litellm_requests_per_second` and `litellm_tokens_per_second` with an
+    # AverageValue target. They coexist with the CPU/memory targets above: the
+    # HPA scales on whichever metric asks for the most replicas. Kubernetes has
+    # no idea what a token is, so a Prometheus Adapter must serve those two
+    # names on custom.metrics.k8s.io from the proxy's counters, grouped by the
+    # scrape target's `pod` label (enable serviceMonitor above):
+    #   litellm_requests_per_second:
+    #     sum(rate(litellm_proxy_total_requests_metric_total{<<.LabelMatchers>>}[1m])) by (<<.GroupBy>>)
+    #   litellm_tokens_per_second:
+    #     sum(rate(litellm_total_tokens_metric_total{<<.LabelMatchers>>}[1m])) by (<<.GroupBy>>)
+    # rate() over [1m] is already per second, so no `* 60`. How fast the HPA
+    # reacts is set by that window, the scrape interval and the HPA sync period
+    # (15s by default), not by the unit: keep serviceMonitor.interval at 15s or
+    # faster so a 1m window holds at least 4 samples. averageValue takes SI
+    # suffixes, so "6M" is six million tokens per second per pod. Tokens are
+    # counted when a response completes, so TPS trails long streams.
+    targetRequestsPerSecond: ""
+    targetTokensPerSecond: ""
   # PodDisruptionBudget for the gateway pods. Set exactly one of
   # `minAvailable` / `maxUnavailable` (minAvailable wins if both are set;
   # enabling without either falls back to `maxUnavailable: 1`). Disabled by

litellm-proxy-extras/litellm_proxy_extras/migrations/20260717000000_cascade_delete_jwt_key_mapping_on_token_delete/migration.sql

@@ -0,0 +1,15 @@
+-- DropForeignKey
+DO $$
+BEGIN
+    IF EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_JWTKeyMapping_token_fkey') THEN
+        ALTER TABLE "LiteLLM_JWTKeyMapping" DROP CONSTRAINT "LiteLLM_JWTKeyMapping_token_fkey";
+    END IF;
+END $$;
+
+-- AddForeignKey
+DO $$
+BEGIN
+    IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_JWTKeyMapping_token_fkey') THEN
+        ALTER TABLE "LiteLLM_JWTKeyMapping" ADD CONSTRAINT "LiteLLM_JWTKeyMapping_token_fkey" FOREIGN KEY ("token") REFERENCES "LiteLLM_VerificationToken"("token") ON DELETE CASCADE ON UPDATE CASCADE;
+    END IF;
+END $$;

litellm-proxy-extras/litellm_proxy_extras/migrations/20260831120000_add_litellm_call_id_spend_logs/migration.sql

@@ -0,0 +1,2 @@
+-- AlterTable
+ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN IF NOT EXISTS "litellm_call_id" TEXT;

litellm-proxy-extras/litellm_proxy_extras/migrations/20260831120001_spend_logs_litellm_call_id_index/migration.sql

@@ -0,0 +1,12 @@
+-- CreateIndex (CONCURRENTLY)
+--
+-- Disclaimer:
+-- - CREATE INDEX CONCURRENTLY cannot run inside a transaction. This migration must stay a
+--   single statement so Prisma Migrate on PostgreSQL can apply it outside a transaction.
+-- - Builds are slower and use more I/O than a blocking CREATE INDEX; if the build is
+--   interrupted, Postgres may leave an INVALID index that must be dropped and recreated.
+-- - Do not edit this file after it has been applied to any database: Prisma checksums
+--   migrations; add a new migration instead.
+-- - Requires PostgreSQL that supports CONCURRENTLY with IF NOT EXISTS (use a new migration
+--   without IF NOT EXISTS if you must support older versions).
+CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_SpendLogs_litellm_call_id_idx" ON "LiteLLM_SpendLogs"("litellm_call_id");

litellm-proxy-extras/litellm_proxy_extras/migrations/20260903230000_add_org_id_to_managed_object_table/migration.sql

@@ -0,0 +1,4 @@
+-- Add org_id column to LiteLLM_ManagedObjectTable
+-- Snapshots the creating key's organization at submission time, like team_id,
+-- so CheckBatchCost can bill organization spend hours later without re-resolving
+ALTER TABLE "LiteLLM_ManagedObjectTable" ADD COLUMN IF NOT EXISTS "org_id" TEXT;

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

@@ -0,0 +1 @@
+ALTER TABLE "LiteLLM_AutoRouterSession" ADD COLUMN IF NOT EXISTS "baseline_models" JSONB NOT NULL DEFAULT '{}';

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

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

litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py

@@ -37,11 +37,13 @@ raised it above the deploy default keeps that larger budget for deploy unless
 the deploy override says otherwise.
 """
 
+import importlib.util
 import math
 import os
 import shutil
 import signal
 import subprocess
+import sys
 from collections.abc import Mapping, Sequence
 from dataclasses import dataclass
 from pathlib import Path
@@ -64,6 +66,7 @@ DEFAULT_PRISMA_BOOTSTRAP_TIMEOUT = 600.0
 DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT = 600.0
 
 BOOTSTRAP_ARG = "--version"
+PRISMA_CONSOLE_SCRIPT = "prisma"
 
 
 @dataclass(frozen=True)
@@ -184,6 +187,28 @@ def _kill_process_group(process: "subprocess.Popen[str]") -> None:
         return
 
 
+def prisma_cli_available() -> bool:
+    """Whether some way of running the Prisma CLI exists: the console script on PATH or the importable package."""
+    if shutil.which(PRISMA_CONSOLE_SCRIPT) is not None:
+        return True
+    return importlib.util.find_spec(PRISMA_CONSOLE_SCRIPT) is not None
+
+
+def resolve_prisma_argv(argv: Sequence[str]) -> tuple[str, ...]:
+    """Route a bare ``prisma`` command through ``python -m prisma`` when the console script is not on PATH.
+
+    The console script and ``python -m prisma`` are the same entry point, but
+    only the module form survives an interpreter whose ``bin`` directory is
+    missing from PATH, which is how the proxy gets started under launchers and
+    init systems. Any other executable name is left untouched.
+    """
+    if not argv or argv[0] != PRISMA_CONSOLE_SCRIPT:
+        return tuple(argv)
+    if shutil.which(PRISMA_CONSOLE_SCRIPT) is not None:
+        return tuple(argv)
+    return (sys.executable, "-m", PRISMA_CONSOLE_SCRIPT, *argv[1:])
+
+
 def run_prisma(
     argv: Sequence[str],
     *,
@@ -200,7 +225,7 @@ def run_prisma(
     text unless ``stdout``/``stderr`` say otherwise.
     """
     with subprocess.Popen(
-        argv,
+        resolve_prisma_argv(argv),
         env=env,
         stdout=stdout,
         stderr=stderr,

litellm-proxy-extras/litellm_proxy_extras/schema.prisma

@@ -282,6 +282,7 @@ model LiteLLM_ObjectPermissionTable {
   mcp_toolsets          String[]       @default([]) // Toolset IDs granted to this key/team/user
   search_tools          String[]       @default([]) // search_tool_name values this key/team/user may call
   mcp_tool_search_enabled Boolean?
+  skills                String[]       @default([]) // Claude Code plugin names granted to this key/team beyond the public (enabled) set
   teams                 LiteLLM_TeamTable[]
   projects              LiteLLM_ProjectTable[]
   verification_tokens   LiteLLM_VerificationToken[]
@@ -492,7 +493,7 @@ model LiteLLM_JWTKeyMapping {
     updated_at        DateTime @default(now()) @updatedAt
     updated_by        String?
 
-    litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token])
+    litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade)
 
     @@unique([jwt_claim_name, jwt_claim_value])
     @@index([jwt_claim_name, jwt_claim_value, is_active])
@@ -658,12 +659,14 @@ model LiteLLM_SpendLogs {
   mcp_namespaced_tool_name String? 
   agent_id            String?
   proxy_server_request Json?     @default("{}")
+  litellm_call_id     String?
   created_at          DateTime   @default(now()) @map("created_at")
   updated_at          DateTime   @default(now()) @updatedAt @map("updated_at")
   @@index([startTime])
   @@index([startTime, request_id])
   @@index([end_user])
   @@index([session_id])
+  @@index([litellm_call_id])
 }
 
 model LiteLLM_BudgetWindowSpend {
@@ -1036,6 +1039,7 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t
   created_at DateTime @default(now())
   created_by String?
   team_id String?
+  org_id String? // creating key's organization at submission time; CheckBatchCost bills org spend against it
   api_key String?
   request_tags Json? @default("[]")
   updated_at DateTime @updatedAt
@@ -1512,6 +1516,7 @@ model LiteLLM_AutoRouterSession {
   classifier_cost       Float    @default(0)
   classifier_cost_recorded_turns Int @default(0)
   tier_turns            Json     @default("{}")
+  baseline_models       Json     @default("{}")
 
   @@id([api_key, session_id, router_name])
   @@index([last_turn_at], map: "idx_autorouter_session_last_turn")

litellm-proxy-extras/litellm_proxy_extras/utils.py

@@ -8,14 +8,10 @@ import tempfile
 import time
 from dataclasses import dataclass, replace
 from pathlib import Path
-from typing import Optional
+from typing import TYPE_CHECKING, Final, Optional
 
 from litellm_proxy_extras import prisma_toolchain
 from litellm_proxy_extras._logging import logger
-from litellm_proxy_extras.replica_identity import (
-    REPLICA_IDENTITY_FULL_ENV_VAR,
-    apply_replica_identity_full,
-)
 from litellm_proxy_extras.prisma_toolchain import (
     PRISMA_COMMAND_TIMEOUT_ENV_VAR,
     PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR,
@@ -23,6 +19,14 @@ from litellm_proxy_extras.prisma_toolchain import (
     prisma_command_timeout,
     prisma_migrate_deploy_timeout,
 )
+from litellm_proxy_extras.replica_identity import (
+    REPLICA_IDENTITY_FULL_ENV_VAR,
+    apply_replica_identity_full,
+)
+
+if TYPE_CHECKING:
+    import psycopg
+    import psycopg.sql
 
 
 def str_to_bool(value: Optional[str]) -> bool:
@@ -46,6 +50,28 @@ def _get_prisma_env() -> dict:
 _MIGRATION_TS_RE = re.compile(r"^(\d{14})_")
 
 _MIGRATION_DEADLOCK_MARKER = "deadlock detected"
+INDEX_REPAIR_ADVISORY_LOCK_KEY: Final = int.from_bytes(b"litellm", "big")
+_TRANSIENT_INDEX_SUFFIX_RE: Final = re.compile(r"_cc(?:new|old)\d*$")
+_INVALID_LITELLM_INDEXES_SQL: Final = (
+    "SELECT n.nspname, c.relname, pg_size_pretty(pg_table_size(t.oid)) "
+    "FROM pg_index i "
+    "JOIN pg_class c ON c.oid = i.indexrelid "
+    "JOIN pg_class t ON t.oid = i.indrelid "
+    "JOIN pg_namespace n ON n.oid = t.relnamespace "
+    "WHERE NOT i.indisvalid "
+    "  AND c.relkind = 'i' "
+    "  AND n.nspname = %s "
+    "  AND t.relname LIKE %s "
+    "  AND NOT EXISTS (SELECT 1 FROM pg_constraint k WHERE k.conindid = i.indexrelid) "
+    "ORDER BY c.relname"
+)
+
+
+@dataclass(frozen=True, slots=True)
+class _InvalidIndex:
+    schema: str
+    name: str
+    table_size: str
 
 MAX_MIGRATE_DEPLOY_ATTEMPTS = 4
 
@@ -624,7 +650,7 @@ class ProxyExtrasDBManager:
     def _strip_prisma_query_params(url: str) -> str:
         """Remove Prisma-specific query params (connection_limit, pool_timeout,
         schema, etc.) from DATABASE_URL so psycopg can parse it."""
-        from urllib.parse import urlparse, urlunparse, parse_qsl, urlencode
+        from urllib.parse import parse_qsl, quote, urlencode, urlparse, urlunparse
 
         parsed = urlparse(url)
         if not parsed.query:
@@ -645,7 +671,7 @@ class ProxyExtrasDBManager:
             "target_session_attrs",
         }
         kept = [(k, v) for k, v in parse_qsl(parsed.query) if k in libpq_params]
-        return urlunparse(parsed._replace(query=urlencode(kept)))
+        return urlunparse(parsed._replace(query=urlencode(kept, quote_via=quote)))
 
     @staticmethod
     def _warn_if_db_ahead_of_head(migrations_dir: str) -> None:
@@ -719,6 +745,95 @@ class ProxyExtrasDBManager:
             ", ".join(sorted_hostile[:5]) + (" ..." if len(sorted_hostile) > 5 else ""),
         )
 
+    @staticmethod
+    def _invalid_litellm_indexes(
+        conn: "psycopg.Connection[tuple[str, str, str]]", schema: str
+    ) -> tuple[_InvalidIndex, ...]:
+        rows: Final = conn.execute(_INVALID_LITELLM_INDEXES_SQL, (schema, "LiteLLM\\_%")).fetchall()
+        return tuple(_InvalidIndex(*row) for row in rows)
+
+    @staticmethod
+    def _index_repair(index: _InvalidIndex) -> tuple["psycopg.sql.Composed", str]:
+        from psycopg import sql
+
+        target: Final = sql.Identifier(index.schema, index.name)
+        if _TRANSIENT_INDEX_SUFFIX_RE.search(index.name):
+            return sql.SQL("DROP INDEX CONCURRENTLY IF EXISTS {}").format(target), "Dropped leftover"
+        return sql.SQL("REINDEX INDEX CONCURRENTLY {}").format(target), "Rebuilt"
+
+    @staticmethod
+    def _repair_index(conn: "psycopg.Connection[tuple[str, str, str]]", index: _InvalidIndex) -> None:
+        import psycopg
+
+        statement, action = ProxyExtrasDBManager._index_repair(index)
+        try:
+            conn.execute(statement)
+        except psycopg.Error as e:
+            logger.warning(
+                "Could not repair invalid index %s.%s, will retry on the next startup. "
+                "If this keeps happening, run `%s` by hand as the index owner. Error: %s",
+                index.schema,
+                index.name,
+                statement.as_string(conn),
+                e,
+            )
+            return
+        logger.info("%s invalid index %s.%s", action, index.schema, index.name)
+
+    @staticmethod
+    def repair_invalid_indexes(lock_timeout: str = "30s") -> bool:
+        """Rebuild LiteLLM indexes an interrupted CREATE INDEX CONCURRENTLY left
+        INVALID (a migration deadlock between replicas is the usual cause; the
+        retried migration skips them because of IF NOT EXISTS). Never raises:
+        returns True when no invalid index remains, False when the repair was
+        skipped or failed and will be retried on the next startup. Looks in the
+        schema DATABASE_URL names, the only URL Prisma migrates through, but
+        connects over DIRECT_URL when set: the session settings, the advisory
+        lock and REINDEX CONCURRENTLY all need one server session, which a
+        transaction pooler does not give."""
+        prisma_url: Final = os.getenv("DATABASE_URL")
+        if not prisma_url:
+            return False
+
+        try:
+            import psycopg
+            from psycopg import sql
+        except ImportError:
+            logger.warning(
+                "psycopg is not installed; skipping the invalid index check. "
+                "Install the litellm[extra_proxy] extra, which includes psycopg."
+            )
+            return False
+
+        schema: Final = ProxyExtrasDBManager._prisma_schema_param(prisma_url) or "public"
+        cleaned_url: Final = ProxyExtrasDBManager._strip_prisma_query_params(os.getenv("DIRECT_URL") or prisma_url)
+        try:
+            with psycopg.connect(cleaned_url, connect_timeout=10, autocommit=True) as conn:
+                conn.execute("SET statement_timeout = 0")
+                conn.execute(sql.SQL("SET lock_timeout = {}").format(sql.Literal(lock_timeout)))
+                found: Final = ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema)
+                if not found:
+                    return True
+                logger.warning(
+                    "Found %d invalid index(es) left by an interrupted CREATE INDEX "
+                    "CONCURRENTLY, rebuilding: %s",
+                    len(found),
+                    ", ".join(f"{index.name} (table size {index.table_size})" for index in found),
+                )
+                lock_row: Final = conn.execute(
+                    "SELECT pg_try_advisory_lock(%s)", (INDEX_REPAIR_ADVISORY_LOCK_KEY,)
+                ).fetchone()
+                if lock_row is None or not lock_row[0]:
+                    logger.info("Another replica is already rebuilding the invalid indexes, skipping")
+                    return False
+                for index in ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema):
+                    ProxyExtrasDBManager._repair_index(conn, index)
+                remaining: Final = ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema)
+        except psycopg.Error as e:
+            logger.warning("Could not check for invalid indexes, will retry on the next startup. Error: %s", e)
+            return False
+        return not remaining
+
     @staticmethod
     def _setup_database_v2(use_migrate: bool) -> bool:
         """
@@ -994,6 +1109,7 @@ class ProxyExtrasDBManager:
             use_migrate=use_migrate, use_v2_resolver=use_v2_resolver
         )
         if migrated:
+            ProxyExtrasDBManager.repair_invalid_indexes()
             ProxyExtrasDBManager.apply_replica_identity_full_if_requested()
         return migrated
 

litellm-proxy-extras/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "litellm-proxy-extras"
-version = "0.4.94"
+version = "0.4.97"
 description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
 readme = "README.md"
 requires-python = ">=3.9"
@@ -26,7 +26,7 @@ required-version = ">=0.10.9"
 module-root = ""
 
 [tool.commitizen]
-version = "0.4.94"
+version = "0.4.97"
 version_files = [
     "pyproject.toml:^version",
     "../pyproject.toml:litellm-proxy-extras==",

litellm-rust/ADDING_A_PROVIDER.md

@@ -1,29 +0,0 @@
-# Adding a provider / route to litellm-rust
-
-Everything for a route lives in `crates/core/src/<route>/`; `crates/core/src/messages` is the reference. A host (the axum gateway, the Python bridge) only calls the route's entrypoint.
-
-1. **Entrypoint** — `mod.rs`: `pub async fn <route>(request) -> CoreResult<Response>`, the Rust equivalent of `litellm.<route>()`, plus a `<route>_stream` variant when the route streams. It is the only thing a host touches.
-2. **Transform contract** — `transformation.rs`: a `…ProviderConfig` trait (URL build + request/response transforms) with types in `types.rs`.
-3. **Provider config** — `crates/core/src/providers/<provider>/<route>/transformation.rs`: implement that trait as a `const <PROVIDER>_<ROUTE>_CONFIG`, mirroring the Python provider tree. Add parity unit tests.
-4. **Prepare + handler** — `prepare.rs` resolves provider/model, credentials, auth headers, and URL, then transforms the request; `handler.rs` performs the provider call through the shared client in `client.rs` and transforms the response.
-
-## Coding standards
-
-Before writing new logic, look for an existing base to extend. When a change is
-“the same behavior for one more provider/endpoint/integration”, the codebase
-almost always already has a shared abstraction for it (for example, provider
-`BaseConfig` transformation classes in `litellm/llms/base_llm/`, shared
-helpers in `litellm_core_utils/`, typed request/response models, or factory
-functions). Find it first with a search, then add the new variant by inheriting
-from or composing that base, overriding only what genuinely differs (model
-name, parameter mapping, or auth).
-
-Never copy an existing implementation and edit it in place, and never hand-roll
-a parallel version of logic a base already provides. If you catch yourself
-writing a second copy of a pattern that exists twice already, stop and extract a
-base instead: put the shared shape in one place and make both call sites thin
-variants of it. The test for a good abstraction is that adding the next provider
-is a few declarative lines, not a new file of duplicated flow. Only diverge from
-the base when behavior is genuinely different, and say so explicitly in the PR.
-
-**Calling:** hosts invoke the core entrypoint — the Python bridge and the `ai-gateway` route service both call `litellm_core::messages::messages`. Never add a provider handler to `ai-gateway`. Register new modules in `lib.rs` / `mod.rs`, then run the commands under "Checks" in [CLAUDE.md](CLAUDE.md).

litellm-rust/AGENTS.md

@@ -1,44 +0,0 @@
-# AGENTS.md
-
-litellm-rust has five crates. A crate is a layer or shared foundation, not a route. Routes (ocr, realtime, chat) and providers (mistral, openai) are modules inside the layers.
-
-## Crates
-
-| Crate | Role |
-|-------|------|
-| litellm-core | The LiteLLM SDK in Rust. One public entrypoint per top-level call (`messages::messages()`), owning types, transforms, provider resolution, auth, and the provider HTTP call. Call it, get a typed response. |
-| litellm-config | Config-loading boundary. Returns resolved core deployment data and optionally delegates loading to Python. |
-| litellm-ai-gateway | The axum server (behind the `server` feature) plus the WebSocket hosts. Translates HTTP/WS to core entrypoints; owns no provider logic and no handlers. |
-| litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. |
-| litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. Owns API registration, domain wiring, and Python exception mapping. |
-
-Dependency direction is acyclic: `litellm-config` depends on `litellm-core`, the gateway depends on both, and `litellm-python-bridge` depends on the domain layers and `litellm-python-interop`. The interop foundation depends on no LiteLLM domain crate.
-
-## Where a route lives
-
-A top-level LiteLLM call is a module under `crates/core/src/<route>/`, shaped like `messages`:
-
-```
-core/src/messages/
-  mod.rs             # pub async fn messages(..) -> CoreResult<..>  (+ messages_stream for SSE)
-  types.rs           # request/response types, MessagesRequest
-  transformation.rs  # the provider template trait
-  prepare.rs         # provider resolution, auth headers, URL
-  handler.rs         # the provider call
-  client.rs          # the shared reqwest client
-```
-
-Handlers never live in `ai-gateway`. `ocr`, `audio_transcription`, and `realtime` are still hosted there from before this rule; they move to `core` as they are touched.
-
-Adding a crate: default to a module. A new crate requires a real trigger: separate artifact (binary/cdylib), proc-macro, shared foundation, or publishable standalone. A new provider or route is none of these.
-
-Adding a crate fails crates/core/tests/workspace_crate_allowlist.rs until you update its allowlist and this file — intentional.
-
-## Style
-
-All Rust in `litellm-rust/` follows the official Rust Style Guide:
-https://doc.rust-lang.org/style-guide/
-
-`rustfmt` implements its formatting by default, so run `cargo fmt` before committing; CI gates every PR on `cargo fmt --check`. Do not hand-format against rustfmt or add a `rustfmt.toml` that diverges from the default style.
-
-Beyond formatting, follow the guide's naming and idiom conventions rustfmt cannot auto-apply: `snake_case` items/functions/modules, `UpperCamelCase` types/traits/variants, `SCREAMING_SNAKE_CASE` constants/statics (acronyms as one word, e.g. `HttpClient`), and the import grouping and item ordering it prescribes. See CLAUDE.md for the detailed version.

litellm-rust/CLAUDE.md

@@ -1,189 +0,0 @@
-# CLAUDE.md
-
-This file defines the rules for Rust work in LiteLLM.
-
-## Provider Coding Standards
-
-Before writing new logic, look for an existing base to extend. When a change is
-“the same behavior for one more provider/endpoint/integration”, the codebase
-almost always already has a shared abstraction for it (for example, provider
-`BaseConfig` transformation classes in `litellm/llms/base_llm/`, shared
-helpers in `litellm_core_utils/`, typed request/response models, or factory
-functions). Find it first with a search, then add the new variant by inheriting
-from or composing that base, overriding only what genuinely differs (model
-name, parameter mapping, or auth).
-
-Never copy an existing implementation and edit it in place, and never hand-roll
-a parallel version of logic a base already provides. If you catch yourself
-writing a second copy of a pattern that exists twice already, stop and extract a
-base instead: put the shared shape in one place and make both call sites thin
-variants of it. The test for a good abstraction is that adding the next provider
-is a few declarative lines, not a new file of duplicated flow. Only diverge from
-the base when behavior is genuinely different, and say so explicitly in the PR.
-
-## Crates (see AGENTS.md)
-
-`litellm-core` **is** the LiteLLM SDK in Rust: it makes the LLM call.
-`litellm-config` is the config-loading boundary and returns resolved core types.
-`litellm-ai-gateway` is an HTTP/WebSocket server in front of it, and
-`litellm-python-bridge` exposes it to the Python SDK. `litellm-python-interop`
-holds domain-neutral PyO3 primitives shared by Python-facing Rust code. A crate
-is a layer or shared foundation, not a route; add modules, not crates.
-
-## Core Boundary
-
-`litellm-core` owns the whole call. The Rust equivalent of `litellm.messages()`
-is `litellm_core::messages::messages(request).await`: you call it, it does the
-provider call, and you get a typed non-streaming response back.
-
-Route-level Rust structure mirrors LiteLLM's Python responsibilities:
-- `core/src/<route>/` owns the route end to end: the public entrypoint fn named
-  after the route in `mod.rs`, the request/response types (`types.rs`), the
-  provider template trait (`transformation.rs`), the provider/auth/URL
-  resolution (`prepare.rs`), the HTTP client (`client.rs`), and the handler that
-  performs the call (`handler.rs`). `core/src/messages` is the reference.
-- `core/src/providers/<provider>/<route>/transformation.rs` owns the
-  provider-specific transform. For Anthropic Messages, this means
-  `core/src/providers/anthropic/messages/transformation.rs`.
-- Handlers live in `core`, never in a host. `ai-gateway` must not contain a
-  route handler that talks to a provider; its axum route reads the HTTP request,
-  picks a deployment, and calls the `core` entrypoint. `python-bridge` marshals
-  Python objects and calls the same entrypoint.
-
-Streaming keeps the same shape: the route entrypoint has a `<route>_stream`
-variant in `core` that returns the upstream response so a host can splice it to
-its own caller; the host still owns no provider logic.
-
-Call-hook and lifecycle instrumentation, including phase timing, usage
-accumulation, and callback payload construction, always lives in `core`.
-Hosts feed observed events into core and dispatch the completed payloads through
-their I/O logger; hosts must not own callback orchestration.
-
-Allowed in `core`:
-- The public entrypoint for a top-level LiteLLM call
-- Request/response transforms and stream chunk normalization
-- Provider resolution, auth header construction, and URL building
-- The provider HTTP call itself, through a shared reused client with connect and
-  request timeouts
-- Shared data types and validation errors
-- Deterministic token/cost helper logic
-
-Not allowed in `core`:
-- Serving HTTP: axum routes, extractors, and transport concerns stay in the host
-- Filesystem access
-- Database access
-- Config file reading and rollout state
-- Logging callbacks, spend writes, or custom callbacks
-- Global mutable runtime state
-
-Env reads in `core` are limited to credential fallback inside a route's
-`prepare.rs` (the `env_lookup` closure), mirroring what the Python SDK does when
-no key is passed. Everything else config-shaped is resolved by the host and
-passed in.
-
-Routes still hosted in `ai-gateway` (`ocr`, `audio_transcription`, `realtime`)
-predate this rule and are being moved into `core` route modules; do not add new
-ones there, and prefer moving one when you touch it.
-
-Python owns rollout state and fallback while Rust is being introduced. Rust
-paths must be off by default until parity tests prove equivalence with Python.
-A new provider/route may instead be implemented rust-only with no Python
-reference; then the Python interface is a thin dispatch that calls Rust with no
-fallback, and you state the rust-only choice explicitly in the PR. Either way
-the Python side stays minimal (it only marshals inputs and calls the Rust
-interface), never add a per-route feature flag, and never push provider
-dispatch into `litellm/main.py`; put it in a thin dispatch class under
-`litellm/llms/<provider>/<route>/`.
-
-## Production Bar
-
-Rust code in this workspace is held to a strict parity and robustness bar from
-the first PR:
-
-- Correctness parity is proven with tests. Do not rely on README claims or
-  manual inspection for a port that mirrors Python behavior.
-- Every provider transform must have unit tests for supported-parameter
-  filtering, request body shape, response normalization, missing/null fields,
-  and bad-input errors.
-- When Rust is exposed through Python, add Python tests that prove disabled,
-  enabled, and unavailable-bridge fallback behavior.
-- Avoid panics on user/provider input. Return typed errors and let the host map
-  them to Python exceptions or HTTP responses.
-- OCR handles documents that often contain personal data. Do not log document
-  contents, base64 payloads, provider response bodies, or secrets.
-- Error messages must be useful but data-minimized. Truncate or sanitize any
-  upstream body before it crosses a host boundary.
-- Treat empty or whitespace-only credentials, URLs, and config values as absent
-  at the host/config resolution layer.
-- Preserve Python output shape intentionally. If a field is always serialized as
-  `null` for Python parity, leave a short comment explaining that parity choice.
-
-## Network I/O Rules
-
-These rules apply to every module that executes network I/O, whether it is a
-`core` route handler or a host such as `ai-gateway`:
-
-- Set connect and full-request timeouts. No unbounded waits.
-- Reuse HTTP clients; do not construct clients per request.
-- Prefer rustls TLS for portable Python wheels and Linux images unless there is
-  a documented reason not to.
-- Add request IDs and structured tracing at the host layer, without logging OCR
-  document contents or secrets.
-- Do not echo raw upstream response bodies to callers. Sanitize and bound them.
-- Avoid `expect`/`unwrap` in server startup and request paths unless the panic is
-  impossible by construction and documented.
-
-## Rust Style Guide
-
-All Rust in `litellm-rust/` follows the official Rust Style Guide:
-https://doc.rust-lang.org/style-guide/
-
-`rustfmt` implements the guide's formatting rules by default, so the mechanical
-side is enforced for you: run `cargo fmt` before committing and CI gates every
-PR on `cargo fmt --check` (see Checks). Do not hand-format against rustfmt or add
-a `rustfmt.toml` that diverges from the default style; the default style *is* the
-guide.
-
-The guide also covers conventions rustfmt cannot auto-apply; follow these too:
-- Naming: `snake_case` for items, functions, and modules; `UpperCamelCase` for
-  types, traits, and enum variants; `SCREAMING_SNAKE_CASE` for constants and
-  statics; acronyms count as one word (`HttpClient`, not `HTTPClient`).
-- Ordering and grouping the guide prescribes: imports grouped std / external /
-  crate-local, derives before other attributes, and consistent item order.
-- Idioms the guide recommends over the formatter fighting you (e.g. prefer
-  restructuring an over-long expression rather than forcing an awkward wrap).
-
-## Constants
-
-Magic numbers and fixed strings go in a crate-level `constants.rs`, never
-hardcoded inline — the Rust mirror of Python's `litellm/constants.py`.
-
-- Each crate that needs them has `src/constants.rs` (declared `mod constants;`);
-  import from it (`use crate::constants::...`). Don't scatter `const` values at
-  the top of feature modules.
-- An env-overridable tunable still lives in `constants.rs` as its `DEFAULT_*`
-  value; the env read (with fallback to that default) happens at the host/config
-  resolution layer, not in `core`/`providers`.
-- Exception: a value that is purely local to one function and has no meaning
-  elsewhere may stay inline, but prefer `constants.rs` when in doubt.
-
-## Checks
-
-Run these before pushing Rust changes. The same checks run in GitHub Actions
-for changes under `litellm-rust/`.
-
-```bash
-cd litellm-rust
-cargo fmt --check
-cargo clippy --workspace --all-targets -- -D warnings
-cargo clippy -p litellm-core --all-targets --features bedrock-auth -- -D warnings
-# the ai-gateway binary + server code is behind the `server` feature
-cargo clippy -p litellm-ai-gateway --all-targets --all-features -- -D warnings
-cargo test --workspace
-cargo test -p litellm-core --features bedrock-auth
-# the `auth`, `routes`, `state` and `realtime` tests only exist under `server`
-cargo test -p litellm-ai-gateway --features server
-```
-
-When a Rust path is exposed through Python, add Python parity tests that compare
-the existing Python output with the Rust-backed output.

litellm-rust/Cargo.toml

@@ -1,6 +1,7 @@
 [workspace]
 members = [
     "crates/core",
+    "crates/token-counter",
     "crates/config",
     "crates/ai-gateway",
     "crates/python-interop",
@@ -15,9 +16,11 @@ license = "MIT"
 repository = "https://github.com/BerriAI/litellm"
 
 [workspace.dependencies]
+bytes = "1"
 tracing = "0.1"
 tracing-subscriber = { version = "0.3", default-features = false, features = ["registry", "std"] }
 litellm-core = { path = "crates/core" }
+litellm-token-counter = { path = "crates/token-counter" }
 litellm-config = { path = "crates/config" }
 litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false }
 litellm-python-interop = { path = "crates/python-interop" }
@@ -39,6 +42,14 @@ tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"]
 tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] }
 futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] }
 base64 = "0.22"
+gcp_auth = "0.12.7"
+azure_core = "1.0.0"
+azure_identity = { version = "1.0.0", features = ["tokio"] }
+moka = { version = "0.12.16", features = ["future"] }
+strum = { version = "0.28.0", features = ["derive"] }
+url = "2.5.8"
+criterion = "0.8.2"
+veil = "0.3.0"
 
 [profile.release]
 opt-level = 3

litellm-rust/README.md

@@ -1,56 +0,0 @@
-# LiteLLM Rust
-
-This workspace contains the staged Rust implementation for LiteLLM.
-
-`litellm-core` is the LiteLLM SDK in Rust: one entrypoint per top-level call
-that makes the LLM call and hands back a typed response, the same shape as
-`litellm.messages()` in Python.
-
-```rust
-let response = litellm_core::messages::messages(MessagesRequest {
-    model: "claude-sonnet-4-5",
-    body,
-    api_key: Some(key),
-    ..
-})
-.await?;
-```
-
-Python continues to own configuration, retries, routing policy, logging,
-callbacks, spend tracking, and customer plugins until each Rust path has parity
-coverage and production evidence.
-
-## Crates
-
-| Crate | Role |
-|-------|------|
-| litellm-core | The SDK. Per-route entrypoints (`messages::messages()`), types, provider transforms (modules under `providers/`), provider resolution, auth, the provider HTTP call, and the router. |
-| litellm-config | Config-loading boundary. Returns resolved deployments and optionally delegates loading to Python. |
-| litellm-ai-gateway | The axum server (behind the `server` feature) and WebSocket hosts. Translates HTTP/WS to core entrypoints; no provider handlers. |
-| litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. |
-| litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. Owns API registration, domain wiring, and Python exception mapping. |
-
-Dependency direction is acyclic: config depends on core, the gateway depends on config and core, and the Python bridge depends on the domain layers and Python interop.
-
-## Layout
-
-```text
-crates/
-  core/           The SDK: route modules + provider transforms.
-    src/messages/   mod.rs (entrypoint), types, transformation, prepare, handler, client
-    src/providers/anthropic/messages/transformation.rs
-  config/         Config loading and resolved deployments.
-  ai-gateway/     Axum server + WebSocket hosts; calls core entrypoints.
-  python-interop/ Domain-neutral PyO3 conversion and GIL primitives.
-  python-bridge/  PyO3 API adapter for Python LiteLLM.
-```
-
-The folder shape follows the Python provider tree:
-`core/src/providers/<provider>/<route>/transformation.rs`. The bridge exposes one
-function per top-level route, mirroring the core entrypoints.
-
-## Checks
-
-Run the commands under "Checks" in [CLAUDE.md](CLAUDE.md) before pushing Rust
-changes. That list is the single source of truth and matches what GitHub Actions
-runs for changes under `litellm-rust/`.

litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md

@@ -1,53 +0,0 @@
-# Provider coding standards (litellm-rust)
-
-Rules for adding or changing an LLM provider/route in `litellm-rust`. `messages` (`core/src/messages`, `ANTHROPIC_MESSAGES_CONFIG`) is the reference: a route is a `core` module with a public entrypoint that makes the call and returns a typed response.
-
-## Provider resolution
-
-1. Always resolve the provider/model first with `get_custom_llm_provider` (`core/src/routing_utils/provider.rs`). Nothing downstream may branch on a raw model string.
-2. Model/provider is resolved once, in `prepare.rs`, and passed down as typed fields. Don't re-resolve or re-parse it in transforms or handlers.
-
-## Transforms and the base config
-
-3. Every route defines a base config trait with `transform_request` + `transform_response` (+ `complete_url`, `supported_params`), living in `core/src/<route>/transformation.rs` (e.g. `AnthropicMessagesProviderConfig`, mirroring `OcrProviderConfig`).
-4. Each provider implements that trait as a `const <PROVIDER>_<ROUTE>_CONFIG` in `core/src/providers/<provider>/<route>/transformation.rs`, mirroring the Python provider tree.
-5. Individual configs implement only the request/response transforms. Shared behavior (param filtering, defaults) stays as trait default methods so future providers inherit existing logic instead of reimplementing it.
-6. Prefer composition: a provider that extends another reuses the base trait's defaults or wraps another config; don't copy transform bodies between providers.
-
-## Boundaries
-
-7. Layers never cross: `core` = the call itself (entrypoint, types, transforms, provider resolution, auth headers, provider HTTP, lifecycle hooks); `ai-gateway` = serving HTTP/WS (routing, extractors, auth of *our* callers, streaming to the client); `python-bridge` = thin PyO3 adapter. Hosts call the core entrypoint; they never build a provider request.
-8. Generic/route files contain zero provider-specific branches. A provider is one module under `core/src/providers/<provider>/<route>/`; a route is a module, never a new crate.
-9. Route entry point stays thin: `core::<route>::<route>()` -> `prepare_*` -> handler (or `CallLifecycle::run_request`, which owns the pre_call -> during_call -> provider call -> success/failure order and phase timing). Axum handlers validate and delegate to a service that calls the entrypoint; no business logic in them.
-10. Constants (URLs, env-var names, API versions, error messages) live in a crate `constants.rs`, never inline. Config-shaped env reads happen at the host/config layer with the `DEFAULT_*` fallback defined in `constants.rs`; the only env read in `core` is the credential fallback in a route's `prepare.rs`.
-
-## Types and errors
-
-11. Typed contracts only: no bare `serde_json::Value` / `String` / `Vec<String>` as a transform input or output. Parse wire bytes into typed structs/enums at the host edge; a `type` discriminator is a typed field, not a raw string.
-12. Model failures as values: return typed `CoreError`, don't panic. No `unwrap`/`expect`/`panic!` on user or provider input.
-13. No mutation: build values in one shot (comprehensions/iterators, `collect`), prefer immutable bindings and owned typed structs over seeding-and-mutating.
-14. Early returns over deep nesting; small focused files over god modules.
-15. Preserve Python output shape intentionally. If a field is always serialized as `null` for parity, keep it and pin it with a test.
-
-## Safety and data minimization
-
-16. Never log request/response bodies, base64 payloads, document contents, or secrets. Truncate and bound any upstream body before it crosses a host boundary.
-17. Treat empty/whitespace credentials, URLs, and config values as absent at the host resolution layer.
-18. Network I/O sets connect + request timeouts (no unbounded waits), reuses a shared HTTP client, and prefers rustls TLS.
-
-## Tests and rollout
-
-19. Every provider transform ships tests for: supported-param filtering, request body shape, response normalization, missing/null fields, bad input, and `*_match_python` fixture parity.
-20. Lifecycle/hook tests cover hook order, success + failure callback payloads, pre-call guardrail blocking before any provider I/O, during-call body mutation, and provider-error mapping.
-21. When a route has a Python reference implementation, the Rust path stays off by default and behind Python parity tests (disabled / enabled-equals-Python / bridge-unavailable fallback) until parity is proven. A new provider/route may instead be implemented rust-only with no Python reference; then the Python interface is a thin dispatch to Rust with no fallback, and tests cover the rust-backed path plus the unavailable-bridge error. State the rust-only choice explicitly in the PR.
-
-## Python bridge (SDK side)
-
-22. A Python -> Rust bridge keeps the Python side minimal: the Python interface only marshals inputs and calls the Rust interface, with no transform, handler, or business logic. Aim for well under 100 lines of interface code per route; if the Python grows past that, the logic belongs in Rust.
-23. Do not bloat `litellm/main.py`. A route's provider dispatch lives in a thin dispatch class under `litellm/llms/<provider>/<route>/` that calls the Rust bridge; `main.py` only instantiates it and calls its sync/async method.
-24. Do not add new feature flags unless explicitly requested. Reuse the existing LiteLLM Rust rollout mechanism (`litellm.rust`); never introduce a per-route env flag such as `LITELLM_USE_RUST_<ROUTE>`.
-
-## Checks before push
-
-25. Run, and keep green, the commands under "Checks" in `litellm-rust/CLAUDE.md`.
-    That list is the single source of truth and matches what GitHub Actions runs.

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

@@ -1,54 +0,0 @@
-# ai-gateway — folder architecture
-
-The Axum server that fronts the Rust gateway. It owns transport + config + auth
-only; deployment selection lives in `core::router`, and the LLM call itself
-(transforms, auth headers, provider HTTP) lives behind a `core` route entrypoint
-such as `litellm_core::messages::messages`. No provider handler lives here.
-
-```
-src/
-  main.rs            # entrypoint: build AppState (router + master key), bind, serve
-  state.rs           # AppState — shared Arc<Router> + master_key
-  auth/              # authentication as an axum extractor — added to handler args
-    mod.rs           #   RequireMasterKey: FromRequestParts, single master key (LITELLM_MASTER_KEY)
-  routes/            # one module per route, all matching the same template
-    AGENTS.md        #   ← the route template (read this before adding a route)
-    mod.rs           #   app(): merges every module's router()
-    health.rs        #   simple route (one file): router() + liveness/readiness
-    realtime/        #   route with logic → axum surface + a no-axum service:
-      mod.rs         #     router() + handler + WS<->events adapter (the axum surface)
-      service.rs     #     business logic (select deployment, call provider) — no axum, testable
-```
-
-## Rules
-
-- **Routes follow one template.** Each route module exposes
-  `pub fn router() -> Router<AppState>`; `routes/mod.rs` only merges them. Simple
-  routes are one file; non-trivial routes are a folder (`handler`/`service`/
-  `transport`). See `routes/AGENTS.md`.
-- **Auth is an extractor.** Add `crate::auth::RequireMasterKey` to a handler's
-  args; it runs during extraction. Never re-implement the check per route.
-- **Handlers are thin.** A handler validates and delegates to its `service`. No
-  business logic, no provider calls, no transforms in handlers.
-- **Services call `core`, they don't reimplement it.** A `service` picks the
-  deployment and calls the `core` route entrypoint. Provider resolution, auth
-  headers, URL building, and the HTTP call are `core`'s job; a service that
-  builds a provider request itself is a bug (`routes/messages/service.rs` is
-  the reference).
-- **State is shared and cheap to clone.** Long-lived handles live behind `Arc` in
-  `state.rs`; read env/config only in `main.rs` when building state.
-
-## Auth (interim)
-
-A single **master key** (`LITELLM_MASTER_KEY`), enforced by the
-`auth::RequireMasterKey` extractor: any caller presenting it as
-`Authorization: Bearer <key>` may invoke the gateway. Fails closed (500) when
-unset; constant-time compare. The server binds `127.0.0.1` by default (`HOST` to
-override). Full per-key auth + budgets/rate-limits are delegated to the Python
-proxy in a later phase. Health routes don't add the extractor (unauthenticated).
-
-## Python interop
-
-Python-backed loading lives in `litellm-config` and is **load-time only**. The
-gateway's `python-config` feature forwards to that crate. The realtime data path
-never takes the GIL.

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

@@ -1,14 +0,0 @@
-# ai-gateway architecture
-
-The Rust ai-gateway does LLM inference (realtime WebSocket). Spend tracking is an
-API callback: it POSTs each finished session to the LiteLLM proxy, which records
-spend and runs the usual callbacks.
-
-```mermaid
-flowchart LR
-  C[client] <--> G[Rust ai-gateway<br/>LLM inference]
-  G <--> O[OpenAI realtime]
-  G -. spend tracking callback .-> P[litellm proxy]
-  F[litellm-config<br/>load-time only] --> G
-  F -. Python backend .-> P
-```

litellm-rust/crates/ai-gateway/Cargo.toml

@@ -13,6 +13,11 @@ name = "litellm-ai-gateway"
 path = "src/main.rs"
 required-features = ["server"]
 
+[[bin]]
+name = "trace-parity-gateway"
+path = "src/bin/trace_parity_gateway.rs"
+required-features = ["trace-parity"]
+
 [dependencies]
 tracing.workspace = true
 litellm-core = { workspace = true, features = ["bedrock-auth"] }

litellm-rust/crates/ai-gateway/Dockerfile

@@ -14,15 +14,20 @@
 # ---- Chef -------------------------------------------------------------------
 # cargo-chef caches the dependency build so only the gateway crate recompiles on
 # a source-only change. python3-dev is present in every rust stage because the
-# `python-config` feature links libpython via pyo3 (even in the cook step).
-FROM rust:1.90-slim-bookworm AS chef
+# `python-config` feature links libpython via pyo3 (even in the cook step), and
+# python3-pip builds the litellm wheel in the builder stage.
+FROM rust:1.98-slim-bookworm AS chef
 ENV PYO3_PYTHON=python3.11
+# rustup reads rust-toolchain.toml from any parent of the working directory, so
+# copying it in is what keeps every cargo call below on the repo's pinned
+# channel rather than on whatever the base image happens to ship.
+COPY rust-toolchain.toml /build/rust-toolchain.toml
+WORKDIR /build/litellm-rust
 RUN apt-get update \
     && apt-get install -y --no-install-recommends \
-        python3 python3-dev pkg-config libssl-dev clang \
+        python3 python3-dev python3-pip pkg-config libssl-dev clang \
     && rm -rf /var/lib/apt/lists/* \
     && cargo install cargo-chef --locked --version 0.1.77
-WORKDIR /build/litellm-rust
 
 # ---- Planner ----------------------------------------------------------------
 # Produce the dependency recipe from the rust workspace manifests + Cargo.lock.
@@ -43,6 +48,19 @@ RUN cargo chef cook --locked --release \
 COPY litellm-rust/ .
 RUN cargo build --locked --release -p litellm-ai-gateway --bin litellm-ai-gateway --features server,python-config
 
+# The root pyproject builds with maturin against litellm-rust/crates/python-bridge,
+# so the wheel is built here, next to the crate sources and the cargo toolchain,
+# and the runtime stage installs the artifact instead of compiling anything.
+# litellm[proxy] pins litellm-enterprise and litellm-proxy-extras to the versions
+# in this repo, and those hit PyPI hours after every version bump merges, so both
+# wheels are built from the repo too instead of being resolved from PyPI.
+COPY pyproject.toml README.md LICENSE /build/
+COPY litellm/ /build/litellm/
+COPY enterprise/ /build/enterprise/
+COPY litellm-proxy-extras/ /build/litellm-proxy-extras/
+RUN pip3 wheel --no-cache-dir --no-deps --wheel-dir /build/dist \
+        /build /build/enterprise /build/litellm-proxy-extras
+
 # ---- Runtime ----------------------------------------------------------------
 # python:3.11-slim-bookworm ships libpython3.11, matching the builder's PyO3
 # 3.11 ABI so the embedded interpreter links and imports cleanly.
@@ -56,11 +74,16 @@ RUN apt-get update \
 WORKDIR /app
 
 # Install litellm (with proxy extras) FROM THIS REPO'S SOURCE so
-# `import litellm.proxy.read_model_list` works — it is not on PyPI yet. Copy the
-# package + packaging metadata, then pip install the proxy extra.
-COPY pyproject.toml README.md LICENSE ./
-COPY litellm/ ./litellm/
-RUN pip install --no-cache-dir ".[proxy]"
+# `import litellm.proxy.read_model_list` works — it is not on PyPI yet. The two
+# sibling wheels come from the builder as well, so the pins in litellm[proxy]
+# resolve against them and never wait on a PyPI publish.
+COPY --from=builder /build/dist/*.whl /tmp/wheels/
+RUN wheel="$(ls /tmp/wheels/litellm-*.whl)" \
+    && pip install --no-cache-dir \
+        /tmp/wheels/litellm_enterprise-*.whl \
+        /tmp/wheels/litellm_proxy_extras-*.whl \
+        "${wheel}[proxy]" \
+    && rm -rf /tmp/wheels
 
 # The compiled gateway binary (pure-Rust realtime hot path; Python is load-time
 # only).

litellm-rust/crates/ai-gateway/Dockerfile.dockerignore

@@ -9,19 +9,28 @@
 # Strategy: ignore everything, then re-include only what the build needs:
 #   - litellm/            (pip install . needs the full package + proxy reader)
 #   - litellm-rust/       (the rust workspace; Cargo.lock + crate sources)
-#   - pyproject.toml / README.md / LICENSE  (packaging metadata for pip install)
+#   - enterprise/         (litellm/proxy/enterprise symlinks into it; maturin walks it)
+#   - litellm-proxy-extras/ (built into a wheel alongside enterprise/ for litellm[proxy])
+#   - pyproject.toml / README.md / LICENSE  (packaging metadata for the wheel build)
+#   - rust-toolchain.toml (the pinned channel every cargo call in the build uses)
 *
 
 # --- re-include the build inputs ---
 !litellm/
 !litellm-rust/
+!enterprise/
+!litellm-proxy-extras/
 !pyproject.toml
+!rust-toolchain.toml
 !README.md
 !LICENSE
 
 # --- prune heavy / irrelevant subpaths back out of the re-included trees ---
 # Rust build artifacts (huge; regenerated in the builder).
 **/target/
+# Committed python distribution artifacts; the wheel build does not read them.
+enterprise/dist/
+litellm-proxy-extras/dist/
 # Python caches and compiled bytecode.
 **/__pycache__/
 **/*.pyc

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

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

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

@@ -1,55 +0,0 @@
-# Realtime gateway benchmark — pool on/off
-
-Measures what the gateway adds over talking to OpenAI's realtime WebSocket
-directly, and what the pre-warmed connection pool removes. See
-`../../src/routes/realtime/README.md` for how the pool works.
-
-## Results
-
-5000 calls / 500 concurrency, gateway at 10 instances, pool ON
-(`REALTIME_POOL_SIZE=64`), upstream OpenAI `gpt-realtime`. Each leg run twice.
-Times in **ms**. Phases per connection: **dial** = TCP+TLS+WS upgrade,
-**session** = upgrade → `session.created` (the phase the pool removes),
-**1st-audio** = `response.create` → first audio delta (OpenAI inference),
-**total** = full wall-clock.
-
-| metric             | Direct OpenAI | Gateway (pool ON) | Overhead (ms) | vs OpenAI  |
-| ------------------ | ------------- | ----------------- | ------------- | ---------- |
-| success rate (%)   | 99.8          | 99.8              | —             | —          |
-| dial p50 (ms)      | 276           | 158               | −118          | **faster** |
-| session p50 (ms)   | 7             | 0                 | −7            | **faster** |
-| 1st-audio p50 (ms) | 440           | 664               | +224          | slower¹    |
-| total p50 (ms)     | 816           | 1010              | +194          | slower¹    |
-| total p95 (ms)     | 2152          | 1970              | −182          | **faster** |
-| total p99 (ms)     | 2692          | 2610              | −82           | **faster** |
-
-The gateway is **faster than direct on 4 of 6 metrics**. The warm pool makes the
-**session phase sub-millisecond** at the median — ~76% of connects hit the pool,
-~70% had session < 1 ms. ¹ The two "slower" rows are not gateway overhead:
-`1st-audio` is OpenAI's own inference time (the gateway only relays it), which ran
-slower during the gateway legs and drags `total p50` with it.
-
-**Pool OFF** (control, `REALTIME_POOL_SIZE=0`): session p50 was **367 ms** — the
-fresh-dial overhead the pool removes.
-
-## Reproduce
-
-The load generator lives in a separate repo:
-**https://github.com/ishaan-berri/litellm-realtime-bench**
-
-```bash
-git clone https://github.com/ishaan-berri/litellm-realtime-bench
-cd litellm-realtime-bench && go build -o wsbench .
-
-# Direct to OpenAI (baseline)
-./wsbench -host api.openai.com -key "$OPENAI_API_KEY" -m gpt-realtime -n 5000 -c 500 -t 60
-
-# Through the gateway — run once with pool ON, once with REALTIME_POOL_SIZE=0
-./wsbench -host <gateway-host> -key "$LITELLM_MASTER_KEY" -m gpt-realtime -n 5000 -c 500 -t 60
-```
-
-Run the gateway with the env stand-in (`OPENAI_REALTIME_MODEL=gpt-realtime`,
-`OPENAI_API_KEY`, `LITELLM_MASTER_KEY`, `REALTIME_POOL_SIZE`, `HOST=0.0.0.0`). At
-500 concurrency over N instances, size the pool to `≈ 500 / N` per instance (64 was
-used here for 10 instances). The bench repo's README covers running 500-concurrency
-legs from a hosted multi-vCPU runner. **Never commit keys — pass them via `-key`.**

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

@@ -269,11 +269,15 @@ fn guardrail_error_to_core_error(error: GuardrailError) -> Error {
 
 fn core_error_kind(error: &Error) -> &'static str {
     match error {
-        Error::Auth(_) => "AuthError",
+        Error::Auth(_)
+        | Error::MissingApiKey { .. }
+        | Error::MissingAzureAiCredentials
+        | Error::MissingAzureDocumentIntelligenceCredentials
+        | Error::MissingReductoApiKey => "AuthError",
         Error::InvalidProvider(_) => "InvalidProvider",
         Error::InvalidRequest(_) => "InvalidRequest",
         Error::InvalidType { .. } => "InvalidType",
-        Error::MissingField(_) => "MissingField",
+        Error::MissingField(_) | Error::MissingDocumentUrl => "MissingField",
         Error::Http { .. } => "HttpError",
         Error::InvalidResponse(_) => "InvalidResponse",
         Error::Network(_) => "NetworkError",

litellm-rust/crates/ai-gateway/src/bin/trace_parity_gateway.rs

@@ -0,0 +1,40 @@
+use std::io::Read;
+
+use serde::Deserialize;
+use serde_json::Value;
+
+#[derive(Deserialize)]
+struct Input {
+    model_alias: String,
+    provider_model: String,
+    api_base: String,
+    body: Value,
+}
+
+#[tokio::main]
+async fn main() {
+    let mut input = String::new();
+    if let Err(error) = std::io::stdin().read_to_string(&mut input) {
+        fail(error);
+    }
+    let input: Input = match serde_json::from_str(&input) {
+        Ok(input) => input,
+        Err(error) => fail(error),
+    };
+    let result = litellm_ai_gateway::trace_parity::traced_messages_request(
+        input.model_alias,
+        input.provider_model,
+        input.api_base,
+        input.body,
+    )
+    .await;
+    match serde_json::to_string(&result) {
+        Ok(result) => println!("{result}"),
+        Err(error) => fail(error),
+    }
+}
+
+fn fail(error: impl std::fmt::Display) -> ! {
+    eprintln!("{error}");
+    std::process::exit(1)
+}

litellm-rust/crates/ai-gateway/src/io/realtime.rs

@@ -15,6 +15,8 @@ use std::time::Duration;
 
 use futures_util::stream::{SplitSink, SplitStream};
 use futures_util::{Sink, SinkExt, Stream, StreamExt};
+use litellm_core::AuthError;
+use litellm_core::auth::error::MissingCredential;
 use litellm_core::error::Error;
 use litellm_core::realtime::transformation::RealtimeProviderConfig;
 use litellm_core::realtime::types::RealtimeEvent;
@@ -32,8 +34,6 @@ use crate::io::tls::connect_upstream;
 /// Environment variable holding the OpenAI API key (last-resort fallback).
 const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY";
 
-const MISSING_KEY_MESSAGE: &str = "Missing OpenAI API Key - a realtime call is being made but no key was passed via params or the OPENAI_API_KEY environment variable";
-
 /// Default **idle** timeout: if neither side sends a frame for this long, the
 /// session is reaped. It resets on any activity, so it does not cap a healthy
 /// (continuously streaming) session — it only frees a stalled one (e.g. a
@@ -59,7 +59,7 @@ pub(crate) fn resolve_api_key(api_key: Option<&str>) -> Result<String, Error> {
                 .ok()
                 .filter(|key| !key.trim().is_empty())
         })
-        .ok_or_else(|| Error::Auth(MISSING_KEY_MESSAGE.to_string()))
+        .ok_or_else(|| Error::from(AuthError::from(MissingCredential::OpenAiRealtimeApiKey)))
 }
 
 /// Open the upstream WebSocket to OpenAI for `(model, api_key, api_base)`.

litellm-rust/crates/ai-gateway/src/io/responses_ws.rs

@@ -1,115 +1,28 @@
-use std::collections::HashMap;
-use std::sync::Arc;
 use std::time::Duration;
 
 use futures_util::stream::{SplitSink, SplitStream};
 use futures_util::{Sink, SinkExt, Stream, StreamExt};
+use litellm_core::AuthError;
 use litellm_core::Error;
+use litellm_core::auth::error::MissingCredential;
 use litellm_core::providers::openai::responses::transformation::OPENAI_RESPONSES_WS_CONFIG;
 use litellm_core::responses::types::ResponsesWsEvent;
 use litellm_core::responses::websocket::ResponsesWebSocketProviderConfig;
-use tokio::net::TcpStream;
-use tokio::sync::Mutex;
 use tokio_tungstenite::tungstenite::Message;
 use tokio_tungstenite::tungstenite::client::IntoClientRequest;
 use tokio_tungstenite::tungstenite::http::HeaderValue;
-use tokio_tungstenite::tungstenite::http::header::{AUTHORIZATION, HeaderName};
-use tokio_tungstenite::{MaybeTlsStream, WebSocketStream};
+use tokio_tungstenite::tungstenite::http::header::AUTHORIZATION;
 
-use crate::io::tls::connect_upstream;
+use litellm_core::responses::websocket::{ResponsesUpstreamWs, connect_upstream};
 
 use crate::constants::{
     DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS, DEFAULT_RESPONSES_WS_IDLE_TIMEOUT_SECS,
 };
 
 const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY";
-const MISSING_KEY_MESSAGE: &str = "Missing OpenAI API Key - a Responses WebSocket call is being made but no key was passed via params or the OPENAI_API_KEY environment variable";
-
-pub type ResponsesUpstreamWs = WebSocketStream<MaybeTlsStream<TcpStream>>;
 type UpstreamTx = SplitSink<ResponsesUpstreamWs, Message>;
 type UpstreamRx = SplitStream<ResponsesUpstreamWs>;
 
-#[derive(Clone)]
-pub struct ResponsesWebSocketConnection {
-    socket: Arc<Mutex<Option<ResponsesUpstreamWs>>>,
-}
-
-impl ResponsesWebSocketConnection {
-    pub async fn connect_url(
-        url: &str,
-        headers: &HashMap<String, String>,
-        timeout: Option<Duration>,
-    ) -> Result<Self, Error> {
-        let mut request = url
-            .into_client_request()
-            .map_err(|error| Error::Network(error.to_string()))?;
-        for (name, value) in headers {
-            let header_name = name
-                .parse::<HeaderName>()
-                .map_err(|error| Error::InvalidRequest(error.to_string()))?;
-            let header_value = HeaderValue::from_str(value)
-                .map_err(|error| Error::InvalidRequest(error.to_string()))?;
-            request.headers_mut().insert(header_name, header_value);
-        }
-        let connect = connect_upstream(request);
-        let result = match timeout {
-            Some(timeout) => tokio::time::timeout(timeout, connect).await.map_err(|_| {
-                Error::Network("Responses WebSocket connection timed out".to_string())
-            })?,
-            None => connect.await,
-        };
-        let (socket, _) = result.map_err(|error| match *error {
-            tokio_tungstenite::tungstenite::Error::Http(response) => Error::Http {
-                status: response.status().as_u16(),
-                body: String::new(),
-            },
-            other => Error::Network(other.to_string()),
-        })?;
-        Ok(Self {
-            socket: Arc::new(Mutex::new(Some(socket))),
-        })
-    }
-
-    pub async fn send_text(&self, text: String) -> Result<(), Error> {
-        let mut socket = self.socket.lock().await;
-        let Some(socket) = socket.as_mut() else {
-            return Err(Error::Network("Responses WebSocket is closed".to_string()));
-        };
-        socket
-            .send(Message::Text(text))
-            .await
-            .map_err(|error| Error::Network(error.to_string()))
-    }
-
-    pub async fn recv_text(&self) -> Result<Option<String>, Error> {
-        let mut socket_guard = self.socket.lock().await;
-        let Some(socket) = socket_guard.as_mut() else {
-            return Ok(None);
-        };
-        match socket.next().await {
-            Some(Ok(Message::Text(text))) => Ok(Some(text)),
-            Some(Ok(Message::Binary(bytes))) => String::from_utf8(bytes.to_vec())
-                .map(Some)
-                .map_err(|error| Error::InvalidResponse(error.to_string())),
-            Some(Ok(Message::Close(_))) | None => Ok(None),
-            Some(Ok(_)) => Ok(None),
-            Some(Err(error)) => Err(Error::Network(error.to_string())),
-        }
-    }
-
-    pub async fn close(&self) -> Result<(), Error> {
-        let mut socket = self.socket.lock().await;
-        if let Some(socket) = socket.as_mut() {
-            socket
-                .close(None)
-                .await
-                .map_err(|error| Error::Network(error.to_string()))?;
-        }
-        *socket = None;
-        Ok(())
-    }
-}
-
 pub(crate) fn resolve_api_key(api_key: Option<&str>) -> Result<String, Error> {
     api_key
         .map(str::trim)
@@ -120,7 +33,7 @@ pub(crate) fn resolve_api_key(api_key: Option<&str>) -> Result<String, Error> {
                 .ok()
                 .filter(|value| !value.trim().is_empty())
         })
-        .ok_or_else(|| Error::Auth(MISSING_KEY_MESSAGE.to_string()))
+        .ok_or_else(|| Error::from(AuthError::from(MissingCredential::OpenAiResponsesApiKey)))
 }
 
 async fn dial_upstream(

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

@@ -1,529 +0,0 @@
-use std::net::IpAddr;
-use std::time::{Duration, Instant};
-
-use base64::Engine;
-use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
-use litellm_core::error::Error;
-use litellm_core::ocr::transformation::OcrProviderConfig;
-use reqwest::Url;
-use serde_json::{Map, Value};
-
-use litellm_core::providers::azure_ai::ocr::transformation::{
-    AZURE_AI_OCR_CONFIG, AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG,
-};
-use litellm_core::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG;
-use litellm_core::providers::reducto::ocr::transformation as reducto;
-use litellm_core::providers::vertex_ai::ocr::transformation as vertex_ai;
-use litellm_core::providers::vertex_ai::ocr::transformation::{
-    VERTEX_AI_DEEPSEEK_OCR_CONFIG, VERTEX_AI_OCR_CONFIG,
-};
-
-use crate::client::http_client;
-
-const ERROR_BODY_MAX_CHARS: usize = 256;
-const AZURE_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS: u64 = 120;
-const DEFAULT_MAX_IMAGE_URL_DOWNLOAD_SIZE_MB: f64 = 50.0;
-const MAX_SAFE_FETCH_REDIRECTS: usize = 10;
-
-pub(super) fn truncate_error_body(body: &str) -> String {
-    if body.chars().count() <= ERROR_BODY_MAX_CHARS {
-        return body.to_string();
-    }
-    let truncated: String = body.chars().take(ERROR_BODY_MAX_CHARS).collect();
-    format!("{truncated}... (truncated)")
-}
-
-#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
-pub(super) fn ocr_provider_config(
-    provider: &str,
-    model: &str,
-) -> Option<&'static dyn OcrProviderConfig> {
-    match provider {
-        "mistral" => Some(&MISTRAL_OCR_CONFIG),
-        "reducto" => reducto::config_for_model(model),
-        "azure_ai" if is_azure_document_intelligence_model(model) => {
-            Some(&AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG)
-        }
-        "azure_ai" => Some(&AZURE_AI_OCR_CONFIG),
-        "vertex_ai" if vertex_ai::is_deepseek_model(model) => Some(&VERTEX_AI_DEEPSEEK_OCR_CONFIG),
-        "vertex_ai" => Some(&VERTEX_AI_OCR_CONFIG),
-        _ => None,
-    }
-}
-
-fn is_azure_document_intelligence_model(model: &str) -> bool {
-    let model = model.to_ascii_lowercase();
-    model.contains("doc-intelligence") || model.contains("documentintelligence")
-}
-
-pub(super) fn string_headers(
-    extra_headers: Option<Map<String, Value>>,
-) -> Result<Vec<(String, String)>, Error> {
-    extra_headers
-        .unwrap_or_default()
-        .into_iter()
-        .map(|(key, value)| {
-            value
-                .as_str()
-                .map(|value| (key.clone(), value.to_string()))
-                .ok_or_else(|| {
-                    Error::InvalidRequest(format!(
-                        "OCR extra_headers.{key} must be a string, got {}",
-                        litellm_core::error::json_type_name(&value)
-                    ))
-                })
-        })
-        .collect()
-}
-
-fn document_url_field(document: &Value) -> Result<Option<(&str, &str)>, Error> {
-    let Some(object) = document.as_object() else {
-        return Ok(None);
-    };
-    let Some(doc_type) = object.get("type").and_then(Value::as_str) else {
-        return Ok(None);
-    };
-    let field = match doc_type {
-        "document_url" => "document_url",
-        "image_url" => "image_url",
-        _ => return Ok(None),
-    };
-    let Some(url) = object.get(field).and_then(Value::as_str) else {
-        return Ok(None);
-    };
-    Ok(Some((field, url)))
-}
-
-fn is_url_requiring_fetch(url: &str) -> bool {
-    !url.starts_with("data:") && (url.starts_with("http://") || url.starts_with("https://"))
-}
-
-fn max_document_download_bytes() -> u64 {
-    let max_size_mb = std::env::var("MAX_IMAGE_URL_DOWNLOAD_SIZE_MB")
-        .ok()
-        .and_then(|value| value.parse::<f64>().ok())
-        .unwrap_or(DEFAULT_MAX_IMAGE_URL_DOWNLOAD_SIZE_MB);
-    (max_size_mb.max(0.0) * 1024.0 * 1024.0) as u64
-}
-
-fn is_blocked_ip(ip: IpAddr) -> bool {
-    match ip {
-        IpAddr::V4(ip) => {
-            ip.is_private()
-                || ip.is_loopback()
-                || ip.is_link_local()
-                || ip.is_broadcast()
-                || ip.is_multicast()
-                || ip.is_unspecified()
-        }
-        IpAddr::V6(ip) => {
-            let first_segment = ip.segments()[0];
-            let is_unique_local = (first_segment & 0xfe00) == 0xfc00;
-            let is_link_local = (first_segment & 0xffc0) == 0xfe80;
-            ip.is_loopback()
-                || ip.is_unspecified()
-                || ip.is_multicast()
-                || is_unique_local
-                || is_link_local
-                || ip
-                    .to_ipv4_mapped()
-                    .or_else(|| ip.to_ipv4())
-                    .map(|v4| is_blocked_ip(IpAddr::V4(v4)))
-                    .unwrap_or(false)
-        }
-    }
-}
-
-fn blocked_url_error(url: &Url) -> Error {
-    Error::InvalidRequest(format!(
-        "OCR document URL rejected by SSRF protection: {url}"
-    ))
-}
-
-async fn validate_safe_fetch_url(url: &Url) -> Result<(), Error> {
-    if !matches!(url.scheme(), "http" | "https") {
-        return Err(blocked_url_error(url));
-    }
-
-    let host = url.host_str().ok_or_else(|| blocked_url_error(url))?;
-    if let Ok(ip) = host.parse::<IpAddr>() {
-        if is_blocked_ip(ip) {
-            return Err(blocked_url_error(url));
-        }
-        return Ok(());
-    }
-
-    let port = url
-        .port_or_known_default()
-        .ok_or_else(|| blocked_url_error(url))?;
-    let addresses = tokio::net::lookup_host((host, port))
-        .await
-        .map_err(|err| Error::Network(err.to_string()))?;
-    let mut saw_address = false;
-    for address in addresses {
-        saw_address = true;
-        if is_blocked_ip(address.ip()) {
-            return Err(blocked_url_error(url));
-        }
-    }
-    if !saw_address {
-        return Err(blocked_url_error(url));
-    }
-    Ok(())
-}
-
-fn redirect_location(response: &reqwest::Response, url: &Url) -> Result<Url, Error> {
-    let location = response
-        .headers()
-        .get(reqwest::header::LOCATION)
-        .and_then(|value| value.to_str().ok())
-        .ok_or_else(|| {
-            Error::InvalidResponse("OCR document redirect missing Location header".to_string())
-        })?;
-    url.join(location)
-        .map_err(|err| Error::InvalidResponse(format!("invalid OCR document redirect: {err}")))
-}
-
-async fn safe_get_document_url(url: &str) -> Result<(Url, reqwest::Response), Error> {
-    let client = reqwest::Client::builder()
-        .redirect(reqwest::redirect::Policy::none())
-        .build()
-        .map_err(|err| Error::Network(err.to_string()))?;
-    let mut current_url = Url::parse(url)
-        .map_err(|err| Error::InvalidRequest(format!("invalid OCR document URL: {err}")))?;
-
-    for _ in 0..MAX_SAFE_FETCH_REDIRECTS {
-        validate_safe_fetch_url(&current_url).await?;
-        let response = client
-            .get(current_url.clone())
-            .send()
-            .await
-            .map_err(|err| Error::Network(err.to_string()))?;
-        if !response.status().is_redirection() {
-            return Ok((current_url, response));
-        }
-        current_url = redirect_location(&response, &current_url)?;
-    }
-
-    Err(Error::InvalidRequest(
-        "Too many redirects while fetching OCR document URL".to_string(),
-    ))
-}
-
-fn enforce_download_size(content_length: u64, max_bytes: u64, url: &Url) -> Result<(), Error> {
-    if max_bytes == 0 {
-        return Err(Error::InvalidRequest(format!(
-            "OCR document URL download is disabled (MAX_IMAGE_URL_DOWNLOAD_SIZE_MB=0). url={url}"
-        )));
-    }
-    if content_length > max_bytes {
-        let size_mb = content_length as f64 / (1024.0 * 1024.0);
-        let max_size_mb = max_bytes as f64 / (1024.0 * 1024.0);
-        return Err(Error::InvalidRequest(format!(
-            "OCR document size ({size_mb:.2}MB) exceeds maximum allowed size ({max_size_mb:.2}MB). url={url}"
-        )));
-    }
-    Ok(())
-}
-
-async fn read_response_with_limit(
-    mut response: reqwest::Response,
-    url: &Url,
-) -> Result<Vec<u8>, Error> {
-    let max_bytes = max_document_download_bytes();
-    if let Some(content_length) = response.content_length() {
-        enforce_download_size(content_length, max_bytes, url)?;
-    } else {
-        enforce_download_size(0, max_bytes, url)?;
-    }
-
-    let mut bytes = Vec::new();
-    let mut bytes_downloaded: u64 = 0;
-    while let Some(chunk) = response
-        .chunk()
-        .await
-        .map_err(|err| Error::Network(err.to_string()))?
-    {
-        bytes_downloaded += chunk.len() as u64;
-        enforce_download_size(bytes_downloaded, max_bytes, url)?;
-        bytes.extend_from_slice(&chunk);
-    }
-    Ok(bytes)
-}
-
-pub(super) async fn convert_document_url_to_data_uri(document: Value) -> Result<Value, Error> {
-    let Some((field, url)) = document_url_field(&document)? else {
-        return Ok(document);
-    };
-    if !is_url_requiring_fetch(url) {
-        return Ok(document);
-    }
-
-    let (final_url, response) = safe_get_document_url(url).await?;
-    let status = response.status();
-    if !status.is_success() {
-        let body = response.text().await.unwrap_or_default();
-        return Err(Error::Http {
-            status: status.as_u16(),
-            body: truncate_error_body(&body),
-        });
-    }
-    let content_type = response
-        .headers()
-        .get(reqwest::header::CONTENT_TYPE)
-        .and_then(|value| value.to_str().ok())
-        .and_then(|value| value.split(';').next())
-        .map(str::trim)
-        .filter(|value| !value.is_empty())
-        .unwrap_or("application/octet-stream")
-        .to_string();
-    let bytes = read_response_with_limit(response, &final_url).await?;
-    let data_uri = format!(
-        "data:{content_type};base64,{}",
-        BASE64_STANDARD.encode(bytes)
-    );
-
-    let mut transformed = document
-        .as_object()
-        .cloned()
-        .ok_or_else(|| Error::InvalidRequest("OCR document must be an object".to_string()))?;
-    transformed.insert(field.to_string(), Value::String(data_uri));
-    Ok(Value::Object(transformed))
-}
-
-fn same_origin(left: &str, right: &str) -> bool {
-    let Ok(left) = reqwest::Url::parse(left) else {
-        return false;
-    };
-    let Ok(right) = reqwest::Url::parse(right) else {
-        return false;
-    };
-    left.scheme() == right.scheme()
-        && left.host_str() == right.host_str()
-        && left.port_or_known_default() == right.port_or_known_default()
-}
-
-fn retry_after_secs(response: &reqwest::Response) -> u64 {
-    response
-        .headers()
-        .get(reqwest::header::RETRY_AFTER)
-        .and_then(|value| value.to_str().ok())
-        .and_then(|value| value.parse::<u64>().ok())
-        .unwrap_or(2)
-}
-
-fn operation_status(response_json: &Value) -> Result<&str, Error> {
-    let status = response_json
-        .get("status")
-        .and_then(Value::as_str)
-        .ok_or(Error::MissingField("status"))?;
-    match status {
-        "succeeded" => Ok("succeeded"),
-        "running" | "notStarted" => Ok("running"),
-        "failed" => {
-            let message = response_json
-                .get("error")
-                .and_then(|error| error.get("message"))
-                .and_then(Value::as_str)
-                .unwrap_or("Unknown error");
-            Err(Error::InvalidResponse(format!(
-                "Azure Document Intelligence analysis failed: {message}"
-            )))
-        }
-        other => Err(Error::InvalidResponse(format!(
-            "Unknown operation status: {other}"
-        ))),
-    }
-}
-
-#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
-pub(super) async fn poll_document_intelligence(
-    operation_url: &str,
-    original_url: &str,
-    headers: &[(String, String)],
-    timeout: Option<Duration>,
-) -> Result<Value, Error> {
-    if !same_origin(operation_url, original_url) {
-        return Err(Error::InvalidResponse(
-            "Azure Document Intelligence: rejected cross-origin polling URL".to_string(),
-        ));
-    }
-
-    let start = Instant::now();
-    let timeout = timeout.unwrap_or(Duration::from_secs(
-        AZURE_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS,
-    ));
-    loop {
-        if start.elapsed() > timeout {
-            return Err(Error::Network(format!(
-                "Azure Document Intelligence operation polling timed out after {} seconds",
-                timeout.as_secs()
-            )));
-        }
-
-        let mut request_builder = http_client().get(operation_url);
-        for (key, value) in headers {
-            if key.eq_ignore_ascii_case("ocp-apim-subscription-key") {
-                request_builder = request_builder.header(key, value);
-            }
-        }
-        let response = request_builder
-            .send()
-            .await
-            .map_err(|err| Error::Network(err.to_string()))?;
-        let retry_after = retry_after_secs(&response);
-        let status = response.status();
-        let text = response
-            .text()
-            .await
-            .map_err(|err| Error::Network(err.to_string()))?;
-        if !status.is_success() {
-            return Err(Error::Http {
-                status: status.as_u16(),
-                body: truncate_error_body(&text),
-            });
-        }
-        let response_json: Value = serde_json::from_str(&text).map_err(|err| {
-            Error::InvalidResponse(format!("invalid Azure DI poll response JSON: {err}"))
-        })?;
-        if operation_status(&response_json)? == "succeeded" {
-            return Ok(response_json);
-        }
-        tokio::time::sleep(Duration::from_secs(retry_after)).await;
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use litellm_core::ocr::transformation::OcrResponseHandling;
-    use serde_json::json;
-
-    use super::*;
-
-    #[test]
-    fn blocks_private_and_metadata_ips() {
-        assert!(is_blocked_ip("127.0.0.1".parse().unwrap()));
-        assert!(is_blocked_ip("10.0.0.1".parse().unwrap()));
-        assert!(is_blocked_ip("169.254.169.254".parse().unwrap()));
-        assert!(is_blocked_ip("::1".parse().unwrap()));
-        assert!(is_blocked_ip("fd00::1".parse().unwrap()));
-        assert!(is_blocked_ip("fe80::1".parse().unwrap()));
-        assert!(is_blocked_ip("::ffff:169.254.169.254".parse().unwrap()));
-        assert!(is_blocked_ip("::ffff:10.0.0.1".parse().unwrap()));
-        assert!(!is_blocked_ip("8.8.8.8".parse().unwrap()));
-        assert!(!is_blocked_ip("::ffff:8.8.8.8".parse().unwrap()));
-    }
-
-    #[tokio::test]
-    async fn convert_document_url_rejects_loopback_fetch() {
-        let error = convert_document_url_to_data_uri(json!({
-            "type": "image_url",
-            "image_url": "http://127.0.0.1/image.png"
-        }))
-        .await
-        .unwrap_err();
-
-        assert!(matches!(
-            error,
-            Error::InvalidRequest(message)
-                if message.contains("SSRF protection")
-        ));
-    }
-
-    #[tokio::test]
-    async fn convert_document_url_leaves_data_uri_untouched() {
-        let document = json!({
-            "type": "image_url",
-            "image_url": "data:image/png;base64,abcd"
-        });
-
-        let transformed = convert_document_url_to_data_uri(document.clone())
-            .await
-            .unwrap();
-
-        assert_eq!(transformed, document);
-    }
-
-    #[test]
-    fn truncate_error_body_passes_short_strings_through() {
-        let body = "Unauthorized";
-        assert_eq!(truncate_error_body(body), "Unauthorized");
-    }
-
-    #[test]
-    fn truncate_error_body_caps_long_payloads() {
-        let body = "x".repeat(306);
-        let truncated = truncate_error_body(&body);
-
-        assert!(truncated.ends_with("... (truncated)"));
-        let prefix_chars = truncated
-            .strip_suffix("... (truncated)")
-            .expect("truncated marker present")
-            .chars()
-            .count();
-        assert_eq!(prefix_chars, 256);
-    }
-
-    #[test]
-    fn truncate_error_body_does_not_split_multibyte_chars() {
-        let body = "é".repeat(266);
-        let truncated = truncate_error_body(&body);
-        assert!(truncated.is_char_boundary(truncated.len()));
-    }
-
-    #[test]
-    fn ocr_dispatch_supports_migrated_providers() {
-        assert!(ocr_provider_config("mistral", "mistral-ocr-latest").is_some());
-        assert!(
-            ocr_provider_config("azure_ai", "pixtral-12b-2409")
-                .expect("azure ai config resolves")
-                .requires_data_uri_document()
-        );
-        assert_eq!(
-            ocr_provider_config("azure_ai", "doc-intelligence/prebuilt-read")
-                .expect("document intelligence config resolves")
-                .response_handling(),
-            OcrResponseHandling::AzureDocumentIntelligencePoll
-        );
-        assert!(
-            ocr_provider_config("vertex_ai", "deepseek-ocr-maas")
-                .expect("vertex deepseek config resolves")
-                .supported_ocr_params()
-                .contains(&"temperature")
-        );
-        assert!(ocr_provider_config("openai", "gpt-4o").is_none());
-    }
-
-    #[test]
-    fn string_headers_accepts_string_values() {
-        let headers = json!({
-            "x-trace-id": "trace-1"
-        })
-        .as_object()
-        .unwrap()
-        .clone();
-
-        assert_eq!(
-            string_headers(Some(headers)).expect("string headers accepted"),
-            vec![("x-trace-id".to_string(), "trace-1".to_string())]
-        );
-    }
-
-    #[test]
-    fn string_headers_rejects_non_string_values() {
-        let headers = json!({
-            "x-retry-count": 3
-        })
-        .as_object()
-        .unwrap()
-        .clone();
-
-        let err = string_headers(Some(headers)).expect_err("non-string header rejected");
-        assert_eq!(
-            err,
-            Error::InvalidRequest(
-                "OCR extra_headers.x-retry-count must be a string, got number".to_string()
-            )
-        );
-    }
-}

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

@@ -1,84 +0,0 @@
-use litellm_core::error::Error;
-use litellm_core::http_utils::http_request;
-use litellm_core::ocr::transformation::OcrResponseHandling;
-use serde_json::Value;
-
-use super::common_utils::{poll_document_intelligence, truncate_error_body};
-use super::hooks::OcrLifecycleHooks;
-use super::types::PreparedOcrRequest;
-use crate::client::http_client;
-
-#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
-pub(crate) async fn execute_ocr_provider_call(
-    request: PreparedOcrRequest,
-    hooks: &OcrLifecycleHooks,
-) -> Result<Value, Error> {
-    let request = hooks.prepare_provider_request(request).await?;
-    let mut request_builder = http_client().post(&request.url).json(&request.body);
-    for (key, value) in &request.upstream_headers {
-        request_builder = request_builder.header(key, value);
-    }
-    if let Some(duration) = request.timeout {
-        request_builder = request_builder.timeout(duration);
-    }
-
-    let response = http_request(request_builder)
-        .await
-        .map_err(|err| Error::Network(err.to_string()))?;
-
-    let status = response.status();
-    if request.config.response_handling() == OcrResponseHandling::AzureDocumentIntelligencePoll
-        && status.as_u16() == 202
-    {
-        let operation_url = response
-            .headers()
-            .get("operation-location")
-            .and_then(|value| value.to_str().ok())
-            .map(str::to_string)
-            .ok_or_else(|| {
-                Error::InvalidResponse(
-                    "Azure Document Intelligence returned 202 but no Operation-Location header found"
-                        .to_string(),
-                )
-            })?;
-        let response_json = poll_document_intelligence(
-            &operation_url,
-            &request.url,
-            &request.upstream_headers,
-            request.timeout,
-        )
-        .await?;
-        return Ok(request
-            .config
-            .transform_ocr_response_with_params(
-                &request.model,
-                response_json,
-                &request.optional_params,
-            )?
-            .into_json());
-    }
-
-    let text = response
-        .text()
-        .await
-        .map_err(|err| Error::Network(err.to_string()))?;
-
-    if !status.is_success() {
-        return Err(Error::Http {
-            status: status.as_u16(),
-            body: truncate_error_body(&text),
-        });
-    }
-
-    let response_json: Value = serde_json::from_str(&text)
-        .map_err(|err| Error::InvalidResponse(format!("invalid OCR response JSON: {err}")))?;
-
-    Ok(request
-        .config
-        .transform_ocr_response_with_params(
-            &request.model,
-            response_json,
-            &request.optional_params,
-        )?
-        .into_json())
-}

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

@@ -1,401 +0,0 @@
-use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming};
-use litellm_core::error::Error;
-use litellm_core::providers::reducto::ocr::transformation::{
-    build_upload_request, extract_document_source, extract_upload_file_id,
-};
-use serde_json::{Map, Value, json};
-use std::future::Future;
-use std::pin::Pin;
-
-use super::common_utils::{convert_document_url_to_data_uri, string_headers, truncate_error_body};
-use super::types::{PreparedOcrRequest, ProviderOcrRequest};
-use crate::client::http_client;
-use crate::integrations::custom_guardrail::{
-    CustomGuardrailRunner, GuardrailContext, GuardrailError, GuardrailRequest,
-};
-use crate::integrations::custom_logger::{
-    CallType, CallbackTiming, CallbackValue, CustomLoggerRunner, LoggingError, ModelCallDetails,
-};
-use crate::integrations::types::{
-    RequestMetadata, StandardLoggingMetadata, StandardLoggingPayload,
-};
-
-pub(crate) struct OcrLifecycleHooks {
-    logger_runner: CustomLoggerRunner,
-    guardrail_runner: CustomGuardrailRunner,
-    request_metadata: RequestMetadata,
-}
-
-type OcrFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
-type OcrLogFuture<'a> = Pin<Box<dyn Future<Output = ()> + Send + 'a>>;
-
-impl OcrLifecycleHooks {
-    pub(crate) fn new(
-        logger_runner: CustomLoggerRunner,
-        guardrail_runner: CustomGuardrailRunner,
-        request_metadata: RequestMetadata,
-    ) -> Self {
-        Self {
-            logger_runner,
-            guardrail_runner,
-            request_metadata,
-        }
-    }
-
-    async fn run_pre_call_guardrails(
-        &self,
-        request: PreparedOcrRequest,
-    ) -> Result<PreparedOcrRequest, Error> {
-        if self.guardrail_runner.is_empty() {
-            return Ok(request);
-        }
-
-        let context = guardrail_context(&self.request_metadata);
-        let guardrail_request = GuardrailRequest::new(json!({
-            "model": request.model,
-            "custom_llm_provider": request.custom_llm_provider,
-            "document": request.document,
-            "optional_params": request.optional_params,
-        }));
-        let (guardrail_request, _) = self
-            .guardrail_runner
-            .run_pre_call(&context, guardrail_request)
-            .await
-            .map_err(guardrail_error_to_core_error)?;
-        let (document, optional_params) = parse_ocr_pre_call_guardrail_request(guardrail_request)?;
-        let optional_params = match &request.config {
-            Ok(config) => config.map_ocr_params(&optional_params),
-            Err(_) => optional_params,
-        };
-        Ok(PreparedOcrRequest {
-            document,
-            optional_params,
-            ..request
-        })
-    }
-
-    pub(crate) async fn prepare_provider_request(
-        &self,
-        request: PreparedOcrRequest,
-    ) -> Result<ProviderOcrRequest, Error> {
-        let config = request.config?;
-        let env_lookup = |key: &str| std::env::var(key).ok();
-        let upstream_headers = config.validate_environment(
-            string_headers(request.extra_headers)?,
-            request.api_key.as_deref(),
-            &env_lookup,
-        )?;
-        let url = config.complete_url(
-            request.api_base.as_deref(),
-            &request.model,
-            &request.optional_params,
-            &env_lookup,
-        )?;
-        let model = request.model.clone();
-        let custom_llm_provider = request.custom_llm_provider.clone();
-        let is_reducto = custom_llm_provider == "reducto";
-        let document = if is_reducto {
-            let guarded_document = self
-                .run_during_call_guardrails(&model, &custom_llm_provider, &url, request.document)
-                .await?;
-            upload_reducto_document(
-                &guarded_document,
-                request.api_base.as_deref(),
-                request.timeout,
-                &upstream_headers,
-            )
-            .await?
-        } else if config.requires_data_uri_document() {
-            convert_document_url_to_data_uri(request.document).await?
-        } else {
-            request.document
-        };
-        let optional_params = request.optional_params;
-        let body = config
-            .transform_ocr_request(&request.model, document, optional_params.clone())?
-            .data;
-        let body = if is_reducto {
-            body
-        } else {
-            self.run_during_call_guardrails(&model, &custom_llm_provider, &url, body)
-                .await?
-        };
-        Ok(ProviderOcrRequest {
-            model,
-            config,
-            url,
-            body,
-            optional_params,
-            upstream_headers,
-            timeout: request.timeout,
-        })
-    }
-
-    async fn run_during_call_guardrails(
-        &self,
-        model: &str,
-        custom_llm_provider: &str,
-        url: &str,
-        body: Value,
-    ) -> Result<Value, Error> {
-        if self.guardrail_runner.is_empty() {
-            return Ok(body);
-        }
-
-        let context = guardrail_context(&self.request_metadata);
-        let guardrail_request = GuardrailRequest::new(json!({
-            "model": model,
-            "custom_llm_provider": custom_llm_provider,
-            "url": url,
-            "body": body,
-        }));
-        let (guardrail_request, _) = self
-            .guardrail_runner
-            .run_during_call(&context, guardrail_request)
-            .await
-            .map_err(guardrail_error_to_core_error)?;
-        parse_ocr_during_call_guardrail_request(guardrail_request)
-    }
-
-    fn standard_logging_payload(
-        &self,
-        context: &CallLifecycleContext,
-        timing: &CallLifecycleTiming,
-    ) -> StandardLoggingPayload {
-        StandardLoggingPayload {
-            id: context.litellm_call_id.clone(),
-            litellm_call_id: context.litellm_call_id.clone(),
-            call_type: context.call_type.clone(),
-            model: context.model.clone(),
-            custom_llm_provider: context.custom_llm_provider.clone(),
-            response_cost: 0.0,
-            prompt_tokens: 0,
-            completion_tokens: 0,
-            total_tokens: 0,
-            start_time: timing.start_time,
-            end_time: timing.end_time,
-            stream: false,
-            metadata: StandardLoggingMetadata {
-                user_api_key_hash: self.request_metadata.user_api_key_hash.clone(),
-                user_api_key_user_id: self.request_metadata.user_api_key_user_id.clone(),
-                user_api_key_team_id: self.request_metadata.user_api_key_team_id.clone(),
-                ..Default::default()
-            },
-            messages: None,
-        }
-    }
-}
-
-async fn upload_reducto_document(
-    document: &Value,
-    api_base: Option<&str>,
-    timeout: Option<std::time::Duration>,
-    upstream_headers: &[(String, String)],
-) -> Result<Value, Error> {
-    let source = extract_document_source(document)?;
-    let Some(authorization) = upstream_headers
-        .iter()
-        .find(|(name, _)| name.eq_ignore_ascii_case("authorization"))
-        .map(|(_, value)| value.as_str())
-    else {
-        return Err(Error::Auth(
-            "Reducto upload requires an Authorization header".to_string(),
-        ));
-    };
-    let Some(upload) = build_upload_request(source, authorization, api_base) else {
-        return Ok(document.clone());
-    };
-    let part = reqwest::multipart::Part::bytes(upload.bytes)
-        .file_name(upload.file_name)
-        .mime_str(&upload.mime_type)
-        .map_err(|error| Error::InvalidRequest(error.to_string()))?;
-    let form = reqwest::multipart::Form::new().part("file", part);
-    let mut request_builder = http_client().post(upload.url).multipart(form);
-    for (name, value) in upstream_headers {
-        if !name.eq_ignore_ascii_case("content-type")
-            && !name.eq_ignore_ascii_case("content-length")
-        {
-            request_builder = request_builder.header(name, value);
-        }
-    }
-    if let Some(timeout) = timeout {
-        request_builder = request_builder.timeout(timeout);
-    }
-    let response = request_builder
-        .send()
-        .await
-        .map_err(|error| Error::Network(error.to_string()))?;
-    let status = response.status();
-    let body = response
-        .text()
-        .await
-        .map_err(|error| Error::Network(error.to_string()))?;
-    if !status.is_success() {
-        return Err(Error::Http {
-            status: status.as_u16(),
-            body: truncate_error_body(&body),
-        });
-    }
-    let response_json: Value = serde_json::from_str(&body).map_err(|error| {
-        Error::InvalidResponse(format!("invalid Reducto upload response JSON: {error}"))
-    })?;
-    let file_id = extract_upload_file_id(&response_json)?;
-    Ok(json!({"type": "document_url", "document_url": file_id}))
-}
-
-impl CallLifecycleHooks<PreparedOcrRequest, PreparedOcrRequest, Value> for OcrLifecycleHooks {
-    type PreCallFuture<'a> = OcrFuture<'a, PreparedOcrRequest>;
-    type DuringCallFuture<'a> = OcrFuture<'a, PreparedOcrRequest>;
-    type SuccessFuture<'a> = OcrLogFuture<'a>;
-    type FailureFuture<'a> = OcrLogFuture<'a>;
-
-    fn async_pre_call_hook<'a>(
-        &'a self,
-        _context: &'a CallLifecycleContext,
-        request: PreparedOcrRequest,
-    ) -> Self::PreCallFuture<'a> {
-        Box::pin(async move { self.run_pre_call_guardrails(request).await })
-    }
-
-    fn async_during_call_hook<'a>(
-        &'a self,
-        _context: &'a CallLifecycleContext,
-        request: PreparedOcrRequest,
-    ) -> Self::DuringCallFuture<'a> {
-        Box::pin(async move { Ok(request) })
-    }
-
-    #[tracing::instrument(
-        name = "success_callback",
-        target = "litellm::function_trace",
-        level = "trace",
-        skip_all
-    )]
-    fn async_log_success_event<'a>(
-        &'a self,
-        context: &'a CallLifecycleContext,
-        response: &'a Value,
-        timing: &'a CallLifecycleTiming,
-    ) -> Self::SuccessFuture<'a> {
-        Box::pin(async move {
-            if self.logger_runner.is_empty() {
-                return;
-            }
-            let response_obj = CallbackValue::new("ocr", response.clone());
-            self.logger_runner
-                .async_log_success_event(
-                    &ModelCallDetails::from_standard_logging_payload(
-                        self.standard_logging_payload(context, timing),
-                    ),
-                    &response_obj,
-                    CallbackTiming::new(timing.start_time, timing.end_time),
-                )
-                .await;
-        })
-    }
-
-    #[tracing::instrument(
-        name = "failure_callback",
-        target = "litellm::function_trace",
-        level = "trace",
-        skip_all
-    )]
-    fn async_log_failure_event<'a>(
-        &'a self,
-        context: &'a CallLifecycleContext,
-        error: &'a Error,
-        timing: &'a CallLifecycleTiming,
-    ) -> Self::FailureFuture<'a> {
-        Box::pin(async move {
-            if self.logger_runner.is_empty() {
-                return;
-            }
-            let logging_error = LoggingError {
-                message: error.to_string(),
-                kind: core_error_kind(error).to_string(),
-            };
-            let response_obj = CallbackValue::new(
-                "error",
-                json!({
-                    "message": logging_error.message,
-                    "kind": logging_error.kind,
-                }),
-            );
-            self.logger_runner
-                .async_log_failure_event(
-                    &ModelCallDetails::from_standard_logging_payload(
-                        self.standard_logging_payload(context, timing),
-                    )
-                    .with_failure_error(logging_error),
-                    Some(&response_obj),
-                    CallbackTiming::new(timing.start_time, timing.end_time),
-                )
-                .await;
-        })
-    }
-}
-
-fn guardrail_context(metadata: &RequestMetadata) -> GuardrailContext {
-    GuardrailContext {
-        call_type: CallType::Ocr,
-        selected_guardrails: Vec::new(),
-        metadata: std::collections::HashMap::new(),
-        user_api_key_hash: metadata.user_api_key_hash.clone(),
-        user_api_key_user_id: metadata.user_api_key_user_id.clone(),
-        user_api_key_team_id: metadata.user_api_key_team_id.clone(),
-        trace_parent: None,
-    }
-}
-
-fn parse_ocr_pre_call_guardrail_request(
-    request: GuardrailRequest,
-) -> Result<(Value, Map<String, Value>), Error> {
-    let Value::Object(mut data) = request.data else {
-        return Err(Error::InvalidRequest(
-            "OCR pre_call guardrail must return an object".to_string(),
-        ));
-    };
-    let document = data.remove("document").ok_or_else(|| {
-        Error::InvalidRequest("OCR pre_call guardrail removed document".to_string())
-    })?;
-    let optional_params = match data.remove("optional_params") {
-        Some(Value::Object(params)) => params,
-        Some(_) => {
-            return Err(Error::InvalidRequest(
-                "OCR pre_call guardrail optional_params must be an object".to_string(),
-            ));
-        }
-        None => Map::new(),
-    };
-    Ok((document, optional_params))
-}
-
-fn parse_ocr_during_call_guardrail_request(request: GuardrailRequest) -> Result<Value, Error> {
-    let Value::Object(mut data) = request.data else {
-        return Err(Error::InvalidRequest(
-            "OCR during_call guardrail must return an object".to_string(),
-        ));
-    };
-    data.remove("body")
-        .ok_or_else(|| Error::InvalidRequest("OCR during_call guardrail removed body".to_string()))
-}
-
-fn guardrail_error_to_core_error(error: GuardrailError) -> Error {
-    Error::InvalidRequest(format!("{}: {}", error.kind, error.message))
-}
-
-fn core_error_kind(error: &Error) -> &'static str {
-    match error {
-        Error::Auth(_) => "AuthError",
-        Error::InvalidProvider(_) => "InvalidProvider",
-        Error::InvalidRequest(_) => "InvalidRequest",
-        Error::InvalidType { .. } => "InvalidType",
-        Error::MissingField(_) => "MissingField",
-        Error::Http { .. } => "HttpError",
-        Error::InvalidResponse(_) => "InvalidResponse",
-        Error::Network(_) => "NetworkError",
-        Error::Connect(_) => "ConnectError",
-        Error::Routing(_) => "RoutingError",
-        Error::Unsupported(_) => "UnsupportedRequest",
-    }
-}

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

@@ -1,174 +1,127 @@
 use litellm_core::Error;
-use litellm_core::call_lifecycle::CallLifecycle;
+use litellm_core::ocr::{
+    OcrClient,
+    wire::{OcrWireRequest, decode_request},
+};
 use serde_json::Value;
 
-mod common_utils;
-mod handler;
-mod hooks;
-mod prepare;
 mod types;
 
 pub use types::OcrRequest;
 
-use handler::execute_ocr_provider_call;
-use prepare::{PreparedOcrCall, prepare_ocr_call};
-
 #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
 pub async fn ocr(request: OcrRequest<'_>) -> Result<Value, Error> {
-    let PreparedOcrCall { request, hooks } = prepare_ocr_call(request);
-    CallLifecycle::default()
-        .run_request(request, &hooks, |request| {
-            execute_ocr_provider_call(request, &hooks)
-        })
+    core_ocr(request).await
+}
+
+async fn core_ocr(request: OcrRequest<'_>) -> Result<Value, Error> {
+    validate_host_hooks(&request)?;
+    let client = OcrClient::new(crate::client::http_client().clone())?;
+    let core_request = decode_request(OcrWireRequest {
+        model: request.model.to_string(),
+        document: request.document,
+        api_key: request.api_key.map(str::to_string),
+        api_base: request.api_base.map(str::to_string),
+        custom_llm_provider: request.custom_llm_provider.map(str::to_string),
+        extra_headers: request.extra_headers,
+        optional_params: request.optional_params,
+        input_sources: Default::default(),
+        timeout_seconds: request.timeout.map(|timeout| timeout.as_secs_f64()),
+    })?;
+    client
+        .perform(core_request)
         .await
+        .map(|response| response.into_json())
+}
+
+fn validate_host_hooks(request: &OcrRequest<'_>) -> Result<(), Error> {
+    if !request.guardrails.is_empty() {
+        return Err(Error::Unsupported(
+            "OCR host guardrails are not wired to the core path",
+        ));
+    }
+    if !request.callbacks.is_empty() {
+        return Err(Error::Unsupported(
+            "OCR host callbacks are not wired to the core path",
+        ));
+    }
+    Ok(())
 }
 
 #[cfg(test)]
 mod tests {
+    use std::sync::Arc;
+
+    use litellm_core::ocr::wire::is_supported_request;
     use serde_json::{Map, json};
-    use tokio::io::{AsyncReadExt, AsyncWriteExt};
-    use tokio::net::{TcpListener, TcpStream};
 
-    use super::{OcrRequest, ocr};
-    use crate::integrations::types::RequestMetadata;
+    use super::{OcrRequest, validate_host_hooks};
+    use crate::integrations::custom_guardrail::{CustomGuardrail, GuardrailEventHook};
+    use crate::integrations::custom_logger::CustomLogger;
 
-    async fn read_http_request(socket: &mut TcpStream) -> String {
-        let mut request = Vec::new();
-        let mut buffer = [0_u8; 1024];
-        let header_end = loop {
-            let n = socket.read(&mut buffer).await.expect("reads request");
-            if n == 0 {
-                break request.len();
-            }
-            request.extend_from_slice(&buffer[..n]);
-            if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") {
-                break position + 4;
-            }
-        };
-        let headers = String::from_utf8_lossy(&request[..header_end]);
-        let content_length = headers
-            .lines()
-            .find_map(|line| {
-                let (name, value) = line.split_once(':')?;
-                name.eq_ignore_ascii_case("content-length")
-                    .then(|| value.trim().parse::<usize>().ok())
-                    .flatten()
-            })
-            .unwrap_or(0);
-        while request.len().saturating_sub(header_end) < content_length {
-            let n = socket.read(&mut buffer).await.expect("reads body");
-            if n == 0 {
-                break;
-            }
-            request.extend_from_slice(&buffer[..n]);
+    struct TestGuardrail;
+
+    impl CustomGuardrail for TestGuardrail {
+        fn guardrail_name(&self) -> &str {
+            "test"
+        }
+
+        fn supported_event_hooks(&self) -> &[GuardrailEventHook] {
+            &[]
         }
-        String::from_utf8(request).expect("request is utf8")
     }
 
-    fn base_ocr_request(model: &str) -> OcrRequest<'_> {
+    struct TestLogger;
+
+    impl CustomLogger for TestLogger {}
+
+    fn request() -> OcrRequest<'static> {
         OcrRequest {
-            model,
-            document: json!({
-                "type": "document_url",
-                "document_url": "https://example.com/doc.pdf"
-            }),
-            api_key: Some("sk-test"),
+            model: "model",
+            document: json!({"type":"image_url","image_url":"data:image/png;base64,YQ=="}),
+            api_key: None,
             api_base: None,
-            custom_llm_provider: None,
+            custom_llm_provider: Some("mistral"),
             extra_headers: None,
             optional_params: Map::new(),
             timeout: None,
             callbacks: Vec::new(),
             guardrails: Vec::new(),
-            request_metadata: RequestMetadata::default(),
+            request_metadata: Default::default(),
             litellm_call_id: None,
         }
     }
 
-    #[tokio::test]
-    async fn reducto_file_upload_then_parse_maps_response() {
-        let listener = TcpListener::bind("127.0.0.1:0")
-            .await
-            .expect("test listener binds");
-        let address = listener.local_addr().expect("listener has local address");
-        let server = tokio::spawn(async move {
-            let (mut upload_socket, _) = listener.accept().await.expect("accepts upload request");
-            let upload_request = read_http_request(&mut upload_socket).await;
-            let upload_body = r#"{"file_id":"reducto://uploaded.pdf"}"#;
-            let upload_response = format!(
-                "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
-                upload_body.len(),
-                upload_body
-            );
-            upload_socket
-                .write_all(upload_response.as_bytes())
-                .await
-                .expect("writes upload response");
+    #[test]
+    fn core_activation_includes_migrated_providers() {
+        assert!(is_supported_request("model", Some("mistral")));
+        assert!(is_supported_request("pixtral-12b", Some("azure_ai")));
+        assert!(is_supported_request(
+            "doc-intelligence/prebuilt-layout",
+            Some("azure_ai")
+        ));
+        assert!(is_supported_request("parse-v3", Some("reducto")));
+        assert!(is_supported_request("mistral-ocr", Some("vertex_ai")));
+        assert!(is_supported_request("deepseek-ocr", Some("vertex_ai")));
+    }
 
-            let (mut parse_socket, _) = listener.accept().await.expect("accepts parse request");
-            let parse_request = read_http_request(&mut parse_socket).await;
-            let parse_body = r#"{"job_id":"job_123","usage":{"num_pages":3,"credits":3},"result":{"chunks":[{"content":"Page 1 block A","blocks":[{"content":"Page 1 block A","bbox":{"page":1},"kind":"text"}]},{"content":"Page 2 block A","blocks":[{"content":"Page 2 block A","bbox":{"page":2},"kind":"table"}]},{"content":"Page 1 block B","blocks":[{"content":"Page 1 block B","bbox":{"page":1},"kind":"text"}]},{"content":"Page 3 block A","blocks":[{"content":"Page 3 block A","bbox":{"page":3},"kind":"figure"}]}]}}"#;
-            let parse_response = format!(
-                "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
-                parse_body.len(),
-                parse_body
-            );
-            parse_socket
-                .write_all(parse_response.as_bytes())
-                .await
-                .expect("writes parse response");
-            (upload_request, parse_request)
-        });
-        let api_base = format!("http://{address}");
-        let mut request = base_ocr_request("reducto/parse-v3");
-        request.api_base = Some(&api_base);
-        request.api_key = None;
-        request.extra_headers = Some(Map::from_iter([
-            ("Authorization".to_string(), json!("Bearer test-key")),
-            ("x-trace-id".to_string(), json!("trace-1")),
-        ]));
-        request.document = json!({
-            "type": "document_url",
-            "document_url": "data:application/pdf;base64,JVBERi0xLjQ="
-        });
-        request.optional_params = Map::from_iter([
-            (
-                "formatting".to_string(),
-                json!({"table_output_format": "html"}),
-            ),
-            ("retrieval".to_string(), json!({"chunk_mode": "section"})),
-            ("settings".to_string(), json!({"ocr_system": "standard"})),
-        ]);
+    #[test]
+    fn core_path_rejects_unwired_guardrails() {
+        let request = OcrRequest {
+            guardrails: vec![Arc::new(TestGuardrail)],
+            ..request()
+        };
+        let error = validate_host_hooks(&request).unwrap_err();
+        assert!(error.to_string().contains("guardrails are not wired"));
+    }
 
-        let response = ocr(request).await.expect("Reducto OCR succeeds");
-
-        assert_eq!(response["pages"].as_array().map(Vec::len), Some(3));
-        assert_eq!(
-            response["pages"][0]["markdown"],
-            "Page 1 block A\n\nPage 1 block B"
-        );
-        assert_eq!(response["pages"][1]["markdown"], "Page 2 block A");
-        assert_eq!(response["pages"][2]["markdown"], "Page 3 block A");
-        assert_eq!(response["usage_info"]["pages_processed"], 3);
-        assert_eq!(response["usage_info"]["credits"], 3);
-        assert_eq!(response["provider_native_response"]["job_id"], "job_123");
-        let (upload_request, parse_request) = server.await.expect("server task completes");
-        assert!(
-            upload_request
-                .to_ascii_lowercase()
-                .contains("authorization: bearer test-key")
-        );
-        assert!(upload_request.contains("application/pdf"));
-        assert!(upload_request.contains("%PDF-1.4"));
-        assert!(upload_request.contains("x-trace-id: trace-1"));
-        assert!(
-            parse_request
-                .to_ascii_lowercase()
-                .contains("authorization: bearer test-key")
-        );
-        assert!(parse_request.contains(r#""input":"reducto://uploaded.pdf""#));
-        assert!(parse_request.contains(r#""table_output_format":"html""#));
-        assert!(parse_request.contains(r#""chunk_mode":"section""#));
-        assert!(parse_request.contains(r#""ocr_system":"standard""#));
+    #[test]
+    fn core_path_rejects_unwired_callbacks() {
+        let request = OcrRequest {
+            callbacks: vec![Arc::new(TestLogger)],
+            ..request()
+        };
+        let error = validate_host_hooks(&request).unwrap_err();
+        assert!(error.to_string().contains("callbacks are not wired"));
     }
 }

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

@@ -1,163 +0,0 @@
-use std::sync::atomic::{AtomicU64, Ordering};
-use std::time::{SystemTime, UNIX_EPOCH};
-
-use litellm_core::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider};
-use serde_json::{Map, Value};
-
-use super::common_utils::ocr_provider_config;
-use super::hooks::OcrLifecycleHooks;
-use super::types::{OcrRequest, PreparedOcrRequest};
-use crate::integrations::custom_guardrail::CustomGuardrailRunner;
-use crate::integrations::custom_logger::CustomLoggerRunner;
-
-pub(crate) struct PreparedOcrCall {
-    pub(crate) request: PreparedOcrRequest,
-    pub(crate) hooks: OcrLifecycleHooks,
-}
-
-#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
-pub(crate) fn prepare_ocr_call(request: OcrRequest<'_>) -> PreparedOcrCall {
-    let call_id = request
-        .litellm_call_id
-        .map(str::to_string)
-        .unwrap_or_else(new_ocr_call_id);
-    let provider_info = get_custom_llm_provider(request.model, request.custom_llm_provider)
-        .unwrap_or(CustomLlmProvider {
-            model: request.model,
-            custom_llm_provider: "mistral",
-        });
-    let model = provider_info.model.to_string();
-    let custom_llm_provider = provider_info.custom_llm_provider.to_string();
-    let config = ocr_provider_config(&custom_llm_provider, &model)
-        .ok_or_else(|| litellm_core::Error::InvalidProvider(custom_llm_provider.clone()))
-        .and_then(|config| {
-            validate_request_format(config, &request.optional_params, &custom_llm_provider)?;
-            Ok(config)
-        });
-    let optional_params = match &config {
-        Ok(config) => {
-            let supported = config.supported_ocr_params();
-            let mut mapped = config.map_ocr_params(
-                &request
-                    .optional_params
-                    .iter()
-                    .filter(|(name, _)| supported.contains(&name.as_str()))
-                    .map(|(name, value)| (name.clone(), value.clone()))
-                    .collect(),
-            );
-            for name in [
-                "vertex_project",
-                "vertex_ai_project",
-                "vertex_location",
-                "vertex_ai_location",
-            ] {
-                if let Some(value) = request.optional_params.get(name) {
-                    mapped.insert(name.to_string(), value.clone());
-                }
-            }
-            mapped
-        }
-        Err(_) => request.optional_params,
-    };
-
-    PreparedOcrCall {
-        request: PreparedOcrRequest {
-            config,
-            model,
-            custom_llm_provider,
-            litellm_call_id: call_id,
-            document: request.document,
-            api_key: request.api_key.map(str::to_string),
-            api_base: request.api_base.map(str::to_string),
-            extra_headers: request.extra_headers,
-            optional_params,
-            timeout: request.timeout,
-        },
-        hooks: OcrLifecycleHooks::new(
-            CustomLoggerRunner::new(request.callbacks),
-            CustomGuardrailRunner::new(request.guardrails),
-            request.request_metadata,
-        ),
-    }
-}
-
-fn validate_request_format(
-    config: &'static dyn litellm_core::ocr::transformation::OcrProviderConfig,
-    optional_params: &Map<String, Value>,
-    provider: &str,
-) -> Result<(), litellm_core::Error> {
-    let Some(format) = optional_params.get("req_format") else {
-        return Ok(());
-    };
-    match format.as_str() {
-        Some("litellm") => Ok(()),
-        Some("native") if config.supported_ocr_params().contains(&"req_format") => Ok(()),
-        Some("native") => Err(litellm_core::Error::InvalidRequest(format!(
-            "`req_format=native` is not supported for provider {provider}"
-        ))),
-        _ => Err(litellm_core::Error::InvalidRequest(format!(
-            "Invalid `req_format`: {format}. Expected `litellm` or `native`"
-        ))),
-    }
-}
-
-fn new_ocr_call_id() -> String {
-    static COUNTER: AtomicU64 = AtomicU64::new(1);
-    let sequence = COUNTER.fetch_add(1, Ordering::Relaxed);
-    let timestamp = SystemTime::now()
-        .duration_since(UNIX_EPOCH)
-        .map(|duration| duration.as_nanos())
-        .unwrap_or(0);
-    format!("ocr-{timestamp}-{sequence}")
-}
-
-#[cfg(test)]
-mod tests {
-    use litellm_core::error::Error;
-    use serde_json::{Map, json};
-
-    use super::{OcrRequest, prepare_ocr_call};
-    use crate::integrations::types::RequestMetadata;
-
-    fn base_ocr_request(model: &str) -> OcrRequest<'_> {
-        OcrRequest {
-            model,
-            document: json!({
-                "type": "document_url",
-                "document_url": "https://example.com/doc.pdf"
-            }),
-            api_key: Some("sk-test"),
-            api_base: None,
-            custom_llm_provider: None,
-            extra_headers: None,
-            optional_params: Map::new(),
-            timeout: None,
-            callbacks: Vec::new(),
-            guardrails: Vec::new(),
-            request_metadata: RequestMetadata::default(),
-            litellm_call_id: None,
-        }
-    }
-
-    fn request_with_format(format: &str) -> OcrRequest<'_> {
-        let mut request = base_ocr_request("mistral/mistral-ocr-latest");
-        request.optional_params = Map::from_iter([("req_format".to_string(), json!(format))]);
-        request
-    }
-
-    #[test]
-    fn native_format_rejected_for_provider_without_support_as_bad_request() {
-        let prepared = prepare_ocr_call(request_with_format("native"));
-        assert!(
-            matches!(prepared.request.config, Err(Error::InvalidRequest(message)) if message.contains("not supported for provider"))
-        );
-    }
-
-    #[test]
-    fn unknown_format_rejected_for_provider_without_support_as_bad_request() {
-        let prepared = prepare_ocr_call(request_with_format("raw"));
-        assert!(
-            matches!(prepared.request.config, Err(Error::InvalidRequest(message)) if message.contains("Invalid `req_format`"))
-        );
-    }
-}

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

@@ -1,8 +1,6 @@
 use std::sync::Arc;
 use std::time::Duration;
 
-use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleRequest};
-use litellm_core::ocr::transformation::OcrProviderConfig;
 use serde_json::{Map, Value};
 
 use crate::integrations::custom_guardrail::CustomGuardrail;
@@ -23,37 +21,3 @@ pub struct OcrRequest<'a> {
     pub request_metadata: RequestMetadata,
     pub litellm_call_id: Option<&'a str>,
 }
-
-pub(crate) struct PreparedOcrRequest {
-    pub(crate) config: Result<&'static dyn OcrProviderConfig, litellm_core::Error>,
-    pub(crate) model: String,
-    pub(crate) custom_llm_provider: String,
-    pub(crate) litellm_call_id: String,
-    pub(crate) document: Value,
-    pub(crate) api_key: Option<String>,
-    pub(crate) api_base: Option<String>,
-    pub(crate) extra_headers: Option<Map<String, Value>>,
-    pub(crate) optional_params: Map<String, Value>,
-    pub(crate) timeout: Option<Duration>,
-}
-
-impl CallLifecycleRequest for PreparedOcrRequest {
-    fn lifecycle_context(&self) -> CallLifecycleContext {
-        CallLifecycleContext::new(
-            "ocr",
-            self.model.clone(),
-            self.custom_llm_provider.clone(),
-            self.litellm_call_id.clone(),
-        )
-    }
-}
-
-pub(crate) struct ProviderOcrRequest {
-    pub(crate) model: String,
-    pub(crate) config: &'static dyn OcrProviderConfig,
-    pub(crate) url: String,
-    pub(crate) body: Value,
-    pub(crate) optional_params: Map<String, Value>,
-    pub(crate) upstream_headers: Vec<(String, String)>,
-    pub(crate) timeout: Option<Duration>,
-}

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

@@ -105,7 +105,11 @@ impl IntoResponse for MessagesRouteError {
                 StatusCode::NOT_FOUND,
                 "no messages deployment is configured for this model".to_string(),
             ),
-            Error::Auth(_) => (
+            Error::Auth(_)
+            | Error::MissingApiKey { .. }
+            | Error::MissingAzureAiCredentials
+            | Error::MissingAzureDocumentIntelligenceCredentials
+            | Error::MissingReductoApiKey => (
                 StatusCode::BAD_GATEWAY,
                 "messages provider authentication failed".to_string(),
             ),
@@ -114,7 +118,8 @@ impl IntoResponse for MessagesRouteError {
             | Error::Connect(_)
             | Error::InvalidResponse(_)
             | Error::InvalidType { .. }
-            | Error::MissingField(_) => (
+            | Error::MissingField(_)
+            | Error::MissingDocumentUrl => (
                 StatusCode::BAD_GATEWAY,
                 "messages provider request failed".to_string(),
             ),