chore(ui): resolve routing forecast merge conflict

.circleci/config.yml

@@ -1,10 +1,33 @@
 version: 2.1
+parameters:
+  run_migration_tests:
+    type: boolean
+    default: false
+  migration_candidate_image:
+    type: string
+    default: ""
+  migration_source_sha:
+    type: string
+    default: ""
 orbs:
   codecov: codecov/codecov@4.0.1
   node: circleci/node@5.1.0 # Add this line to declare the node orb
   win: circleci/windows@5.0 # Add Windows orb
 
 commands:
+  checkout_migration_source:
+    steps:
+      - run:
+          name: Select the requested migration test revision
+          environment:
+            MIGRATION_SOURCE_SHA: << pipeline.parameters.migration_source_sha >>
+          command: |
+            revision="${MIGRATION_SOURCE_SHA:-$CIRCLE_SHA1}"
+            [[ "$revision" =~ ^[0-9a-f]{40}$ ]] || exit 1
+            git init
+            git remote add origin https://github.com/BerriAI/litellm.git
+            git fetch --depth 1 origin "$revision"
+            git checkout --detach FETCH_HEAD
   skip_if_unrelated_changes:
     parameters:
       category:
@@ -1485,7 +1508,7 @@ jobs:
       - run:
           name: Run tests
           command: |
-            uv run --no-sync python -m pytest -vv tests/local_testing/test_basic_python_version.py -k "not v2_resolver"
+            uv run --no-sync python -m pytest -vv tests/local_testing/test_basic_python_version.py -k "not legacy_resolver"
 
   installing_litellm_on_python_3_13:
     docker:
@@ -1509,7 +1532,7 @@ jobs:
       - run:
           name: Run tests
           command: |
-            uv run --no-sync python -m pytest -v tests/local_testing/test_basic_python_version.py -k "not v2_resolver"
+            uv run --no-sync python -m pytest -v tests/local_testing/test_basic_python_version.py -k "not legacy_resolver"
 
   installing_litellm_on_python_v2_migration_resolver:
     docker:
@@ -1538,10 +1561,11 @@ jobs:
           url: tcp://localhost:5432
           timeout: "60"
       - run:
-          name: Run v2 migration resolver proxy smoke test
+          name: Run both migration resolvers against Postgres
           command: |
             uv run --no-sync python -m pytest -vv \
-              tests/local_testing/test_basic_python_version.py::test_litellm_proxy_server_config_no_general_settings_v2_resolver
+              tests/local_testing/test_basic_python_version.py::test_litellm_proxy_server_config_no_general_settings \
+              tests/local_testing/test_basic_python_version.py::test_litellm_proxy_server_config_no_general_settings_legacy_resolver
 
   helm_chart_testing:
     machine:
@@ -1650,6 +1674,7 @@ jobs:
           command: |
             docker run -d \
               -p 4001:4000 \
+              -e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \
               -e DATABASE_URL="postgresql://postgres:postgres@host.docker.internal:5432/litellm_test" \
               -e LITELLM_MASTER_KEY="sk-1234" \
               --name schema-seed \
@@ -1670,6 +1695,7 @@ jobs:
           command: |
             docker run -d \
               -p 4000:4000 \
+              -e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \
               -e DATABASE_URL="postgresql://postgres:postgres@host.docker.internal:5432/litellm_test" \
               -e DEFAULT_NUM_WORKERS_LITELLM_PROXY=1 \
               -e DISABLE_SCHEMA_UPDATE="True" \
@@ -1744,7 +1770,9 @@ jobs:
           command: |
             docker run -d \
               -p 4000:4000 \
+              -e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \
               -e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \
+              -e LITELLM_MASTER_KEY="sk-1234" \
               -e USE_PRISMA_MIGRATE=True \
               -e FAKE_OPENAI_API_BASE=http://host.docker.internal:8190 \
               -e AZURE_API_KEY=$AZURE_API_KEY \
@@ -1839,7 +1867,9 @@ jobs:
           command: |
             docker run -d \
               -p 4000:4000 \
+              -e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \
               -e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \
+              -e LITELLM_MASTER_KEY="sk-1234" \
               -e AZURE_API_KEY=$AZURE_API_KEY \
               -e AZURE_API_BASE=$AZURE_API_BASE \
               -e AZURE_API_VERSION="2024-05-01-preview" \
@@ -1927,6 +1957,7 @@ jobs:
           command: |
             docker run -d \
               -p 4000:4000 \
+              -e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \
               -e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \
               -e REDIS_HOST=$REDIS_HOST \
               -e REDIS_PASSWORD=$REDIS_PASSWORD \
@@ -1987,6 +2018,7 @@ jobs:
           command: |
             docker run -d \
               -p 4000:4000 \
+              -e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \
               -e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \
               -e REDIS_HOST=$REDIS_HOST \
               -e REDIS_PASSWORD=$REDIS_PASSWORD \
@@ -2064,6 +2096,7 @@ jobs:
           command: |
             docker run -d \
               -p 4000:4000 \
+              -e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \
               -e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \
               -e REDIS_HOST=host.docker.internal \
               -e REDIS_PORT=6379 \
@@ -2146,6 +2179,7 @@ jobs:
           command: |
             docker run -d \
               -p 4000:4000 \
+              -e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \
               -e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \
               -e REDIS_HOST=$REDIS_HOST \
               -e REDIS_PASSWORD=$REDIS_PASSWORD \
@@ -2168,6 +2202,7 @@ jobs:
           command: |
             docker run -d \
               -p 4001:4001 \
+              -e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \
               -e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \
               -e REDIS_HOST=$REDIS_HOST \
               -e REDIS_PASSWORD=$REDIS_PASSWORD \
@@ -2245,6 +2280,7 @@ jobs:
             docker run -d \
               --restart on-failure \
               -p 4000:4000 \
+              -e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \
               -e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \
               -e STORE_MODEL_IN_DB="True" \
               -e LITELLM_MASTER_KEY="sk-1234" \
@@ -2319,6 +2355,7 @@ jobs:
           command: |
             docker run -d \
               -p 4000:4000 \
+              -e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \
               -e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \
               -e REDIS_HOST=$REDIS_HOST \
               -e REDIS_PASSWORD=$REDIS_PASSWORD \
@@ -2401,6 +2438,7 @@ jobs:
           command: |
             docker run -d \
               -p 4000:4000 \
+              -e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \
               -e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \
               -e LITELLM_MASTER_KEY="sk-1234" \
               -e OPENAI_API_KEY=$OPENAI_API_KEY \
@@ -2492,6 +2530,7 @@ jobs:
           command: |
             docker run -d \
               -p 4000:4000 \
+              -e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \
               -e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \
               -e LITELLM_MASTER_KEY="sk-1234" \
               -e ANTHROPIC_API_KEY=$ANTHROPIC_API_KEY \
@@ -2673,6 +2712,7 @@ jobs:
           name: Start LiteLLM proxy
           environment:
             LITELLM_MASTER_KEY: "sk-1234"
+            LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY: "true"
             MOCK_LLM_URL: "http://127.0.0.1:8090/v1"
             DISABLE_SCHEMA_UPDATE: "true"
             SERVER_ROOT_PATH: ""
@@ -2816,6 +2856,7 @@ jobs:
           name: Start LiteLLM proxy under a server root path
           environment:
             LITELLM_MASTER_KEY: "sk-1234"
+            LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY: "true"
             MOCK_LLM_URL: "http://127.0.0.1:8090/v1"
             DISABLE_SCHEMA_UPDATE: "true"
           # Output flows to this step's own log, so a boot crash is visible here
@@ -2854,20 +2895,41 @@ jobs:
           destination: e2e-server-root-path-playwright-report
 
   build_docker_database_image:
+    parameters:
+      migration_qualification:
+        type: boolean
+        default: false
     machine:
       image: ubuntu-2204:2024.04.1
     resource_class: large
     working_directory: ~/project
     steps:
-      - checkout
-      - skip_if_unrelated_changes
+      - when:
+          condition: << parameters.migration_qualification >>
+          steps:
+            - checkout_migration_source
+      - unless:
+          condition: << parameters.migration_qualification >>
+          steps:
+            - checkout
+            - skip_if_unrelated_changes
 
       - run:
           name: Build Docker image
+          environment:
+            MIGRATION_CANDIDATE_IMAGE: << pipeline.parameters.migration_candidate_image >>
           command: |
-            docker build \
-              -t litellm-docker-database:ci \
-              -f docker/Dockerfile.database .
+            if [ -n "$MIGRATION_CANDIDATE_IMAGE" ]; then
+              [[ "$MIGRATION_CANDIDATE_IMAGE" =~ ^ghcr.io/berriai/[a-z0-9._/-]+@sha256:[0-9a-f]{64}$ ]] || exit 1
+              docker pull "$MIGRATION_CANDIDATE_IMAGE"
+              docker tag "$MIGRATION_CANDIDATE_IMAGE" litellm-docker-database:ci
+            else
+              docker build \
+                --label org.opencontainers.image.revision="$(git rev-parse HEAD)" \
+                -t litellm-docker-database:ci \
+                -f docker/Dockerfile.database .
+            fi
+            python3 .circleci/scripts/run_migration_tests.py record-image
 
       - run:
           name: Save Docker image to workspace root
@@ -2878,6 +2940,78 @@ jobs:
           root: .
           paths:
             - litellm-docker-database.tar.zst
+            - migration-image.json
+
+  migration_startup_tests:
+    parameters:
+      suite:
+        type: enum
+        enum: [startup, recovery, legacy]
+    machine:
+      image: ubuntu-2204:2024.04.1
+    resource_class: large
+    working_directory: ~/project
+    environment:
+      LITELLM_MIGRATION_TESTS: "1"
+      LITELLM_MIGRATION_TEST_IMAGE: litellm-docker-database:ci
+      MIGRATION_TEST_ADMIN_URL: postgresql://postgres:postgres@127.0.0.1:5432/postgres
+      MIGRATION_TEST_CONTAINER_ADMIN_URL: postgresql://postgres:postgres@host.docker.internal:5432/postgres
+      MIGRATION_TEST_OUTPUT: /tmp/migration-results
+      PYTHONPATH: tests/e2e
+    steps:
+      - checkout_migration_source
+      - install_uv
+      - install_rust
+      - restore_cache:
+          keys:
+            - v1-uv-cache-{{ checksum "uv.lock" }}
+      - run:
+          name: Install test dependencies
+          command: uv sync --frozen --all-groups --all-extras --python 3.12
+      - attach_workspace:
+          at: ~/project
+      - run:
+          name: Load the shared candidate and start PostgreSQL
+          command: |
+            zstd -d litellm-docker-database.tar.zst --stdout | docker load
+            docker run -d --name migration-postgres \
+              -e POSTGRES_USER=postgres -e POSTGRES_PASSWORD=postgres \
+              -p 5432:5432 \
+              postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5
+      - wait_for_service:
+          url: tcp://localhost:5432
+          timeout: "60"
+      - run:
+          name: Run migration startup regressions
+          environment:
+            MIGRATION_TEST_SUITE: << parameters.suite >>
+            MIGRATION_CANDIDATE_IMAGE: << pipeline.parameters.migration_candidate_image >>
+          command: |
+            mkdir -p /tmp/migration-results
+            uv run --no-sync python .circleci/scripts/run_migration_tests.py
+          no_output_timeout: 15m
+      - store_test_results:
+          path: /tmp/migration-results/junit
+      - run:
+          name: Package migration diagnostics
+          when: always
+          command: |
+            mkdir -p /tmp/migration-artifacts
+            if [ -d /tmp/migration-results ]; then
+              tar -czf /tmp/migration-artifacts/diagnostics.tar.gz -C /tmp/migration-results .
+              if [ -f /tmp/migration-results/verdict.json ]; then
+                cp /tmp/migration-results/verdict.json /tmp/migration-artifacts/verdict.json
+              fi
+            fi
+      - store_artifacts:
+          path: /tmp/migration-artifacts
+          destination: migration-results
+      - run:
+          name: Remove migration test containers
+          when: always
+          command: |
+            docker ps -aq --filter label=litellm-migration-test=true | xargs -r docker rm -f
+            docker rm -f migration-postgres || true
 
   test_bad_database_url:
     machine:
@@ -2901,6 +3035,7 @@ jobs:
           command: |
             docker run --name my-app \
               -p 4000:4000 \
+              -e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \
               -e DEFAULT_NUM_WORKERS_LITELLM_PROXY=1 \
               -e DATABASE_URL="postgresql://wrong:wrong@wrong:5432/wrong" \
               myapp:latest \
@@ -3008,8 +3143,61 @@ jobs:
       - store_artifacts:
           path: test-results
 
+  unit:
+    machine:
+      image: ubuntu-2204:2024.04.1
+    resource_class: large
+    working_directory: ~/project
+    steps:
+      - setup_litellm_test_deps
+      - run:
+          name: Generate Prisma client
+          command: uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
+      - run:
+          name: Run unit tests
+          command: |
+            mkdir -p test-results/unit
+            mapfile -t files < <(find tests/unit -name 'test_*.py' | sort)
+            if [ "${#files[@]}" -eq 0 ]; then echo "tests/unit holds no test_*.py files; nothing to run"; exit 0; fi
+            set +e
+            LITELLM_LOCAL_MODEL_COST_MAP=True uv run --no-sync pytest "${files[@]}" -p no:rerunfailures -p no:pytest-retry --timeout=90 -n 4 --dist=loadscope --tb=short --junitxml=test-results/unit/junit.xml
+            status=$?
+            set -e
+            if [ "$status" -eq 5 ]; then echo "pytest collected no tests from tests/unit; passing"; exit 0; fi
+            exit "$status"
+      - store_test_results:
+          path: test-results
+      - store_artifacts:
+          path: test-results
+
 workflows:
+  migration_startup:
+    when: << pipeline.parameters.run_migration_tests >>
+    jobs: &migration_jobs
+      - build_docker_database_image:
+          migration_qualification: true
+      - migration_startup_tests:
+          name: migration-startup
+          suite: startup
+          requires: [build_docker_database_image]
+      - migration_startup_tests:
+          name: migration-recovery
+          suite: recovery
+          requires: [build_docker_database_image]
+      - migration_startup_tests:
+          name: migration-legacy-and-pooling
+          suite: legacy
+          requires: [build_docker_database_image]
+  migration_startup_scheduled:
+    triggers:
+      - schedule:
+          cron: "17 0,6,12,18 * * *"
+          filters:
+            branches:
+              only: litellm_internal_staging
+    jobs: *migration_jobs
   integration:
+    unless: << pipeline.parameters.run_migration_tests >>
     jobs:
       - integration_contracts:
           name: integration-<< matrix.suite >>
@@ -3022,6 +3210,7 @@ workflows:
                 - main
                 - /litellm_.*/
   build_and_test:
+    unless: << pipeline.parameters.run_migration_tests >>
     jobs:
       - using_litellm_on_windows:
           filters: &main_branches
@@ -3029,6 +3218,8 @@ workflows:
               only:
                 - main
                 - /litellm_.*/
+      - unit:
+          filters: *main_branches
       - provider_replay_harness
       - base_sdk_install:
           filters: *main_branches

.circleci/scripts/run_migration_tests.py

@@ -0,0 +1,113 @@
+from __future__ import annotations
+
+import json
+import os
+import re
+import subprocess
+import sys
+from pathlib import Path
+from typing import Final
+from xml.etree import ElementTree
+
+SUITES: Final = {
+    "startup": (("test_startup.py",), 12),
+    "recovery": (("test_recovery.py",), 15),
+    "legacy": (("test_legacy.py", "test_pooling.py"), 11),
+}
+
+
+def successful_junit(path: Path, expected: int, exit_code: int) -> bool:
+    if exit_code != 0 or not path.is_file():
+        return False
+    try:
+        root: Final = ElementTree.parse(path).getroot()
+    except ElementTree.ParseError:
+        return False
+    cases: Final = tuple(root.iter("testcase"))
+    identities: Final = frozenset((case.get("classname"), case.get("name")) for case in cases)
+    return len(cases) == len(identities) == expected and all(
+        not any(case.find(tag) is not None for tag in ("failure", "error", "skipped")) for case in cases
+    )
+
+
+def output(*command: str) -> str:
+    return subprocess.check_output(command, text=True, timeout=90).strip()
+
+
+def record_image() -> None:
+    source: Final = output("git", "rev-parse", "HEAD")
+    image: Final = output("docker", "image", "inspect", "litellm-docker-database:ci", "--format", "{{.Id}}")
+    revision: Final = output(
+        "docker",
+        "image",
+        "inspect",
+        "litellm-docker-database:ci",
+        "--format",
+        '{{index .Config.Labels "org.opencontainers.image.revision"}}',
+    )
+    assert re.fullmatch(r"[0-9a-f]{40}", source), "Invalid source revision"
+    assert revision == source, "Candidate image revision differs from the tested source"
+    Path("migration-image.json").write_text(
+        json.dumps(
+            {
+                "source_sha": source,
+                "image_id": image,
+                "candidate_image": os.environ.get("MIGRATION_CANDIDATE_IMAGE", ""),
+            }
+        )
+    )
+
+
+def main() -> int:
+    suite: Final = os.environ["MIGRATION_TEST_SUITE"]
+    files, expected = SUITES[suite]
+    metadata: Final = json.loads(Path("migration-image.json").read_text())
+    assert metadata["source_sha"] == output("git", "rev-parse", "HEAD"), "Image and test source revisions differ"
+    assert metadata["image_id"] == output(
+        "docker", "image", "inspect", os.environ["LITELLM_MIGRATION_TEST_IMAGE"], "--format", "{{.Id}}"
+    ), "Loaded image differs from the build output"
+    assert metadata["candidate_image"] == os.environ.get("MIGRATION_CANDIDATE_IMAGE", ""), "Wrong release candidate"
+    destination: Final = Path(os.environ["MIGRATION_TEST_OUTPUT"])
+    junit: Final = destination / "junit" / "results.xml"
+    junit.parent.mkdir(parents=True, exist_ok=True)
+    result: Final = subprocess.run(
+        (
+            sys.executable,
+            "-m",
+            "pytest",
+            *(f"tests/e2e/migrations/{name}" for name in files),
+            "-vv",
+            "--tb=short",
+            "--durations=10",
+            f"--junitxml={junit}",
+            "-o",
+            "addopts=",
+            "--reruns=0",
+        ),
+        check=False,
+        timeout=1200,
+    )
+    passed: Final = successful_junit(junit, expected, result.returncode)
+    (destination / "verdict.json").write_text(
+        json.dumps(
+            {
+                **metadata,
+                "suite": suite,
+                "expected_cases": expected,
+                "passed": passed,
+                "pytest_exit_code": result.returncode,
+                "test_revision": metadata["source_sha"],
+                "workflow_id": os.environ.get("CIRCLE_WORKFLOW_ID", ""),
+                "job_number": os.environ.get("CIRCLE_BUILD_NUM", ""),
+            },
+            indent=2,
+        )
+    )
+    return 0 if passed else 1
+
+
+if __name__ == "__main__":
+    if sys.argv[1:] == ["record-image"]:
+        record_image()
+    else:
+        raise SystemExit(main())

.env.example

@@ -26,6 +26,7 @@ NOVITA_API_KEY = ""
 INFINITY_API_KEY = ""
 
 # Development Configs
-LITELLM_MASTER_KEY = "sk-1234"
+# Generate one with: echo "LITELLM_MASTER_KEY=sk-$(openssl rand -hex 32)"
+LITELLM_MASTER_KEY = ""
 DATABASE_URL = "postgresql://llmproxy:dbpassword9090@db:5432/litellm"
 STORE_MODEL_IN_DB = "True"

.github/e2e-stack/assert_tests_ran.py

@@ -1,3 +1,5 @@
+import os
+import re
 import sys
 import xml.etree.ElementTree as ET
 from pathlib import Path
@@ -15,6 +17,7 @@ def main() -> int:
         _ = sys.stdout.write("::error::could not read the test execution report\n")
         return 1
     cases: Final = tuple(report.iter("testcase"))
+    expected_count: Final = os.environ.get("E2E_REQUIRED_TEST_COUNT")
     passed: Final = frozenset(
         case.get("file") for case in cases if all(case.find(tag) is None for tag in ("skipped", "failure", "error"))
     )
@@ -35,9 +38,22 @@ def main() -> int:
         skipped: Final = sum(case.get("file") == path and case.find("skipped") is not None for case in cases)
         _ = sys.stdout.write(f"{path}: {collected} collected, {skipped} skipped\n")
         for case in cases:
-            if case.get("file") != path or all(case.find(tag) is None for tag in ("failure", "error")):
+            if case.get("file") != path or all(case.find(tag) is None for tag in ("failure", "error", "skipped")):
                 continue
-            _ = sys.stdout.write(f"  failed: {case.get('classname', '')}::{case.get('name', '')}\n")
+            outcome = "skipped" if case.find("skipped") is not None else "failed"
+            _ = sys.stdout.write(f"  {outcome}: {case.get('classname', '')}::{case.get('name', '')}\n")
+            for prop in case.findall("./properties/property"):
+                name = prop.get("name", "")
+                value = prop.get("value", "")
+                if name in ("oauth_failure_phase", "oauth_exception_type", "oauth_frame") and re.fullmatch(
+                    r"[A-Za-z0-9_.:<>-]{1,240}", value
+                ):
+                    _ = sys.stdout.write(f"    {name}: {value}\n")
+    if expected_count is not None and (
+        len(cases) != int(expected_count) or any(case.find("skipped") is not None for case in cases)
+    ):
+        _ = sys.stdout.write("::error::required test count was not met or a required case was skipped\n")
+        return 1
     if (
         selected
         and not missing

.github/e2e-stack/redact_output.py

@@ -0,0 +1,83 @@
+import argparse
+import os
+import sys
+from functools import reduce
+from pathlib import Path
+from typing import Final
+from xml.sax.saxutils import escape
+
+from pydantic import JsonValue, TypeAdapter, ValidationError
+from secrets_to_env import MIN_MASKED_LENGTH
+
+REDACTED: Final = "***"
+json_adapter: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
+
+
+def string_leaves(node: JsonValue) -> tuple[str, ...]:
+    match node:
+        case str():
+            return (node,)
+        case list():
+            return tuple(leaf for child in node for leaf in string_leaves(child))
+        case dict():
+            return tuple(leaf for child in node.values() for leaf in string_leaves(child))
+    return ()
+
+
+def field_lines(value: str) -> tuple[str, ...]:
+    try:
+        return tuple(line for leaf in string_leaves(json_adapter.validate_json(value)) for line in leaf.splitlines())
+    except ValidationError:
+        return ()
+
+
+def masked_values(values_files: tuple[Path, ...]) -> tuple[str, ...]:
+    values: Final = frozenset(
+        line.split("=", 1)[1].strip().strip("'")
+        for path in values_files
+        for line in path.read_text().splitlines()
+        if "=" in line
+    )
+    texts: Final = frozenset(text for value in values for text in (value, *field_lines(value)))
+    renderings: Final = frozenset(
+        rendering
+        for text in texts
+        if len(text) >= MIN_MASKED_LENGTH
+        for rendering in (text, escape(text), escape(text, {'"': """}))
+    )
+    return tuple(sorted(renderings, key=lambda rendering: (-len(rendering), rendering)))
+
+
+def redact(text: str, values: tuple[str, ...]) -> str:
+    return reduce(lambda redacted, value: redacted.replace(value, REDACTED), values, text)
+
+
+def write_redacted(source: Path, out_dir: Path, values: tuple[str, ...]) -> None:
+    target: Final = out_dir / source.name
+    with os.fdopen(os.open(target, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW, 0o600), "w") as handle:
+        _ = handle.write(redact(source.read_text(errors="replace"), values))
+
+
+def main() -> int:
+    parser: Final = argparse.ArgumentParser()
+    _ = parser.add_argument("--values", action="append", type=Path, required=True)
+    _ = parser.add_argument("--out", type=Path, required=True)
+    _ = parser.add_argument("files", nargs="*", type=Path)
+    args: Final = parser.parse_args()
+    values_files: Final = tuple(args.values)
+    out_dir: Final[Path] = args.out
+    sources: Final = tuple(args.files)
+    try:
+        values: Final = masked_values(values_files)
+        out_dir.mkdir(mode=0o700, exist_ok=True)
+        for source in sources:
+            write_redacted(source, out_dir, values)
+    except OSError as error:
+        _ = sys.stderr.write(f"could not redact {error.filename}\n")
+        return 1
+    _ = sys.stdout.write(f"redacted {len(sources)} file(s) into {out_dir}\n")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

.github/e2e-stack/select_tests.py

@@ -4,7 +4,8 @@ from typing import Final
 
 SELECTABLE: Final = re.compile(r"^tests/e2e/([A-Za-z0-9_.-]+/)*test_[A-Za-z0-9_.-]+\.py$")
 UNSUPPORTED: Final = re.compile(
-    r"^tests/e2e/(ui|claude_code|load)/"
+    r"^tests/e2e/(ui|claude_code|load|migrations)/"
+    r"|^tests/e2e/mcp/test_mcp_oauth_happy_path_e2e\.py$"
     r"|^tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e\.py$"
     r"|^tests/e2e/batches/test_managed_files_enforcement_e2e\.py$"
     r"|^tests/e2e/guardrails/test_presidio_masking_e2e\.py$"

.github/e2e-stack/up.sh

@@ -143,7 +143,7 @@ env "${SERVER_ENV[@]}" uv run --no-sync python migrations/run.py >"${LOGS_DIR}/m
 
 start_server() {
   local name="$1"; shift
-  env "${SERVER_ENV[@]}" "$@" >"${LOGS_DIR}/${name}.log" 2>&1 &
+  env -u AWS_ROLE_NAME "${SERVER_ENV[@]}" "$@" >"${LOGS_DIR}/${name}.log" 2>&1 &
   echo $! > "${PIDS_DIR}/${name}.pid"
 }
 

.github/template.yaml

@@ -21,7 +21,7 @@ Parameters:
   WorkerConfigParameter:
     Type: String
     Description: Sample environment variable
-    Default: '{"model": null, "alias": null, "api_base": null, "api_version": "2023-07-01-preview", "debug": false, "temperature": null, "max_tokens": null, "request_timeout": 600, "max_budget": null, "telemetry": true, "drop_params": false, "add_function_to_prompt": false, "headers": null, "save": false, "config": null, "use_queue": false}'
+    Default: '{"model": null, "alias": null, "api_base": null, "api_version": "2023-07-01-preview", "debug": false, "temperature": null, "max_tokens": null, "request_timeout": 600, "max_budget": null, "drop_params": false, "add_function_to_prompt": false, "headers": null, "save": false, "config": null, "use_queue": false}'
 
 Resources:
   MyUrlFunctionPermissions:

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

@@ -165,33 +165,41 @@ jobs:
           DIST: ${{ inputs.dist }}
           COVERAGE_CORE: sysmon
         run: |
-          if [ "${WORKERS}" = "0" ]; then
-            uv run --no-sync pytest ${TEST_PATH:?} \
-              --tb=short -vv \
-              --maxfail="${MAX_FAILURES}" \
-              --reruns "${RERUNS}" \
-              --reruns-delay 1 \
-              --timeout="${TEST_TIMEOUT_SECONDS}" \
-              --rerun-except "from pytest-timeout" \
-              --durations=20 \
-              --cov=./litellm --cov=./enterprise/litellm_enterprise \
-              --cov-report=xml:coverage.xml \
-              --cov-config=pyproject.toml
-          else
-            uv run --no-sync pytest ${TEST_PATH:?} \
-              --tb=short -vv \
-              --maxfail="${MAX_FAILURES}" \
-              -n "${WORKERS}" \
-              --reruns "${RERUNS}" \
-              --reruns-delay 1 \
-              --timeout="${TEST_TIMEOUT_SECONDS}" \
-              --rerun-except "from pytest-timeout" \
-              --dist="${DIST}" \
-              --durations=20 \
-              --cov=./litellm --cov=./enterprise/litellm_enterprise \
-              --cov-report=xml:coverage.xml \
-              --cov-config=pyproject.toml
+          found_path=false
+          for path in ${TEST_PATH}; do
+            if [ -e "${path%%::*}" ]; then
+              found_path=true
+              break
+            fi
+          done
+          if [ "$found_path" = false ]; then
+            echo "No path in TEST_PATH exists (${TEST_PATH}); nothing to run"
+            exit 0
           fi
+          xdist_args=()
+          if [ "${WORKERS}" != "0" ]; then
+            xdist_args=(-n "${WORKERS}" --dist="${DIST}")
+          fi
+          set +e
+          uv run --no-sync pytest ${TEST_PATH:?} \
+            --tb=short -vv \
+            --maxfail="${MAX_FAILURES}" \
+            "${xdist_args[@]}" \
+            --reruns "${RERUNS}" \
+            --reruns-delay 1 \
+            --timeout="${TEST_TIMEOUT_SECONDS}" \
+            --rerun-except "from pytest-timeout" \
+            --durations=20 \
+            --cov=./litellm --cov=./enterprise/litellm_enterprise \
+            --cov-report=xml:coverage.xml \
+            --cov-config=pyproject.toml
+          status=$?
+          set -e
+          if [ "$status" -eq 5 ]; then
+            echo "pytest collected no tests from ${TEST_PATH}; passing"
+            exit 0
+          fi
+          exit "$status"
 
       - name: Save coverage report
         if: always() && steps.changes.outputs.decision != 'skip'

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

@@ -175,6 +175,8 @@ jobs:
         env:
           TESTS: ${{ needs.detect.outputs.tests }}
           E2E_FIXTURE_MODE: live
+          E2E_PROVIDER_EDGE_HOST_REACHABLE: '1'
+          COLUMNS: '400'
         run: |
           umask 077
           read -r -a test_files <<< "${TESTS}"
@@ -189,6 +191,7 @@ jobs:
             uv run --no-sync python .github/e2e-stack/assert_tests_ran.py "${report}" "${test_files[@]}"
             verified=$?
             set -e
+            grep -E '^(FAILED|ERROR) ' "${log}" || true
             grep -E '^=+ .* in [0-9.]+s( \([0-9:]+\))? =+$' "${log}" | tail -n 1
             echo "::endgroup::"
             if [ "${status}" = "5" ]; then
@@ -206,6 +209,24 @@ jobs:
             echo "pass ${pass} of 3 passed"
           done
 
+      - name: Redact the pytest output
+        if: always() && steps.boot.outcome == 'success'
+        run: |
+          umask 077
+          shopt -s nullglob
+          uv run --no-sync python .github/e2e-stack/redact_output.py \
+            --values tests/e2e/.env --values "${RUNNER_TEMP}/litellm-e2e-stack/stack.env" \
+            --out "${RUNNER_TEMP}/e2e-redacted" "${RUNNER_TEMP}"/e2e-pass-*.log "${RUNNER_TEMP}"/e2e-pass-*.xml
+
+      - name: Keep the redacted pytest output
+        if: always() && steps.boot.outcome == 'success'
+        uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1
+        with:
+          name: e2e-changed-pytest-output-${{ github.run_attempt }}
+          path: ${{ runner.temp }}/e2e-redacted
+          retention-days: 14
+          if-no-files-found: ignore
+
       - name: Stop the stack
         if: always() && steps.boot.outcome != 'skipped'
         run: bash .github/e2e-stack/down.sh
@@ -214,7 +235,7 @@ jobs:
         if: always()
         run: |
           rm -f tests/e2e/.env "${RUNNER_TEMP}/e2e-boot.log" "${RUNNER_TEMP}"/e2e-pass-*.log "${RUNNER_TEMP}"/e2e-pass-*.xml
-          rm -rf "${RUNNER_TEMP}/litellm-e2e-stack"
+          rm -rf "${RUNNER_TEMP}/litellm-e2e-stack" "${RUNNER_TEMP}/e2e-redacted"
 
   gate:
     name: e2e-changed-tests

.github/workflows/test-mcp-oauth-e2e.yml

@@ -0,0 +1,180 @@
+name: MCP OAuth happy path
+
+on:
+  pull_request:
+    paths:
+      - '.github/workflows/test-mcp-oauth-e2e.yml'
+      - '.github/e2e-stack/**'
+      - 'tests/e2e/*.py'
+      - 'tests/e2e/pytest.ini'
+      - 'tests/e2e/idp_realm.json'
+      - 'tests/e2e/mcp/**'
+      - 'litellm/experimental_mcp_client/**'
+      - 'litellm/proxy/_experimental/mcp_server/**'
+      - 'litellm/proxy/auth/**'
+      - 'litellm/proxy/management_endpoints/mcp_management_endpoints.py'
+      - 'litellm/proxy/_types.py'
+      - 'litellm/types/mcp_server/mcp_server_manager.py'
+      - 'litellm/proxy/management_endpoints/*sso*.py'
+      - 'litellm/proxy/management_endpoints/sso/**'
+      - 'litellm/proxy/common_utils/encrypt_decrypt_utils.py'
+      - 'litellm/proxy/proxy_server.py'
+      - 'litellm/proxy/schema.prisma'
+      - 'ui/litellm-dashboard/src/app/connect/**'
+      - 'ui/litellm-dashboard/src/app/mcp/oauth/**'
+      - 'pyproject.toml'
+      - 'uv.lock'
+  workflow_dispatch:
+
+permissions: {}
+
+concurrency:
+  group: mcp-oauth-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  oauth:
+    if: github.event_name == 'workflow_dispatch' || github.event.pull_request.head.repo.full_name == github.repository
+    runs-on: ubuntu-latest
+    environment: e2e-changed
+    timeout-minutes: 45
+    permissions:
+      contents: read
+      id-token: write
+    services:
+      postgres:
+        image: postgres:16.6
+        env:
+          POSTGRES_USER: litellm
+          POSTGRES_PASSWORD: dbpassword9090
+          POSTGRES_DB: litellm
+        ports:
+          - 5432:5432
+        options: >-
+          --health-cmd "pg_isready -U litellm"
+          --health-interval 5s
+          --health-timeout 5s
+          --health-retries 10
+    env:
+      DATABASE_HOST: 127.0.0.1
+      DATABASE_PORT: '5432'
+      DATABASE_USER: litellm
+      DATABASE_PASSWORD: dbpassword9090
+      DATABASE_NAME: litellm
+      DATABASE_URL: postgresql://litellm:dbpassword9090@127.0.0.1:5432/litellm
+      E2E_KEYCLOAK_URL: http://127.0.0.1:8081
+      E2E_KEYCLOAK_ADMIN_USER: admin
+      E2E_KEYCLOAK_ADMIN_PASSWORD: e2e-ephemeral-idp-not-a-secret
+      E2E_FIXTURE_MODE: live
+      E2E_PROVIDER_CACHE: '0'
+      E2E_MCP_OAUTH_LIVE: '1'
+      E2E_REQUIRED_TEST_COUNT: '4'
+    steps:
+      - name: Checkout the tested source
+        uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
+        with:
+          ref: ${{ github.event.pull_request.head.sha || github.sha }}
+          persist-credentials: false
+
+      - name: Require and materialize the upstream login
+        env:
+          STORAGE_STATE: ${{ secrets.E2E_LINEAR_STORAGE_STATE_B64 }}
+        run: |
+          umask 077
+          python3 - <<'PY'
+          import base64
+          import json
+          import os
+          import secrets
+          from pathlib import Path
+          encoded = os.environ.get("STORAGE_STATE", "")
+          if not encoded:
+              raise SystemExit("E2E_LINEAR_STORAGE_STATE_B64 is required; capture and provision a test-account login")
+          state = json.loads(base64.b64decode(encoded, validate=True))
+          if not isinstance(state, dict) or not state.get("cookies"):
+              raise SystemExit("The captured login must contain browser cookies")
+          directory = Path(os.environ["RUNNER_TEMP"]) / "mcp-oauth-private"
+          directory.mkdir(mode=0o700)
+          path = directory / "linear-state.json"
+          path.write_text(json.dumps(state))
+          with open(os.environ["GITHUB_ENV"], "a") as output:
+              output.write(f"E2E_LINEAR_STORAGE_STATE={path}\n")
+              for name in ("LITELLM_MASTER_KEY", "LITELLM_SALT_KEY"):
+                  value = "sk-e2e-" + secrets.token_hex(24)
+                  print(f"::add-mask::{value}")
+                  output.write(f"{name}={value}\n")
+          PY
+
+      - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
+        with:
+          python-version: '3.13'
+      - uses: ./.github/actions/setup-uv-with-retries
+        with:
+          version: '0.10.9'
+      - uses: ./.github/actions/cache-cargo-build
+      - name: Install the frozen E2E environment
+        run: |
+          .github/scripts/uv_sync_with_retries.sh --frozen --extra proxy --extra proxy-runtime --extra extra_proxy --group ci --group proxy-dev --group e2e-dev
+          uv run --no-sync python scripts/prisma_generate_if_needed.py
+          uv run --no-sync playwright install --with-deps chromium
+
+      - name: Configure license access
+        id: aws
+        uses: aws-actions/configure-aws-credentials@e7f100cf4c008499ea8adda475de1042d6975c7b # v6.2.0
+        with:
+          role-to-assume: ${{ vars.E2E_AWS_ROLE_TO_ASSUME }}
+          aws-region: us-east-1
+          role-session-name: mcp-oauth-${{ github.run_id }}
+          role-duration-seconds: 900
+          output-env-credentials: false
+          output-credentials: true
+      - name: Load the E2E license
+        env:
+          AWS_ACCESS_KEY_ID: ${{ steps.aws.outputs.aws-access-key-id }}
+          AWS_SECRET_ACCESS_KEY: ${{ steps.aws.outputs.aws-secret-access-key }}
+          AWS_SESSION_TOKEN: ${{ steps.aws.outputs.aws-session-token }}
+          AWS_DEFAULT_REGION: us-east-1
+        run: |
+          license="$(aws secretsmanager get-secret-value --secret-id litellm-e2e-changed-license --query SecretString --output text)"
+          test -n "${license}"
+          echo "::add-mask::${license}"
+          echo "LITELLM_LICENSE=${license}" >> "${GITHUB_ENV}"
+
+      - uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020 # v4.4.0
+        with:
+          node-version-file: ui/litellm-dashboard/.nvmrc
+      - name: Build the gateway consent UI at the tested commit
+        run: |
+          cd ui/litellm-dashboard
+          ../../scripts/with_dashboard_node.sh npm ci
+          ../../scripts/with_dashboard_node.sh npm run build
+          mkdir -p ../../litellm/proxy/_experimental/out
+          cp -r out/. ../../litellm/proxy/_experimental/out/
+          find ../../litellm/proxy/_experimental/out -name '*.html' ! -name index.html | while read -r page; do
+            mkdir -p "${page%.html}"
+            mv "${page}" "${page%.html}/index.html"
+          done
+
+      - name: Prepare the isolated database and IdP
+        run: |
+          umask 077
+          bash .github/e2e-stack/start-idp.sh
+          uv run --no-sync python migrations/run.py > "${RUNNER_TEMP}/mcp-oauth-private/migrations.log" 2>&1
+
+      - name: Run every required OAuth variant without retries
+        run: |
+          umask 077
+          uv run --no-sync pytest -c tests/e2e/pytest.ini tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py \
+            --rootdir=. --reruns 0 --tb=short -o junit_family=xunit1 \
+            --junitxml="${RUNNER_TEMP}/mcp-oauth-private/results.xml" \
+            > "${RUNNER_TEMP}/mcp-oauth-private/pytest.log" 2>&1
+      - name: Report JUnit results and reject skipped or missing cases
+        if: always()
+        run: |
+          uv run --no-sync python .github/e2e-stack/assert_tests_ran.py \
+            "${RUNNER_TEMP}/mcp-oauth-private/results.xml" tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py
+      - name: Remove private login and logs
+        if: always()
+        run: |
+          docker rm -f e2e-keycloak >/dev/null 2>&1 || true
+          rm -rf "${RUNNER_TEMP}/mcp-oauth-private"

.github/workflows/test-rust.yml

@@ -120,6 +120,20 @@ jobs:
 
       - run: cargo test --workspace --doc --locked
 
+      - name: Test token counter feature combinations
+        run: |
+          for features in '' fast huggingface tiktoken fast,huggingface fast,tiktoken huggingface,tiktoken fast,huggingface,tiktoken; do
+            cargo test -p litellm-token-counter --locked --no-default-features --features "$features"
+            cargo check -p litellm-python-bridge --locked --no-default-features --features "abi3${features:+,$features}"
+          done
+
+      - name: Test secret manager feature combinations
+        run: |
+          cargo test -p litellm-auth-gcp --locked --no-default-features
+          for features in '' aws google aws,google; do
+            cargo test -p litellm-secrets --locked --no-default-features --features "$features"
+          done
+
   rust-wheel:
     runs-on: ubuntu-latest
     timeout-minutes: 30

.github/workflows/test-unit.yml

@@ -51,7 +51,7 @@ jobs:
         include:
           - shard: mcp-integration
             artifact-name: mcp-integration
-            test-path: "tests/mcp_tests"
+            test-path: "tests/mcp_tests tests/test_litellm/experimental_mcp_client"
             workers: 2
             reruns: 0
             timeout-minutes: 20
@@ -113,7 +113,6 @@ jobs:
               tests/test_litellm/compression
               tests/test_litellm/containers
               tests/test_litellm/endpoints
-              tests/test_litellm/experimental_mcp_client
               tests/test_litellm/models
               tests/test_litellm/repositories
               tests/test_litellm/images

CONTRIBUTING.md

@@ -268,10 +268,13 @@ If you want to build the Docker image yourself:
 # Build using the non-root Dockerfile
 docker build -f docker/Dockerfile.non_root -t litellm_dev .
 
+# Generate a master key. Requests send it as the bearer token
+export LITELLM_MASTER_KEY="sk-$(openssl rand -hex 32)"
+
 # Run with your config
 docker run \
     -v $(pwd)/proxy_config.yaml:/app/config.yaml \
-    -e LITELLM_MASTER_KEY="sk-1234" \
+    -e LITELLM_MASTER_KEY \
     -p 4000:4000 \
     litellm_dev \
     --config /app/config.yaml --detailed_debug

README.md

@@ -168,7 +168,7 @@ from a2a.utils.constants import TransportProtocol
 from uuid import uuid4
 
 base_url = "http://localhost:4000/a2a/my-agent"  # LiteLLM proxy + agent name
-headers = {"Authorization": "Bearer sk-1234"}    # LiteLLM Virtual Key
+headers = {"Authorization": "Bearer <your-master-key>"}    # LiteLLM master key or a virtual key
 
 async with httpx.AsyncClient(headers=headers, timeout=60.0) as http_client:
     resolver = A2ACardResolver(httpx_client=http_client, base_url=base_url)
@@ -233,7 +233,7 @@ async with stdio_client(server_params) as (read, write):
 
 ```bash
 curl -X POST 'http://0.0.0.0:4000/v1/chat/completions' \
-  -H 'Authorization: Bearer sk-1234' \
+  -H 'Authorization: Bearer <your-master-key>' \
   -H 'Content-Type: application/json' \
   -d '{
     "model": "gpt-4o",
@@ -255,7 +255,7 @@ curl -X POST 'http://0.0.0.0:4000/v1/chat/completions' \
     "LiteLLM": {
       "url": "http://localhost:4000/mcp/",
       "headers": {
-        "x-litellm-api-key": "Bearer sk-1234"
+        "x-litellm-api-key": "Bearer <your-master-key>"
       }
     }
   }

backend/routes/allowlist.py

@@ -51,6 +51,7 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = (
     "/cache_settings",
     "/coordination_redis/",
     "/cost_tracking",
+    "/cost_optimization/",
     "/cost/",
     "/credentials",
     "/credential",

cookbook/livekit_agent_sdk/config.example.yaml

@@ -15,7 +15,6 @@ model_list:
 
 litellm_settings:
   drop_params: True
-  telemetry: False
 
 general_settings:
   master_key: sk-1234  # Change this to a secure key

cookbook/misc/config.yaml

@@ -55,7 +55,6 @@ litellm_settings:
   # budget_duration: 30d
   num_retries: 5
   request_timeout: 600
-  telemetry: False
   context_window_fallbacks: [{"gpt-3.5-turbo": ["gpt-3.5-turbo-large"]}]
 
 general_settings: 

docker/.env.example

@@ -3,7 +3,8 @@
 # YOU MUST CHANGE THESE BEFORE GOING INTO PRODUCTION
 ############
 
-LITELLM_MASTER_KEY="sk-1234"
+# Generate one with: echo "LITELLM_MASTER_KEY=sk-$(openssl rand -hex 32)"
+LITELLM_MASTER_KEY=""
 
 ############
 # Database - You can change these to any PostgreSQL database that has logical replication enabled.

helm/litellm/templates/backend/deployment.yaml

@@ -7,6 +7,9 @@ metadata:
     {{- include "litellm.commonLabels" . | nindent 4 }}
     app.kubernetes.io/component: backend
 spec:
+  {{- if and (not .Values.backend.hpa.enabled) (not (kindIs "invalid" .Values.backend.replicaCount)) }}
+  replicas: {{ .Values.backend.replicaCount }}
+  {{- end }}
   {{- with .Values.backend.strategy }}
   strategy:
     {{- toYaml . | nindent 4 }}

helm/litellm/templates/gateway/deployment.yaml

@@ -7,6 +7,9 @@ metadata:
     {{- include "litellm.commonLabels" . | nindent 4 }}
     app.kubernetes.io/component: gateway
 spec:
+  {{- if and (not .Values.gateway.hpa.enabled) (not (kindIs "invalid" .Values.gateway.replicaCount)) }}
+  replicas: {{ .Values.gateway.replicaCount }}
+  {{- end }}
   {{- with .Values.gateway.strategy }}
   strategy:
     {{- toYaml . | nindent 4 }}

helm/litellm/templates/ui/deployment.yaml

@@ -7,6 +7,9 @@ metadata:
     {{- include "litellm.commonLabels" . | nindent 4 }}
     app.kubernetes.io/component: ui
 spec:
+  {{- if and (not .Values.ui.hpa.enabled) (not (kindIs "invalid" .Values.ui.replicaCount)) }}
+  replicas: {{ .Values.ui.replicaCount }}
+  {{- end }}
   {{- with .Values.ui.strategy }}
   strategy:
     {{- toYaml . | nindent 4 }}

helm/litellm/tests/replica_count_tests.yaml

@@ -0,0 +1,100 @@
+suite: test fixed replica count when HPA is disabled
+templates:
+  - gateway/deployment.yaml
+  - gateway/configmap.yaml
+  - backend/deployment.yaml
+  - ui/deployment.yaml
+values:
+  - ./values/required.yaml
+tests:
+  - it: gateway renders replicaCount into spec.replicas when its HPA is disabled
+    template: gateway/deployment.yaml
+    set:
+      gateway.hpa.enabled: false
+      gateway.replicaCount: 3
+    asserts:
+      - isKind:
+          of: Deployment
+      - equal:
+          path: spec.replicas
+          value: 3
+
+  - it: backend renders replicaCount into spec.replicas when its HPA is disabled
+    template: backend/deployment.yaml
+    set:
+      backend.hpa.enabled: false
+      backend.replicaCount: 2
+    asserts:
+      - equal:
+          path: spec.replicas
+          value: 2
+
+  - it: ui renders replicaCount into spec.replicas when its HPA is disabled
+    template: ui/deployment.yaml
+    set:
+      ui.hpa.enabled: false
+      ui.replicaCount: 2
+    asserts:
+      - equal:
+          path: spec.replicas
+          value: 2
+
+  - it: replicaCount 0 scales the gateway to zero instead of being treated as unset
+    template: gateway/deployment.yaml
+    set:
+      gateway.hpa.enabled: false
+      gateway.replicaCount: 0
+    asserts:
+      - equal:
+          path: spec.replicas
+          value: 0
+
+  - it: a component with HPA disabled but no replicaCount set keeps omitting spec.replicas, so upgrades do not reset a hand-scaled Deployment
+    set:
+      gateway.hpa.enabled: false
+      backend.hpa.enabled: false
+      ui.hpa.enabled: false
+    asserts:
+      - notExists:
+          path: spec.replicas
+        template: gateway/deployment.yaml
+      - notExists:
+          path: spec.replicas
+        template: backend/deployment.yaml
+      - notExists:
+          path: spec.replicas
+        template: ui/deployment.yaml
+
+  - it: every component omits spec.replicas when its HPA is enabled, so the autoscaler owns the count
+    set:
+      gateway.hpa.enabled: true
+      gateway.replicaCount: 3
+      backend.hpa.enabled: true
+      backend.replicaCount: 3
+      ui.hpa.enabled: true
+      ui.replicaCount: 3
+    asserts:
+      - notExists:
+          path: spec.replicas
+        template: gateway/deployment.yaml
+      - notExists:
+          path: spec.replicas
+        template: backend/deployment.yaml
+      - notExists:
+          path: spec.replicas
+        template: ui/deployment.yaml
+
+  - it: a component with HPA disabled renders replicas while a sibling with HPA enabled does not
+    set:
+      gateway.hpa.enabled: false
+      gateway.replicaCount: 4
+      backend.hpa.enabled: true
+      backend.replicaCount: 4
+    asserts:
+      - equal:
+          path: spec.replicas
+          value: 4
+        template: gateway/deployment.yaml
+      - notExists:
+          path: spec.replicas
+        template: backend/deployment.yaml

helm/litellm/values.yaml

@@ -397,6 +397,11 @@ gateway:
   #   failureThreshold: 30
   #   periodSeconds: 10
   startupProbe: {}
+  # Optional fixed pod count, rendered into the Deployment's spec.replicas only
+  # when hpa.enabled is false. Unset by default so an existing Deployment keeps
+  # its current count; with the HPA on, the autoscaler owns the count, e.g.:
+  #   replicaCount: 3
+  replicaCount:
   hpa:
     enabled: true
     minReplicas: 1
@@ -524,6 +529,8 @@ backend:
   strategy: {}
   # Optional startupProbe; same shape as gateway.startupProbe. Empty by default.
   startupProbe: {}
+  # Same semantics as gateway.replicaCount.
+  replicaCount:
   hpa:
     enabled: true
     minReplicas: 1
@@ -590,6 +597,8 @@ ui:
   strategy: {}
   # Optional startupProbe; same shape as gateway.startupProbe. Empty by default.
   startupProbe: {}
+  # Same semantics as gateway.replicaCount.
+  replicaCount:
   hpa:
     enabled: false
     minReplicas: 1

litellm-proxy-extras/litellm_proxy_extras/migration_lock.py

@@ -0,0 +1,89 @@
+import random
+import time
+from collections.abc import Generator, Mapping
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Final
+from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
+
+from litellm_proxy_extras._logging import logger
+from litellm_proxy_extras.prisma_toolchain import MIGRATION_LOCK_TIMEOUT_ENV_VAR, migration_lock_timeout
+
+MIGRATION_LOCK_KEY: Final = int.from_bytes(b"llm_mig2", "big")
+
+if TYPE_CHECKING:
+    import psycopg
+
+
+def migration_environment(environment: Mapping[str, str]) -> Mapping[str, str]:
+    database_url: Final = environment.get("DATABASE_URL")
+    direct_url: Final = environment.get("DIRECT_URL")
+    if not database_url or not direct_url:
+        return environment
+    schema: Final = next((value for key, value in parse_qsl(urlsplit(database_url).query) if key == "schema"), "public")
+    direct: Final = urlsplit(direct_url)
+    parameters: Final = tuple((key, value) for key, value in parse_qsl(direct.query) if key != "schema")
+    return {
+        **environment,
+        "DATABASE_URL": urlunsplit(direct._replace(query=urlencode((*parameters, ("schema", schema))))),
+    }
+
+
+@dataclass(frozen=True, slots=True)
+class _LockResult:
+    acquired: bool
+
+
+def _try_lock(connection: "psycopg.Connection[tuple[object, ...]]", key: int = MIGRATION_LOCK_KEY) -> bool:
+    from psycopg.rows import class_row
+
+    with connection.cursor(row_factory=class_row(_LockResult)) as cursor:
+        row: Final = cursor.execute("SELECT pg_try_advisory_xact_lock(%s) AS acquired", (key,)).fetchone()
+    return row is not None and row.acquired
+
+
+@dataclass(frozen=True, slots=True)
+class MigrationCoordinator:
+    connection: "psycopg.Connection[tuple[object, ...]]"
+
+    def check_connection(self) -> None:
+        self.connection.execute("SELECT 1")
+
+    def acquire_prisma_lock(self) -> None:
+        deadline: Final = time.monotonic() + migration_lock_timeout()
+        while time.monotonic() < deadline:
+            if _try_lock(self.connection, 72707369):
+                return
+            time.sleep(min(random.uniform(0.5, 1.5), max(0.0, deadline - time.monotonic())))
+        raise RuntimeError(
+            "Timed out waiting for Prisma's lock to recover migration history. LiteLLM startup has stopped. "
+            "Another migration or a pooled database session may still hold the lock. Check the database lock holder. "
+            "When using a transaction pooler, configure DIRECT_URL to reach the same database without the pooler."
+        )
+
+
+@contextmanager
+def migration_lock(database_url: str) -> Generator[MigrationCoordinator, None, None]:
+    import psycopg
+
+    wait_seconds: Final = migration_lock_timeout()
+    deadline: Final = time.monotonic() + wait_seconds
+    try:
+        with psycopg.connect(database_url, connect_timeout=10, autocommit=True) as connection:
+            coordinator: Final = MigrationCoordinator(connection)
+            logger.info("Waiting for the v2 migration coordinator lock (up to %ss)", wait_seconds)
+            while time.monotonic() < deadline:
+                with connection.transaction():
+                    if _try_lock(connection):
+                        logger.info("Acquired the v2 migration coordinator lock")
+
+                        yield coordinator
+                        coordinator.check_connection()
+                        return
+                time.sleep(min(random.uniform(0.5, 1.5), max(0.0, deadline - time.monotonic())))
+    except psycopg.Error as exc:
+        raise RuntimeError(f"Lost or could not establish v2 migration coordination with the database: {exc}") from exc
+    raise RuntimeError(
+        f"Timed out waiting for another v2 migration resolver after {wait_seconds}s. "
+        f"Check the running migration or increase {MIGRATION_LOCK_TIMEOUT_ENV_VAR}."
+    )

litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py

@@ -0,0 +1,158 @@
+import hashlib
+import subprocess
+from collections.abc import Mapping
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Final
+from uuid import uuid4
+
+from litellm_proxy_extras import prisma_toolchain
+from litellm_proxy_extras._logging import logger
+from litellm_proxy_extras.migration_lock import MigrationCoordinator
+
+if TYPE_CHECKING:
+    import psycopg
+
+
+@dataclass(frozen=True, slots=True)
+class MigrationProgress:
+    checksum: str
+    applied_steps_count: int
+    logs: str
+    id: str = ""
+    finished: bool = False
+
+    def confirms_completion(self, script: bytes) -> bool:
+        return (
+            self.applied_steps_count == 1
+            and not self.logs.strip()
+            and self.checksum == hashlib.sha256(script).hexdigest()
+        )
+
+
+def _migration_records(
+    connection: "psycopg.Connection[tuple[object, ...]]", schema: str, migration: Path
+) -> tuple[MigrationProgress, ...]:
+    from psycopg import sql
+    from psycopg.rows import class_row
+
+    with connection.cursor(row_factory=class_row(MigrationProgress)) as cursor:
+        records: Final = cursor.execute(
+            sql.SQL(
+                "SELECT id, checksum, applied_steps_count, coalesce(logs, '') AS logs, "
+                "finished_at IS NOT NULL AS finished FROM {} "
+                "WHERE migration_name = %s AND rolled_back_at IS NULL"
+            ).format(sql.Identifier(schema, "_prisma_migrations")),
+            (migration.parent.name,),
+        ).fetchall()
+    return tuple(records)
+
+
+def recover_completed_migration(coordinator: MigrationCoordinator, schema: str, migration: Path) -> bool:
+    """Finish a proven successful row without erasing its durable completion evidence.
+
+    The caller commits this checkpoint before running another Prisma command.
+    """
+    from psycopg import sql
+
+    coordinator.acquire_prisma_lock()
+    records: Final = _migration_records(coordinator.connection, schema, migration)
+    unfinished: Final = tuple(record for record in records if not record.finished)
+    script: Final = migration.read_bytes()
+    if not unfinished:
+        return any(record.checksum == hashlib.sha256(script).hexdigest() for record in records)
+    if len(unfinished) != 1 or not unfinished[0].confirms_completion(script):
+        return False
+    progress: Final = unfinished[0]
+    result: Final = coordinator.connection.execute(
+        sql.SQL(
+            "UPDATE {} SET finished_at = current_timestamp "
+            "WHERE id = %s AND checksum = %s AND applied_steps_count = 1 "
+            "AND finished_at IS NULL AND rolled_back_at IS NULL AND coalesce(logs, '') = %s"
+        ).format(sql.Identifier(schema, "_prisma_migrations")),
+        (progress.id, progress.checksum, progress.logs),
+    )
+    if result.rowcount != 1:
+        raise RuntimeError("Could not complete the confirmed migration history row; retry startup.")
+    logger.info("Completed migration %s using its successful SQL step and matching checksum", migration.parent.name)
+    return True
+
+
+def migration_files(directory: Path) -> tuple[tuple[str, str], ...]:
+    return tuple(
+        (path.parent.name, hashlib.sha256(path.read_bytes()).hexdigest())
+        for path in sorted((directory / "migrations").glob("*/migration.sql"))
+    )
+
+
+def baseline_current_schema(
+    coordinator: MigrationCoordinator,
+    schema: str,
+    migrations_dir: Path,
+    prisma_command: str,
+    prisma_env: Mapping[str, str],
+) -> None:
+    from psycopg import sql
+
+    packaged_dir: Final = Path(__file__).parent
+    migrations: Final = migration_files(migrations_dir)
+    if (
+        not migrations
+        or migrations != migration_files(packaged_dir)
+        or (migrations_dir / "schema.prisma").read_bytes() != (packaged_dir / "schema.prisma").read_bytes()
+    ):
+        raise RuntimeError("Cannot automatically baseline an existing database with custom migration history.")
+
+    coordinator.acquire_prisma_lock()
+    existing: Final = coordinator.connection.execute(
+        "SELECT to_regclass(%s)", (sql.Identifier(schema, "_prisma_migrations").as_string(coordinator.connection),)
+    ).fetchone()
+    if existing is not None and existing[0] is not None:
+        return
+    try:
+        prisma_toolchain.run_prisma(
+            (
+                prisma_command,
+                "migrate",
+                "diff",
+                "--from-schema-datasource",
+                str(migrations_dir / "schema.prisma"),
+                "--to-schema-datamodel",
+                str(migrations_dir / "schema.prisma"),
+                "--exit-code",
+            ),
+            timeout=prisma_toolchain.prisma_command_timeout(),
+            env=prisma_env,
+        )
+    except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as exc:
+        raise RuntimeError(
+            "Cannot automatically baseline this database: its schema has not been verified to match this build. "
+            "Establish the existing migration history before retrying. No schema reconciliation was performed. "
+            "If using a transaction pooler, configure DIRECT_URL to reach the same database without the pooler. "
+            f"Schema verification detail: {exc.stderr}"
+        ) from exc
+
+    coordinator.check_connection()
+    ledger: Final = sql.Identifier(schema, "_prisma_migrations")
+    coordinator.connection.execute(
+        sql.SQL(
+            "CREATE TABLE {} (id varchar(36) PRIMARY KEY NOT NULL, checksum varchar(64) NOT NULL, "
+            "finished_at timestamptz, migration_name varchar(255) NOT NULL, logs text, rolled_back_at timestamptz, "
+            "started_at timestamptz NOT NULL DEFAULT now(), applied_steps_count integer NOT NULL DEFAULT 0)"
+        ).format(ledger)
+    )
+    with coordinator.connection.cursor() as cursor:
+        cursor.executemany(
+            sql.SQL(
+                "INSERT INTO {} (id, checksum, migration_name, logs, started_at, finished_at) "
+                "VALUES (%s, %s, %s, '', current_timestamp, current_timestamp)"
+            ).format(ledger),
+            tuple((str(uuid4()), checksum, name) for name, checksum in migrations),
+        )
+    logger.warning(
+        "Legacy migration history was missing. The existing Prisma schema matches this build; "
+        "adopted %s packaged migrations as a baseline. No schema changes were applied, and "
+        "historical data backfills were not replayed or verified. Continuing startup; "
+        "review any feature-specific backfill requirements.",
+        len(migrations),
+    )

litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py

@@ -59,6 +59,7 @@ except ImportError:
 PRISMA_COMMAND_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_COMMAND_TIMEOUT"
 PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_BOOTSTRAP_TIMEOUT"
 PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_MIGRATE_DEPLOY_TIMEOUT"
+MIGRATION_LOCK_TIMEOUT_ENV_VAR = "LITELLM_MIGRATION_LOCK_TIMEOUT"
 NODEENV_CACHE_DIR_ENV_VAR = "PRISMA_NODEENV_CACHE_DIR"
 
 DEFAULT_PRISMA_COMMAND_TIMEOUT = 60.0
@@ -106,6 +107,10 @@ def prisma_command_timeout() -> float:
     )
 
 
+def migration_lock_timeout() -> float:
+    return _timeout_from_env(MIGRATION_LOCK_TIMEOUT_ENV_VAR, 600.0)
+
+
 def prisma_bootstrap_timeout() -> float:
     """Seconds the one-time Node toolchain install may run for."""
     return _timeout_from_env(

litellm-proxy-extras/litellm_proxy_extras/utils.py

@@ -6,6 +6,7 @@ import shutil
 import subprocess
 import tempfile
 import time
+from collections.abc import Callable
 from dataclasses import dataclass, replace
 from pathlib import Path
 from typing import TYPE_CHECKING, Final, Optional
@@ -78,15 +79,10 @@ MAX_MIGRATE_DEPLOY_ATTEMPTS = 4
 
 @dataclass(frozen=True)
 class _MigrateAttemptBudget:
-    """Retries left, and the recoveries already run.
-
-    A recovery that lands something new costs nothing, so a database full of
-    objects `prisma db push` created works through them one per pass. Anything
-    that made no progress spends an attempt, so a stuck run still gives up.
-    """
+    """Independent bounds for failed attempts and Prisma lock contention."""
 
     attempts_left: int
-    recoveries: frozenset[str] = frozenset()
+    contention_seconds_left: float = 600.0
 
     @property
     def exhausted(self) -> bool:
@@ -99,10 +95,14 @@ class _MigrateAttemptBudget:
     def spend(self) -> "_MigrateAttemptBudget":
         return replace(self, attempts_left=self.attempts_left - 1)
 
-    def after_recovery(self, recovery: str) -> "_MigrateAttemptBudget":
-        if recovery in self.recoveries:
-            return self.spend()
-        return replace(self, recoveries=self.recoveries | {recovery})
+    def after_contention(self, elapsed: float) -> "_MigrateAttemptBudget":
+        remaining: Final = self.contention_seconds_left - elapsed
+        if remaining <= 0:
+            raise RuntimeError(
+                "Timed out waiting for Prisma's migration advisory lock. Check the running migration "
+                "or increase LITELLM_MIGRATION_LOCK_TIMEOUT."
+            )
+        return replace(self, contention_seconds_left=remaining)
 
 
 _SPEND_LOGS_ALTER_RE = re.compile(r'^ALTER\s+TABLE\s+"LiteLLM_SpendLogs"\s', re.IGNORECASE)
@@ -836,12 +836,51 @@ class ProxyExtrasDBManager:
 
     @staticmethod
     def _setup_database_v2(use_migrate: bool) -> bool:
+        if not use_migrate:
+            return ProxyExtrasDBManager._run_database_v2(False)
+        from litellm_proxy_extras.migration_lock import migration_environment, migration_lock
+        from litellm_proxy_extras.migration_recovery import baseline_current_schema, recover_completed_migration
+
+        database_url: Final = os.environ.get("DATABASE_URL")
+        if not database_url:
+            raise RuntimeError("DATABASE_URL is required for v2 migrations")
+        lock_url: Final = ProxyExtrasDBManager._strip_prisma_query_params(os.environ.get("DIRECT_URL") or database_url)
+        schema: Final = ProxyExtrasDBManager._prisma_schema_param(database_url) or "public"
+
+        def recover_completed(name: str) -> bool:
+            if Path(name).name != name or "\\" in name:
+                return False
+            migration: Final = Path(os.getcwd()) / "migrations" / name / "migration.sql"
+            if not migration.is_file():
+                return False
+            with migration_lock(lock_url) as coordinator:
+                return recover_completed_migration(coordinator, schema, migration)
+
+        def baseline_existing(migrations_dir: str) -> None:
+            with migration_lock(lock_url) as coordinator:
+                baseline_current_schema(
+                    coordinator,
+                    schema,
+                    Path(migrations_dir),
+                    _get_prisma_command(),
+                    migration_environment(_get_prisma_env()),
+                )
+
+        while not ProxyExtrasDBManager._run_database_v2(True, recover_completed, baseline_existing):
+            continue
+        return True
+
+    @staticmethod
+    def _run_database_v2(
+        use_migrate: bool,
+        recover_completed: Callable[[str], bool] = lambda name: False,
+        baseline_existing: "Callable[[str], None] | None" = None,
+    ) -> bool:
         """
         v2 migration resolver (opt-in via --use_v2_migration_resolver).
 
-        Runs `prisma migrate deploy` and handles standard recovery paths
-        (P3005 baseline, P3009/P3018 idempotent errors, deadlocks against a
-        concurrent migrate deploy). Critically, it does
+        Runs `prisma migrate deploy`, baselines verified existing schemas,
+        and recovers confirmed SQL completion or reported deadlocks. It does
         NOT call `_resolve_all_migrations` — the diff-and-force recovery that
         caused schema thrashing when two LiteLLM versions contended for the
         same DB during rolling deploys.
@@ -850,10 +889,9 @@ class ProxyExtrasDBManager:
         is logged as a warning, not a fatal error — users whose DBs got into
         weird shapes from the old thrashing should still be able to start.
 
-        The retry budget only counts attempts that made no progress: see
-        _MigrateAttemptBudget.
+        False requests a committed recovery checkpoint and another deploy
+        pass. True means every pending migration is complete.
         """
-        schema_path = ProxyExtrasDBManager._get_prisma_dir() + "/schema.prisma"
         migrations_dir = ProxyExtrasDBManager._get_prisma_dir()
 
         if not use_migrate:
@@ -886,14 +924,22 @@ class ProxyExtrasDBManager:
         original_dir = os.getcwd()
         os.chdir(migrations_dir)
         deploy_timeout = prisma_migrate_deploy_timeout()
-        budget = _MigrateAttemptBudget(attempts_left=MAX_MIGRATE_DEPLOY_ATTEMPTS)
+        from litellm_proxy_extras.migration_lock import migration_environment, migration_lock_timeout
+
+        migration_env: Final = migration_environment(_get_prisma_env())
+
+        budget = _MigrateAttemptBudget(
+            attempts_left=MAX_MIGRATE_DEPLOY_ATTEMPTS,
+            contention_seconds_left=migration_lock_timeout(),
+        )
         try:
             while not budget.exhausted:
+                attempt_started = time.monotonic()
                 try:
                     result = prisma_toolchain.run_prisma(
                         [_get_prisma_command(), "migrate", "deploy"],
                         timeout=deploy_timeout,
-                        env=_get_prisma_env(),
+                        env=migration_env,
                     )
                     logger.info(f"prisma migrate deploy stdout: {result.stdout}")
                     return True
@@ -909,8 +955,16 @@ class ProxyExtrasDBManager:
                     next_budget = budget.spend()
 
                 except subprocess.CalledProcessError as e:
+                    if "P3005" in (e.stderr or "") and baseline_existing is not None:
+                        baseline_existing(migrations_dir)
+                        return False
+                    failed_migration = ProxyExtrasDBManager._v2_failed_migration_name(e.stderr or "")
+                    if failed_migration and recover_completed(failed_migration):
+                        return False
                     next_budget = ProxyExtrasDBManager._budget_after_deploy_failure(
-                        e, budget, schema_path
+                        e,
+                        budget,
+                        time.monotonic() - attempt_started,
                     )
 
                 if next_budget.attempts_left < budget.attempts_left:
@@ -919,19 +973,41 @@ class ProxyExtrasDBManager:
 
             raise RuntimeError(
                 f"Database migration failed after {MAX_MIGRATE_DEPLOY_ATTEMPTS} "
-                "attempts that made no progress (timeouts, deadlock retries, or a "
-                "recovery that had already run once). Check database connectivity, "
+                "attempts that made no progress (timeouts or deadlock retries). Check database connectivity, "
                 "load, and _prisma_migrations ledger state, and raise "
                 f"{PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR} if the attempts timed out."
             )
         finally:
             os.chdir(original_dir)
 
+    @staticmethod
+    def _v2_failed_migration_name(stderr: str) -> "str | None":
+        if "P3009" in stderr:
+            match = re.search(r"`(\d+_[^`\r\n]+)`", stderr)
+            return match.group(1) if match else None
+        if "P3018" in stderr:
+            match = re.search(r"Migration name: (\d+_[^\r\n]+)", stderr)
+            return match.group(1) if match else None
+        return None
+
+    @staticmethod
+    def _v2_roll_back_migration_best_effort(migration_name: str) -> None:
+        from litellm_proxy_extras.migration_lock import migration_environment
+
+        try:
+            prisma_toolchain.run_prisma(
+                [_get_prisma_command(), "migrate", "resolve", "--rolled-back", migration_name],
+                timeout=prisma_command_timeout(),
+                env=migration_environment(_get_prisma_env()),
+            )
+        except (subprocess.CalledProcessError, subprocess.TimeoutExpired):
+            pass
+
     @staticmethod
     def _budget_after_deploy_failure(
         error: subprocess.CalledProcessError,
         budget: "_MigrateAttemptBudget",
-        schema_path: str,
+        attempt_seconds: float = 0.0,
     ) -> "_MigrateAttemptBudget":
         """Recover from one failed `prisma migrate deploy`, and price the pass.
 
@@ -940,37 +1016,35 @@ class ProxyExtrasDBManager:
         """
         stderr = error.stderr or ""
 
-        if "P3005" in stderr and "database schema is not empty" in stderr:
-            logger.info("Schema exists but no migrations ledger — creating baseline")
-            if ProxyExtrasDBManager._create_baseline_migration(schema_path):
-                return budget.after_recovery("baseline")
-            return budget.spend()
-
         if "P3009" in stderr:
-            migration_match = re.search(r"`(\d+_\S+?)`", stderr)
-            if migration_match and ProxyExtrasDBManager._is_idempotent_error(stderr):
-                name = migration_match.group(1)
-                logger.info(
-                    f"Migration {name} failed idempotently — marking applied and retrying"
-                )
-                ProxyExtrasDBManager._mark_migration_applied(name)
-                return budget.after_recovery(f"resolved:{name}")
-            if migration_match:
-                migration_name = migration_match.group(1)
+            migration_name = ProxyExtrasDBManager._v2_failed_migration_name(stderr)
+            if migration_name:
                 ledger_logs = ProxyExtrasDBManager._failed_migration_logs(migration_name)
-                if ledger_logs is not None and (
-                    ledger_logs == "" or _MIGRATION_DEADLOCK_MARKER in ledger_logs
-                ):
+                if ledger_logs and _MIGRATION_DEADLOCK_MARKER in ledger_logs:
                     logger.info(
                         "Migration %s failed in a concurrent migrate deploy "
                         "deadlock race, rolling its ledger row back and retrying",
                         migration_name,
                     )
-                    ProxyExtrasDBManager._roll_back_migration_best_effort(migration_name)
+                    ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name)
                     return budget.spend()
             raise RuntimeError(
-                "Database migration failed and cannot be auto-recovered. "
-                f"Manual intervention required.\n\nPrisma error:\n{stderr}"
+                "Migration completion could not be verified. LiteLLM startup has stopped.\n\n"
+                f"Prisma migration history (migration name and start time):\n{stderr}\n\n"
+                "A migration has a start record but no successful completion record. "
+                "LiteLLM cannot determine whether its SQL committed from this record alone. "
+                "Startup stopped to avoid repeating or skipping database changes.\n\n"
+                "Before resolving, stop other migration runners and inspect _prisma_migrations, "
+                "the named migration.sql from this build, database logs, and the actual database objects and data. "
+                "Use the same database and this build's schema and migration files for recovery:\n"
+                "- Only after verifying every migration change is present, run "
+                "prisma migrate resolve --applied <migration_name>, then retry startup.\n"
+                "- Only after verifying no migration changes remain (or fully undoing partial changes), run "
+                "prisma migrate resolve --rolled-back <migration_name>, then retry startup. "
+                "This command updates history; it does not undo SQL.\n"
+                "Replace <migration_name> with the reported name. If the outcome remains uncertain, "
+                "leave migration history unchanged and contact your database administrator. "
+                "Repeated restarts alone will not resolve this state."
             ) from error
 
         if "P3018" in stderr:
@@ -981,25 +1055,14 @@ class ProxyExtrasDBManager:
                     f"and retry.\n\nPrisma error:\n{stderr}"
                 ) from error
 
-            migration_match = re.search(r"Migration name: (\d+_\S+)", stderr)
-            if migration_match and ProxyExtrasDBManager._is_idempotent_error(stderr):
-                name = migration_match.group(1)
+            migration_name = ProxyExtrasDBManager._v2_failed_migration_name(stderr)
+            if migration_name and _MIGRATION_DEADLOCK_MARKER in stderr:
                 logger.info(
-                    f"Migration {name} SQL hit idempotent error — marking applied and retrying"
-                )
-                ProxyExtrasDBManager._mark_migration_applied(name)
-                return budget.after_recovery(f"resolved:{name}")
-
-            if migration_match and _MIGRATION_DEADLOCK_MARKER in stderr:
-                logger.info(
-                    "Migration %s deadlocked against a concurrent "
-                    "migrate deploy, rolling its ledger row back "
-                    "and retrying",
-                    migration_match.group(1),
-                )
-                ProxyExtrasDBManager._roll_back_migration_best_effort(
-                    migration_match.group(1)
+                    "Migration %s deadlocked against a concurrent migrate deploy, "
+                    "rolling its ledger row back and retrying",
+                    migration_name,
                 )
+                ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name)
                 return budget.spend()
 
             raise RuntimeError(
@@ -1009,19 +1072,17 @@ class ProxyExtrasDBManager:
 
         if _MIGRATION_DEADLOCK_MARKER in stderr:
             logger.info(
-                "prisma migrate deploy attempt %s deadlocked against "
-                "a concurrent migrate deploy, retrying",
+                "prisma migrate deploy attempt %s deadlocked against a concurrent migrate deploy, retrying",
                 budget.attempt_number,
             )
             return budget.spend()
 
         if "P1002" in stderr and "advisory lock" in stderr:
             logger.info(
-                "prisma migrate deploy attempt %s timed out waiting for "
-                "the advisory lock a concurrent migrate deploy holds, retrying",
-                budget.attempt_number,
+                "Waiting for the advisory lock held by another Prisma migration; "
+                "contention does not spend a migration failure attempt"
             )
-            return budget.spend()
+            return budget.after_contention(attempt_seconds)
 
         raise RuntimeError(
             "Database migration failed and cannot be auto-recovered. "

litellm-proxy-extras/tests/test_setup_database_fail_fast.py

@@ -6,7 +6,8 @@ The v2 resolver is opt-in via `--use_v2_migration_resolver` / the
 """
 
 import subprocess
-from unittest.mock import patch
+from types import SimpleNamespace
+from unittest.mock import MagicMock, Mock, patch
 
 import pytest
 
@@ -31,12 +32,7 @@ def _fake_migrate_deploy_failure(returncode: int, stderr: str):
 
 def test_v2_p3018_permission_error_raises_runtime_error(monkeypatch, tmp_path):
     """v2: a permission failure during migrate deploy raises RuntimeError."""
-    monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
-    monkeypatch.setattr(
-        ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
-    )
-    monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
-    (tmp_path / "schema.prisma").write_text("// stub")
+    _stub_v2_env(monkeypatch, tmp_path)
 
     stderr = (
         "Error: P3018\nMigration name: 20250326162113_baseline\n"
@@ -49,19 +45,14 @@ def test_v2_p3018_permission_error_raises_runtime_error(monkeypatch, tmp_path):
 
 def test_v2_non_idempotent_p3009_raises_runtime_error(monkeypatch, tmp_path):
     """v2: a non-idempotent migration failure raises (no silent recovery)."""
-    monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
-    monkeypatch.setattr(
-        ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
-    )
-    monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
-    (tmp_path / "schema.prisma").write_text("// stub")
+    _stub_v2_env(monkeypatch, tmp_path)
 
     stderr = (
         "Error: P3009\nMigration `20260101000000_genuinely_broken` failed\n"
         'Reason: syntax error at or near "BRKN" LINE 42'
     )
     with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)):
-        with pytest.raises(RuntimeError, match="cannot be auto-recovered"):
+        with pytest.raises(RuntimeError, match="Migration completion could not be verified"):
             ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
 
 
@@ -135,8 +126,7 @@ def test_v1_default_still_calls_resolve_all_migrations(monkeypatch, tmp_path):
 def test_v2_db_push_wraps_subprocess_error_as_runtime_error(monkeypatch, tmp_path):
     """v2: a failing `prisma db push` must raise RuntimeError, not leak
     CalledProcessError past proxy_cli.py's `except RuntimeError`."""
-    monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
-    (tmp_path / "schema.prisma").write_text("// stub")
+    monkeypatch.setenv("LITELLM_MIGRATION_DIR", str(tmp_path))
 
     stderr = "db push error"
     with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)):
@@ -153,8 +143,7 @@ def test_v2_warn_ahead_of_head_swallows_db_errors(monkeypatch, tmp_path):
     import psycopg
 
     monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
-    monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
-    (tmp_path / "schema.prisma").write_text("// stub")
+    monkeypatch.setenv("LITELLM_MIGRATION_DIR", str(tmp_path))
 
     class _FakeConn:
         def __enter__(self):
@@ -176,70 +165,28 @@ def test_v2_warn_ahead_of_head_swallows_db_errors(monkeypatch, tmp_path):
     ProxyExtrasDBManager._warn_if_db_ahead_of_head(str(tmp_path))
 
 
-def test_v2_resolve_specific_migration_failure_raises_runtime_error(
-    monkeypatch, tmp_path
-):
-    """If marking a migration as applied fails inside P3009 idempotent
-    recovery, the subprocess error must be re-raised as RuntimeError so
-    proxy_cli.py catches it cleanly (instead of leaking CalledProcessError)."""
-    monkeypatch.setattr(
-        ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
-    )
-    monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
-    (tmp_path / "schema.prisma").write_text("// stub")
-    monkeypatch.setattr(
-        ProxyExtrasDBManager, "_roll_back_migration", lambda *a, **kw: None
-    )
-
-    # First call: migrate deploy -> P3009 idempotent error.
-    # Recovery path tries _resolve_specific_migration; that also raises.
-    def _failing_resolve(*a, **kw):
-        raise subprocess.CalledProcessError(
-            returncode=1,
-            cmd="prisma migrate resolve --applied",
-            stderr="resolve failed",
-            output="",
-        )
-
-    monkeypatch.setattr(
-        ProxyExtrasDBManager, "_resolve_specific_migration", _failing_resolve
-    )
-
-    stderr = (
-        "Error: P3009\nMigration `20260101000000_some_migration` failed\n"
-        "relation already exists"
-    )
-    with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)):
-        with pytest.raises(
-            RuntimeError, match="Failed to mark migration .* as applied"
-        ):
+def test_v2_duplicate_object_p3009_is_not_marked_applied(monkeypatch, tmp_path):
+    _stub_v2_env(monkeypatch, tmp_path, ledger_logs="relation already exists")
+    stderr = "Error: P3009\nMigration `20260101000000_some_migration` failed\nrelation already exists"
+    with patch(
+        "litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)
+    ) as run:
+        with pytest.raises(RuntimeError, match="Migration completion could not be verified"):
             ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
+    assert tuple(call.args[0][1:] for call in run.call_args_list if "migrate" in call.args[0]) == (
+        ["migrate", "deploy"],
+    )
 
 
 def test_v2_does_not_call_resolve_all_migrations(monkeypatch, tmp_path):
-    """v2 must never call _resolve_all_migrations — that's the bug it fixes."""
-    monkeypatch.setattr(
-        ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
+    _stub_v2_env(monkeypatch, tmp_path)
+    run = Mock(side_effect=_succeed_after(0, ""))
+    monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", run)
+
+    assert ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) is True
+    assert tuple(call.args[0][1:] for call in run.call_args_list if "migrate" in call.args[0]) == (
+        ["migrate", "deploy"],
     )
-    monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
-    (tmp_path / "schema.prisma").write_text("// stub")
-
-    class FakeResult:
-        stdout = "Applied migration.\n"
-        stderr = ""
-
-    monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", lambda *a, **kw: FakeResult())
-
-    resolve_called = {"n": 0}
-    monkeypatch.setattr(
-        ProxyExtrasDBManager,
-        "_resolve_all_migrations",
-        lambda *a, **kw: resolve_called.__setitem__("n", resolve_called["n"] + 1),
-    )
-
-    ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
-    assert ok is True
-    assert resolve_called["n"] == 0, "v2 must not invoke the diff-and-force recovery"
 
 
 _DEADLOCK_P3018_STDERR = (
@@ -250,14 +197,34 @@ _DEADLOCK_P3018_STDERR = (
 )
 
 
-def _stub_v2_env(monkeypatch, tmp_path):
+def _stub_v2_env(monkeypatch, tmp_path, ledger_logs=""):
+    import psycopg
+
     monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
-    monkeypatch.setattr(
-        ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
-    )
-    monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
-    (tmp_path / "schema.prisma").write_text("// stub")
+    monkeypatch.delenv("DIRECT_URL", raising=False)
+    monkeypatch.setenv("LITELLM_MIGRATION_DIR", str(tmp_path))
     monkeypatch.setattr("time.sleep", lambda _: None)
+    connection = MagicMock()
+    connection.__enter__.return_value = connection
+    cursor = connection.cursor.return_value.__enter__.return_value
+    cursor.execute.return_value = cursor
+    cursor.fetchone.return_value = SimpleNamespace(acquired=True)
+    cursor.fetchall.return_value = []
+    empty = MagicMock()
+    empty.fetchall.return_value = []
+    empty.fetchone.return_value = None
+    ledger = MagicMock()
+    ledger.fetchone.return_value = (ledger_logs,)
+
+    def execute(query, *args, **kwargs):
+        if "SELECT logs FROM" in str(query):
+            if ledger_logs is None:
+                raise psycopg.OperationalError("ledger is unavailable")
+            return ledger
+        return empty
+
+    connection.execute.side_effect = execute
+    monkeypatch.setattr("psycopg.connect", lambda *args, **kwargs: connection)
 
 
 def _succeed_after(failures: int, stderr: str):
@@ -272,9 +239,7 @@ def _succeed_after(failures: int, stderr: str):
             return _OkResult()
         calls["n"] += 1
         if calls["n"] <= failures:
-            raise subprocess.CalledProcessError(
-                returncode=1, cmd=args[0], stderr=stderr, output=""
-            )
+            raise subprocess.CalledProcessError(returncode=1, cmd=args[0], stderr=stderr, output="")
         return _OkResult()
 
     return _run
@@ -285,28 +250,21 @@ def test_v2_p3018_deadlock_rolls_back_and_retries(monkeypatch, tmp_path):
     instance rolls the ledger row back and retries instead of dying."""
     _stub_v2_env(monkeypatch, tmp_path)
 
-    rolled_back = []
-    monkeypatch.setattr(
-        ProxyExtrasDBManager,
-        "_roll_back_migration",
-        lambda name: rolled_back.append(name),
-    )
-    monkeypatch.setattr(
-        ProxyExtrasDBManager,
-        "_resolve_specific_migration",
-        lambda name: pytest.fail("a deadlocked migration must never be marked applied"),
-    )
-    monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, _DEADLOCK_P3018_STDERR))
+    run = Mock(side_effect=_succeed_after(1, _DEADLOCK_P3018_STDERR))
+    monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", run)
 
     ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
     assert ok is True
-    assert rolled_back == ["20260415120000_health_check_latest_per_model_index"]
+    assert tuple(call.args[0][1:] for call in run.call_args_list if "migrate" in call.args[0]) == (
+        ["migrate", "deploy"],
+        ["migrate", "resolve", "--rolled-back", "20260415120000_health_check_latest_per_model_index"],
+        ["migrate", "deploy"],
+    )
 
 
 def test_v2_p3018_persistent_deadlock_exhausts_attempts(monkeypatch, tmp_path):
     """v2: a deadlock on every attempt still fails after the retry budget."""
     _stub_v2_env(monkeypatch, tmp_path)
-    monkeypatch.setattr(ProxyExtrasDBManager, "_roll_back_migration", lambda name: None)
 
     with patch(
         "litellm_proxy_extras.prisma_toolchain.run_prisma",
@@ -319,7 +277,7 @@ def test_v2_p3018_persistent_deadlock_exhausts_attempts(monkeypatch, tmp_path):
 def test_v2_p3009_deadlocked_ledger_row_rolls_back_and_retries(monkeypatch, tmp_path):
     """v2: the surviving instance sees the victim's failed ledger row as P3009.
     When that row's logs show a deadlock, roll it back and retry."""
-    _stub_v2_env(monkeypatch, tmp_path)
+    _stub_v2_env(monkeypatch, tmp_path, ledger_logs="ERROR: deadlock detected\nDETAIL: Process 72 waits for ShareLock")
 
     stderr = (
         "Error: P3009\n"
@@ -327,61 +285,39 @@ def test_v2_p3009_deadlocked_ledger_row_rolls_back_and_retries(monkeypatch, tmp_
         "The `20260415120000_health_check_latest_per_model_index` migration "
         "started at 2026-09-01 18:46:13 UTC failed"
     )
-    monkeypatch.setattr(
-        ProxyExtrasDBManager,
-        "_failed_migration_logs",
-        lambda name: "ERROR: deadlock detected\nDETAIL: Process 72 waits for ShareLock",
-    )
-    rolled_back = []
-    monkeypatch.setattr(
-        ProxyExtrasDBManager,
-        "_roll_back_migration",
-        lambda name: rolled_back.append(name),
-    )
-    monkeypatch.setattr(
-        ProxyExtrasDBManager,
-        "_resolve_specific_migration",
-        lambda name: pytest.fail("a deadlocked migration must never be marked applied"),
-    )
-    monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, stderr))
+    run = Mock(side_effect=_succeed_after(1, stderr))
+    monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", run)
 
     ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
     assert ok is True
-    assert rolled_back == ["20260415120000_health_check_latest_per_model_index"]
+    assert tuple(call.args[0][1:] for call in run.call_args_list if "migrate" in call.args[0]) == (
+        ["migrate", "deploy"],
+        ["migrate", "resolve", "--rolled-back", "20260415120000_health_check_latest_per_model_index"],
+        ["migrate", "deploy"],
+    )
 
 
-def test_v2_p3009_empty_ledger_logs_rolls_back_and_retries(monkeypatch, tmp_path):
-    """v2: empty failed ledger logs mean a concurrent deploy moved it on."""
+def test_v2_p3009_empty_ledger_logs_do_not_prove_completion(monkeypatch, tmp_path):
     _stub_v2_env(monkeypatch, tmp_path)
-
     stderr = (
         "Error: P3009\n"
         "migrate found failed migrations in the target database\n"
         "The `20260415120000_health_check_latest_per_model_index` migration "
         "started at 2026-09-01 18:46:13 UTC failed"
     )
-    monkeypatch.setattr(ProxyExtrasDBManager, "_failed_migration_logs", lambda name: "")
-    rolled_back = []
-    monkeypatch.setattr(
-        ProxyExtrasDBManager,
-        "_roll_back_migration",
-        lambda name: rolled_back.append(name),
+    with patch(
+        "litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)
+    ) as run:
+        with pytest.raises(RuntimeError, match="Migration completion could not be verified"):
+            ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
+    assert tuple(call.args[0][1:] for call in run.call_args_list if "migrate" in call.args[0]) == (
+        ["migrate", "deploy"],
     )
-    monkeypatch.setattr(
-        ProxyExtrasDBManager,
-        "_resolve_specific_migration",
-        lambda name: pytest.fail("a deadlocked migration must never be marked applied"),
-    )
-    monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, stderr))
-
-    ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
-    assert ok is True
-    assert rolled_back == ["20260415120000_health_check_latest_per_model_index"]
 
 
 def test_v2_p3009_unreadable_ledger_still_raises(monkeypatch, tmp_path):
     """v2: an unreadable ledger cannot establish that P3009 was a deadlock."""
-    _stub_v2_env(monkeypatch, tmp_path)
+    _stub_v2_env(monkeypatch, tmp_path, ledger_logs=None)
 
     stderr = (
         "Error: P3009\n"
@@ -389,21 +325,15 @@ def test_v2_p3009_unreadable_ledger_still_raises(monkeypatch, tmp_path):
         "The `20260415120000_health_check_latest_per_model_index` migration "
         "started at 2026-09-01 18:46:13 UTC failed"
     )
-    monkeypatch.setattr(ProxyExtrasDBManager, "_failed_migration_logs", lambda name: None)
-    monkeypatch.setattr(
-        ProxyExtrasDBManager,
-        "_roll_back_migration",
-        lambda name: pytest.fail("an unreadable ledger must not trigger a retry"),
-    )
     monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, stderr))
 
-    with pytest.raises(RuntimeError, match="cannot be auto-recovered"):
+    with pytest.raises(RuntimeError, match="Migration completion could not be verified"):
         ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
 
 
 def test_v2_p3009_non_deadlock_ledger_row_still_raises(monkeypatch, tmp_path):
     """v2: a failed ledger row whose logs show a real SQL error stays fatal."""
-    _stub_v2_env(monkeypatch, tmp_path)
+    _stub_v2_env(monkeypatch, tmp_path, ledger_logs='ERROR: syntax error at or near "BRKN"')
 
     stderr = (
         "Error: P3009\n"
@@ -411,14 +341,9 @@ def test_v2_p3009_non_deadlock_ledger_row_still_raises(monkeypatch, tmp_path):
         "The `20260101000000_genuinely_broken` migration started at "
         "2026-09-01 18:46:13 UTC failed"
     )
-    monkeypatch.setattr(
-        ProxyExtrasDBManager,
-        "_failed_migration_logs",
-        lambda name: 'ERROR: syntax error at or near "BRKN"',
-    )
 
     with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)):
-        with pytest.raises(RuntimeError, match="cannot be auto-recovered"):
+        with pytest.raises(RuntimeError, match="Migration completion could not be verified"):
             ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
 
 

litellm-rust/Cargo.toml

@@ -14,9 +14,14 @@ litellm-host = { path = "crates/host" }
 litellm-callbacks-legacy-python = { path = "crates/callbacks-legacy-python" }
 litellm-framing = { path = "crates/framer" }
 litellm-auth = { path = "crates/auth" }
+litellm-auth-types = { path = "crates/auth-types" }
 litellm-auth-aws = { path = "crates/auth-aws" }
 litellm-auth-azure = { path = "crates/auth-azure" }
 litellm-auth-gcp = { path = "crates/auth-gcp" }
+litellm-secrets = { path = "crates/secrets" }
+litellm-secrets-types = { path = "crates/secrets-types" }
+litellm-secrets-aws = { path = "crates/secrets-aws" }
+litellm-secrets-google = { path = "crates/secrets-google" }
 litellm-http = { path = "crates/http" }
 litellm-llms = { path = "crates/llms" }
 litellm-types = { path = "crates/types" }
@@ -24,10 +29,15 @@ litellm-core-utils = { path = "crates/core-utils" }
 litellm-cache = { path = "crates/cache" }
 litellm-cache-memory = { path = "crates/cache-memory" }
 litellm-token-counter = { path = "crates/token-counter" }
+litellm-token-counter-fast = { path = "crates/token-counter-fast" }
+litellm-token-counter-huggingface = { path = "crates/token-counter-huggingface" }
+litellm-token-counter-tiktoken = { path = "crates/token-counter-tiktoken" }
 litellm-host-python = { path = "crates/host-python" }
 
 bytes = "1"
 http = "1"
+google-cloud-auth = { version = "1.16.0", default-features = false }
+jsonwebtoken = { version = "11.1.0", default-features = false }
 hyper-util = { version = "0.1.20", default-features = false, features = ["client-proxy"] }
 proptest = "1.7.0"
 pyo3 = "0.29.2"
@@ -45,6 +55,8 @@ serde_with = { version = "=3.16.1", default-features = false, features = ["std",
 sha2 = "0.10"
 subtle = "2"
 thiserror = "2.0"
+tokenizers = { version = "0.23.1", default-features = false, features = ["onig"] }
+tiktoken-rs = "0.12.0"
 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"] }

litellm-rust/crates/auth-aws/Cargo.toml

@@ -6,7 +6,7 @@ license.workspace = true
 repository.workspace = true
 
 [dependencies]
-litellm-auth.workspace = true
+litellm-auth-types.workspace = true
 litellm-http.workspace = true
 
 moka = { workspace = true, features = ["sync"] }

litellm-rust/crates/auth-aws/src/constants.rs

@@ -3,6 +3,8 @@ pub const AWS_SECRET_ACCESS_KEY: &str = "AWS_SECRET_ACCESS_KEY";
 pub const AWS_SESSION_TOKEN: &str = "AWS_SESSION_TOKEN";
 pub const AWS_REGION_NAME: &str = "AWS_REGION_NAME";
 pub const AWS_REGION: &str = "AWS_REGION";
+pub const AWS_DEFAULT_REGION: &str = "AWS_DEFAULT_REGION";
+pub const AWS_BEDROCK_RUNTIME_ENDPOINT: &str = "AWS_BEDROCK_RUNTIME_ENDPOINT";
 pub const AWS_SESSION_NAME: &str = "AWS_SESSION_NAME";
 pub const AWS_PROFILE_NAME: &str = "AWS_PROFILE_NAME";
 pub const AWS_ROLE_NAME: &str = "AWS_ROLE_NAME";

litellm-rust/crates/auth-aws/src/error.rs

@@ -22,7 +22,7 @@ pub enum Error {
     AwsMissingWebIdentityCredentials,
 }
 
-impl From<Error> for litellm_auth::Error {
+impl From<Error> for litellm_auth_types::Error {
     fn from(error: Error) -> Self {
         Self::ProviderAuthentication(error.to_string())
     }
@@ -34,11 +34,11 @@ mod tests {
 
     #[test]
     fn converts_to_shared_auth_error_without_losing_context() {
-        let error = litellm_auth::Error::from(Error::AwsProfile("profile not found".into()));
+        let error = litellm_auth_types::Error::from(Error::AwsProfile("profile not found".into()));
 
         assert_eq!(
             error,
-            litellm_auth::Error::ProviderAuthentication(
+            litellm_auth_types::Error::ProviderAuthentication(
                 "AWS profile credentials failed: profile not found".into()
             )
         );

litellm-rust/crates/auth-azure/Cargo.toml

@@ -6,7 +6,7 @@ license.workspace = true
 repository.workspace = true
 
 [dependencies]
-litellm-auth.workspace = true
+litellm-auth-types.workspace = true
 
 moka.workspace = true
 serde_json.workspace = true

litellm-rust/crates/auth-azure/src/credential_provider_cache.rs

@@ -4,7 +4,7 @@ use std::sync::Arc;
 use azure_core::credentials::TokenCredential;
 use moka::future::Cache;
 
-use litellm_auth::Error;
+use litellm_auth_types::Error;
 
 #[derive(Clone, Debug, PartialEq, Eq, Hash)]
 pub(crate) struct AzureCredentialProviderCacheKey {

litellm-rust/crates/auth-azure/src/native.rs

@@ -12,8 +12,8 @@ use azure_identity::{
 };
 use sha2::{Digest, Sha256};
 
-use litellm_auth::Error;
-use litellm_auth::{InputSource, ResolvedCredential, SecretValue, Sourced};
+use litellm_auth_types::Error;
+use litellm_auth_types::{InputSource, ResolvedCredential, SecretValue, Sourced};
 
 use super::credential_provider_cache::{
     AzureCredentialProviderCache, AzureCredentialProviderCacheKey,
@@ -484,7 +484,7 @@ mod tests {
     use azure_core::{Bytes, Result};
 
     use super::{NativeAzureRequest, NativeAzureTokenAcquirer, ValidatedAzureRequest};
-    use litellm_auth::{InputSource, SecretValue, Sourced};
+    use litellm_auth_types::{InputSource, SecretValue, Sourced};
 
     fn deployment<T>(value: T) -> Sourced<T> {
         Sourced::new(value, InputSource::Deployment)
@@ -649,7 +649,7 @@ mod tests {
 
         assert!(matches!(
             error,
-            litellm_auth::Error::MixedAzureCredentialSources
+            litellm_auth_types::Error::MixedAzureCredentialSources
         ));
     }
 
@@ -679,7 +679,10 @@ mod tests {
                 authority,
             ))
             .unwrap_err();
-            assert!(matches!(error, litellm_auth::Error::InvalidAzureAuthority));
+            assert!(matches!(
+                error,
+                litellm_auth_types::Error::InvalidAzureAuthority
+            ));
         }
     }
 }

litellm-rust/crates/auth-azure/src/resolve.rs

@@ -1,5 +1,5 @@
-use litellm_auth::Error;
-use litellm_auth::{
+use litellm_auth_types::Error;
+use litellm_auth_types::{
     CredentialFileRef, CredentialLookup, CredentialRef, InputSource, ResolvedCredential,
     SecretValue, Sourced, TokenProviderHandle,
 };
@@ -451,9 +451,9 @@ mod tests {
     };
     use crate::native::ValidatedAzureRequest;
     use crate::types::AzureAuthInputs;
-    use litellm_auth::Error;
-    use litellm_auth::ResolvedCredential;
-    use litellm_auth::{
+    use litellm_auth_types::Error;
+    use litellm_auth_types::ResolvedCredential;
+    use litellm_auth_types::{
         CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialRef,
         CredentialResolver, CredentialResolverHandle, InputSource, SecretValue, Sourced,
     };
@@ -661,8 +661,8 @@ mod tests {
     #[derive(Debug)]
     struct CallerToken(&'static str);
 
-    impl litellm_auth::TokenProvider for CallerToken {
-        fn acquire(&self) -> litellm_auth::TokenFuture<'_> {
+    impl litellm_auth_types::TokenProvider for CallerToken {
+        fn acquire(&self) -> litellm_auth_types::TokenFuture<'_> {
             Box::pin(async move {
                 Ok(ResolvedCredential::AccessToken {
                     token: SecretValue::new(self.0),
@@ -675,7 +675,7 @@ mod tests {
     fn caller_inputs(token: &'static str) -> AzureAuthInputs {
         let params = json!({"azure_ad_token": "static-token"});
         AzureAuthInputs {
-            azure_ad_token_provider: Some(litellm_auth::TokenProviderHandle::new(Arc::new(
+            azure_ad_token_provider: Some(litellm_auth_types::TokenProviderHandle::new(Arc::new(
                 CallerToken(token),
             ))),
             ..AzureAuthInputs::from_optional_params(params.as_object().unwrap()).unwrap()

litellm-rust/crates/auth-azure/src/types.rs

@@ -1,6 +1,6 @@
 use std::collections::BTreeMap;
 
-use litellm_auth::{
+use litellm_auth_types::{
     CredentialResolverHandle, Error, InputSource, SecretValue, Sourced, TokenProviderHandle,
 };
 use serde_json::{Map, Value};
@@ -126,7 +126,7 @@ fn source_for(sources: &BTreeMap<String, InputSource>, name: &str) -> InputSourc
 mod tests {
     use std::collections::BTreeMap;
 
-    use litellm_auth::{InputSource, Sourced};
+    use litellm_auth_types::{InputSource, Sourced};
     use serde_json::json;
 
     use super::{AzureAuthInputs, AzureCredentialType, ConfigValue};

litellm-rust/crates/auth-gcp/Cargo.toml

@@ -5,8 +5,11 @@ edition.workspace = true
 license.workspace = true
 repository.workspace = true
 
+[features]
+google-sdk = ["dep:google-cloud-auth", "dep:http"]
+
 [dependencies]
-litellm-auth.workspace = true
+litellm-auth-types.workspace = true
 
 moka.workspace = true
 serde_json.workspace = true
@@ -14,3 +17,5 @@ sha2.workspace = true
 tokio.workspace = true
 
 gcp_auth = "0.12.7"
+google-cloud-auth = { workspace = true, optional = true }
+http = { workspace = true, optional = true }

litellm-rust/crates/auth-gcp/src/lib.rs

@@ -1,13 +1,18 @@
 use std::{collections::BTreeMap, future::Future, path::Path, pin::Pin, sync::Arc};
 
 use gcp_auth::{CustomServiceAccount, TokenProvider};
-use litellm_auth::{
+use litellm_auth_types::{
     CredentialPlacement, Error, InputSource, SecretValue, Sourced, http::apply_credential,
 };
 use moka::future::Cache;
 use serde_json::{Map, Value};
 use sha2::{Digest, Sha256};
 
+#[cfg(feature = "google-sdk")]
+mod sdk;
+#[cfg(feature = "google-sdk")]
+pub use sdk::GoogleCredentials;
+
 const CLOUD_PLATFORM_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform";
 const GOOGLE_OAUTH_TOKEN_ENDPOINT: &str = "https://oauth2.googleapis.com/token";
 const GOOGLE_APPLICATION_CREDENTIALS_ENV: &str = "GOOGLE_APPLICATION_CREDENTIALS";
@@ -26,19 +31,31 @@ pub struct VertexConfig {
 }
 
 impl VertexConfig {
+    pub fn new(
+        credentials: Option<Sourced<SecretValue>>,
+        project_id: Option<String>,
+        location: Option<String>,
+    ) -> Self {
+        Self {
+            credentials: credentials.filter(|value| !value.value().expose().trim().is_empty()),
+            project_id: project_id.filter(|value| !value.trim().is_empty()),
+            location: location.filter(|value| !value.trim().is_empty()),
+        }
+    }
+
     pub fn from_sourced_optional_params(
         params: &Map<String, Value>,
         sources: &BTreeMap<String, InputSource>,
     ) -> Result<Self, Error> {
-        Ok(Self {
-            credentials: optional_credentials(
+        Ok(Self::new(
+            optional_credentials(
                 params,
                 sources,
                 &["vertex_credentials", "vertex_ai_credentials"],
             )?,
-            project_id: optional_string(params, &["vertex_project", "vertex_ai_project"])?,
-            location: optional_string(params, &["vertex_location", "vertex_ai_location"])?,
-        })
+            optional_string(params, &["vertex_project", "vertex_ai_project"])?,
+            optional_string(params, &["vertex_location", "vertex_ai_location"])?,
+        ))
     }
 
     pub fn or_configured(self, project_id: Option<&str>, location: Option<&str>) -> Self {
@@ -469,6 +486,39 @@ mod tests {
         assert_eq!(config.location(), Some("alias-location"));
     }
 
+    #[test]
+    fn typed_config_preserves_source_and_empty_value_fallback() {
+        let configured = VertexConfig::new(
+            Some(Sourced::new(
+                SecretValue::new("inline-json"),
+                InputSource::Request,
+            )),
+            Some("project".into()),
+            Some("location".into()),
+        );
+        assert!(matches!(
+            credential_source(&configured, &|_| Some("environment-json".into())),
+            CredentialSource::Inline(value) if value.expose() == "inline-json"
+        ));
+        let empty = VertexConfig::new(
+            Some(Sourced::new(SecretValue::new(" "), InputSource::Request)),
+            Some(" ".into()),
+            Some(" ".into()),
+        );
+        assert!(matches!(
+            credential_source(&empty, &|_| None),
+            CredentialSource::Adc
+        ));
+        assert_eq!(
+            get_vertex_ai_project(&empty, &|_| Some("env-project".into())).as_deref(),
+            Some("env-project")
+        );
+        assert_eq!(
+            get_vertex_ai_location(&empty, &|_| Some("env-location".into())).as_deref(),
+            Some("env-location")
+        );
+    }
+
     #[test]
     fn project_and_location_prefer_input_then_environment() {
         let configured =

litellm-rust/crates/auth-gcp/src/sdk.rs

@@ -0,0 +1,106 @@
+use std::sync::Arc;
+
+use google_cloud_auth::credentials::{CacheableResource, CredentialsProvider, EntityTag};
+use google_cloud_auth::errors::CredentialsError;
+use http::{Extensions, HeaderMap, HeaderName, HeaderValue};
+use litellm_auth_types::Error;
+
+use crate::{VertexAuth, VertexConfig};
+
+type EnvironmentLookup = dyn Fn(&str) -> Option<String> + Send + Sync;
+
+pub struct GoogleCredentials {
+    auth: VertexAuth,
+    config: VertexConfig,
+    environment: Arc<EnvironmentLookup>,
+}
+
+impl GoogleCredentials {
+    pub fn new(config: VertexConfig, environment: Arc<EnvironmentLookup>) -> Self {
+        Self {
+            auth: VertexAuth::default(),
+            config,
+            environment,
+        }
+    }
+
+    pub async fn request_headers(&self) -> Result<HeaderMap, Error> {
+        let response = self
+            .auth
+            .validate_environment(Vec::new(), None, &self.config, &|name| {
+                (self.environment)(name)
+            })
+            .await?;
+        response
+            .headers
+            .into_iter()
+            .map(|(key, value)| {
+                let name =
+                    HeaderName::from_bytes(key.as_bytes()).map_err(|_| Error::InvalidHeader)?;
+                let value = HeaderValue::from_str(&value).map_err(|_| Error::InvalidHeader)?;
+                Ok((name, value))
+            })
+            .collect()
+    }
+}
+
+impl CredentialsProvider for GoogleCredentials {
+    async fn headers(
+        &self,
+        _: Extensions,
+    ) -> Result<CacheableResource<HeaderMap>, CredentialsError> {
+        self.request_headers()
+            .await
+            .map(|data| CacheableResource::New {
+                entity_tag: EntityTag::new(),
+                data,
+            })
+            .map_err(|_| CredentialsError::from_msg(false, "Google authentication failed"))
+    }
+
+    async fn universe_domain(&self) -> Option<String> {
+        None
+    }
+}
+
+impl std::fmt::Debug for GoogleCredentials {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("GoogleCredentials").finish_non_exhaustive()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn sdk_and_http_credentials_share_token_resolution_and_redaction() {
+        let credentials = GoogleCredentials::new(
+            VertexConfig::new(None, Some("project".into()), None),
+            Arc::new(|name| (name == "VERTEX_AI_API_KEY").then(|| "private-token".into())),
+        );
+        let direct = credentials.request_headers().await.unwrap();
+        let CacheableResource::New { data, .. } =
+            credentials.headers(Extensions::new()).await.unwrap()
+        else {
+            panic!("first request did not return headers");
+        };
+        assert_eq!(direct, data);
+        assert_eq!(data[http::header::AUTHORIZATION], "Bearer private-token");
+        assert!(!format!("{credentials:?}").contains("private-token"));
+    }
+
+    #[tokio::test]
+    async fn invalid_token_headers_return_a_redacted_sdk_error() {
+        let credentials = GoogleCredentials::new(
+            VertexConfig::new(None, Some("project".into()), None),
+            Arc::new(|name| (name == "VERTEX_AI_API_KEY").then(|| "private\nvalue".into())),
+        );
+        assert_eq!(
+            credentials.request_headers().await.unwrap_err(),
+            Error::InvalidHeader
+        );
+        let error = credentials.headers(Extensions::new()).await.unwrap_err();
+        assert!(!format!("{error:?}").contains("private"));
+    }
+}

litellm-rust/crates/auth-types/Cargo.toml

@@ -0,0 +1,15 @@
+[package]
+name = "litellm-auth-types"
+version = "0.1.0"
+edition.workspace = true
+license.workspace = true
+repository.workspace = true
+
+[dependencies]
+serde.workspace = true
+subtle.workspace = true
+thiserror.workspace = true
+veil.workspace = true
+
+[dev-dependencies]
+tokio.workspace = true

litellm-rust/crates/auth-types/src/credential.rs

@@ -5,9 +5,7 @@ use std::sync::Arc;
 
 use veil::Redact;
 
-use crate::Error;
-
-use super::{ResolvedCredential, SecretValue, TokenProviderHandle};
+use crate::{Error, ResolvedCredential, SecretValue, TokenProviderHandle};
 
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub enum CredentialFileRef {

litellm-rust/crates/auth-types/src/http.rs

@@ -40,9 +40,6 @@ pub fn apply_credential(
     )
 }
 
-/// How the upstream call is authenticated. API-key strategies become headers
-/// in `prepare`; SigV4 covers the serialized body, so it is applied where the
-/// outbound request is built.
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub enum RequestAuth {
     Header {

litellm-rust/crates/auth-types/src/lib.rs

@@ -0,0 +1,57 @@
+#![forbid(unsafe_code)]
+
+mod credential;
+mod error;
+pub mod http;
+mod policy;
+mod secret;
+mod token;
+
+use serde::{Deserialize, Serialize};
+
+#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, Hash, PartialEq, Serialize)]
+#[serde(rename_all = "snake_case")]
+pub enum InputSource {
+    Request,
+    #[default]
+    Deployment,
+    Environment,
+}
+
+#[derive(Clone, Debug, Default, Eq, PartialEq)]
+pub struct Sourced<T> {
+    value: T,
+    source: InputSource,
+}
+
+impl<T> Sourced<T> {
+    pub fn new(value: T, source: InputSource) -> Self {
+        Self { value, source }
+    }
+
+    pub fn value(&self) -> &T {
+        &self.value
+    }
+
+    pub fn source(&self) -> InputSource {
+        self.source
+    }
+
+    pub fn into_value(self) -> T {
+        self.value
+    }
+
+    pub fn map<U>(self, map: impl FnOnce(T) -> U) -> Sourced<U> {
+        Sourced::new(map(self.value), self.source)
+    }
+}
+
+pub use credential::{
+    CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialPlan,
+    CredentialPlanResolution, CredentialRef, CredentialResolver, CredentialResolverHandle,
+};
+pub use error::Error;
+pub use http::{CredentialPlacement, RequestAuth};
+pub use policy::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy};
+pub use secret::SecretValue;
+pub use token::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle};

litellm-rust/crates/auth-types/src/policy.rs

@@ -1,7 +1,5 @@
-use crate::Error;
-
-use super::http::apply_credential;
-use super::{CredentialPlacement, ResolvedCredential};
+use crate::http::apply_credential;
+use crate::{CredentialPlacement, Error, ResolvedCredential};
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub enum CredentialPlanKind {

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

@@ -5,9 +5,7 @@ use std::time::SystemTime;
 
 use veil::Redact;
 
-use crate::Error;
-
-use super::secret::SecretValue;
+use crate::{Error, SecretValue};
 
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub enum ResolvedCredential {

litellm-rust/crates/auth/Cargo.toml

@@ -5,11 +5,14 @@ edition.workspace = true
 license.workspace = true
 repository.workspace = true
 
-[dependencies]
-serde.workspace = true
-subtle.workspace = true
-thiserror.workspace = true
-veil.workspace = true
+[features]
+default = []
+aws = ["dep:litellm-auth-aws"]
+azure = ["dep:litellm-auth-azure"]
+gcp = ["dep:litellm-auth-gcp"]
 
-[dev-dependencies]
-tokio.workspace = true
+[dependencies]
+litellm-auth-types.workspace = true
+litellm-auth-aws = { workspace = true, optional = true }
+litellm-auth-azure = { workspace = true, optional = true }
+litellm-auth-gcp = { workspace = true, optional = true }

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

@@ -1,55 +1,10 @@
-mod credential;
-mod error;
-pub mod http;
-mod policy;
-mod secret;
-mod token;
+#![forbid(unsafe_code)]
 
-use serde::{Deserialize, Serialize};
+pub use litellm_auth_types::*;
 
-#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, Hash, PartialEq, Serialize)]
-#[serde(rename_all = "snake_case")]
-pub enum InputSource {
-    Request,
-    #[default]
-    Deployment,
-    Environment,
-}
-
-#[derive(Clone, Debug, Default, Eq, PartialEq)]
-pub struct Sourced<T> {
-    value: T,
-    source: InputSource,
-}
-
-impl<T> Sourced<T> {
-    pub fn new(value: T, source: InputSource) -> Self {
-        Self { value, source }
-    }
-
-    pub fn value(&self) -> &T {
-        &self.value
-    }
-
-    pub fn source(&self) -> InputSource {
-        self.source
-    }
-
-    pub fn into_value(self) -> T {
-        self.value
-    }
-
-    pub fn map<U>(self, map: impl FnOnce(T) -> U) -> Sourced<U> {
-        Sourced::new(map(self.value), self.source)
-    }
-}
-
-pub use credential::{
-    CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialPlan,
-    CredentialPlanResolution, CredentialRef, CredentialResolver, CredentialResolverHandle,
-};
-pub use error::Error;
-pub use http::{CredentialPlacement, RequestAuth};
-pub use policy::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy};
-pub use secret::SecretValue;
-pub use token::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle};
+#[cfg(feature = "aws")]
+pub use litellm_auth_aws as aws;
+#[cfg(feature = "azure")]
+pub use litellm_auth_azure as azure;
+#[cfg(feature = "gcp")]
+pub use litellm_auth_gcp as gcp;

litellm-rust/crates/auth/tests/facade.rs

@@ -0,0 +1,33 @@
+use litellm_auth::{
+    CredentialPlacement, CredentialPlanKind, CredentialRule, ExistingHeaderBehavior,
+    ProviderAuthPolicy, ResolvedCredential, SecretValue,
+};
+
+const RULES: &[CredentialRule] = &[CredentialRule {
+    kind: CredentialPlanKind::Static,
+    placement: CredentialPlacement::Header("x-api-key"),
+}];
+
+#[test]
+fn facade_applies_shared_auth_policy() {
+    let policy = ProviderAuthPolicy {
+        rules: RULES,
+        accepted_existing_headers: &["x-api-key"],
+        existing_header_behavior: ExistingHeaderBehavior::Preserve,
+        scope: None,
+        audience: None,
+    };
+
+    let headers = policy
+        .apply(
+            Vec::new(),
+            CredentialPlanKind::Static,
+            &ResolvedCredential::Static(SecretValue::new("secret")),
+        )
+        .expect("facade policy applies");
+
+    assert_eq!(
+        headers,
+        vec![("x-api-key".to_string(), "secret".to_string())]
+    );
+}

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

@@ -10,10 +10,13 @@ name = "_native"
 crate-type = ["cdylib"]
 
 [features]
-default = ["abi3"]
+default = ["abi3", "fast"]
 abi3 = ["pyo3/abi3-py310"]
 extension-module = ["pyo3/extension-module"]
 panic-test = []
+fast = ["litellm-token-counter/fast"]
+huggingface = ["litellm-token-counter/huggingface"]
+tiktoken = ["litellm-token-counter/tiktoken"]
 
 [dependencies]
 bytes.workspace = true
@@ -26,7 +29,7 @@ litellm-http.workspace = true
 litellm-llms.workspace = true
 litellm-types.workspace = true
 litellm-host-python.workspace = true
-litellm-token-counter.workspace = true
+litellm-token-counter = { path = "../token-counter", default-features = false }
 pyo3.workspace = true
 pyo3-async-runtimes.workspace = true
 serde_json.workspace = true

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

@@ -1,6 +1,11 @@
-use std::{num::NonZero, sync::Arc, thread::available_parallelism};
+use std::sync::Arc;
 
-use litellm_host_python::{release_gil, run_async};
+#[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))]
+use std::{num::NonZero, thread::available_parallelism};
+
+#[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))]
+use litellm_host_python::release_gil;
+use litellm_host_python::run_async;
 use litellm_token_counter::{
     CountableRequest, Error, InputTokenCount, TokenCounter as CoreTokenCounter,
 };
@@ -28,17 +33,66 @@ pub(crate) struct TokenCounter {
 impl TokenCounter {
     #[new]
     fn new(py: Python<'_>, tokenizer_json: &str) -> PyResult<Self> {
-        Self::load(py, || CoreTokenCounter::from_json(tokenizer_json))
+        #[cfg(feature = "fast")]
+        {
+            Self::load(py, || CoreTokenCounter::from_json_fast(tokenizer_json))
+        }
+        #[cfg(all(not(feature = "fast"), feature = "huggingface"))]
+        {
+            Self::load(py, || CoreTokenCounter::from_json(tokenizer_json))
+        }
+        #[cfg(not(any(feature = "fast", feature = "huggingface")))]
+        {
+            let _ = (py, tokenizer_json);
+            Err(RustBridgeDeclined::new_err(
+                "tokenizer backend requires the fast or huggingface feature",
+            ))
+        }
     }
 
     #[staticmethod]
     fn from_cl100k_ranks(py: Python<'_>, rank_file: &str) -> PyResult<Self> {
-        Self::load(py, || CoreTokenCounter::from_cl100k_ranks(rank_file))
+        #[cfg(feature = "fast")]
+        {
+            Self::load(py, || CoreTokenCounter::from_cl100k_ranks(rank_file))
+        }
+        #[cfg(not(feature = "fast"))]
+        {
+            let _ = (py, rank_file);
+            Err(RustBridgeDeclined::new_err(
+                "tokenizer backend requires the fast feature",
+            ))
+        }
     }
 
     #[staticmethod]
     fn from_o200k_ranks(py: Python<'_>, rank_file: &str) -> PyResult<Self> {
-        Self::load(py, || CoreTokenCounter::from_o200k_ranks(rank_file))
+        #[cfg(feature = "fast")]
+        {
+            Self::load(py, || CoreTokenCounter::from_o200k_ranks(rank_file))
+        }
+        #[cfg(not(feature = "fast"))]
+        {
+            let _ = (py, rank_file);
+            Err(RustBridgeDeclined::new_err(
+                "tokenizer backend requires the fast feature",
+            ))
+        }
+    }
+
+    #[staticmethod]
+    fn from_tiktoken(py: Python<'_>, encoding: &str) -> PyResult<Self> {
+        #[cfg(feature = "tiktoken")]
+        {
+            Self::load(py, || CoreTokenCounter::from_tiktoken(encoding))
+        }
+        #[cfg(not(feature = "tiktoken"))]
+        {
+            let _ = (py, encoding);
+            Err(RustBridgeDeclined::new_err(
+                "tokenizer backend requires the tiktoken feature",
+            ))
+        }
     }
 
     fn acount_request<'py>(&self, py: Python<'py>, body: &[u8]) -> PyResult<Bound<'py, PyAny>> {
@@ -62,6 +116,7 @@ impl TokenCounter {
 }
 
 impl TokenCounter {
+    #[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))]
     fn load(
         py: Python<'_>,
         load: impl FnOnce() -> Result<CoreTokenCounter, Error> + Send,
@@ -74,6 +129,7 @@ impl TokenCounter {
     }
 }
 
+#[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))]
 fn encode_parallelism() -> usize {
     available_parallelism().map_or(1, NonZero::get)
 }
@@ -86,7 +142,10 @@ fn count_body(counter: &CoreTokenCounter, body: &[u8]) -> Result<InputTokenCount
 fn token_count_error_to_pyerr(error: Error) -> PyErr {
     let message = error.to_string();
     match error {
-        Error::Load(_) | Error::Ranks(_) | Error::UnicodeClasses => PyValueError::new_err(message),
+        Error::Load(_)
+        | Error::Ranks(_)
+        | Error::UnicodeClasses
+        | Error::UnsupportedTokenizer(_) => PyValueError::new_err(message),
         Error::RequestParse(_)
         | Error::MissingInput
         | Error::FloatText

litellm-rust/crates/secrets-aws/Cargo.toml

@@ -0,0 +1,24 @@
+[package]
+name = "litellm-secrets-aws"
+version = "0.1.0"
+edition.workspace = true
+license.workspace = true
+repository.workspace = true
+
+[dependencies]
+litellm-auth-aws.workspace = true
+litellm-secrets-types.workspace = true
+litellm-core-utils.workspace = true
+serde_json.workspace = true
+thiserror.workspace = true
+tracing = "0.1"
+veil.workspace = true
+aws-sdk-kms = "1.120.0"
+aws-sdk-secretsmanager = "1.117.0"
+aws-credential-types = "1.3.0"
+
+[dev-dependencies]
+base64.workspace = true
+rstest.workspace = true
+tokio.workspace = true
+wiremock = "0.6.5"

litellm-rust/crates/secrets-aws/src/auth.rs

@@ -0,0 +1,79 @@
+use std::sync::Arc;
+
+use aws_credential_types::provider::{ProvideCredentials, error::CredentialsError, future};
+use litellm_auth_aws::{
+    AwsAuthConfig,
+    constants::{AWS_DEFAULT_REGION, AWS_REGION, AWS_REGION_NAME},
+    resolve_credentials,
+};
+use litellm_core_utils::settings::Lookup;
+use litellm_secrets_types::KeyManagementSettings;
+
+use crate::Error;
+
+#[derive(Clone)]
+pub(crate) struct Credentials {
+    config: AwsAuthConfig,
+    environment: Arc<dyn Lookup + Send + Sync>,
+}
+
+impl Credentials {
+    pub(crate) fn new(
+        settings: &KeyManagementSettings,
+        environment: Arc<dyn Lookup + Send + Sync>,
+    ) -> Self {
+        Self {
+            config: AwsAuthConfig {
+                region_name: region(settings, environment.as_ref()).ok(),
+                role_name: settings.aws_role_name.clone(),
+                session_name: settings.aws_session_name.clone(),
+                external_id: settings
+                    .aws_external_id
+                    .as_ref()
+                    .map(|v| v.expose().to_owned()),
+                profile_name: settings.aws_profile_name.clone(),
+                web_identity_token: settings
+                    .aws_web_identity_token
+                    .as_ref()
+                    .map(|v| v.expose().to_owned()),
+                sts_endpoint: settings.aws_sts_endpoint.clone(),
+                ..Default::default()
+            },
+            environment,
+        }
+    }
+}
+
+impl ProvideCredentials for Credentials {
+    fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a>
+    where
+        Self: 'a,
+    {
+        future::ProvideCredentials::new(async {
+            resolve_credentials(self.config.clone(), &|name| self.environment.get(name))
+                .await
+                .map_err(|_| {
+                    CredentialsError::provider_error("secret manager authentication failed")
+                })
+        })
+    }
+}
+
+pub(crate) fn region(
+    settings: &KeyManagementSettings,
+    environment: &dyn Lookup,
+) -> Result<String, Error> {
+    settings
+        .aws_region_name
+        .clone()
+        .or_else(|| environment.get(AWS_REGION_NAME))
+        .or_else(|| environment.get(AWS_REGION))
+        .or_else(|| environment.get(AWS_DEFAULT_REGION))
+        .ok_or(Error::MissingRegion)
+}
+
+impl std::fmt::Debug for Credentials {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("Credentials").finish_non_exhaustive()
+    }
+}

litellm-rust/crates/secrets-aws/src/error.rs

@@ -0,0 +1,31 @@
+use aws_sdk_secretsmanager::error::SdkError;
+
+#[derive(thiserror::Error, veil::Redact)]
+pub enum Error {
+    #[error("AWS authentication failed")]
+    Auth(#[from] #[redact] litellm_auth_aws::Error),
+    #[error("AWS region is not configured")]
+    MissingRegion,
+    #[error("KMS response has no plaintext")]
+    MissingPlaintext,
+    #[error("AWS request timed out")]
+    Timeout,
+    #[error("AWS KMS decrypt failed")]
+    Decrypt(#[from] #[redact] Box<SdkError<aws_sdk_kms::operation::decrypt::DecryptError>>),
+    #[error("AWS Secrets Manager read failed")]
+    Read(#[from] #[redact] Box<SdkError<aws_sdk_secretsmanager::operation::get_secret_value::GetSecretValueError>>),
+    #[error("AWS Secrets Manager create failed")]
+    Create(#[from] #[redact] Box<SdkError<aws_sdk_secretsmanager::operation::create_secret::CreateSecretError>>),
+    #[error("AWS Secrets Manager update failed")]
+    Put(#[from] #[redact] Box<SdkError<aws_sdk_secretsmanager::operation::put_secret_value::PutSecretValueError>>),
+    #[error("AWS Secrets Manager delete failed")]
+    Delete(#[from] #[redact] Box<SdkError<aws_sdk_secretsmanager::operation::delete_secret::DeleteSecretError>>),
+    #[error("AWS Secrets Manager replication failed")]
+    Replicate(#[from] #[redact] Box<SdkError<aws_sdk_secretsmanager::operation::replicate_secret_to_regions::ReplicateSecretToRegionsError>>),
+    #[error("AWS Secrets Manager response has no string payload")]
+    MissingString,
+    #[error("primary secret is not a JSON object")]
+    PrimarySecret,
+    #[error(transparent)]
+    Operation(#[from] litellm_secrets_types::Error),
+}

litellm-rust/crates/secrets-aws/src/kms.rs

@@ -0,0 +1,63 @@
+use litellm_auth_aws::constants::AWS_REGION_NAME;
+use std::sync::Arc;
+
+use aws_sdk_kms::{
+    Client,
+    config::{BehaviorVersion, Region},
+    primitives::Blob,
+};
+use litellm_core_utils::settings::Lookup;
+use litellm_secrets_types::KeyManagementSettings;
+
+use crate::{Error, auth};
+
+#[derive(Clone)]
+pub struct AwsKms {
+    client: Client,
+}
+
+impl AwsKms {
+    pub fn new(client: Client) -> Self {
+        Self { client }
+    }
+
+    pub async fn decrypt(&self, ciphertext: Vec<u8>) -> Result<Vec<u8>, Error> {
+        let response = self
+            .client
+            .decrypt()
+            .ciphertext_blob(Blob::new(ciphertext))
+            .send()
+            .await
+            .map_err(|error| Error::Decrypt(Box::new(error)))?;
+        Ok(response
+            .plaintext
+            .ok_or(Error::MissingPlaintext)?
+            .into_inner())
+    }
+}
+
+pub fn validate_environment(environment: &dyn Lookup) -> Result<(), Error> {
+    environment
+        .get(AWS_REGION_NAME)
+        .map(|_| ())
+        .ok_or(Error::MissingRegion)
+}
+
+pub fn load_aws_kms(
+    use_aws_kms: Option<bool>,
+    settings: &KeyManagementSettings,
+    environment: Arc<dyn Lookup + Send + Sync>,
+) -> Result<Option<AwsKms>, Error> {
+    if use_aws_kms != Some(true) {
+        return Ok(None);
+    }
+    if settings.aws_region_name.is_none() {
+        validate_environment(environment.as_ref())?;
+    }
+    let config = aws_sdk_kms::Config::builder()
+        .behavior_version(BehaviorVersion::latest())
+        .region(Region::new(auth::region(settings, environment.as_ref())?))
+        .credentials_provider(auth::Credentials::new(settings, environment))
+        .build();
+    Ok(Some(AwsKms::new(Client::from_conf(config))))
+}

litellm-rust/crates/secrets-aws/src/lib.rs

@@ -0,0 +1,10 @@
+#![forbid(unsafe_code)]
+
+mod auth;
+mod error;
+pub mod kms;
+pub mod secret_manager;
+
+pub use error::Error;
+pub use kms::{AwsKms, load_aws_kms};
+pub use secret_manager::{AwsSecretWriteSettings, AwsSecretsManagerV2, RotationResponse};

litellm-rust/crates/secrets-aws/src/secret_manager.rs

@@ -0,0 +1,287 @@
+use litellm_auth_aws::constants::AWS_BEDROCK_RUNTIME_ENDPOINT;
+use std::{collections::BTreeMap, sync::Arc};
+
+use aws_sdk_secretsmanager::{
+    Client,
+    config::{BehaviorVersion, Region},
+    operation::{
+        create_secret::CreateSecretOutput, delete_secret::DeleteSecretOutput,
+        put_secret_value::PutSecretValueOutput,
+        replicate_secret_to_regions::ReplicateSecretToRegionsOutput,
+    },
+    types::{ReplicaRegionType, Tag},
+};
+use litellm_auth_aws::constants::{
+    AWS_ACCESS_KEY_ID, AWS_REGION, AWS_REGION_NAME, AWS_SECRET_ACCESS_KEY,
+};
+use litellm_core_utils::settings::Lookup;
+use litellm_secrets_types::{
+    BaseSecretManager, KeyManagementSettings, Secret, SecretValue, async_rotate_secret,
+};
+use serde_json::Value;
+
+use crate::{Error, auth};
+
+#[derive(Clone)]
+pub struct AwsSecretsManagerV2 {
+    client: Client,
+    write_settings: AwsSecretWriteSettings,
+}
+
+#[derive(Clone, Debug, Default)]
+pub struct AwsSecretWriteSettings {
+    pub kms_key_id: Option<String>,
+    pub tags: Option<BTreeMap<String, String>>,
+    pub replica_regions: Option<Vec<String>>,
+}
+
+impl From<&KeyManagementSettings> for AwsSecretWriteSettings {
+    fn from(settings: &KeyManagementSettings) -> Self {
+        Self {
+            kms_key_id: settings.kms_key_id.clone(),
+            tags: settings.tags.clone(),
+            replica_regions: settings.replica_regions.clone(),
+        }
+    }
+}
+
+#[derive(Debug)]
+pub enum RotationResponse {
+    Created(CreateSecretOutput),
+    Updated(PutSecretValueOutput),
+}
+
+impl AwsSecretsManagerV2 {
+    pub fn new(client: Client, write_settings: AwsSecretWriteSettings) -> Self {
+        Self {
+            client,
+            write_settings,
+        }
+    }
+
+    pub fn load_aws_secret_manager(
+        use_aws_secret_manager: Option<bool>,
+        settings: KeyManagementSettings,
+        environment: Arc<dyn Lookup + Send + Sync>,
+    ) -> Result<Option<Self>, Error> {
+        if use_aws_secret_manager != Some(true) {
+            return Ok(None);
+        }
+        let builder = aws_sdk_secretsmanager::Config::builder()
+            .behavior_version(BehaviorVersion::latest())
+            .region(Region::new(auth::region(&settings, environment.as_ref())?))
+            .credentials_provider(auth::Credentials::new(&settings, environment.clone()));
+        let config = match environment.get(AWS_BEDROCK_RUNTIME_ENDPOINT) {
+            Some(url) => builder
+                .endpoint_url(url.replace("bedrock-runtime", "secretsmanager"))
+                .build(),
+            None => builder.build(),
+        };
+        Ok(Some(Self::new(
+            Client::from_conf(config),
+            (&settings).into(),
+        )))
+    }
+
+    pub async fn read_secret_for_resolver(
+        &self,
+        name: &str,
+        primary_name: Option<&str>,
+        environment: &(dyn Lookup + Sync),
+    ) -> Result<Option<Secret>, Error> {
+        if bootstrap_key(name) {
+            return Ok(environment
+                .get(name)
+                .map(SecretValue::new)
+                .map(Secret::String));
+        }
+        match primary_name.filter(|name| !name.is_empty()) {
+            None => self
+                .async_read_secret(name)
+                .await
+                .map(|value| value.map(Secret::String)),
+            Some(primary) => {
+                let value = if bootstrap_key(primary) {
+                    environment.get(primary).map(SecretValue::new)
+                } else {
+                    self.async_read_secret(primary).await?
+                };
+                let Some(value) = value else {
+                    return Ok(None);
+                };
+                let object: Value =
+                    serde_json::from_str(value.expose()).map_err(|_| Error::PrimarySecret)?;
+                let object = object.as_object().ok_or(Error::PrimarySecret)?;
+                Ok(object.get(name).cloned().map(Secret::from_json))
+            }
+        }
+    }
+
+    pub async fn async_read_secret(&self, name: &str) -> Result<Option<SecretValue>, Error> {
+        match self.client.get_secret_value().secret_id(name).send().await {
+            Ok(response) => response
+                .secret_string
+                .map(SecretValue::new)
+                .map(Some)
+                .ok_or(Error::MissingString),
+            Err(error)
+                if matches!(
+                    &error,
+                    aws_sdk_secretsmanager::error::SdkError::TimeoutError(_)
+                ) || matches!(&error, aws_sdk_secretsmanager::error::SdkError::DispatchFailure(failure) if failure.is_timeout()) =>
+            {
+                Err(Error::Timeout)
+            }
+            Err(error)
+                if error
+                    .as_service_error()
+                    .is_some_and(|error| error.is_resource_not_found_exception()) =>
+            {
+                Ok(None)
+            }
+            Err(error) => Err(Error::Read(Box::new(error))),
+        }
+    }
+
+    pub async fn async_write_secret(
+        &self,
+        name: &str,
+        value: &SecretValue,
+        description: Option<&str>,
+    ) -> Result<CreateSecretOutput, Error> {
+        let response = self
+            .client
+            .create_secret()
+            .name(name)
+            .secret_string(value.expose())
+            .set_description(description.filter(|v| !v.is_empty()).map(str::to_owned))
+            .set_kms_key_id(
+                self.write_settings
+                    .kms_key_id
+                    .clone()
+                    .filter(|v| !v.is_empty()),
+            )
+            .set_tags(self.write_settings.tags.as_ref().map(|tags| {
+                tags.iter()
+                    .map(|(key, value)| Tag::builder().key(key).value(value).build())
+                    .collect()
+            }))
+            .send()
+            .await
+            .map_err(|error| Error::Create(Box::new(error)))?;
+        if let Some(regions) = &self.write_settings.replica_regions
+            && !regions.is_empty()
+            && self.async_replicate_secret(name, regions).await.is_err()
+        {
+            tracing::warn!("secret created but replication failed");
+        }
+        Ok(response)
+    }
+
+    pub async fn async_replicate_secret(
+        &self,
+        name: &str,
+        regions: &[String],
+    ) -> Result<Option<ReplicateSecretToRegionsOutput>, Error> {
+        if regions.is_empty() {
+            return Ok(None);
+        }
+        self.client
+            .replicate_secret_to_regions()
+            .secret_id(name)
+            .set_add_replica_regions(Some(
+                regions
+                    .iter()
+                    .map(|region| ReplicaRegionType::builder().region(region).build())
+                    .collect(),
+            ))
+            .send()
+            .await
+            .map(Some)
+            .map_err(|error| Error::Replicate(Box::new(error)))
+    }
+
+    pub async fn async_put_secret_value(
+        &self,
+        name: &str,
+        value: &SecretValue,
+    ) -> Result<PutSecretValueOutput, Error> {
+        self.client
+            .put_secret_value()
+            .secret_id(name)
+            .secret_string(value.expose())
+            .send()
+            .await
+            .map_err(|error| Error::Put(Box::new(error)))
+    }
+
+    pub async fn async_delete_secret(
+        &self,
+        name: &str,
+        recovery_window_in_days: i64,
+    ) -> Result<DeleteSecretOutput, Error> {
+        self.client
+            .delete_secret()
+            .secret_id(name)
+            .recovery_window_in_days(recovery_window_in_days)
+            .send()
+            .await
+            .map_err(|error| Error::Delete(Box::new(error)))
+    }
+
+    pub async fn async_rotate_secret(
+        &self,
+        current_name: &str,
+        new_name: &str,
+        value: &SecretValue,
+    ) -> Result<RotationResponse, Error> {
+        if current_name == new_name {
+            return self
+                .async_put_secret_value(current_name, value)
+                .await
+                .map(RotationResponse::Updated);
+        }
+        async_rotate_secret(self, current_name, new_name, value)
+            .await
+            .map(RotationResponse::Created)
+    }
+}
+
+impl BaseSecretManager for AwsSecretsManagerV2 {
+    type Error = Error;
+    type WriteResponse = CreateSecretOutput;
+    type DeleteResponse = DeleteSecretOutput;
+
+    async fn async_read_secret(&self, name: &str) -> Result<Option<SecretValue>, Error> {
+        self.async_read_secret(name).await
+    }
+
+    async fn async_write_secret(
+        &self,
+        name: &str,
+        value: &SecretValue,
+        description: Option<&str>,
+    ) -> Result<CreateSecretOutput, Error> {
+        self.async_write_secret(name, value, description).await
+    }
+
+    async fn async_delete_secret(
+        &self,
+        name: &str,
+        recovery_window_in_days: i64,
+    ) -> Result<DeleteSecretOutput, Error> {
+        self.async_delete_secret(name, recovery_window_in_days)
+            .await
+    }
+}
+
+fn bootstrap_key(name: &str) -> bool {
+    matches!(
+        name,
+        AWS_ACCESS_KEY_ID
+            | AWS_SECRET_ACCESS_KEY
+            | AWS_REGION_NAME
+            | AWS_REGION
+            | AWS_BEDROCK_RUNTIME_ENDPOINT
+    )
+}

litellm-rust/crates/secrets-aws/tests/kms.rs

@@ -0,0 +1,59 @@
+use aws_sdk_kms::{
+    Client,
+    config::{BehaviorVersion, Credentials, Region, retry::RetryConfig},
+};
+use base64::{Engine, engine::general_purpose::STANDARD};
+use litellm_secrets_aws::AwsKms;
+use wiremock::{
+    Mock, MockServer, ResponseTemplate,
+    matchers::{body_json, header},
+};
+
+#[tokio::test]
+async fn kms_decrypt_calls_the_sdk_without_applying_lookup_policy() {
+    let server = MockServer::start().await;
+    let plaintext = "  private-value\n";
+    Mock::given(header("x-amz-target", "TrentService.Decrypt"))
+        .and(body_json(
+            serde_json::json!({"CiphertextBlob": STANDARD.encode("encrypted")}),
+        ))
+        .respond_with(
+            ResponseTemplate::new(200)
+                .set_body_json(serde_json::json!({"Plaintext": STANDARD.encode(plaintext)})),
+        )
+        .expect(1)
+        .mount(&server)
+        .await;
+    let client = Client::from_conf(
+        aws_sdk_kms::Config::builder()
+            .behavior_version(BehaviorVersion::latest())
+            .region(Region::new("us-east-1"))
+            .credentials_provider(Credentials::new("test", "test", None, None, "test"))
+            .endpoint_url(server.uri())
+            .retry_config(RetryConfig::disabled())
+            .build(),
+    );
+    let manager = AwsKms::new(client);
+    assert_eq!(
+        manager.decrypt(b"encrypted".to_vec()).await.unwrap(),
+        plaintext.as_bytes()
+    );
+}
+
+#[test]
+fn disabled_kms_loader_does_not_require_environment_configuration() {
+    use litellm_secrets_aws::load_aws_kms;
+    use litellm_secrets_types::KeyManagementSettings;
+    use std::sync::Arc;
+    for enabled in [None, Some(false)] {
+        assert!(
+            load_aws_kms(
+                enabled,
+                &KeyManagementSettings::default(),
+                Arc::new(|_: &str| None)
+            )
+            .unwrap()
+            .is_none()
+        );
+    }
+}

litellm-rust/crates/secrets-aws/tests/secret_manager.rs

@@ -0,0 +1,312 @@
+use std::sync::atomic::{AtomicUsize, Ordering};
+
+use aws_sdk_secretsmanager::{
+    Client,
+    config::{BehaviorVersion, Credentials, Region, retry::RetryConfig},
+};
+use litellm_secrets_aws::{AwsSecretsManagerV2, Error, RotationResponse};
+use litellm_secrets_types::{KeyManagementSettings, SecretValue};
+use serde_json::json;
+use wiremock::{
+    Mock, MockServer, ResponseTemplate,
+    matchers::{body_partial_json, header},
+};
+
+fn manager(server: &MockServer, settings: KeyManagementSettings) -> AwsSecretsManagerV2 {
+    let client = Client::from_conf(
+        aws_sdk_secretsmanager::Config::builder()
+            .behavior_version(BehaviorVersion::latest())
+            .region(Region::new("us-east-1"))
+            .credentials_provider(Credentials::new("test", "test", None, None, "test"))
+            .endpoint_url(server.uri())
+            .retry_config(RetryConfig::disabled())
+            .build(),
+    );
+    AwsSecretsManagerV2::new(client, (&settings).into())
+}
+
+#[rstest::rstest]
+#[case::string_value("KEY", Some("value"))]
+#[case::missing_value("missing", None)]
+#[case::non_string_value("BOOL", None)]
+#[tokio::test]
+async fn primary_lookup_preserves_read_semantics(
+    #[case] name: &str,
+    #[case] expected: Option<&str>,
+) {
+    let server = MockServer::start().await;
+    Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue"))
+        .and(body_partial_json(json!({"SecretId":"primary"})))
+        .respond_with(
+            ResponseTemplate::new(200).set_body_json(
+                json!({"SecretString":json!({"KEY":"value", "BOOL":true}).to_string()}),
+            ),
+        )
+        .expect(1)
+        .mount(&server)
+        .await;
+    let manager = manager(&server, KeyManagementSettings::default());
+    assert_eq!(
+        manager
+            .read_secret_for_resolver(name, Some("primary"), &|_: &str| None)
+            .await
+            .unwrap()
+            .and_then(|v| v.as_str().map(str::to_owned))
+            .as_deref(),
+        expected
+    );
+}
+
+#[rstest::rstest]
+#[case::access_key("AWS_ACCESS_KEY_ID")]
+#[case::secret_access_key("AWS_SECRET_ACCESS_KEY")]
+#[case::region_name("AWS_REGION_NAME")]
+#[case::region("AWS_REGION")]
+#[case::bedrock_endpoint("AWS_BEDROCK_RUNTIME_ENDPOINT")]
+#[tokio::test]
+async fn bootstrap_keys_bypass_primary_lookup(#[case] name: &str) {
+    let server = MockServer::start().await;
+    let manager = manager(&server, KeyManagementSettings::default());
+    assert_eq!(
+        manager
+            .read_secret_for_resolver(name, Some("primary"), &|_: &str| Some("bootstrap".into()))
+            .await
+            .unwrap()
+            .unwrap()
+            .as_str()
+            .unwrap(),
+        "bootstrap"
+    );
+}
+
+#[tokio::test]
+async fn failed_read_returns_none_but_invalid_primary_json_is_an_error() {
+    let server = MockServer::start().await;
+    Mock::given(body_partial_json(json!({"SecretId":"missing"})))
+        .respond_with(
+            ResponseTemplate::new(400).set_body_json(json!({"__type":"ResourceNotFoundException"})),
+        )
+        .mount(&server)
+        .await;
+    Mock::given(body_partial_json(json!({"SecretId":"invalid"})))
+        .respond_with(ResponseTemplate::new(200).set_body_json(json!({"SecretString":"not-json"})))
+        .mount(&server)
+        .await;
+    let manager = manager(&server, KeyManagementSettings::default());
+    assert!(
+        manager
+            .async_read_secret("missing")
+            .await
+            .unwrap()
+            .is_none()
+    );
+    assert!(matches!(
+        manager
+            .read_secret_for_resolver("KEY", Some("invalid"), &|_: &str| None)
+            .await,
+        Err(Error::PrimarySecret)
+    ));
+}
+
+#[tokio::test]
+async fn same_name_rotation_uses_put_and_returns_its_response() {
+    let server = MockServer::start().await;
+    Mock::given(header("x-amz-target", "secretsmanager.PutSecretValue"))
+        .and(body_partial_json(
+            json!({"SecretId":"key", "SecretString":"replacement"}),
+        ))
+        .respond_with(
+            ResponseTemplate::new(200).set_body_json(json!({"Name":"key", "VersionId":"version"})),
+        )
+        .expect(1)
+        .mount(&server)
+        .await;
+    let response = manager(&server, KeyManagementSettings::default())
+        .async_rotate_secret("key", "key", &SecretValue::new("replacement"))
+        .await
+        .unwrap();
+    match response {
+        RotationResponse::Updated(output) => assert_eq!(output.version_id(), Some("version")),
+        _ => panic!("rotation created a second secret"),
+    }
+    assert_eq!(server.received_requests().await.unwrap().len(), 1);
+}
+
+#[tokio::test]
+async fn renamed_rotation_reads_creates_verifies_then_deletes() {
+    let server = MockServer::start().await;
+    let step = AtomicUsize::new(0);
+    Mock::given(wiremock::matchers::method("POST"))
+        .respond_with(move |request: &wiremock::Request| {
+            let body: serde_json::Value = request.body_json().unwrap();
+            let action = request
+                .headers
+                .get("x-amz-target")
+                .unwrap()
+                .to_str()
+                .unwrap();
+            match step.fetch_add(1, Ordering::SeqCst) {
+                0 => {
+                    assert_eq!(action, "secretsmanager.GetSecretValue");
+                    assert_eq!(body["SecretId"], "old");
+                    ResponseTemplate::new(200).set_body_json(json!({"SecretString":"old-value"}))
+                }
+                1 => {
+                    assert_eq!(action, "secretsmanager.CreateSecret");
+                    assert_eq!(body["Name"], "new");
+                    assert_eq!(body["Description"], "Rotated from old");
+                    assert_eq!(body["SecretString"], "replacement");
+                    ResponseTemplate::new(200).set_body_json(json!({"Name":"new"}))
+                }
+                2 => {
+                    assert_eq!(action, "secretsmanager.GetSecretValue");
+                    assert_eq!(body["SecretId"], "new");
+                    ResponseTemplate::new(200).set_body_json(json!({"SecretString":"replacement"}))
+                }
+                3 => {
+                    assert_eq!(action, "secretsmanager.DeleteSecret");
+                    assert_eq!(body["SecretId"], "old");
+                    assert_eq!(body["RecoveryWindowInDays"], 7);
+                    ResponseTemplate::new(200).set_body_json(json!({"Name":"old"}))
+                }
+                _ => panic!("unexpected request"),
+            }
+        })
+        .expect(4)
+        .mount(&server)
+        .await;
+    assert!(matches!(
+        manager(&server, KeyManagementSettings::default())
+            .async_rotate_secret("old", "new", &SecretValue::new("replacement"))
+            .await
+            .unwrap(),
+        RotationResponse::Created(_)
+    ));
+}
+
+#[tokio::test]
+async fn creation_passes_tags_and_kms_and_survives_replication_failure() {
+    let server = MockServer::start().await;
+    Mock::given(header("x-amz-target", "secretsmanager.CreateSecret"))
+        .and(body_partial_json(json!({"Name":"key", "SecretString":"value", "KmsKeyId":"kms-key", "Tags":[{"Key":"stage", "Value":"test"}]})))
+        .respond_with(ResponseTemplate::new(200).set_body_json(json!({"Name":"key"}))).expect(1).mount(&server).await;
+    Mock::given(header(
+        "x-amz-target",
+        "secretsmanager.ReplicateSecretToRegions",
+    ))
+    .and(body_partial_json(
+        json!({"SecretId":"key", "AddReplicaRegions":[{"Region":"replica-region"}]}),
+    ))
+    .respond_with(
+        ResponseTemplate::new(400).set_body_json(json!({"__type":"InvalidRequestException"})),
+    )
+    .expect(1)
+    .mount(&server)
+    .await;
+    let settings = KeyManagementSettings {
+        kms_key_id: Some("kms-key".into()),
+        tags: Some(std::collections::BTreeMap::from([(
+            "stage".into(),
+            "test".into(),
+        )])),
+        replica_regions: Some(vec!["replica-region".into()]),
+        ..Default::default()
+    };
+    let manager = manager(&server, settings);
+    assert_eq!(
+        manager
+            .async_write_secret("key", &SecretValue::new("value"), None)
+            .await
+            .unwrap()
+            .name(),
+        Some("key")
+    );
+    assert!(
+        manager
+            .async_replicate_secret("key", &[])
+            .await
+            .unwrap()
+            .is_none()
+    );
+}
+
+#[tokio::test]
+async fn credential_failures_are_not_swallowed_as_missing_secrets() {
+    use aws_credential_types::provider::{ProvideCredentials, error::CredentialsError, future};
+    #[derive(Debug)]
+    struct FailedCredentials;
+    impl ProvideCredentials for FailedCredentials {
+        fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a>
+        where
+            Self: 'a,
+        {
+            future::ProvideCredentials::ready(Err(CredentialsError::provider_error(
+                "private-auth-detail",
+            )))
+        }
+    }
+    let server = MockServer::start().await;
+    let config = aws_sdk_secretsmanager::Config::builder()
+        .behavior_version(BehaviorVersion::latest())
+        .region(Region::new("us-east-1"))
+        .credentials_provider(FailedCredentials)
+        .endpoint_url(server.uri())
+        .retry_config(RetryConfig::disabled())
+        .build();
+    let manager = AwsSecretsManagerV2::new(Client::from_conf(config), Default::default());
+    let error = manager.async_read_secret("key").await.unwrap_err();
+    assert!(!format!("{error:?}").contains("private-auth-detail"));
+    assert!(matches!(error, Error::Read(_)));
+    assert!(server.received_requests().await.unwrap().is_empty());
+}
+
+#[tokio::test]
+async fn read_timeout_is_an_error_and_cannot_be_mistaken_for_missing() {
+    use std::time::Duration;
+    let server = MockServer::start().await;
+    Mock::given(wiremock::matchers::method("POST"))
+        .respond_with(
+            ResponseTemplate::new(200)
+                .set_delay(Duration::from_secs(1))
+                .set_body_json(json!({"SecretString":"late"})),
+        )
+        .mount(&server)
+        .await;
+    let config = aws_sdk_secretsmanager::Config::builder()
+        .behavior_version(BehaviorVersion::latest())
+        .region(Region::new("us-east-1"))
+        .credentials_provider(Credentials::new("test", "test", None, None, "test"))
+        .endpoint_url(server.uri())
+        .retry_config(RetryConfig::disabled())
+        .timeout_config(
+            aws_sdk_secretsmanager::config::timeout::TimeoutConfig::builder()
+                .operation_timeout(Duration::from_millis(30))
+                .build(),
+        )
+        .build();
+    let manager = AwsSecretsManagerV2::new(Client::from_conf(config), Default::default());
+    assert!(matches!(
+        manager.async_read_secret("key").await,
+        Err(Error::Timeout)
+    ));
+}
+
+#[rstest::rstest]
+#[case::denied(400, "AccessDeniedException")]
+#[case::throttled(400, "ThrottlingException")]
+#[case::unavailable(503, "ServiceUnavailableException")]
+#[tokio::test]
+async fn service_failures_remain_errors(#[case] status: u16, #[case] code: &str) {
+    let server = MockServer::start().await;
+    Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue"))
+        .respond_with(ResponseTemplate::new(status).set_body_json(json!({"__type":code})))
+        .expect(1)
+        .mount(&server)
+        .await;
+    assert!(matches!(
+        manager(&server, KeyManagementSettings::default())
+            .async_read_secret("key")
+            .await,
+        Err(Error::Read(_))
+    ));
+}

litellm-rust/crates/secrets-google/Cargo.toml

@@ -0,0 +1,28 @@
+[package]
+name = "litellm-secrets-google"
+version = "0.1.0"
+edition.workspace = true
+license.workspace = true
+repository.workspace = true
+
+[dependencies]
+litellm-auth-gcp = { workspace = true, features = ["google-sdk"] }
+litellm-secrets-types.workspace = true
+litellm-auth-types.workspace = true
+litellm-core-utils.workspace = true
+base64.workspace = true
+serde_json.workspace = true
+thiserror.workspace = true
+moka.workspace = true
+veil.workspace = true
+google-cloud-kms-v1 = "1.14.0"
+google-cloud-gax = { version = "1.14.0", default-features = false }
+percent-encoding = "2.3"
+serde.workspace = true
+reqwest.workspace = true
+
+[dev-dependencies]
+google-cloud-auth.workspace = true
+rstest.workspace = true
+tokio.workspace = true
+wiremock = "0.6.5"

litellm-rust/crates/secrets-google/src/auth.rs

@@ -0,0 +1,21 @@
+use std::sync::Arc;
+
+use litellm_auth_gcp::{GoogleCredentials, VertexConfig};
+use litellm_auth_types::{InputSource, Sourced};
+use litellm_core_utils::settings::Lookup;
+use litellm_secrets_types::SecretValue;
+
+pub(crate) fn credentials(
+    project: Option<String>,
+    credentials: Option<SecretValue>,
+    environment: Arc<dyn Lookup + Send + Sync>,
+) -> GoogleCredentials {
+    GoogleCredentials::new(
+        VertexConfig::new(
+            credentials.map(|value| Sourced::new(value, InputSource::Environment)),
+            project,
+            None,
+        ),
+        Arc::new(move |name| environment.get(name)),
+    )
+}

litellm-rust/crates/secrets-google/src/error.rs

@@ -0,0 +1,43 @@
+#[derive(thiserror::Error, veil::Redact)]
+pub enum Error {
+    #[error("Google KMS client configuration failed")]
+    Client(
+        #[from]
+        #[redact]
+        google_cloud_gax::client_builder::Error,
+    ),
+    #[error("Google authentication failed")]
+    Auth(
+        #[from]
+        #[redact]
+        litellm_auth_types::Error,
+    ),
+    #[error("Google KMS request failed")]
+    Kms(
+        #[from]
+        #[redact]
+        google_cloud_gax::error::Error,
+    ),
+    #[error("Google Secret Manager HTTP request failed")]
+    Http(
+        #[from]
+        #[redact]
+        reqwest::Error,
+    ),
+    #[error("Google Secret Manager returned HTTP {0}")]
+    Status(u16),
+    #[error("Google Secret Manager returned no payload")]
+    MissingPayload,
+    #[error("required environment variable is missing: {0}")]
+    MissingEnvironment(&'static str),
+    #[error("invalid refresh interval")]
+    RefreshInterval,
+    #[error("payload is not valid base64")]
+    Base64(#[from] base64::DecodeError),
+    #[error("decrypted value is not UTF-8")]
+    Utf8,
+    #[error("invalid Google Secret Manager endpoint")]
+    Endpoint,
+    #[error("Google Secret Manager requires an enterprise license")]
+    EnterpriseRequired,
+}

litellm-rust/crates/secrets-google/src/kms.rs

@@ -0,0 +1,67 @@
+use std::sync::Arc;
+
+use google_cloud_kms_v1::client::KeyManagementService;
+use litellm_core_utils::settings::Lookup;
+use litellm_secrets_types::SecretValue;
+
+use crate::{Error, auth};
+
+const GOOGLE_APPLICATION_CREDENTIALS: &str = "GOOGLE_APPLICATION_CREDENTIALS";
+const GOOGLE_KMS_RESOURCE_NAME: &str = "GOOGLE_KMS_RESOURCE_NAME";
+
+#[derive(Clone)]
+pub struct GoogleKms {
+    client: KeyManagementService,
+    resource_name: String,
+}
+
+impl GoogleKms {
+    pub fn new(client: KeyManagementService, resource_name: String) -> Self {
+        Self {
+            client,
+            resource_name,
+        }
+    }
+
+    pub async fn decrypt(&self, ciphertext: Vec<u8>) -> Result<Vec<u8>, Error> {
+        let response = self
+            .client
+            .decrypt()
+            .set_name(&self.resource_name)
+            .set_ciphertext(ciphertext)
+            .send()
+            .await?;
+        Ok(response.plaintext.to_vec())
+    }
+}
+
+pub fn validate_environment(environment: &dyn Lookup) -> Result<(), Error> {
+    for key in [GOOGLE_APPLICATION_CREDENTIALS, GOOGLE_KMS_RESOURCE_NAME] {
+        if environment.get(key).is_none() {
+            return Err(Error::MissingEnvironment(key));
+        }
+    }
+    Ok(())
+}
+
+pub async fn load_google_kms(
+    use_google_kms: Option<bool>,
+    environment: Arc<dyn Lookup + Send + Sync>,
+) -> Result<Option<GoogleKms>, Error> {
+    if use_google_kms != Some(true) {
+        return Ok(None);
+    }
+    validate_environment(environment.as_ref())?;
+    let credentials = environment
+        .get(GOOGLE_APPLICATION_CREDENTIALS)
+        .ok_or(Error::MissingEnvironment(GOOGLE_APPLICATION_CREDENTIALS))?;
+    let resource_name = environment
+        .get(GOOGLE_KMS_RESOURCE_NAME)
+        .ok_or(Error::MissingEnvironment(GOOGLE_KMS_RESOURCE_NAME))?;
+    let credentials = auth::credentials(None, Some(SecretValue::new(credentials)), environment);
+    let client = KeyManagementService::builder()
+        .with_credentials(credentials)
+        .build()
+        .await?;
+    Ok(Some(GoogleKms::new(client, resource_name)))
+}

litellm-rust/crates/secrets-google/src/lib.rs

@@ -0,0 +1,10 @@
+#![forbid(unsafe_code)]
+
+mod auth;
+mod error;
+pub mod kms;
+pub mod secret_manager;
+
+pub use error::Error;
+pub use kms::{GoogleKms, load_google_kms};
+pub use secret_manager::GoogleSecretManager;

litellm-rust/crates/secrets-google/src/secret_manager.rs

@@ -0,0 +1,157 @@
+use std::{sync::Arc, time::Duration};
+
+use base64::{Engine, engine::general_purpose::STANDARD};
+use litellm_core_utils::settings::Lookup;
+use litellm_secrets_types::{Secret, SecretValue};
+use moka::future::Cache;
+use serde::Deserialize;
+
+use litellm_auth_gcp::GoogleCredentials;
+
+use crate::{Error, auth};
+
+const GOOGLE_SECRET_MANAGER_PROJECT_ID: &str = "GOOGLE_SECRET_MANAGER_PROJECT_ID";
+const GOOGLE_SECRET_MANAGER_REFRESH_INTERVAL: &str = "GOOGLE_SECRET_MANAGER_REFRESH_INTERVAL";
+const SECRET_MANAGER_REFRESH_INTERVAL: &str = "SECRET_MANAGER_REFRESH_INTERVAL";
+const GOOGLE_SECRET_MANAGER_ALWAYS_READ_SECRET_MANAGER: &str =
+    "GOOGLE_SECRET_MANAGER_ALWAYS_READ_SECRET_MANAGER";
+const GCS_PATH_SERVICE_ACCOUNT: &str = "GCS_PATH_SERVICE_ACCOUNT";
+const DEFAULT_REFRESH_INTERVAL: Duration = Duration::from_secs(86400);
+const DEFAULT_CACHE_TTL: Duration = Duration::from_secs(600);
+const CACHE_CAPACITY: u64 = 200;
+
+#[derive(Clone)]
+pub struct GoogleSecretManager {
+    client: reqwest::Client,
+    credentials: Arc<GoogleCredentials>,
+    endpoint: reqwest::Url,
+    project: String,
+    cache: Cache<String, SecretValue>,
+    always_read: bool,
+}
+
+#[derive(Deserialize)]
+struct Response {
+    payload: Option<Payload>,
+}
+
+#[derive(Deserialize)]
+struct Payload {
+    data: Option<String>,
+}
+
+impl GoogleSecretManager {
+    pub fn with_client(
+        client: reqwest::Client,
+        endpoint: reqwest::Url,
+        project: String,
+        environment: Arc<dyn Lookup + Send + Sync>,
+        refresh_interval: Option<Duration>,
+        always_read: bool,
+    ) -> Result<Self, Error> {
+        let credentials = auth::credentials(
+            Some(project.clone()),
+            environment
+                .get(GCS_PATH_SERVICE_ACCOUNT)
+                .map(SecretValue::new),
+            environment,
+        );
+        let ttl = refresh_interval
+            .filter(|ttl| !ttl.is_zero())
+            .unwrap_or(DEFAULT_CACHE_TTL);
+        let cache = Cache::builder()
+            .max_capacity(CACHE_CAPACITY)
+            .time_to_live(ttl)
+            .build();
+        Ok(Self {
+            client,
+            credentials: Arc::new(credentials),
+            endpoint,
+            project,
+            cache,
+            always_read,
+        })
+    }
+
+    pub fn new(
+        environment: Arc<dyn Lookup + Send + Sync>,
+        enterprise_enabled: bool,
+    ) -> Result<Self, Error> {
+        if !enterprise_enabled {
+            return Err(Error::EnterpriseRequired);
+        }
+        let project = environment
+            .get(GOOGLE_SECRET_MANAGER_PROJECT_ID)
+            .ok_or(Error::MissingEnvironment(GOOGLE_SECRET_MANAGER_PROJECT_ID))?;
+        let ttl = environment
+            .get(GOOGLE_SECRET_MANAGER_REFRESH_INTERVAL)
+            .filter(|v| !v.is_empty())
+            .map(|v| v.parse::<i64>().map_err(|_| Error::RefreshInterval))
+            .transpose()?
+            .unwrap_or(
+                environment
+                    .get(SECRET_MANAGER_REFRESH_INTERVAL)
+                    .map(|v| v.parse::<i64>().map_err(|_| Error::RefreshInterval))
+                    .transpose()?
+                    .unwrap_or(DEFAULT_REFRESH_INTERVAL.as_secs() as i64),
+            );
+        let always_read = environment
+            .get(GOOGLE_SECRET_MANAGER_ALWAYS_READ_SECRET_MANAGER)
+            .is_some_and(|v| v.eq_ignore_ascii_case("true"));
+        Self::with_client(
+            reqwest::Client::new(),
+            reqwest::Url::parse("https://secretmanager.googleapis.com").expect("static URL"),
+            project,
+            environment,
+            Some(if ttl < 0 {
+                Duration::from_nanos(1)
+            } else {
+                Duration::from_secs(ttl as u64)
+            }),
+            always_read,
+        )
+    }
+
+    pub async fn get_secret_from_google_secret_manager(
+        &self,
+        name: &str,
+    ) -> Result<Option<Secret>, Error> {
+        if !self.always_read
+            && let Some(cached) = self.cache.get(name).await
+        {
+            return Ok(Some(Secret::String(cached)));
+        }
+        let url = self
+            .endpoint
+            .join(&format!(
+                "/v1/projects/{}/secrets/{}/versions/latest:access",
+                percent_encoding::utf8_percent_encode(
+                    &self.project,
+                    percent_encoding::NON_ALPHANUMERIC
+                ),
+                percent_encoding::utf8_percent_encode(name, percent_encoding::NON_ALPHANUMERIC)
+            ))
+            .map_err(|_| Error::Endpoint)?;
+        let response = self
+            .client
+            .get(url)
+            .headers(self.credentials.request_headers().await?)
+            .send()
+            .await?;
+        if response.status() == reqwest::StatusCode::NOT_FOUND {
+            return Ok(None);
+        }
+        if response.status() != reqwest::StatusCode::OK {
+            return Err(Error::Status(response.status().as_u16()));
+        }
+        let response: Response = response.json().await?;
+        let Some(data) = response.payload.and_then(|payload| payload.data) else {
+            return Err(Error::MissingPayload);
+        };
+        let bytes = STANDARD.decode(data)?;
+        let plaintext = String::from_utf8(bytes).map_err(|_| Error::Utf8)?;
+        let value = SecretValue::new(plaintext);
+        self.cache.insert(name.to_owned(), value.clone()).await;
+        Ok(Some(Secret::String(value)))
+    }
+}

litellm-rust/crates/secrets-google/tests/kms.rs

@@ -0,0 +1,49 @@
+use base64::{Engine, engine::general_purpose::STANDARD};
+use google_cloud_kms_v1::client::KeyManagementService;
+use litellm_secrets_google::GoogleKms;
+use wiremock::{
+    Mock, MockServer, ResponseTemplate,
+    matchers::{body_json, path},
+};
+
+#[tokio::test]
+async fn google_kms_decrypts_using_the_configured_resource() {
+    let server = MockServer::start().await;
+    let resource = "projects/project/locations/global/keyRings/ring/cryptoKeys/key";
+    Mock::given(path(format!("/v1/{resource}:decrypt")))
+        .and(body_json(
+            serde_json::json!({"ciphertext":STANDARD.encode("encrypted")}),
+        ))
+        .respond_with(
+            ResponseTemplate::new(200)
+                .set_body_json(serde_json::json!({"plaintext":STANDARD.encode(" value\n")})),
+        )
+        .expect(1)
+        .mount(&server)
+        .await;
+    let client = KeyManagementService::builder()
+        .with_endpoint(server.uri())
+        .with_credentials(google_cloud_auth::credentials::anonymous::Builder::new().build())
+        .with_retry_policy(google_cloud_gax::retry_policy::NeverRetry)
+        .build()
+        .await
+        .unwrap();
+    let manager = GoogleKms::new(client, resource.into());
+    assert_eq!(
+        manager.decrypt(b"encrypted".to_vec()).await.unwrap(),
+        b" value\n"
+    );
+}
+
+#[tokio::test]
+async fn disabled_google_kms_loader_does_not_require_environment_configuration() {
+    use std::sync::Arc;
+    for enabled in [None, Some(false)] {
+        assert!(
+            litellm_secrets_google::load_google_kms(enabled, Arc::new(|_: &str| None))
+                .await
+                .unwrap()
+                .is_none()
+        );
+    }
+}

litellm-rust/crates/secrets-google/tests/secret_manager.rs

@@ -0,0 +1,188 @@
+use std::{sync::Arc, time::Duration};
+
+use base64::{Engine, engine::general_purpose::STANDARD};
+use litellm_secrets_google::{Error, GoogleSecretManager};
+
+use wiremock::{
+    Mock, MockServer, ResponseTemplate,
+    matchers::{header, path},
+};
+
+fn manager(server: &MockServer, always_read: bool, ttl: Duration) -> GoogleSecretManager {
+    GoogleSecretManager::with_client(
+        reqwest::Client::new(),
+        server.uri().parse().unwrap(),
+        "project".into(),
+        Arc::new(|name: &str| (name == "VERTEX_AI_API_KEY").then(|| "token".into())),
+        Some(ttl),
+        always_read,
+    )
+    .unwrap()
+}
+
+#[rstest::rstest]
+#[case::nonempty("private-value")]
+#[case::empty("")]
+#[tokio::test]
+async fn successful_reads_use_auth_latest_version_and_cache_including_empty_values(
+    #[case] value: &str,
+) {
+    let server = MockServer::start().await;
+    Mock::given(path(
+        "/v1/projects/project/secrets/key/versions/latest:access",
+    ))
+    .and(header("authorization", "Bearer token"))
+    .respond_with(
+        ResponseTemplate::new(200)
+            .set_body_json(serde_json::json!({"payload":{"data":STANDARD.encode(value)}})),
+    )
+    .expect(1)
+    .mount(&server)
+    .await;
+    let manager = manager(&server, false, Duration::from_secs(60));
+    for _ in 0..2 {
+        assert_eq!(
+            manager
+                .get_secret_from_google_secret_manager("key")
+                .await
+                .unwrap()
+                .unwrap()
+                .as_str()
+                .unwrap(),
+            value
+        );
+    }
+}
+
+#[rstest::rstest]
+#[case::not_found(404, serde_json::json!({}))]
+#[case::unauthorized(401, serde_json::json!({}))]
+#[case::forbidden(403, serde_json::json!({}))]
+#[case::throttled(429, serde_json::json!({}))]
+#[case::unavailable(503, serde_json::json!({}))]
+#[case::missing_payload(200, serde_json::json!({"payload":{}}))]
+#[case::invalid_base64(200, serde_json::json!({"payload":{"data":"%%%"}}))]
+#[tokio::test]
+async fn failed_or_missing_reads_are_not_cached(
+    #[case] status: u16,
+    #[case] body: serde_json::Value,
+) {
+    let server = MockServer::start().await;
+    let manager = manager(&server, false, Duration::from_secs(60));
+    let failing = Mock::given(path(
+        "/v1/projects/project/secrets/key/versions/latest:access",
+    ))
+    .respond_with(ResponseTemplate::new(status).set_body_json(body))
+    .expect(1)
+    .mount_as_scoped(&server)
+    .await;
+    let result = manager.get_secret_from_google_secret_manager("key").await;
+    match status {
+        404 => assert_eq!(result.unwrap(), None),
+        200 => assert!(matches!(
+            result,
+            Err(Error::MissingPayload | Error::Base64(_))
+        )),
+        status => assert!(matches!(result, Err(Error::Status(actual)) if actual == status)),
+    }
+    drop(failing);
+    Mock::given(path(
+        "/v1/projects/project/secrets/key/versions/latest:access",
+    ))
+    .respond_with(
+        ResponseTemplate::new(200)
+            .set_body_json(serde_json::json!({"payload":{"data":STANDARD.encode("recovered")}})),
+    )
+    .expect(1)
+    .mount(&server)
+    .await;
+    for _ in 0..2 {
+        assert_eq!(
+            manager
+                .get_secret_from_google_secret_manager("key")
+                .await
+                .unwrap()
+                .unwrap()
+                .as_str(),
+            Some("recovered")
+        );
+    }
+}
+
+#[rstest::rstest]
+#[case::always_read(true, Duration::from_secs(60))]
+#[case::expired_cache(false, Duration::from_millis(1))]
+#[tokio::test]
+async fn always_read_and_expired_cache_fetch_again(
+    #[case] always_read: bool,
+    #[case] ttl: Duration,
+) {
+    let server = MockServer::start().await;
+    Mock::given(path(
+        "/v1/projects/project/secrets/key/versions/latest:access",
+    ))
+    .respond_with(
+        ResponseTemplate::new(200)
+            .set_body_json(serde_json::json!({"payload":{"data":STANDARD.encode("value")}})),
+    )
+    .expect(2)
+    .mount(&server)
+    .await;
+    let manager = manager(&server, always_read, ttl);
+    for _ in 0..2 {
+        tokio::time::sleep(Duration::from_millis(5)).await;
+        assert!(
+            manager
+                .get_secret_from_google_secret_manager("key")
+                .await
+                .unwrap()
+                .is_some()
+        );
+    }
+}
+
+#[test]
+fn google_manager_requires_host_license_and_project_configuration() {
+    assert!(matches!(
+        GoogleSecretManager::new(Arc::new(|_: &str| None), false),
+        Err(Error::EnterpriseRequired)
+    ));
+    assert!(matches!(
+        GoogleSecretManager::new(Arc::new(|_: &str| None), true),
+        Err(Error::MissingEnvironment(
+            "GOOGLE_SECRET_MANAGER_PROJECT_ID"
+        ))
+    ));
+}
+
+#[rstest::rstest]
+#[case("true")]
+#[case("null")]
+#[case("\"text\"")]
+#[case("{\"key\":1}")]
+#[tokio::test]
+async fn cache_preserves_raw_values(#[case] raw: &str) {
+    let server = MockServer::start().await;
+    Mock::given(path(
+        "/v1/projects/project/secrets/key/versions/latest:access",
+    ))
+    .respond_with(
+        ResponseTemplate::new(200)
+            .set_body_json(serde_json::json!({"payload":{"data":STANDARD.encode(raw)}})),
+    )
+    .expect(1)
+    .mount(&server)
+    .await;
+    let manager = manager(&server, false, Duration::from_secs(60));
+    for _ in 0..2 {
+        assert_eq!(
+            manager
+                .get_secret_from_google_secret_manager("key")
+                .await
+                .unwrap()
+                .unwrap()
+                .as_str(),
+            Some(raw)
+        );
+    }
+}

litellm-rust/crates/secrets-types/Cargo.toml

@@ -0,0 +1,17 @@
+[package]
+name = "litellm-secrets-types"
+version = "0.1.0"
+edition.workspace = true
+license.workspace = true
+repository.workspace = true
+
+[dependencies]
+litellm-auth-types.workspace = true
+serde.workspace = true
+serde_json.workspace = true
+thiserror.workspace = true
+veil.workspace = true
+
+[dev-dependencies]
+rstest.workspace = true
+tokio.workspace = true

litellm-rust/crates/secrets-types/src/base_secret_manager.rs

@@ -0,0 +1,58 @@
+use crate::{Error, SecretValue};
+
+pub fn validate_secret_name(name: &str) -> Result<(), Error> {
+    if name.split('/').any(|segment| segment == "..")
+        || name
+            .chars()
+            .any(|c| c.is_control() || matches!(c, '\u{2028}' | '\u{2029}'))
+    {
+        return Err(Error::UnsafeSecretName);
+    }
+    Ok(())
+}
+
+#[expect(
+    async_fn_in_trait,
+    reason = "closed backend dispatch does not require Send bounds on generic rotation"
+)]
+pub trait BaseSecretManager {
+    type Error: From<Error>;
+    type WriteResponse;
+    type DeleteResponse;
+
+    async fn async_read_secret(&self, name: &str) -> Result<Option<SecretValue>, Self::Error>;
+    async fn async_write_secret(
+        &self,
+        name: &str,
+        value: &SecretValue,
+        description: Option<&str>,
+    ) -> Result<Self::WriteResponse, Self::Error>;
+    async fn async_delete_secret(
+        &self,
+        name: &str,
+        recovery_window_in_days: i64,
+    ) -> Result<Self::DeleteResponse, Self::Error>;
+}
+
+pub async fn async_rotate_secret<M: BaseSecretManager>(
+    manager: &M,
+    current_name: &str,
+    new_name: &str,
+    value: &SecretValue,
+) -> Result<M::WriteResponse, M::Error> {
+    if manager.async_read_secret(current_name).await?.is_none() {
+        return Err(Error::CurrentSecretMissing.into());
+    }
+    let response = manager
+        .async_write_secret(
+            new_name,
+            value,
+            Some(&format!("Rotated from {current_name}")),
+        )
+        .await?;
+    if manager.async_read_secret(new_name).await?.is_none() {
+        return Err(Error::NewSecretMissing.into());
+    }
+    manager.async_delete_secret(current_name, 7).await?;
+    Ok(response)
+}

litellm-rust/crates/secrets-types/src/config.rs

@@ -0,0 +1,92 @@
+use std::collections::BTreeMap;
+
+use serde::{Deserialize, Serialize};
+
+use crate::SecretValue;
+
+#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
+#[serde(rename_all = "snake_case")]
+pub enum KeyManagementSystem {
+    GoogleKms,
+    AzureKeyVault,
+    AwsSecretManager,
+    GoogleSecretManager,
+    HashicorpVault,
+    Cyberark,
+    Local,
+    AwsKms,
+    Custom,
+}
+
+#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
+#[serde(rename_all = "snake_case")]
+pub enum AccessMode {
+    #[default]
+    ReadOnly,
+    WriteOnly,
+    ReadAndWrite,
+}
+
+impl AccessMode {
+    pub fn readable(self) -> bool {
+        matches!(self, Self::ReadOnly | Self::ReadAndWrite)
+    }
+}
+
+#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)]
+#[serde(default)]
+pub struct KeyManagementSettings {
+    pub hosted_keys: Option<Vec<String>>,
+    pub store_virtual_keys: Option<bool>,
+    pub prefix_for_stored_virtual_keys: String,
+    pub access_mode: AccessMode,
+    pub primary_secret_name: Option<String>,
+    pub description: Option<String>,
+    pub tags: Option<BTreeMap<String, String>>,
+    pub kms_key_id: Option<String>,
+    pub custom_secret_manager: Option<String>,
+    pub aws_region_name: Option<String>,
+    pub aws_role_name: Option<String>,
+    pub aws_session_name: Option<String>,
+    #[serde(serialize_with = "serialize_secret")]
+    pub aws_external_id: Option<SecretValue>,
+    pub aws_profile_name: Option<String>,
+    #[serde(serialize_with = "serialize_secret")]
+    pub aws_web_identity_token: Option<SecretValue>,
+    pub aws_sts_endpoint: Option<String>,
+    pub replica_regions: Option<Vec<String>>,
+}
+
+impl Default for KeyManagementSettings {
+    fn default() -> Self {
+        Self {
+            hosted_keys: None,
+            store_virtual_keys: Some(false),
+            prefix_for_stored_virtual_keys: "litellm/".into(),
+            access_mode: AccessMode::ReadOnly,
+            primary_secret_name: None,
+            description: None,
+            tags: None,
+            kms_key_id: None,
+            custom_secret_manager: None,
+            aws_region_name: None,
+            aws_role_name: None,
+            aws_session_name: None,
+            aws_external_id: None,
+            aws_profile_name: None,
+            aws_web_identity_token: None,
+            aws_sts_endpoint: None,
+            replica_regions: None,
+        }
+    }
+}
+
+fn serialize_secret<S: serde::Serializer>(
+    value: &Option<SecretValue>,
+    serializer: S,
+) -> Result<S::Ok, S::Error> {
+    value
+        .as_ref()
+        .map(SecretValue::expose)
+        .serialize(serializer)
+}

litellm-rust/crates/secrets-types/src/error.rs

@@ -0,0 +1,9 @@
+#[derive(Debug, thiserror::Error, PartialEq, Eq)]
+pub enum Error {
+    #[error("secret name contains an unsafe path segment or control character")]
+    UnsafeSecretName,
+    #[error("current secret was not found")]
+    CurrentSecretMissing,
+    #[error("new secret could not be verified")]
+    NewSecretMissing,
+}

litellm-rust/crates/secrets-types/src/lib.rs

@@ -0,0 +1,12 @@
+#![forbid(unsafe_code)]
+
+mod base_secret_manager;
+mod config;
+mod error;
+mod value;
+
+pub use base_secret_manager::{BaseSecretManager, async_rotate_secret, validate_secret_name};
+pub use config::{AccessMode, KeyManagementSettings, KeyManagementSystem};
+pub use error::Error;
+pub use litellm_auth_types::SecretValue;
+pub use value::Secret;

litellm-rust/crates/secrets-types/src/value.rs

@@ -0,0 +1,31 @@
+use crate::SecretValue;
+
+#[derive(Clone, PartialEq, Eq, veil::Redact)]
+pub enum Secret {
+    String(SecretValue),
+    Bool(#[redact] bool),
+    Json(#[redact] serde_json::Value),
+}
+
+impl From<SecretValue> for Secret {
+    fn from(value: SecretValue) -> Self {
+        Self::String(value)
+    }
+}
+
+impl Secret {
+    pub fn from_json(value: serde_json::Value) -> Self {
+        match value {
+            serde_json::Value::String(value) => Self::String(SecretValue::new(value)),
+            serde_json::Value::Bool(value) => Self::Bool(value),
+            value => Self::Json(value),
+        }
+    }
+
+    pub fn as_str(&self) -> Option<&str> {
+        match self {
+            Self::String(value) => Some(value.expose()),
+            Self::Bool(_) | Self::Json(_) => None,
+        }
+    }
+}

litellm-rust/crates/secrets-types/tests/config.rs

@@ -0,0 +1,60 @@
+use litellm_secrets_types::{
+    AccessMode, KeyManagementSettings, KeyManagementSystem, Secret, SecretValue,
+};
+use serde_json::json;
+
+#[test]
+fn config_preserves_defaults_nulls_and_serialized_names() {
+    let empty: KeyManagementSettings = serde_json::from_value(json!({})).unwrap();
+    assert_eq!(empty, KeyManagementSettings::default());
+    assert_eq!(empty.access_mode, AccessMode::ReadOnly);
+    assert_eq!(empty.store_virtual_keys, Some(false));
+    assert_eq!(empty.prefix_for_stored_virtual_keys, "litellm/");
+    let configured: KeyManagementSettings = serde_json::from_value(json!({
+        "hosted_keys": [], "store_virtual_keys": null, "access_mode": "write_only",
+        "aws_web_identity_token": "private-token", "aws_external_id": "private-id",
+        "tags": {"stage": "test"}, "replica_regions": ["test-region"]
+    }))
+    .unwrap();
+    assert!(!configured.access_mode.readable());
+    assert_eq!(configured.store_virtual_keys, None);
+    assert_eq!(configured.hosted_keys.as_deref(), Some([].as_slice()));
+    assert!(!format!("{configured:?}").contains("private-"));
+    let serialized = serde_json::to_value(&configured).unwrap();
+    assert_eq!(serialized["access_mode"], "write_only");
+    assert_eq!(serialized["aws_web_identity_token"], "private-token");
+    assert_eq!(
+        serde_json::from_value::<KeyManagementSettings>(serialized).unwrap(),
+        configured
+    );
+}
+
+#[rstest::rstest]
+#[case::aws_kms("aws_kms", KeyManagementSystem::AwsKms)]
+#[case::aws_secret_manager("aws_secret_manager", KeyManagementSystem::AwsSecretManager)]
+#[case::google_kms("google_kms", KeyManagementSystem::GoogleKms)]
+#[case::google_secret_manager("google_secret_manager", KeyManagementSystem::GoogleSecretManager)]
+#[case::azure_key_vault("azure_key_vault", KeyManagementSystem::AzureKeyVault)]
+#[case::hashicorp_vault("hashicorp_vault", KeyManagementSystem::HashicorpVault)]
+#[case::cyberark("cyberark", KeyManagementSystem::Cyberark)]
+#[case::custom("custom", KeyManagementSystem::Custom)]
+#[case::local("local", KeyManagementSystem::Local)]
+fn key_management_system_serialization_round_trips(
+    #[case] name: &str,
+    #[case] system: KeyManagementSystem,
+) {
+    assert_eq!(
+        serde_json::from_value::<KeyManagementSystem>(json!(name)).unwrap(),
+        system
+    );
+    assert_eq!(serde_json::to_value(system).unwrap(), name);
+}
+
+#[test]
+fn secret_debug_never_exposes_values() {
+    assert!(
+        !format!("{:?}", Secret::String(SecretValue::new("sensitive-value")))
+            .contains("sensitive-value")
+    );
+    assert!(!format!("{:?}", Secret::Bool(true)).contains("true"));
+}

litellm-rust/crates/secrets-types/tests/rotation.rs

@@ -0,0 +1,105 @@
+use std::sync::atomic::{AtomicUsize, Ordering};
+
+use litellm_secrets_types::{
+    BaseSecretManager, Error, SecretValue, async_rotate_secret, validate_secret_name,
+};
+
+struct Manager {
+    step: AtomicUsize,
+    absent_at: Option<usize>,
+}
+
+impl BaseSecretManager for Manager {
+    type Error = Error;
+    type WriteResponse = &'static str;
+    type DeleteResponse = ();
+
+    async fn async_read_secret(&self, name: &str) -> Result<Option<SecretValue>, Error> {
+        let step = self.step.fetch_add(1, Ordering::SeqCst);
+        assert_eq!(name, if step == 0 { "old" } else { "new" });
+        Ok((self.absent_at != Some(step)).then(|| SecretValue::new("value")))
+    }
+
+    async fn async_write_secret(
+        &self,
+        name: &str,
+        value: &SecretValue,
+        description: Option<&str>,
+    ) -> Result<Self::WriteResponse, Error> {
+        assert_eq!(self.step.fetch_add(1, Ordering::SeqCst), 1);
+        assert_eq!(name, "new");
+        assert_eq!(value.expose(), "replacement");
+        assert_eq!(description, Some("Rotated from old"));
+        Ok("provider-response")
+    }
+
+    async fn async_delete_secret(
+        &self,
+        name: &str,
+        recovery_window_in_days: i64,
+    ) -> Result<(), Error> {
+        assert_eq!(self.step.fetch_add(1, Ordering::SeqCst), 3);
+        assert_eq!(name, "old");
+        assert_eq!(recovery_window_in_days, 7);
+        Ok(())
+    }
+}
+
+#[tokio::test]
+async fn rotation_verifies_before_deleting_and_returns_provider_response() {
+    let manager = Manager {
+        step: AtomicUsize::new(0),
+        absent_at: None,
+    };
+    assert_eq!(
+        async_rotate_secret(&manager, "old", "new", &SecretValue::new("replacement"))
+            .await
+            .unwrap(),
+        "provider-response"
+    );
+    assert_eq!(manager.step.load(Ordering::SeqCst), 4);
+}
+
+#[rstest::rstest]
+#[case::current_secret_missing(0, Error::CurrentSecretMissing, 1)]
+#[case::new_secret_missing(2, Error::NewSecretMissing, 3)]
+#[tokio::test]
+async fn missing_old_or_new_value_stops_rotation_before_deletion(
+    #[case] absent_at: usize,
+    #[case] expected: Error,
+    #[case] calls: usize,
+) {
+    let manager = Manager {
+        step: AtomicUsize::new(0),
+        absent_at: Some(absent_at),
+    };
+    assert_eq!(
+        async_rotate_secret(&manager, "old", "new", &SecretValue::new("replacement"))
+            .await
+            .unwrap_err(),
+        expected
+    );
+    assert_eq!(manager.step.load(Ordering::SeqCst), calls);
+}
+
+#[rstest::rstest]
+#[case::parent("..")]
+#[case::parent_prefix("../x")]
+#[case::parent_segment("x/../y")]
+#[case::parent_suffix("x/..")]
+#[case::line_feed("line\n")]
+#[case::next_line("\u{85}")]
+#[case::line_separator("\u{2028}")]
+#[case::paragraph_separator("\u{2029}")]
+fn names_reject_path_traversal_and_control_characters(#[case] name: &str) {
+    assert_eq!(validate_secret_name(name), Err(Error::UnsafeSecretName));
+}
+
+#[rstest::rstest]
+#[case::embedded_double_dot("release-1.0..2")]
+#[case::path_separator("folder/key")]
+#[case::empty("")]
+#[case::three_dots("...")]
+fn names_allow_safe_values(#[case] name: &str) {
+    assert_eq!(validate_secret_name(name), Ok(()));
+}

litellm-rust/crates/secrets/Cargo.toml

@@ -0,0 +1,34 @@
+[package]
+name = "litellm-secrets"
+version = "0.1.0"
+edition.workspace = true
+license.workspace = true
+repository.workspace = true
+
+[features]
+default = []
+aws = ["dep:litellm-secrets-aws"]
+google = ["dep:litellm-secrets-google"]
+
+[dependencies]
+litellm-secrets-types.workspace = true
+litellm-secrets-aws = { workspace = true, optional = true }
+litellm-secrets-google = { workspace = true, optional = true }
+litellm-core-utils.workspace = true
+base64.workspace = true
+serde.workspace = true
+strum.workspace = true
+jsonwebtoken.workspace = true
+serde_json.workspace = true
+thiserror.workspace = true
+reqwest.workspace = true
+moka.workspace = true
+tokio = { workspace = true, features = ["fs"] }
+
+[dev-dependencies]
+rstest.workspace = true
+wiremock = "0.6.5"
+tempfile = "3"
+aws-sdk-kms = "1.120.0"
+google-cloud-kms-v1 = "1.14.0"
+google-cloud-auth.workspace = true

litellm-rust/crates/secrets/README.md

@@ -0,0 +1,11 @@
+# Secret resolution
+
+Construct `SecretManagerState::new(backend, settings)` for a configured manager or use `SecretManagerState::default()` for environment lookups. The configured backend determines its provider identity. Write-only settings and names excluded by `hosted_keys` use the environment directly. `secret_manager_would_be_consulted` follows the same routing decision as resolution
+
+`get_secret` returns `Ok(Some(value))` for a found value, `Ok(None)` when no source contains the value, and `Err(error)` when lookup fails. For managed names, resolution checks the manager, then the environment, then the caller's default. An empty string, `false`, or an explicitly stored JSON null is a found value
+
+Backend failures propagate by default. To allow fallback during a backend failure, construct the resolver with `.with_failure_policy(FailurePolicy::EnvironmentFallback)`. It then tries the environment and default, in that order. If neither exists, the original error is returned. This policy applies to manager lookups. Explicit OIDC references retain their own authentication errors and never fall back to environment secrets under the reference name
+
+`get_secret` preserves value types. `get_secret_str` accepts a string default and rejects boolean or JSON values with `Error::TypeMismatch`. `get_secret_bool` accepts a boolean default and converts strings containing `true` or `false`, ignoring surrounding whitespace and ASCII case. Other strings and JSON values produce `Error::TypeMismatch`. Conversion failures never activate fallback or replace a found value with the default
+
+Provider payloads remain strings unless explicitly selecting a field from an AWS primary JSON secret. Google caches only successfully decoded string payloads, so reads have identical values and types before and after caching. Confirmed absence and failed reads are not cached. AWS resource-not-found responses and Google HTTP 404 responses indicate absence. Other provider errors remain errors, and successful responses without the required payload are malformed responses rather than missing secrets

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

@@ -0,0 +1,33 @@
+#[derive(Debug, thiserror::Error)]
+pub enum Error {
+    #[error("encrypted environment value is missing")]
+    MissingCiphertext,
+    #[error("ciphertext is not valid base64 for the configured manager")]
+    InvalidCiphertext,
+    #[error("decrypted value is not UTF-8")]
+    Utf8,
+    #[error("unsupported OIDC provider or missing build feature")]
+    UnsupportedOidc,
+    #[error("OIDC reference requires a provider and audience")]
+    InvalidOidc,
+    #[error("OIDC environment variable is missing")]
+    MissingEnvironment,
+    #[error("OIDC request failed")]
+    OidcHttp,
+    #[error("OIDC provider returned HTTP {0}")]
+    OidcStatus(u16),
+    #[error("OIDC response is invalid")]
+    OidcResponse,
+    #[error("OIDC file path must be absolute and within the credential allowlist")]
+    UnsafeOidcPath,
+    #[error("OIDC file could not be read")]
+    OidcFile,
+    #[error("secret cannot be converted to {expected}")]
+    TypeMismatch { expected: &'static str },
+    #[cfg(feature = "aws")]
+    #[error(transparent)]
+    Aws(#[from] litellm_secrets_aws::Error),
+    #[cfg(feature = "google")]
+    #[error(transparent)]
+    Google(#[from] litellm_secrets_google::Error),
+}

litellm-rust/crates/secrets/src/handler.rs

@@ -0,0 +1,117 @@
+use litellm_core_utils::settings::Lookup;
+
+use crate::{Error, KeyManagementSettings, KeyManagementSystem, Secret, SecretValue};
+
+#[derive(Clone)]
+pub enum SecretManager {
+    Local,
+    #[cfg(feature = "aws")]
+    AwsKms(crate::aws::AwsKms),
+    #[cfg(feature = "aws")]
+    AwsSecretsManagerV2(crate::aws::AwsSecretsManagerV2),
+    #[cfg(feature = "google")]
+    GoogleKms(crate::google::GoogleKms),
+    #[cfg(feature = "google")]
+    GoogleSecretManager(crate::google::GoogleSecretManager),
+}
+
+impl SecretManager {
+    pub fn system(&self) -> KeyManagementSystem {
+        match self {
+            Self::Local => KeyManagementSystem::Local,
+            #[cfg(feature = "aws")]
+            Self::AwsKms(_) => KeyManagementSystem::AwsKms,
+            #[cfg(feature = "aws")]
+            Self::AwsSecretsManagerV2(_) => KeyManagementSystem::AwsSecretManager,
+            #[cfg(feature = "google")]
+            Self::GoogleKms(_) => KeyManagementSystem::GoogleKms,
+            #[cfg(feature = "google")]
+            Self::GoogleSecretManager(_) => KeyManagementSystem::GoogleSecretManager,
+        }
+    }
+}
+
+pub async fn get_secret_from_manager(
+    client: &SecretManager,
+    secret_name: &str,
+    _settings: &KeyManagementSettings,
+    environment: &(dyn Lookup + Send + Sync),
+) -> Result<Option<Secret>, Error> {
+    match client {
+        SecretManager::Local => Ok(environment
+            .get(secret_name)
+            .map(SecretValue::new)
+            .map(Secret::String)),
+        #[cfg(feature = "aws")]
+        SecretManager::AwsKms(client) => {
+            let ciphertext = environment
+                .get(secret_name)
+                .ok_or(Error::MissingCiphertext)?;
+            let plaintext = client
+                .decrypt(decode_ciphertext(&ciphertext, Base64Mode::Permissive)?)
+                .await?;
+            let value = String::from_utf8(plaintext).map_err(|_| Error::Utf8)?;
+            Ok(Some(Secret::String(SecretValue::new(value.trim()))))
+        }
+        #[cfg(feature = "google")]
+        SecretManager::GoogleKms(client) => {
+            let ciphertext = environment
+                .get(secret_name)
+                .ok_or(Error::MissingCiphertext)?;
+            let plaintext = client
+                .decrypt(decode_ciphertext(&ciphertext, Base64Mode::Canonical)?)
+                .await?;
+            let value = String::from_utf8(plaintext).map_err(|_| Error::Utf8)?;
+            Ok(Some(Secret::String(SecretValue::new(value))))
+        }
+        #[cfg(feature = "aws")]
+        SecretManager::AwsSecretsManagerV2(client) => client
+            .read_secret_for_resolver(
+                secret_name,
+                _settings.primary_secret_name.as_deref(),
+                environment,
+            )
+            .await
+            .map_err(Error::from),
+        #[cfg(feature = "google")]
+        SecretManager::GoogleSecretManager(client) => client
+            .get_secret_from_google_secret_manager(secret_name)
+            .await
+            .map_err(Error::from),
+    }
+}
+
+#[cfg(any(feature = "aws", feature = "google"))]
+#[derive(Clone, Copy)]
+enum Base64Mode {
+    #[cfg(feature = "google")]
+    Canonical,
+    #[cfg(feature = "aws")]
+    Permissive,
+}
+
+#[cfg(any(feature = "aws", feature = "google"))]
+fn decode_ciphertext(value: &str, mode: Base64Mode) -> Result<Vec<u8>, Error> {
+    use base64::{Engine, engine::general_purpose::STANDARD};
+    let canonical = match mode {
+        #[cfg(feature = "google")]
+        Base64Mode::Canonical => true,
+        #[cfg(feature = "aws")]
+        Base64Mode::Permissive => false,
+    };
+    let encoded = if canonical {
+        value.to_owned()
+    } else {
+        value
+            .chars()
+            .filter(|c| c.is_ascii_alphanumeric() || matches!(c, '+' | '/' | '='))
+            .collect()
+    };
+    let ciphertext = STANDARD
+        .decode(&encoded)
+        .map_err(|_| Error::InvalidCiphertext)?;
+    if canonical && STANDARD.encode(&ciphertext) != encoded {
+        return Err(Error::InvalidCiphertext);
+    }
+    Ok(ciphertext)
+}

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

@@ -0,0 +1,21 @@
+#![forbid(unsafe_code)]
+
+mod error;
+mod handler;
+mod oidc;
+mod resolver;
+mod state;
+
+pub use error::Error;
+pub use handler::{SecretManager, get_secret_from_manager};
+pub use litellm_secrets_types::{
+    AccessMode, KeyManagementSettings, KeyManagementSystem, Secret, SecretValue,
+};
+pub use oidc::{OidcProvider, OidcReference, OidcResolver};
+pub use resolver::{FailurePolicy, SecretResolver};
+pub use state::{SecretManagerState, secret_manager_would_be_consulted};
+
+#[cfg(feature = "aws")]
+pub use litellm_secrets_aws as aws;
+#[cfg(feature = "google")]
+pub use litellm_secrets_google as google;

litellm-rust/crates/secrets/src/oidc.rs

@@ -0,0 +1,269 @@
+use std::{
+    path::Path,
+    time::{Duration, SystemTime, UNIX_EPOCH},
+};
+
+use jsonwebtoken::dangerous::insecure_decode_claims;
+use litellm_core_utils::settings::Lookup;
+use moka::future::Cache;
+use serde::Deserialize;
+
+use crate::{Error, SecretValue};
+
+const GOOGLE_TOKEN_MAX_TTL: Duration = Duration::from_secs(3540);
+const GITHUB_TOKEN_TTL: Duration = Duration::from_secs(295);
+const TOKEN_EXPIRY_MARGIN_SECONDS: f64 = 60.0;
+const CIRCLE_OIDC_TOKEN: &str = "CIRCLE_OIDC_TOKEN";
+const CIRCLE_OIDC_TOKEN_V2: &str = "CIRCLE_OIDC_TOKEN_V2";
+const AZURE_FEDERATED_TOKEN_FILE: &str = "AZURE_FEDERATED_TOKEN_FILE";
+const ACTIONS_ID_TOKEN_REQUEST_URL: &str = "ACTIONS_ID_TOKEN_REQUEST_URL";
+const ACTIONS_ID_TOKEN_REQUEST_TOKEN: &str = "ACTIONS_ID_TOKEN_REQUEST_TOKEN";
+const OIDC_ALLOWED_CREDENTIAL_DIRS: &str = "LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS";
+const DEFAULT_CREDENTIAL_DIRS: &str = "/var/run/secrets,/run/secrets";
+
+#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq, strum::EnumString, strum::AsRefStr)]
+#[strum(serialize_all = "snake_case")]
+pub enum OidcProvider {
+    Google,
+    #[strum(serialize = "circleci")]
+    CircleCi,
+    #[strum(serialize = "circleci_v2")]
+    CircleCiV2,
+    Github,
+    Azure,
+    File,
+    Env,
+    EnvPath,
+}
+
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub struct OidcReference<'a> {
+    pub provider: OidcProvider,
+    pub audience: &'a str,
+}
+
+impl<'a> TryFrom<&'a str> for OidcReference<'a> {
+    type Error = Error;
+
+    fn try_from(reference: &'a str) -> Result<Self, Error> {
+        let (provider, audience) = reference
+            .strip_prefix("oidc/")
+            .and_then(|body| body.split_once('/'))
+            .ok_or(Error::InvalidOidc)?;
+        Ok(Self {
+            provider: provider.parse().map_err(|_| Error::UnsupportedOidc)?,
+            audience,
+        })
+    }
+}
+
+#[derive(Deserialize)]
+struct OidcTokenClaims {
+    exp: Option<NumericDate>,
+}
+
+#[derive(Deserialize)]
+#[serde(untagged)]
+enum NumericDate {
+    Number(f64),
+    String(String),
+}
+
+impl NumericDate {
+    fn seconds(self) -> Option<f64> {
+        match self {
+            Self::Number(value) => Some(value),
+            Self::String(value) => value.parse().ok(),
+        }
+        .filter(|value| value.is_finite())
+    }
+}
+
+pub struct OidcResolver {
+    client: reqwest::Client,
+    google_identity_endpoint: reqwest::Url,
+    cache: Cache<String, (SecretValue, SystemTime)>,
+    clock: fn() -> SystemTime,
+}
+
+impl Default for OidcResolver {
+    fn default() -> Self {
+        let client = reqwest::Client::builder()
+            .timeout(Duration::from_secs(600))
+            .connect_timeout(Duration::from_secs(5))
+            .build()
+            .expect("HTTP client configuration");
+        Self::new(
+            client,
+            reqwest::Url::parse("http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity").expect("static URL"),
+        )
+    }
+}
+
+impl OidcResolver {
+    pub fn new(client: reqwest::Client, google_identity_endpoint: reqwest::Url) -> Self {
+        Self {
+            client,
+            google_identity_endpoint,
+            cache: Cache::builder()
+                .max_capacity(200)
+                .time_to_live(GOOGLE_TOKEN_MAX_TTL)
+                .build(),
+            clock: SystemTime::now,
+        }
+    }
+
+    pub fn with_clock(self, clock: fn() -> SystemTime) -> Self {
+        Self { clock, ..self }
+    }
+
+    pub async fn resolve(
+        &self,
+        reference: &str,
+        environment: &(dyn Lookup + Send + Sync),
+    ) -> Result<Option<SecretValue>, Error> {
+        let OidcReference { provider, audience } = reference.try_into()?;
+        match provider {
+            OidcProvider::CircleCi => required_env(environment, CIRCLE_OIDC_TOKEN)
+                .map(SecretValue::new)
+                .map(Some),
+            OidcProvider::CircleCiV2 => required_env(environment, CIRCLE_OIDC_TOKEN_V2)
+                .map(SecretValue::new)
+                .map(Some),
+            OidcProvider::Env => required_env(environment, audience)
+                .map(SecretValue::new)
+                .map(Some),
+            OidcProvider::EnvPath => read_file(&required_env(environment, audience)?)
+                .await
+                .map(Some),
+            OidcProvider::File => read_allowed_file(audience, environment).await.map(Some),
+            OidcProvider::Azure => {
+                if let Some(path) = environment.get(AZURE_FEDERATED_TOKEN_FILE) {
+                    return read_file(&path).await.map(Some);
+                }
+                Err(Error::UnsupportedOidc)
+            }
+            OidcProvider::Github => {
+                let url = required_env(environment, ACTIONS_ID_TOKEN_REQUEST_URL)?;
+                let authorization = required_env(environment, ACTIONS_ID_TOKEN_REQUEST_TOKEN)?;
+                if let Some(value) = self.cached(reference).await {
+                    return Ok(Some(value));
+                }
+                let response = self
+                    .client
+                    .get(url)
+                    .query(&[("audience", audience)])
+                    .bearer_auth(authorization)
+                    .header("Accept", "application/json; api-version=2.0")
+                    .send()
+                    .await
+                    .map_err(|_| Error::OidcHttp)?;
+                if response.status() != reqwest::StatusCode::OK {
+                    return Err(Error::OidcStatus(response.status().as_u16()));
+                }
+                #[derive(Deserialize)]
+                struct Token {
+                    value: Option<SecretValue>,
+                }
+                let token: Token = response.json().await.map_err(|_| Error::OidcResponse)?;
+                if let Some(value) = &token.value {
+                    self.cache
+                        .insert(
+                            reference.to_owned(),
+                            (value.clone(), (self.clock)() + GITHUB_TOKEN_TTL),
+                        )
+                        .await;
+                }
+                Ok(token.value)
+            }
+            OidcProvider::Google => {
+                if !cfg!(feature = "google") {
+                    return Err(Error::UnsupportedOidc);
+                }
+                if let Some(value) = self.cached(reference).await {
+                    return Ok(Some(value));
+                }
+                let response = self
+                    .client
+                    .get(self.google_identity_endpoint.clone())
+                    .query(&[("audience", audience)])
+                    .header("Metadata-Flavor", "Google")
+                    .send()
+                    .await
+                    .map_err(|_| Error::OidcHttp)?;
+                if response.status() != reqwest::StatusCode::OK {
+                    return Err(Error::OidcStatus(response.status().as_u16()));
+                }
+                let token = response.text().await.map_err(|_| Error::OidcResponse)?;
+                let now = (self.clock)();
+                let ttl = oidc_token_cache_ttl(&token, now, GOOGLE_TOKEN_MAX_TTL);
+                let value = SecretValue::new(token);
+                if let Some(ttl) = ttl.filter(|ttl| !ttl.is_zero()) {
+                    self.cache
+                        .insert(reference.to_owned(), (value.clone(), now + ttl))
+                        .await;
+                }
+                Ok(Some(value))
+            }
+        }
+    }
+
+    async fn cached(&self, reference: &str) -> Option<SecretValue> {
+        self.cache
+            .get(reference)
+            .await
+            .and_then(|(value, expires)| ((self.clock)() < expires).then_some(value))
+    }
+}
+
+fn required_env(environment: &dyn Lookup, name: &str) -> Result<String, Error> {
+    environment.get(name).ok_or(Error::MissingEnvironment)
+}
+
+async fn read_file(path: &str) -> Result<SecretValue, Error> {
+    tokio::fs::read_to_string(path)
+        .await
+        .map(|value| SecretValue::new(value.replace("\r\n", "\n").replace('\r', "\n")))
+        .map_err(|_| Error::OidcFile)
+}
+
+async fn read_allowed_file(
+    path: &str,
+    environment: &(dyn Lookup + Sync),
+) -> Result<SecretValue, Error> {
+    if !Path::new(path).is_absolute() {
+        return Err(Error::UnsafeOidcPath);
+    }
+    let resolved = tokio::fs::canonicalize(path)
+        .await
+        .map_err(|_| Error::OidcFile)?;
+    let allowed = environment
+        .get(OIDC_ALLOWED_CREDENTIAL_DIRS)
+        .filter(|v| !v.is_empty())
+        .unwrap_or_else(|| DEFAULT_CREDENTIAL_DIRS.into());
+    for directory in allowed.split(',').map(str::trim).filter(|d| !d.is_empty()) {
+        if let Ok(directory) = tokio::fs::canonicalize(directory).await
+            && resolved.starts_with(directory)
+        {
+            return tokio::fs::read_to_string(&resolved)
+                .await
+                .map(|value| SecretValue::new(value.replace("\r\n", "\n").replace('\r', "\n")))
+                .map_err(|_| Error::OidcFile);
+        }
+    }
+    Err(Error::UnsafeOidcPath)
+}
+
+fn oidc_token_cache_ttl(token: &str, now: SystemTime, max_ttl: Duration) -> Option<Duration> {
+    let fallback = Some(max_ttl);
+    let Ok(claims) = insecure_decode_claims::<OidcTokenClaims>(token) else {
+        return fallback;
+    };
+    let Some(exp) = claims.exp.and_then(NumericDate::seconds) else {
+        return fallback;
+    };
+    let seconds = exp.trunc()
+        - now.duration_since(UNIX_EPOCH).ok()?.as_secs() as f64
+        - TOKEN_EXPIRY_MARGIN_SECONDS;
+    (seconds > 0.0).then(|| Duration::from_secs_f64(seconds.min(max_ttl.as_secs_f64())))
+}

litellm-rust/crates/secrets/src/resolver.rs

@@ -0,0 +1,135 @@
+use std::sync::Arc;
+
+use litellm_core_utils::settings::{Lookup, ProcessEnvironment};
+
+use crate::state::{LookupTarget, normalize_secret_name};
+use crate::{Error, OidcResolver, Secret, SecretManagerState, SecretValue};
+
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
+pub enum FailurePolicy {
+    #[default]
+    Propagate,
+    EnvironmentFallback,
+}
+
+pub struct SecretResolver {
+    state: Arc<SecretManagerState>,
+    environment: Arc<dyn Lookup + Send + Sync>,
+    oidc: OidcResolver,
+    failure_policy: FailurePolicy,
+}
+
+impl Default for SecretResolver {
+    fn default() -> Self {
+        Self::new(
+            Arc::new(SecretManagerState::default()),
+            Arc::new(ProcessEnvironment),
+            OidcResolver::default(),
+        )
+    }
+}
+
+impl SecretResolver {
+    pub fn new(
+        state: Arc<SecretManagerState>,
+        environment: Arc<dyn Lookup + Send + Sync>,
+        oidc: OidcResolver,
+    ) -> Self {
+        Self {
+            state,
+            environment,
+            oidc,
+            failure_policy: FailurePolicy::default(),
+        }
+    }
+
+    pub fn with_failure_policy(self, failure_policy: FailurePolicy) -> Self {
+        Self {
+            failure_policy,
+            ..self
+        }
+    }
+
+    pub async fn get_secret(
+        &self,
+        name: &str,
+        default_value: Option<Secret>,
+    ) -> Result<Option<Secret>, Error> {
+        let name = normalize_secret_name(name);
+        if name.starts_with("oidc/") {
+            return self
+                .oidc
+                .resolve(name, self.environment.as_ref())
+                .await
+                .map(|value| value.map(Secret::String).or(default_value));
+        }
+        let LookupTarget::Manager { backend, settings } = self.state.lookup_target(name) else {
+            return Ok(self.environment_secret(name).or(default_value));
+        };
+        match crate::get_secret_from_manager(backend, name, settings, self.environment.as_ref())
+            .await
+        {
+            Ok(value) => Ok(value
+                .or_else(|| self.environment_secret(name))
+                .or(default_value)),
+            Err(error) => match self.failure_policy {
+                FailurePolicy::Propagate => Err(error),
+                FailurePolicy::EnvironmentFallback => self
+                    .environment_secret(name)
+                    .or(default_value)
+                    .map(Some)
+                    .ok_or(error),
+            },
+        }
+    }
+
+    fn environment_secret(&self, name: &str) -> Option<Secret> {
+        self.environment
+            .get(name)
+            .map(SecretValue::new)
+            .map(Secret::String)
+    }
+
+    pub async fn get_secret_str(
+        &self,
+        name: &str,
+        default_value: Option<SecretValue>,
+    ) -> Result<Option<SecretValue>, Error> {
+        match self
+            .get_secret(name, default_value.map(Secret::String))
+            .await?
+        {
+            Some(Secret::String(value)) => Ok(Some(value)),
+            None => Ok(None),
+            Some(Secret::Bool(_) | Secret::Json(_)) => {
+                Err(Error::TypeMismatch { expected: "string" })
+            }
+        }
+    }
+
+    pub async fn get_secret_bool(
+        &self,
+        name: &str,
+        default_value: Option<bool>,
+    ) -> Result<Option<bool>, Error> {
+        match self
+            .get_secret(name, default_value.map(Secret::Bool))
+            .await?
+        {
+            Some(Secret::Bool(value)) => Ok(Some(value)),
+            Some(Secret::String(value)) => {
+                match value.expose().trim().to_ascii_lowercase().as_str() {
+                    "true" => Ok(Some(true)),
+                    "false" => Ok(Some(false)),
+                    _ => Err(Error::TypeMismatch {
+                        expected: "boolean",
+                    }),
+                }
+            }
+            Some(Secret::Json(_)) => Err(Error::TypeMismatch {
+                expected: "boolean",
+            }),
+            None => Ok(None),
+        }
+    }
+}

litellm-rust/crates/secrets/src/state.rs

@@ -0,0 +1,59 @@
+use crate::{KeyManagementSettings, KeyManagementSystem, SecretManager};
+
+pub(crate) enum LookupTarget<'a> {
+    Environment,
+    Manager {
+        backend: &'a SecretManager,
+        settings: &'a KeyManagementSettings,
+    },
+}
+
+pub(crate) fn normalize_secret_name(name: &str) -> &str {
+    name.strip_prefix("os.environ/").unwrap_or(name)
+}
+
+#[derive(Clone, Default)]
+pub struct SecretManagerState {
+    manager: Option<(SecretManager, KeyManagementSettings)>,
+}
+
+impl SecretManagerState {
+    pub fn new(backend: SecretManager, settings: KeyManagementSettings) -> Self {
+        Self {
+            manager: Some((backend, settings)),
+        }
+    }
+
+    pub fn system(&self) -> Option<KeyManagementSystem> {
+        self.backend().map(SecretManager::system)
+    }
+
+    pub fn settings(&self) -> Option<&KeyManagementSettings> {
+        self.manager.as_ref().map(|(_, settings)| settings)
+    }
+
+    pub fn backend(&self) -> Option<&SecretManager> {
+        self.manager.as_ref().map(|(backend, _)| backend)
+    }
+
+    pub(crate) fn lookup_target(&self, name: &str) -> LookupTarget<'_> {
+        match &self.manager {
+            Some((backend, settings))
+                if backend.system() != KeyManagementSystem::Local
+                    && settings.access_mode.readable()
+                    && settings
+                        .hosted_keys
+                        .as_ref()
+                        .is_none_or(|keys| keys.iter().any(|key| key == name)) =>
+            {
+                LookupTarget::Manager { backend, settings }
+            }
+            _ => LookupTarget::Environment,
+        }
+    }
+}
+
+pub fn secret_manager_would_be_consulted(state: &SecretManagerState, name: &str) -> bool {
+    let name = normalize_secret_name(name);
+    !name.starts_with("oidc/") && matches!(state.lookup_target(name), LookupTarget::Manager { .. })
+}

litellm-rust/crates/secrets/tests/handler.rs

@@ -0,0 +1,107 @@
+#[cfg(feature = "aws")]
+#[tokio::test]
+async fn aws_handler_reads_ciphertext_decodes_trims_and_redacts() {
+    use aws_sdk_kms::{
+        Client,
+        config::{BehaviorVersion, Credentials, Region},
+    };
+    use base64::{Engine, engine::general_purpose::STANDARD};
+    use litellm_secrets::{
+        Error, KeyManagementSettings, SecretManager, aws::AwsKms, get_secret_from_manager,
+    };
+    use wiremock::{Mock, MockServer, ResponseTemplate, matchers::body_json};
+
+    let server = MockServer::start().await;
+    Mock::given(body_json(
+        serde_json::json!({"CiphertextBlob": STANDARD.encode("encrypted")}),
+    ))
+    .respond_with(
+        ResponseTemplate::new(200)
+            .set_body_json(serde_json::json!({"Plaintext":STANDARD.encode(" value\n")})),
+    )
+    .expect(1)
+    .mount(&server)
+    .await;
+    let client = Client::from_conf(
+        aws_sdk_kms::Config::builder()
+            .behavior_version(BehaviorVersion::latest())
+            .region(Region::new("us-east-1"))
+            .credentials_provider(Credentials::new("test", "test", None, None, "test"))
+            .endpoint_url(server.uri())
+            .build(),
+    );
+    let manager = SecretManager::AwsKms(AwsKms::new(client));
+    let settings = KeyManagementSettings::default();
+    let value = get_secret_from_manager(&manager, "KEY", &settings, &|name: &str| {
+        assert_eq!(name, "KEY");
+        Some(format!(" {}\n", STANDARD.encode("encrypted")))
+    })
+    .await
+    .unwrap()
+    .unwrap();
+    assert_eq!(value.as_str(), Some("value"));
+    assert!(!format!("{value:?}").contains("value"));
+    assert!(matches!(
+        get_secret_from_manager(&manager, "KEY", &settings, &|_: &str| None).await,
+        Err(Error::MissingCiphertext)
+    ));
+    assert!(matches!(
+        get_secret_from_manager(&manager, "KEY", &settings, &|_: &str| Some("abc".into())).await,
+        Err(Error::InvalidCiphertext)
+    ));
+}
+
+#[cfg(feature = "google")]
+#[tokio::test]
+async fn google_handler_requires_canonical_base64_and_preserves_plaintext_whitespace() {
+    use base64::{Engine, engine::general_purpose::STANDARD};
+    use google_cloud_kms_v1::client::KeyManagementService;
+    use litellm_secrets::{
+        Error, KeyManagementSettings, SecretManager, get_secret_from_manager, google::GoogleKms,
+    };
+    use wiremock::{
+        Mock, MockServer, ResponseTemplate,
+        matchers::{body_json, path},
+    };
+
+    let server = MockServer::start().await;
+    let resource = "projects/project/locations/global/keyRings/ring/cryptoKeys/key";
+    Mock::given(path(format!("/v1/{resource}:decrypt")))
+        .and(body_json(
+            serde_json::json!({"ciphertext":STANDARD.encode("encrypted")}),
+        ))
+        .respond_with(
+            ResponseTemplate::new(200)
+                .set_body_json(serde_json::json!({"plaintext":STANDARD.encode(" value\n")})),
+        )
+        .expect(1)
+        .mount(&server)
+        .await;
+    let client = KeyManagementService::builder()
+        .with_endpoint(server.uri())
+        .with_credentials(google_cloud_auth::credentials::anonymous::Builder::new().build())
+        .build()
+        .await
+        .unwrap();
+    let manager = SecretManager::GoogleKms(GoogleKms::new(client, resource.into()));
+    let settings = KeyManagementSettings::default();
+    let value = get_secret_from_manager(&manager, "KEY", &settings, &|_: &str| {
+        Some(STANDARD.encode("encrypted"))
+    })
+    .await
+    .unwrap()
+    .unwrap();
+    assert_eq!(value.as_str(), Some(" value\n"));
+    assert!(matches!(
+        get_secret_from_manager(&manager, "KEY", &settings, &|_: &str| Some(format!(
+            " {}",
+            STANDARD.encode("encrypted")
+        )))
+        .await,
+        Err(Error::InvalidCiphertext)
+    ));
+    assert!(matches!(
+        get_secret_from_manager(&manager, "KEY", &settings, &|_: &str| None).await,
+        Err(Error::MissingCiphertext)
+    ));
+}

litellm-rust/crates/secrets/tests/oidc.rs

@@ -0,0 +1,295 @@
+use std::{collections::BTreeMap, sync::Arc};
+
+use litellm_core_utils::settings::Lookup;
+use litellm_secrets::{Error, OidcResolver, Secret, SecretManagerState, SecretResolver};
+use wiremock::{
+    Mock, MockServer, ResponseTemplate,
+    matchers::{header, method, path, query_param},
+};
+
+fn environment(pairs: &[(&str, &str)]) -> Arc<dyn Lookup + Send + Sync> {
+    let values: BTreeMap<String, String> = pairs
+        .iter()
+        .map(|(k, v)| (k.to_string(), v.to_string()))
+        .collect();
+    Arc::new(move |name: &str| values.get(name).cloned())
+}
+
+#[rstest::rstest]
+#[case::environment("oidc/env/TOKEN", "true")]
+#[case::circleci("oidc/circleci/audience", "circle")]
+#[case::circleci_v2("oidc/circleci_v2/audience", "circle-v2")]
+#[tokio::test]
+async fn environment_sources_resolve_expected_value(
+    #[case] reference: &str,
+    #[case] expected: &str,
+) {
+    let env = environment(&[
+        ("TOKEN", "true"),
+        ("CIRCLE_OIDC_TOKEN", "circle"),
+        ("CIRCLE_OIDC_TOKEN_V2", "circle-v2"),
+    ]);
+    assert_eq!(
+        OidcResolver::default()
+            .resolve(reference, env.as_ref())
+            .await
+            .unwrap()
+            .unwrap()
+            .expose(),
+        expected
+    );
+}
+
+#[tokio::test]
+async fn environment_sources_bypass_boolean_conversion_and_defaults() {
+    let env = environment(&[("TOKEN", "true")]);
+    let oidc = OidcResolver::default();
+    let resolver = SecretResolver::new(Arc::new(SecretManagerState::default()), env, oidc);
+    assert_eq!(
+        resolver
+            .get_secret_str("os.environ/oidc/env/TOKEN", None)
+            .await
+            .unwrap()
+            .unwrap()
+            .expose(),
+        "true"
+    );
+    assert_eq!(
+        resolver
+            .get_secret_bool("oidc/env/TOKEN", None)
+            .await
+            .unwrap(),
+        Some(true)
+    );
+    assert!(matches!(
+        resolver
+            .get_secret("oidc/env/MISSING", Some(Secret::Bool(true)))
+            .await,
+        Err(Error::MissingEnvironment)
+    ));
+    assert!(matches!(
+        resolver.get_secret("oidc/invalid", None).await,
+        Err(Error::InvalidOidc)
+    ));
+}
+
+#[tokio::test]
+async fn github_requests_are_authenticated_cached_and_revalidate_environment() {
+    let server = MockServer::start().await;
+    Mock::given(method("GET"))
+        .and(path("/token"))
+        .and(query_param("audience", "https://service/oidc/path"))
+        .and(header("authorization", "Bearer request-token"))
+        .and(header("accept", "application/json; api-version=2.0"))
+        .respond_with(
+            ResponseTemplate::new(200).set_body_json(serde_json::json!({"value":"identity-token"})),
+        )
+        .expect(1)
+        .mount(&server)
+        .await;
+    let env = environment(&[
+        (
+            "ACTIONS_ID_TOKEN_REQUEST_URL",
+            &format!("{}/token", server.uri()),
+        ),
+        ("ACTIONS_ID_TOKEN_REQUEST_TOKEN", "request-token"),
+    ]);
+    let oidc = OidcResolver::default();
+    for _ in 0..2 {
+        assert_eq!(
+            oidc.resolve("oidc/github/https://service/oidc/path", env.as_ref())
+                .await
+                .unwrap()
+                .unwrap()
+                .expose(),
+            "identity-token"
+        );
+    }
+    assert!(matches!(
+        oidc.resolve(
+            "oidc/github/https://service/oidc/path",
+            environment(&[]).as_ref()
+        )
+        .await,
+        Err(Error::MissingEnvironment)
+    ));
+}
+
+#[tokio::test]
+async fn file_allowlist_resolves_symlinks_while_environment_paths_remain_explicit() {
+    let allowed = tempfile::tempdir().unwrap();
+    let outside = tempfile::tempdir().unwrap();
+    let token = allowed.path().join("token");
+    let private = outside.path().join("private");
+    std::fs::write(&token, "token\r\n").unwrap();
+    std::fs::write(&private, "outside").unwrap();
+    let env = environment(&[
+        (
+            "LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS",
+            allowed.path().to_str().unwrap(),
+        ),
+        ("PATH_TOKEN", private.to_str().unwrap()),
+        ("AZURE_FEDERATED_TOKEN_FILE", token.to_str().unwrap()),
+    ]);
+    let oidc = OidcResolver::default();
+    assert_eq!(
+        oidc.resolve(&format!("oidc/file/{}", token.display()), env.as_ref())
+            .await
+            .unwrap()
+            .unwrap()
+            .expose(),
+        "token\n"
+    );
+    assert!(matches!(
+        oidc.resolve("oidc/file/relative", env.as_ref()).await,
+        Err(Error::UnsafeOidcPath)
+    ));
+    assert!(matches!(
+        oidc.resolve(&format!("oidc/file/{}", private.display()), env.as_ref())
+            .await,
+        Err(Error::UnsafeOidcPath)
+    ));
+    assert_eq!(
+        oidc.resolve("oidc/env_path/PATH_TOKEN", env.as_ref())
+            .await
+            .unwrap()
+            .unwrap()
+            .expose(),
+        "outside"
+    );
+    assert_eq!(
+        oidc.resolve("oidc/azure/scope", env.as_ref())
+            .await
+            .unwrap()
+            .unwrap()
+            .expose(),
+        "token\n"
+    );
+    #[cfg(unix)]
+    {
+        let link = allowed.path().join("link");
+        std::os::unix::fs::symlink(&private, &link).unwrap();
+        assert!(matches!(
+            oidc.resolve(&format!("oidc/file/{}", link.display()), env.as_ref())
+                .await,
+            Err(Error::UnsafeOidcPath)
+        ));
+    }
+}
+
+#[cfg(feature = "google")]
+#[rstest::rstest]
+#[case::at_refresh_boundary(serde_json::json!(1060), 2)]
+#[case::beyond_refresh_boundary(serde_json::json!(1061), 1)]
+#[case::already_expired(serde_json::json!(999), 2)]
+#[case::string_expiry(serde_json::json!("999"), 2)]
+#[case::fractional_expiry(serde_json::json!(1060.9), 2)]
+#[case::negative_expiry(serde_json::json!(-1), 2)]
+#[case::null_expiry(serde_json::Value::Null, 1)]
+#[case::unreadable_expiry(serde_json::json!("invalid"), 1)]
+#[case::nonfinite_expiry(serde_json::json!("NaN"), 1)]
+#[tokio::test]
+async fn google_expiry_caps_cache_and_preserves_audience(
+    #[case] expiry: serde_json::Value,
+    #[case] calls: u64,
+) {
+    use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD};
+    use std::time::{Duration, SystemTime, UNIX_EPOCH};
+    fn now() -> SystemTime {
+        UNIX_EPOCH + Duration::from_secs(1000)
+    }
+    let server = MockServer::start().await;
+    let token = format!(
+        "{}.{}.signature",
+        URL_SAFE_NO_PAD.encode(serde_json::json!({"alg":"RS256","typ":"JWT"}).to_string()),
+        URL_SAFE_NO_PAD.encode(serde_json::json!({"exp":expiry}).to_string())
+    );
+    Mock::given(method("GET"))
+        .and(header("metadata-flavor", "Google"))
+        .and(query_param("audience", "https://service/oidc/path"))
+        .respond_with(ResponseTemplate::new(200).set_body_string(&token))
+        .expect(calls)
+        .mount(&server)
+        .await;
+    let oidc =
+        OidcResolver::new(reqwest::Client::new(), server.uri().parse().unwrap()).with_clock(now);
+    for _ in 0..2 {
+        assert_eq!(
+            oidc.resolve(
+                "oidc/google/https://service/oidc/path",
+                environment(&[]).as_ref()
+            )
+            .await
+            .unwrap()
+            .unwrap()
+            .expose(),
+            token
+        );
+    }
+}
+
+#[cfg(not(feature = "google"))]
+#[tokio::test]
+async fn google_oidc_requires_its_build_feature() {
+    assert!(matches!(
+        OidcResolver::default()
+            .resolve("oidc/google/audience", environment(&[]).as_ref())
+            .await,
+        Err(Error::UnsupportedOidc)
+    ));
+}
+
+#[tokio::test]
+async fn azure_oidc_without_a_token_file_requires_an_unimplemented_backend() {
+    assert!(matches!(
+        OidcResolver::default()
+            .resolve("oidc/azure/scope", environment(&[]).as_ref())
+            .await,
+        Err(Error::UnsupportedOidc)
+    ));
+}
+
+#[rstest::rstest]
+#[case::missing_prefix("env/TOKEN", false)]
+#[case::missing_audience_separator("oidc/env", false)]
+#[case::unknown_provider("oidc/unknown/TOKEN", true)]
+#[tokio::test]
+async fn invalid_references_fail_before_environment_lookup(
+    #[case] reference: &str,
+    #[case] unsupported: bool,
+) {
+    let error = OidcResolver::default()
+        .resolve(reference, &|_: &str| {
+            panic!("invalid reference reached environment lookup")
+        })
+        .await
+        .unwrap_err();
+    assert!(matches!(error, Error::UnsupportedOidc) == unsupported);
+    assert!(matches!(error, Error::InvalidOidc) != unsupported);
+}
+
+#[cfg(feature = "google")]
+#[rstest::rstest]
+#[case::opaque("opaque-token")]
+#[case::missing_expiry("header.e30.signature")]
+#[tokio::test]
+async fn unreadable_expiry_keeps_python_cache_fallback(#[case] token: &str) {
+    let server = MockServer::start().await;
+    Mock::given(method("GET"))
+        .respond_with(ResponseTemplate::new(200).set_body_string(token))
+        .expect(1)
+        .mount(&server)
+        .await;
+    let resolver = OidcResolver::new(reqwest::Client::new(), server.uri().parse().unwrap());
+    for _ in 0..2 {
+        assert_eq!(
+            resolver
+                .resolve("oidc/google/audience", environment(&[]).as_ref())
+                .await
+                .unwrap()
+                .unwrap()
+                .expose(),
+            token,
+        );
+    }
+}

litellm-rust/crates/secrets/tests/resolution.rs

@@ -0,0 +1,368 @@
+use std::sync::Arc;
+
+use litellm_secrets::{
+    Error, KeyManagementSettings, OidcResolver, Secret, SecretManager, SecretManagerState,
+    SecretResolver, SecretValue, secret_manager_would_be_consulted,
+};
+
+fn resolver(value: Option<&str>, configured: bool) -> SecretResolver {
+    let state = if configured {
+        SecretManagerState::new(SecretManager::Local, KeyManagementSettings::default())
+    } else {
+        SecretManagerState::default()
+    };
+    let value = value.map(str::to_owned);
+    SecretResolver::new(
+        Arc::new(state),
+        Arc::new(move |_: &str| value.clone()),
+        OidcResolver::default(),
+    )
+}
+
+#[rstest::rstest]
+#[case("true", Some(true))]
+#[case(" FALSE ", Some(false))]
+#[case("(True)", None)]
+#[case("False # comment", None)]
+#[case("1", None)]
+#[case("secret", None)]
+#[tokio::test]
+async fn conversion_is_explicit_and_independent_of_manager_configuration(
+    #[case] input: &str,
+    #[case] boolean: Option<bool>,
+    #[values(false, true)] configured: bool,
+) {
+    let resolver = resolver(Some(input), configured);
+    assert_eq!(
+        resolver.get_secret("key", None).await.unwrap(),
+        Some(Secret::String(SecretValue::new(input)))
+    );
+    assert_eq!(
+        resolver
+            .get_secret_str("key", None)
+            .await
+            .unwrap()
+            .unwrap()
+            .expose(),
+        input
+    );
+    match boolean {
+        Some(value) => assert_eq!(
+            resolver.get_secret_bool("key", None).await.unwrap(),
+            Some(value)
+        ),
+        None => assert!(matches!(
+            resolver.get_secret_bool("key", Some(true)).await,
+            Err(Error::TypeMismatch {
+                expected: "boolean"
+            })
+        )),
+    }
+}
+
+#[rstest::rstest]
+#[tokio::test]
+async fn defaults_apply_only_to_absence(#[values(false, true)] configured: bool) {
+    let missing = resolver(None, configured);
+    assert_eq!(missing.get_secret("key", None).await.unwrap(), None);
+    assert_eq!(
+        missing.get_secret_bool("key", Some(false)).await.unwrap(),
+        Some(false)
+    );
+    assert_eq!(
+        missing
+            .get_secret_str("key", Some(SecretValue::new("default")))
+            .await
+            .unwrap()
+            .unwrap()
+            .expose(),
+        "default"
+    );
+    for value in [
+        Secret::Bool(false),
+        Secret::from_json(serde_json::json!({"key":1})),
+        Secret::from_json(serde_json::Value::Null),
+    ] {
+        assert_eq!(
+            missing
+                .get_secret("key", Some(value.clone()))
+                .await
+                .unwrap(),
+            Some(value)
+        );
+    }
+    assert_eq!(
+        resolver(Some(""), configured)
+            .get_secret_str("key", Some(SecretValue::new("default")))
+            .await
+            .unwrap()
+            .unwrap()
+            .expose(),
+        ""
+    );
+}
+
+#[tokio::test]
+async fn prefix_is_removed_once_and_local_manager_is_not_consulted() {
+    let state = SecretManagerState::new(SecretManager::Local, KeyManagementSettings::default());
+    assert_eq!(
+        state.system(),
+        Some(litellm_secrets::KeyManagementSystem::Local)
+    );
+    assert!(!secret_manager_would_be_consulted(
+        &state,
+        "os.environ/os.environ/KEY"
+    ));
+    let resolver = SecretResolver::new(
+        Arc::new(state),
+        Arc::new(|name: &str| (name == "os.environ/KEY").then(|| "value".into())),
+        OidcResolver::default(),
+    );
+    assert_eq!(
+        resolver
+            .get_secret_str("os.environ/os.environ/KEY", None)
+            .await
+            .unwrap()
+            .unwrap()
+            .expose(),
+        "value"
+    );
+}
+
+#[tokio::test]
+async fn resolver_future_can_run_on_a_tokio_worker() {
+    let resolver = resolver(Some("worker-value"), false);
+    let result = tokio::spawn(async move { resolver.get_secret_str("KEY", None).await })
+        .await
+        .unwrap()
+        .unwrap();
+    assert_eq!(result.unwrap().expose(), "worker-value");
+}
+
+#[cfg(feature = "aws")]
+mod aws {
+    use super::*;
+    use litellm_secrets::{AccessMode, FailurePolicy, aws::AwsSecretsManagerV2};
+    use wiremock::{Mock, MockServer, ResponseTemplate, matchers::method};
+
+    fn state(server: &MockServer, settings: KeyManagementSettings) -> SecretManagerState {
+        let endpoint = server.uri();
+        let environment = Arc::new(move |name: &str| match name {
+            "AWS_REGION_NAME" => Some("us-east-1".into()),
+            "AWS_ACCESS_KEY_ID" | "AWS_SECRET_ACCESS_KEY" => Some("test".into()),
+            "AWS_BEDROCK_RUNTIME_ENDPOINT" => Some(endpoint.clone()),
+            _ => None,
+        });
+        let manager =
+            AwsSecretsManagerV2::load_aws_secret_manager(Some(true), settings.clone(), environment)
+                .unwrap()
+                .unwrap();
+        SecretManagerState::new(SecretManager::AwsSecretsManagerV2(manager), settings)
+    }
+
+    #[rstest::rstest]
+    #[case::missing(400, serde_json::json!({"__type":"ResourceNotFoundException"}), false)]
+    #[case::denied(400, serde_json::json!({"__type":"AccessDeniedException"}), true)]
+    #[case::malformed(200, serde_json::json!({}), true)]
+    #[tokio::test]
+    async fn failure_policy_preserves_errors_and_fallback_precedence(
+        #[case] status: u16,
+        #[case] body: serde_json::Value,
+        #[case] fails: bool,
+        #[values(FailurePolicy::Propagate, FailurePolicy::EnvironmentFallback)]
+        policy: FailurePolicy,
+        #[values(None, Some("environment"))] environment: Option<&'static str>,
+        #[values(None, Some("default"))] default: Option<&str>,
+    ) {
+        let server = MockServer::start().await;
+        Mock::given(method("POST"))
+            .respond_with(ResponseTemplate::new(status).set_body_json(body))
+            .expect(1)
+            .mount(&server)
+            .await;
+        let resolver = SecretResolver::new(
+            Arc::new(state(&server, KeyManagementSettings::default())),
+            Arc::new(move |_: &str| environment.map(str::to_owned)),
+            OidcResolver::default(),
+        )
+        .with_failure_policy(policy);
+        let result = resolver
+            .get_secret_str("KEY", default.map(SecretValue::new))
+            .await;
+        let fallback = environment.or(default);
+        if fails && (policy == FailurePolicy::Propagate || fallback.is_none()) {
+            assert!(matches!(result, Err(Error::Aws(_))));
+        } else {
+            assert_eq!(result.unwrap().as_ref().map(SecretValue::expose), fallback);
+        }
+    }
+
+    #[rstest::rstest]
+    #[case::boolean(serde_json::json!(false))]
+    #[case::object(serde_json::json!({"key":1}))]
+    #[case::null(serde_json::Value::Null)]
+    #[case::string(serde_json::json!("true"))]
+    #[tokio::test]
+    async fn typed_values_survive_resolution_and_accessors_reject_wrong_types(
+        #[case] value: serde_json::Value,
+    ) {
+        let server = MockServer::start().await;
+        Mock::given(method("POST"))
+            .respond_with(ResponseTemplate::new(200).set_body_json(
+                serde_json::json!({"SecretString":serde_json::json!({"KEY":value}).to_string()}),
+            ))
+            .expect(3)
+            .mount(&server)
+            .await;
+        let settings = KeyManagementSettings {
+            primary_secret_name: Some("primary".into()),
+            ..Default::default()
+        };
+        let resolver = SecretResolver::new(
+            Arc::new(state(&server, settings)),
+            Arc::new(|_: &str| Some("fallback".into())),
+            OidcResolver::default(),
+        );
+        assert_eq!(
+            resolver
+                .get_secret("KEY", Some(Secret::Bool(true)))
+                .await
+                .unwrap(),
+            Some(Secret::from_json(value.clone()))
+        );
+        match &value {
+            serde_json::Value::String(text) => assert_eq!(
+                resolver
+                    .get_secret_str("KEY", None)
+                    .await
+                    .unwrap()
+                    .unwrap()
+                    .expose(),
+                text
+            ),
+            _ => assert!(matches!(
+                resolver.get_secret_str("KEY", None).await,
+                Err(Error::TypeMismatch { expected: "string" })
+            )),
+        }
+        match value {
+            serde_json::Value::Bool(boolean) => assert_eq!(
+                resolver.get_secret_bool("KEY", None).await.unwrap(),
+                Some(boolean)
+            ),
+            serde_json::Value::String(_) => assert_eq!(
+                resolver.get_secret_bool("KEY", None).await.unwrap(),
+                Some(true)
+            ),
+            _ => assert!(matches!(
+                resolver.get_secret_bool("KEY", None).await,
+                Err(Error::TypeMismatch {
+                    expected: "boolean"
+                })
+            )),
+        }
+    }
+
+    #[rstest::rstest]
+    #[tokio::test]
+    async fn gating_prediction_matches_actual_lookup(
+        #[values(AccessMode::ReadOnly, AccessMode::WriteOnly, AccessMode::ReadAndWrite)]
+        access_mode: AccessMode,
+        #[values(None, Some(vec![]), Some(vec!["KEY".into()]))] hosted_keys: Option<Vec<String>>,
+        #[values("os.environ/KEY", "os.environ/oidc/env/KEY")] name: &str,
+    ) {
+        let server = MockServer::start().await;
+        let expected = name == "os.environ/KEY"
+            && access_mode.readable()
+            && hosted_keys
+                .as_ref()
+                .is_none_or(|keys| keys.iter().any(|key| key == "KEY"));
+        Mock::given(method("POST"))
+            .respond_with(
+                ResponseTemplate::new(200)
+                    .set_body_json(serde_json::json!({"SecretString":"remote"})),
+            )
+            .expect(u64::from(expected))
+            .mount(&server)
+            .await;
+        let state = state(
+            &server,
+            KeyManagementSettings {
+                access_mode,
+                hosted_keys,
+                ..Default::default()
+            },
+        );
+        assert!(state.backend().is_some());
+        assert_eq!(state.settings().unwrap().access_mode, access_mode);
+        assert_eq!(secret_manager_would_be_consulted(&state, name), expected);
+        let resolver = SecretResolver::new(
+            Arc::new(state),
+            Arc::new(|_: &str| Some("environment".into())),
+            OidcResolver::default(),
+        );
+        assert_eq!(
+            resolver
+                .get_secret_str(name, None)
+                .await
+                .unwrap()
+                .unwrap()
+                .expose(),
+            if expected { "remote" } else { "environment" }
+        );
+    }
+}
+
+#[cfg(feature = "google")]
+#[rstest::rstest]
+#[case::missing(404)]
+#[case::failure(503)]
+#[tokio::test]
+async fn google_resolver_distinguishes_absence_from_failure(#[case] status: u16) {
+    use litellm_secrets::{FailurePolicy, google::GoogleSecretManager};
+    use wiremock::{Mock, MockServer, ResponseTemplate, matchers::method};
+    let server = MockServer::start().await;
+    Mock::given(method("GET"))
+        .respond_with(ResponseTemplate::new(status))
+        .expect(2)
+        .mount(&server)
+        .await;
+    let environment: Arc<dyn litellm_core_utils::settings::Lookup + Send + Sync> =
+        Arc::new(|name: &str| match name {
+            "VERTEX_AI_API_KEY" => Some("token".into()),
+            "KEY" => Some("environment".into()),
+            _ => None,
+        });
+    let manager = GoogleSecretManager::with_client(
+        reqwest::Client::new(),
+        server.uri().parse().unwrap(),
+        "project".into(),
+        environment.clone(),
+        None,
+        false,
+    )
+    .unwrap();
+    let state = SecretManagerState::new(
+        SecretManager::GoogleSecretManager(manager),
+        KeyManagementSettings::default(),
+    );
+    let resolver = SecretResolver::new(Arc::new(state), environment, OidcResolver::default());
+    let result = resolver.get_secret_str("KEY", None).await;
+    if status == 404 {
+        assert_eq!(result.unwrap().unwrap().expose(), "environment");
+    } else {
+        assert!(
+            matches!(result, Err(Error::Google(litellm_secrets::google::Error::Status(actual))) if actual == status)
+        );
+    }
+    assert_eq!(
+        resolver
+            .with_failure_policy(FailurePolicy::EnvironmentFallback)
+            .get_secret_str("KEY", None)
+            .await
+            .unwrap()
+            .unwrap()
+            .expose(),
+        "environment"
+    );
+}

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

@@ -0,0 +1,19 @@
+[package]
+name = "litellm-token-counter-fast"
+version = "0.1.0"
+edition.workspace = true
+license.workspace = true
+repository.workspace = true
+
+[dependencies]
+base64.workspace = true
+rustc-hash = "2.1.3"
+thiserror.workspace = true
+tokenizers.workspace = true
+unicode-normalization-alignments = "0.1.12"
+
+[dev-dependencies]
+rand.workspace = true
+rstest.workspace = true
+serde.workspace = true
+serde_json.workspace = true

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

@@ -473,13 +473,13 @@ mod tests {
             _ => unreachable!(),
         }
         assert!(ByteLevelCounter::detect(&anthropic_tokenizer).is_none());
-        let counter = crate::TokenCounter::from_json(
+        let counter = crate::FastTokenizer::from_json(
             &anthropic_tokenizer.to_string(false).expect("serialize"),
         )
         .expect("load");
         for text in ["", "Hello WORLD!  ＡＢ fi Ⅳ", "<EOT> stop"] {
             assert_eq!(
-                counter.count_text(text).expect("count"),
+                counter.count_tokens(text).expect("count"),
                 reference_count(&anthropic_tokenizer, text)
             );
         }
@@ -545,7 +545,7 @@ mod tests {
                 .rstrip(rstrip)])
             .expect("add token");
         let fast = ByteLevelCounter::detect(&anthropic_tokenizer).expect("supported");
-        let counter = crate::TokenCounter::from_json(
+        let counter = crate::FastTokenizer::from_json(
             &anthropic_tokenizer.to_string(false).expect("serialize"),
         )
         .expect("load");
@@ -557,7 +557,7 @@ mod tests {
         ] {
             assert_eq!(fast.count(&anthropic_tokenizer, text), None);
             assert_eq!(
-                counter.count_text(text).expect("count"),
+                counter.count_tokens(text).expect("count"),
                 reference_count(&anthropic_tokenizer, text)
             );
         }
@@ -571,17 +571,17 @@ mod tests {
         assert_eq!(fast.count(&tokenizer, "hello"), None);
         assert!(tokenizer.encode_fast("hello", true).is_err());
         let counter =
-            crate::TokenCounter::from_json(&tokenizer.to_string(false).expect("serialize"))
+            crate::FastTokenizer::from_json(&tokenizer.to_string(false).expect("serialize"))
                 .expect("load");
         assert!(matches!(
-            counter.count_text("hello"),
+            counter.count_tokens("hello"),
             Err(crate::Error::Encode(_))
         ));
     }
 
     #[rstest]
     fn shared_counter_matches_encoder_across_threads(anthropic_tokenizer: Tokenizer) {
-        let counter = crate::TokenCounter::from_json(
+        let counter = crate::FastTokenizer::from_json(
             &anthropic_tokenizer.to_string(false).expect("serialize"),
         )
         .expect("load");
@@ -598,7 +598,7 @@ mod tests {
                 scope.spawn(move || {
                     for _ in 0..100 {
                         for (text, count) in inputs.iter().zip(expected) {
-                            assert_eq!(counter.count_text(text).expect("count"), count);
+                            assert_eq!(counter.count_tokens(text).expect("count"), count);
                         }
                     }
                 });
@@ -614,10 +614,10 @@ mod tests {
         let fast = ByteLevelCounter::detect(&anthropic_tokenizer).expect("supported");
         assert_eq!(reference_count(&anthropic_tokenizer, "ABCD EFGH"), 1);
         assert_eq!(fast.count(&anthropic_tokenizer, "ABCD EFGH"), None);
-        let counter = crate::TokenCounter::from_json(
+        let counter = crate::FastTokenizer::from_json(
             &anthropic_tokenizer.to_string(false).expect("serialize"),
         )
         .expect("load");
-        assert_eq!(counter.count_text("ABCD EFGH").expect("count"), 1);
+        assert_eq!(counter.count_tokens("ABCD EFGH").expect("count"), 1);
     }
 }

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

@@ -0,0 +1,13 @@
+use thiserror::Error as ThisError;
+
+#[derive(Debug, ThisError)]
+pub enum Error {
+    #[error("failed to load tokenizer: {0}")]
+    Load(#[source] tokenizers::Error),
+    #[error("failed to load tokenizer: tiktoken rank file: {0}")]
+    Ranks(String),
+    #[error("failed to load tokenizer: Unicode character classes are unavailable")]
+    UnicodeClasses,
+    #[error("tokenization failed: {0}")]
+    Encode(#[source] tokenizers::Error),
+}

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

@@ -0,0 +1,70 @@
+#![forbid(unsafe_code)]
+
+mod byte_level;
+mod cl100k;
+mod error;
+mod o200k;
+mod scanner;
+mod tiktoken;
+mod unicode_classes;
+
+use byte_level::ByteLevelCounter;
+use scanner::{SplitPattern, TiktokenCounter};
+
+pub use error::Error;
+
+enum Encoder {
+    HuggingFace {
+        tokenizer: Box<tokenizers::Tokenizer>,
+        byte_level: Option<ByteLevelCounter>,
+    },
+    Tiktoken(TiktokenCounter),
+}
+
+pub struct FastTokenizer(Encoder);
+
+impl FastTokenizer {
+    pub fn from_json(json: &str) -> Result<Self, Error> {
+        let tokenizer = json.parse::<tokenizers::Tokenizer>().map_err(Error::Load)?;
+        let byte_level = ByteLevelCounter::detect(&tokenizer);
+        Ok(Self(Encoder::HuggingFace {
+            tokenizer: Box::new(tokenizer),
+            byte_level,
+        }))
+    }
+
+    pub fn from_cl100k_ranks(ranks: &str) -> Result<Self, Error> {
+        Self::from_ranks(SplitPattern::Cl100k, ranks)
+    }
+
+    pub fn from_o200k_ranks(ranks: &str) -> Result<Self, Error> {
+        Self::from_ranks(SplitPattern::O200k, ranks)
+    }
+
+    fn from_ranks(split: SplitPattern, ranks: &str) -> Result<Self, Error> {
+        TiktokenCounter::from_ranks(split, ranks)
+            .map(Encoder::Tiktoken)
+            .map(Self)
+    }
+
+    pub fn count_tokens(&self, text: &str) -> Result<usize, Error> {
+        match &self.0 {
+            Encoder::Tiktoken(counter) => Ok(counter.count(text)),
+            Encoder::HuggingFace {
+                tokenizer,
+                byte_level,
+            } => {
+                if let Some(count) = byte_level
+                    .as_ref()
+                    .and_then(|counter| counter.count(tokenizer, text))
+                {
+                    return Ok(count);
+                }
+                tokenizer
+                    .encode_fast(text, true)
+                    .map(|encoding| encoding.len())
+                    .map_err(Error::Encode)
+            }
+        }
+    }
+}

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

@@ -0,0 +1,240 @@
+//! tiktoken's byte-level BPE: a rank file of `base64(token) rank` lines and
+//! the merge loop that turns one regex piece into tokens. The merge order is
+//! tiktoken's (lowest rank first, leftmost pair on ties) so the token count is
+//! identical, but pairs are tracked in a heap so a long piece costs
+//! `O(n log n)` instead of tiktoken's `O(n^2)`.
+
+use std::cmp::Reverse;
+use std::collections::BinaryHeap;
+
+use base64::Engine;
+use base64::engine::general_purpose::STANDARD;
+use rustc_hash::FxHashMap;
+
+use crate::Error;
+
+type Rank = u32;
+
+const NO_RANK: Rank = Rank::MAX;
+const END: usize = usize::MAX;
+
+pub(super) struct MergeRanks(FxHashMap<Box<[u8]>, Rank>);
+
+impl MergeRanks {
+    pub(super) fn parse(text: &str) -> Result<Self, Error> {
+        let ranks = text
+            .lines()
+            .filter(|line| !line.is_empty())
+            .map(parse_line)
+            .collect::<Result<FxHashMap<_, _>, _>>()?;
+        if let Some(byte) = (0..=u8::MAX).find(|byte| !ranks.contains_key(&[*byte][..])) {
+            return Err(Error::Ranks(format!("byte 0x{byte:02X} has no token")));
+        }
+        Ok(Self(ranks))
+    }
+
+    fn rank(&self, bytes: &[u8]) -> Rank {
+        self.0.get(bytes).copied().unwrap_or(NO_RANK)
+    }
+
+    /// Token count of one regex piece, as `encode_ordinary` would produce.
+    pub(super) fn count_piece(&self, piece: &[u8], scratch: &mut MergeScratch) -> usize {
+        if piece.len() < 2 || self.0.contains_key(piece) {
+            return 1;
+        }
+        scratch.reset(piece.len());
+        for start in 0..piece.len() - 1 {
+            scratch.set_rank(start, self.rank(&piece[start..start + 2]));
+        }
+        let mut parts = piece.len();
+        while let Some(Reverse((rank, start))) = scratch.heap.pop() {
+            if scratch.next[start] == END || scratch.rank[start] != rank {
+                continue;
+            }
+            let merged = scratch.next[start];
+            let after = scratch.next[merged];
+            scratch.next[merged] = END;
+            scratch.next[start] = after;
+            parts -= 1;
+            if after < piece.len() {
+                scratch.prev[after] = start;
+                scratch.set_rank(start, self.rank(&piece[start..scratch.end(after)]));
+            } else {
+                scratch.rank[start] = NO_RANK;
+            }
+            let before = scratch.prev[start];
+            if before != END {
+                scratch.set_rank(before, self.rank(&piece[before..scratch.end(start)]));
+            }
+        }
+        parts
+    }
+}
+
+fn parse_line(line: &str) -> Result<(Box<[u8]>, Rank), Error> {
+    let (token, rank) = line
+        .split_once(' ')
+        .ok_or_else(|| Error::Ranks(format!("line without a rank: {line:?}")))?;
+    let bytes = STANDARD
+        .decode(token)
+        .map_err(|error| Error::Ranks(format!("token is not base64: {error}")))?;
+    let rank = rank
+        .parse()
+        .map_err(|error| Error::Ranks(format!("rank is not an integer: {error}")))?;
+    if rank == NO_RANK {
+        return Err(Error::Ranks(format!("rank {rank} is reserved")));
+    }
+    Ok((bytes.into_boxed_slice(), rank))
+}
+
+/// Buffers reused across the pieces of one text. Parts are addressed by the
+/// byte offset they start at, which also gives the leftmost-pair tie break.
+#[derive(Default)]
+pub(super) struct MergeScratch {
+    next: Vec<usize>,
+    prev: Vec<usize>,
+    rank: Vec<Rank>,
+    heap: BinaryHeap<Reverse<(Rank, usize)>>,
+}
+
+impl MergeScratch {
+    fn reset(&mut self, len: usize) {
+        self.next.clear();
+        self.next.extend(1..=len);
+        self.prev.clear();
+        self.prev.push(END);
+        self.prev.extend(0..len - 1);
+        self.rank.clear();
+        self.rank.resize(len, NO_RANK);
+        self.heap.clear();
+    }
+
+    fn end(&self, start: usize) -> usize {
+        self.next[start]
+    }
+
+    fn set_rank(&mut self, start: usize, rank: Rank) {
+        self.rank[start] = rank;
+        if rank != NO_RANK {
+            self.heap.push(Reverse((rank, start)));
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use rand::rngs::StdRng;
+    use rand::{Rng, SeedableRng};
+
+    use super::*;
+
+    fn ranks() -> MergeRanks {
+        let path = concat!(
+            env!("CARGO_MANIFEST_DIR"),
+            "/../../../litellm/litellm_core_utils/tokenizers/9b5ad71b2ce5302211f9c61530b329a4922fc6a4"
+        );
+        MergeRanks::parse(&std::fs::read_to_string(path).expect("cl100k rank file is in the repo"))
+            .expect("rank file parses")
+    }
+
+    /// tiktoken's `_byte_pair_merge`, transcribed, as the reference.
+    fn reference_count(ranks: &MergeRanks, piece: &[u8]) -> usize {
+        if piece.len() < 2 || ranks.0.contains_key(piece) {
+            return 1;
+        }
+        let mut parts: Vec<(usize, Rank)> = (0..piece.len() - 1)
+            .map(|index| (index, ranks.rank(&piece[index..index + 2])))
+            .chain([(piece.len() - 1, NO_RANK), (piece.len(), NO_RANK)])
+            .collect();
+        let get_rank = |parts: &[(usize, Rank)], index: usize| {
+            if index + 3 < parts.len() {
+                ranks.rank(&piece[parts[index].0..parts[index + 3].0])
+            } else {
+                NO_RANK
+            }
+        };
+        loop {
+            let Some(index) = parts[..parts.len() - 1]
+                .iter()
+                .enumerate()
+                .filter(|(_, (_, rank))| *rank != NO_RANK)
+                .min_by_key(|(index, (_, rank))| (*rank, *index))
+                .map(|(index, _)| index)
+            else {
+                return parts.len() - 1;
+            };
+            if index > 0 {
+                parts[index - 1].1 = get_rank(&parts, index - 1);
+            }
+            parts[index].1 = get_rank(&parts, index);
+            parts.remove(index + 1);
+        }
+    }
+
+    #[test]
+    fn every_byte_is_a_token() {
+        let ranks = ranks();
+        assert_eq!(ranks.0.len(), 100_256);
+        assert!((0..=u8::MAX).all(|byte| ranks.rank(&[byte]) != NO_RANK));
+    }
+
+    #[test]
+    fn heap_merge_matches_tiktokens_merge_loop() {
+        let ranks = ranks();
+        let mut scratch = MergeScratch::default();
+        let mut rng = StdRng::seed_from_u64(99);
+        let alphabet = b" abcdeorstn.,'\n\xc3\xa9\xe2\x82\xac0123";
+        for _ in 0..20_000 {
+            let piece: Vec<u8> = (0..rng.gen_range(1..24))
+                .map(|_| alphabet[rng.gen_range(0..alphabet.len())])
+                .collect();
+            assert_eq!(
+                ranks.count_piece(&piece, &mut scratch),
+                reference_count(&ranks, &piece),
+                "piece {:?}",
+                String::from_utf8_lossy(&piece)
+            );
+        }
+    }
+
+    #[test]
+    fn long_repeated_runs_cost_close_to_linear() {
+        let ranks = ranks();
+        let mut scratch = MergeScratch::default();
+        let mut time = |len: usize| {
+            let piece = vec![b' '; len];
+            let started = std::time::Instant::now();
+            assert!(ranks.count_piece(&piece, &mut scratch) > 0);
+            started.elapsed()
+        };
+        let small = (0..3).map(|_| time(1 << 14)).min().unwrap();
+        let large = time(1 << 18);
+        assert!(
+            large < small * 64,
+            "{small:?} for 2^14 bytes, {large:?} for 2^18"
+        );
+    }
+
+    #[test]
+    fn reserved_merge_rank_is_rejected() {
+        let bytes = (0..=u8::MAX)
+            .map(|byte| format!("{} {byte}\n", STANDARD.encode([byte])))
+            .collect::<String>();
+        let rank_file = format!("{bytes}{} {NO_RANK}\n", STANDARD.encode(b"ab"));
+        assert!(matches!(
+            MergeRanks::parse(&rank_file),
+            Err(Error::Ranks(_))
+        ));
+        let valid_rank_file = format!("{bytes}{} {}\n", STANDARD.encode(b"ab"), NO_RANK - 1);
+        let ranks = MergeRanks::parse(&valid_rank_file).unwrap();
+        assert_eq!(ranks.count_piece(b"aab", &mut MergeScratch::default()), 2);
+    }
+
+    #[test]
+    fn malformed_rank_files_are_rejected() {
+        assert!(MergeRanks::parse("IQ==").is_err());
+        assert!(MergeRanks::parse("IQ== x").is_err());
+        assert!(MergeRanks::parse("!!! 1").is_err());
+        assert!(MergeRanks::parse("IQ== 1").is_err());
+    }
+}